From a38f4ef3b7cff206a74e298085f336fde41f778b Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 29 Apr 2026 17:56:35 -0700 Subject: [PATCH 01/15] update ww singleton to have time and angle dependence --- mcdc/object_/technique.py | 74 +++++++++++++++++++++++++++++---------- 1 file changed, 55 insertions(+), 19 deletions(-) diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index e2adfb80..e2df61ca 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -1,7 +1,7 @@ from numpy.typing import NDArray import numpy as np from typing import Annotated -from mcdc.constant import INF +from mcdc.constant import INF, PI, PI_HALF from mcdc.object_.base import ObjectSingleton from mcdc.object_.mesh import MeshBase, MeshUniform from mcdc.print_ import print_error @@ -82,9 +82,17 @@ class WeightWindows(ObjectSingleton): active: bool + # time + time_bounds: NDArray[np.float64] + Nt: int # energy energy_bounds: NDArray[np.float64] Ne: int + # angle + mu_bounds: NDArray[np.float64] + Nmu: int + azi_bounds: NDArray[np.float64] + Na: int # space mesh: MeshBase Nx: int @@ -92,29 +100,41 @@ class WeightWindows(ObjectSingleton): Nz: int # arrays of ww params - lower_weights: Annotated[NDArray[np.float64], ("Ne", "Nx", "Ny", "Nz")] - target_weights: Annotated[NDArray[np.float64], ("Ne", "Nx", "Ny", "Nz")] - upper_weights: Annotated[NDArray[np.float64], ("Ne", "Nx", "Ny", "Nz")] + lower_weights: Annotated[NDArray[np.float64], ("Nt", "Ne", "Nmu", "Na", "Nx", "Ny", "Nz")] + target_weights: Annotated[NDArray[np.float64], ("Nt", "Ne", "Nmu", "Na", "Nx", "Ny", "Nz")] + upper_weights: Annotated[NDArray[np.float64], ("Nt", "Ne", "Nmu", "Na", "Nx", "Ny", "Nz")] def __init__(self): self.active = False - self.energy_bounds = np.array([0.0, 1.0]) + self.time_bounds = np.array([0.0, INF]) + self.Nt = 1 + self.azi_bounds = np.array([-PI, PI]) + self.Na = 1 + self.mu_bounds = np.array([-1.0, 1.0]) + self.Nmu = 1 + self.energy_bounds = np.array([-0.5, INF]) self.Ne = 1 self.mesh_ID = -1 # skirt around having to create a MeshBase instance self.Nx, self.Ny, self.Nz = 1, 1, 1 self.Nt = 1 - shape = (self.Ne, self.Nx, self.Ny, self.Nz) + shape = (self.Nt, self.Ne, self.Nmu, self.Na, self.Nx, self.Ny, self.Nz) self.lower_weights = np.array([1.0]).reshape(*shape) self.target_weights = np.array([1.0]).reshape(*shape) self.upper_weights = np.array([1.0]).reshape(*shape) - def __call__(self, weight_windows, mesh=None, energy=None): + def __call__(self, weight_windows, mesh=None, energy=None, mu=None, azimuthal=None, time=None): # fill in defaults if mesh is None: mesh = MeshUniform() if energy is None: # usable for both groups and max energy energy = np.array([-0.5, INF]) + if mu is None: + mu = np.array([-1.0, 1.0]) + if azimuthal is None: + azimuthal = np.array([-PI, PI]) + if time is None: + time = np.array([0, INF]) # get mesh size match mesh.label: @@ -130,31 +150,38 @@ def __call__(self, weight_windows, mesh=None, energy=None): print_error( f"{type(mesh).__name__} is not supported for weight windows" ) - # validate energy as strictly increasing - if not (np.diff(energy) > 0).all(): - print_error("Energy bounds must be strictly increasing") - # get energy size - if len(energy.shape) != 1: - print_error( - f"Invalid shape for energy; expected 1D got {len(energy.shape)}D" - ) + # validate energy and get size + self.__check_array(energy, "Energy") ne = energy.shape[0] - 1 + # validate mu and get size + self.__check_array(mu, "Mu") + nmu = mu.shape[0] - 1 + # validate azimuthal and get size + self.__check_array(azimuthal, "Azimuthal") + na = azimuthal.shape[0] - 1 + # validate time and get size + self.__check_array(time, "Time") + nt = time.shape[0] - 1 # check correct shape - mesh_shape = (nx, ny, nz) + expected_shape = (nt, ne, nmu, na, nx, ny, nz, 3) ww_shape = weight_windows.shape - expected_shape = (ne, *mesh_shape, 3) if ww_shape != expected_shape: print_error( f"Weight window array has shape {ww_shape}, but expected {expected_shape}" ) self.active = True + self.time_bounds = time + self.Nt = nt self.energy_bounds = energy self.Ne = ne + self.mu_bounds = mu + self.Nmu = nmu + self.azi_bounds = azimuthal + self.Na = na self.mesh = mesh - self.Nx, self.Ny, self.Nz = mesh_shape - shape = (self.Ne, *mesh_shape) + self.Nx, self.Ny, self.Nz = (nx, ny, nz) self.lower_weights = weight_windows[..., 0] self.target_weights = weight_windows[..., 1] self.upper_weights = weight_windows[..., 2] @@ -173,6 +200,15 @@ def __call__(self, weight_windows, mesh=None, energy=None): "Target weight can not be greater than the upper bound weight for any weight window" ) + @staticmethod + def __check_array(array: NDArray[np.float64], name: str): + if not (np.diff(array) > 0).all(): + print_error(f"{name} bounds must be strictly increasing") + if len(array.shape) != 1: + print_error( + f"Invalid shape for {name} bounds; expected 1D got {len(array.shape)}D" + ) + # ====================================================================================== # Population control From 7004dd6584a4e338594a610689a4740bd974f780 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 29 Apr 2026 17:57:14 -0700 Subject: [PATCH 02/15] add 7d accessor for ww --- mcdc/code_factory/numba_objects_generator.py | 37 ++++++ mcdc/mcdc_get/weight_windows.py | 126 ++++++++++++++++--- mcdc/mcdc_set/weight_windows.py | 126 ++++++++++++++++--- 3 files changed, 259 insertions(+), 30 deletions(-) diff --git a/mcdc/code_factory/numba_objects_generator.py b/mcdc/code_factory/numba_objects_generator.py index c466b92d..28ab1b34 100644 --- a/mcdc/code_factory/numba_objects_generator.py +++ b/mcdc/code_factory/numba_objects_generator.py @@ -1021,6 +1021,14 @@ def generate_mcdc_access(targets): text_setter += _accessor_4d_element( object_name, attribute_name, shape[1], shape[2], shape[3], True ) + elif len(shape) == 7: + text_getter += _accessor_7d_element( + object_name, attribute_name, shape[1], shape[2], shape[3], shape[4], shape[5], shape[6] + ) + + text_setter += _accessor_7d_element( + object_name, attribute_name, shape[1], shape[2], shape[3], shape[4], shape[5], shape[6], True + ) text_getter += _accessor_chunk(object_name, attribute_name) text_setter += _accessor_chunk(object_name, attribute_name, True) @@ -1184,6 +1192,35 @@ def _accessor_4d_element( return text +def _accessor_7d_element( + object_name, + attribute_name, + stride_2, + stride_3, + stride_4, + stride_5, + stride_6, + stride_7, + setter=False, +): + text = f"@njit\n" + if setter: + text += f"def {attribute_name}(index_1, index_2, index_3, index_4, index_5, index_6, index_7, {object_name}, data, value):\n" + else: + text += f"def {attribute_name}(index_1, index_2, index_3, index_4, index_5, index_6, index_7, {object_name}, data):\n" + text += f' offset = {object_name}["{attribute_name}_offset"]\n' + text += f' stride_2 = {object_name}["{stride_2}"]\n' + text += f' stride_3 = {object_name}["{stride_3}"]\n' + text += f' stride_4 = {object_name}["{stride_4}"]\n' + text += f' stride_5 = {object_name}["{stride_5}"]\n' + text += f' stride_6 = {object_name}["{stride_6}"]\n' + text += f' stride_7 = {object_name}["{stride_7}"]\n' + if setter: + text += f" data[offset + index_1 * stride_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_4 * stride_5 * stride_6 * stride_7 + index_5 * stride_6 * stride_7 + index_6 * stride_7 + index_7] = value\n\n\n" + else: + text += f" return data[offset + index_1 * stride_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_4 * stride_5 * stride_6 * stride_7 + index_5 * stride_6 * stride_7 + index_6 * stride_7 + index_7]\n\n\n" + return text + # ====================================================================================== # Misc. # ====================================================================================== diff --git a/mcdc/mcdc_get/weight_windows.py b/mcdc/mcdc_get/weight_windows.py index b1b7cbec..f15b61d2 100644 --- a/mcdc/mcdc_get/weight_windows.py +++ b/mcdc/mcdc_get/weight_windows.py @@ -3,6 +3,35 @@ from numba import njit +@njit +def time_bounds(index, weight_windows, data): + offset = weight_windows["time_bounds_offset"] + return data[offset + index] + + +@njit +def time_bounds_all(weight_windows, data): + start = weight_windows["time_bounds_offset"] + size = weight_windows["time_bounds_length"] + end = start + size + return data[start:end] + + +@njit +def time_bounds_last(weight_windows, data): + start = weight_windows["time_bounds_offset"] + size = weight_windows["time_bounds_length"] + end = start + size + return data[end - 1] + + +@njit +def time_bounds_chunk(start, length, weight_windows, data): + start += weight_windows["time_bounds_offset"] + end = start + length + return data[start:end] + + @njit def energy_bounds(index, weight_windows, data): offset = weight_windows["energy_bounds_offset"] @@ -33,12 +62,73 @@ def energy_bounds_chunk(start, length, weight_windows, data): @njit -def lower_weights(index_1, index_2, index_3, index_4, weight_windows, data): +def mu_bounds(index, weight_windows, data): + offset = weight_windows["mu_bounds_offset"] + return data[offset + index] + + +@njit +def mu_bounds_all(weight_windows, data): + start = weight_windows["mu_bounds_offset"] + size = weight_windows["mu_bounds_length"] + end = start + size + return data[start:end] + + +@njit +def mu_bounds_last(weight_windows, data): + start = weight_windows["mu_bounds_offset"] + size = weight_windows["mu_bounds_length"] + end = start + size + return data[end - 1] + + +@njit +def mu_bounds_chunk(start, length, weight_windows, data): + start += weight_windows["mu_bounds_offset"] + end = start + length + return data[start:end] + + +@njit +def azi_bounds(index, weight_windows, data): + offset = weight_windows["azi_bounds_offset"] + return data[offset + index] + + +@njit +def azi_bounds_all(weight_windows, data): + start = weight_windows["azi_bounds_offset"] + size = weight_windows["azi_bounds_length"] + end = start + size + return data[start:end] + + +@njit +def azi_bounds_last(weight_windows, data): + start = weight_windows["azi_bounds_offset"] + size = weight_windows["azi_bounds_length"] + end = start + size + return data[end - 1] + + +@njit +def azi_bounds_chunk(start, length, weight_windows, data): + start += weight_windows["azi_bounds_offset"] + end = start + length + return data[start:end] + + +@njit +def lower_weights(index_1, index_2, index_3, index_4, index_5, index_6, index_7, weight_windows, data): offset = weight_windows["lower_weights_offset"] - stride_2 = weight_windows["Nx"] - stride_3 = weight_windows["Ny"] - stride_4 = weight_windows["Nz"] - return data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4] + stride_2 = weight_windows["Ne"] + stride_3 = weight_windows["Nmu"] + stride_4 = weight_windows["Na"] + stride_5 = weight_windows["Nx"] + stride_6 = weight_windows["Ny"] + stride_7 = weight_windows["Nz"] + return data[offset + index_1 * stride_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_4 * stride_5 * stride_6 * stride_7 + index_5 * stride_6 * stride_7 + index_6 * stride_7 + index_7] @njit @@ -49,12 +139,15 @@ def lower_weights_chunk(start, length, weight_windows, data): @njit -def target_weights(index_1, index_2, index_3, index_4, weight_windows, data): +def target_weights(index_1, index_2, index_3, index_4, index_5, index_6, index_7, weight_windows, data): offset = weight_windows["target_weights_offset"] - stride_2 = weight_windows["Nx"] - stride_3 = weight_windows["Ny"] - stride_4 = weight_windows["Nz"] - return data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4] + stride_2 = weight_windows["Ne"] + stride_3 = weight_windows["Nmu"] + stride_4 = weight_windows["Na"] + stride_5 = weight_windows["Nx"] + stride_6 = weight_windows["Ny"] + stride_7 = weight_windows["Nz"] + return data[offset + index_1 * stride_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_4 * stride_5 * stride_6 * stride_7 + index_5 * stride_6 * stride_7 + index_6 * stride_7 + index_7] @njit @@ -65,12 +158,15 @@ def target_weights_chunk(start, length, weight_windows, data): @njit -def upper_weights(index_1, index_2, index_3, index_4, weight_windows, data): +def upper_weights(index_1, index_2, index_3, index_4, index_5, index_6, index_7, weight_windows, data): offset = weight_windows["upper_weights_offset"] - stride_2 = weight_windows["Nx"] - stride_3 = weight_windows["Ny"] - stride_4 = weight_windows["Nz"] - return data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4] + stride_2 = weight_windows["Ne"] + stride_3 = weight_windows["Nmu"] + stride_4 = weight_windows["Na"] + stride_5 = weight_windows["Nx"] + stride_6 = weight_windows["Ny"] + stride_7 = weight_windows["Nz"] + return data[offset + index_1 * stride_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_4 * stride_5 * stride_6 * stride_7 + index_5 * stride_6 * stride_7 + index_6 * stride_7 + index_7] @njit diff --git a/mcdc/mcdc_set/weight_windows.py b/mcdc/mcdc_set/weight_windows.py index e6a2ab8a..5a77acea 100644 --- a/mcdc/mcdc_set/weight_windows.py +++ b/mcdc/mcdc_set/weight_windows.py @@ -3,6 +3,35 @@ from numba import njit +@njit +def time_bounds(index, weight_windows, data, value): + offset = weight_windows["time_bounds_offset"] + data[offset + index] = value + + +@njit +def time_bounds_all(weight_windows, data, value): + start = weight_windows["time_bounds_offset"] + size = weight_windows["time_bounds_length"] + end = start + size + data[start:end] = value + + +@njit +def time_bounds_last(weight_windows, data, value): + start = weight_windows["time_bounds_offset"] + size = weight_windows["time_bounds_length"] + end = start + size + data[end - 1] = value + + +@njit +def time_bounds_chunk(start, length, weight_windows, data, value): + start += weight_windows["time_bounds_offset"] + end = start + length + data[start:end] = value + + @njit def energy_bounds(index, weight_windows, data, value): offset = weight_windows["energy_bounds_offset"] @@ -33,12 +62,73 @@ def energy_bounds_chunk(start, length, weight_windows, data, value): @njit -def lower_weights(index_1, index_2, index_3, index_4, weight_windows, data, value): +def mu_bounds(index, weight_windows, data, value): + offset = weight_windows["mu_bounds_offset"] + data[offset + index] = value + + +@njit +def mu_bounds_all(weight_windows, data, value): + start = weight_windows["mu_bounds_offset"] + size = weight_windows["mu_bounds_length"] + end = start + size + data[start:end] = value + + +@njit +def mu_bounds_last(weight_windows, data, value): + start = weight_windows["mu_bounds_offset"] + size = weight_windows["mu_bounds_length"] + end = start + size + data[end - 1] = value + + +@njit +def mu_bounds_chunk(start, length, weight_windows, data, value): + start += weight_windows["mu_bounds_offset"] + end = start + length + data[start:end] = value + + +@njit +def azi_bounds(index, weight_windows, data, value): + offset = weight_windows["azi_bounds_offset"] + data[offset + index] = value + + +@njit +def azi_bounds_all(weight_windows, data, value): + start = weight_windows["azi_bounds_offset"] + size = weight_windows["azi_bounds_length"] + end = start + size + data[start:end] = value + + +@njit +def azi_bounds_last(weight_windows, data, value): + start = weight_windows["azi_bounds_offset"] + size = weight_windows["azi_bounds_length"] + end = start + size + data[end - 1] = value + + +@njit +def azi_bounds_chunk(start, length, weight_windows, data, value): + start += weight_windows["azi_bounds_offset"] + end = start + length + data[start:end] = value + + +@njit +def lower_weights(index_1, index_2, index_3, index_4, index_5, index_6, index_7, weight_windows, data, value): offset = weight_windows["lower_weights_offset"] - stride_2 = weight_windows["Nx"] - stride_3 = weight_windows["Ny"] - stride_4 = weight_windows["Nz"] - data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4] = value + stride_2 = weight_windows["Ne"] + stride_3 = weight_windows["Nmu"] + stride_4 = weight_windows["Na"] + stride_5 = weight_windows["Nx"] + stride_6 = weight_windows["Ny"] + stride_7 = weight_windows["Nz"] + data[offset + index_1 * stride_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_4 * stride_5 * stride_6 * stride_7 + index_5 * stride_6 * stride_7 + index_6 * stride_7 + index_7] = value @njit @@ -49,12 +139,15 @@ def lower_weights_chunk(start, length, weight_windows, data, value): @njit -def target_weights(index_1, index_2, index_3, index_4, weight_windows, data, value): +def target_weights(index_1, index_2, index_3, index_4, index_5, index_6, index_7, weight_windows, data, value): offset = weight_windows["target_weights_offset"] - stride_2 = weight_windows["Nx"] - stride_3 = weight_windows["Ny"] - stride_4 = weight_windows["Nz"] - data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4] = value + stride_2 = weight_windows["Ne"] + stride_3 = weight_windows["Nmu"] + stride_4 = weight_windows["Na"] + stride_5 = weight_windows["Nx"] + stride_6 = weight_windows["Ny"] + stride_7 = weight_windows["Nz"] + data[offset + index_1 * stride_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_4 * stride_5 * stride_6 * stride_7 + index_5 * stride_6 * stride_7 + index_6 * stride_7 + index_7] = value @njit @@ -65,12 +158,15 @@ def target_weights_chunk(start, length, weight_windows, data, value): @njit -def upper_weights(index_1, index_2, index_3, index_4, weight_windows, data, value): +def upper_weights(index_1, index_2, index_3, index_4, index_5, index_6, index_7, weight_windows, data, value): offset = weight_windows["upper_weights_offset"] - stride_2 = weight_windows["Nx"] - stride_3 = weight_windows["Ny"] - stride_4 = weight_windows["Nz"] - data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4] = value + stride_2 = weight_windows["Ne"] + stride_3 = weight_windows["Nmu"] + stride_4 = weight_windows["Na"] + stride_5 = weight_windows["Nx"] + stride_6 = weight_windows["Ny"] + stride_7 = weight_windows["Nz"] + data[offset + index_1 * stride_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_4 * stride_5 * stride_6 * stride_7 + index_5 * stride_6 * stride_7 + index_6 * stride_7 + index_7] = value @njit From ff8219ec762d77df57d2b4dcf3076327bc4e2aaf Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 29 Apr 2026 17:57:38 -0700 Subject: [PATCH 03/15] add numba types for ww with new indices --- mcdc/numba_types.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/mcdc/numba_types.py b/mcdc/numba_types.py index 056ce177..d4190c1f 100644 --- a/mcdc/numba_types.py +++ b/mcdc/numba_types.py @@ -577,9 +577,18 @@ weight_windows = into_dtype([ ('active', bool), + ('time_bounds_offset', int64), + ('time_bounds_length', int64), + ('Nt', int64), ('energy_bounds_offset', int64), ('energy_bounds_length', int64), ('Ne', int64), + ('mu_bounds_offset', int64), + ('mu_bounds_length', int64), + ('Nmu', int64), + ('azi_bounds_offset', int64), + ('azi_bounds_length', int64), + ('Na', int64), ('mesh_ID', int64), ('Nx', int64), ('Ny', int64), From f770c874f928fcd5090ac6f106911f97cdfb506e Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 29 Apr 2026 17:58:06 -0700 Subject: [PATCH 04/15] update get indices method for ww, quick regression test fix --- mcdc/transport/technique.py | 15 ++++++++++++++- test/regression/basic_weight_windows/input.py | 2 +- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index 5de63799..05f8c2fc 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -148,6 +148,10 @@ def get_ww_indices(particle_container, ww_obj, simulation, data): """ particle = particle_container[0] + # get time index + time_bounds = ww_get.time_bounds_all(ww_obj, data) + it = util.find_bin(particle["t"], time_bounds) + # get energy index energy_bounds = ww_get.energy_bounds_all(ww_obj, data) if simulation["settings"]["neutron_multigroup_mode"]: @@ -156,11 +160,20 @@ def get_ww_indices(particle_container, ww_obj, simulation, data): energy = particle["E"] ie = util.find_bin(energy, energy_bounds) + # get mu index + mu_bounds = ww_get.mu_bounds_all(ww_obj, data) + imu = util.find_bin(particle["uz"], mu_bounds) + + # get azimuthal index + azi_bounds = ww_get.azi_bounds_all(ww_obj, data) + azimuthal = np.arctan2(particle["uy"], particle["ux"]) + ia = util.find_bin(azimuthal, azi_bounds) + # get spatial index mesh = simulation["meshes"][ww_obj["mesh_ID"]] idx, idy, idz = get_mesh_indices(particle_container, mesh, simulation, data) - return (ie, idx, idy, idz) + return (it, ie, imu, ia, idx, idy, idz) @njit diff --git a/test/regression/basic_weight_windows/input.py b/test/regression/basic_weight_windows/input.py index 38d3e980..84c4735c 100644 --- a/test/regression/basic_weight_windows/input.py +++ b/test/regression/basic_weight_windows/input.py @@ -53,7 +53,7 @@ mcdc.Tally(mesh=mesh, scores=["flux"]) # Weight windows -ww_array = np.ones((1, 20, 20, 1, 3)) +ww_array = np.ones((1, 1, 1, 1, 20, 20, 1, 3)) # Actual bounds are set to arbitrary numbers ww_array[..., 0] = 0.55 # Forces roulette on split particles from 1.0 ww_array[..., 1] = 0.7 # arbitrary in the middle From 90cfb2a65adba583d41e4eb600c8f29c3fb884c3 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 29 Apr 2026 17:58:26 -0700 Subject: [PATCH 05/15] updated unit tests for ww with new indices --- .../transport/technique/weight_windows.py | 90 ++++++++++++------- 1 file changed, 56 insertions(+), 34 deletions(-) diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py index 6ad4eb44..3313b0c2 100644 --- a/test/unit/transport/technique/weight_windows.py +++ b/test/unit/transport/technique/weight_windows.py @@ -42,12 +42,15 @@ def make_ww_model_params(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): mesh, N = make_mesh() Ne = 1 + Nt = 1 + Nmu = 1 + Na = 1 if mess_up_size: - ww_array = np.ones((Ne, N, N, 4, 3)) + ww_array = np.ones((Nt, Ne, Nmu, Na, N, N, 4, 3)) else: # global assign for simplicity - ww_array = np.ones((Ne, N, N, N, 3)) + ww_array = np.ones((Nt, Ne, Nmu, Na, N, N, N, 3)) ww_array[..., 0] = lower ww_array[..., 1] = target ww_array[..., 2] = upper @@ -64,20 +67,29 @@ def make_ww_model_distinct(): mesh, N = make_mesh() energy = np.linspace(0.0, 6.0, 7) Ne = 6 + time = np.linspace(0.0, 4.0, 5) + Nt = 4 + mu = np.linspace(-1.0, 1.0, 5) + Nmu = 4 + azimuthal = np.linspace(-np.pi, np.pi, 5) + Na = 4 - ww_array = np.empty((Ne, N, N, N, 3)) + ww_array = np.empty((Nt, Ne, Nmu, Na, N, N, N, 3)) # value at index is related to index, easy to predict during later test - for e in range(Ne): - for i in range(N): - for j in range(N): - for k in range(N): - val = 1000 * e + 100 * i + 10 * j + k + 1 - ww_array[e, i, j, k, 0] = val - ww_array[e, i, j, k, 1] = 10000 + val - ww_array[e, i, j, k, 2] = 20000 + val - - mcdc.simulation.weight_windows(ww_array, mesh=mesh, energy=energy) + for t in range(Nt): + for e in range(Ne): + for m in range(Nmu): + for a in range(Na): + for i in range(N): + for j in range(N): + for k in range(N): + val = 1_000_000*t + 100_000*e + 10_000*m + 1_000*a + 100 * i + 10 * j + k + 1 + ww_array[t, e, m, a, i, j, k, 0] = val + ww_array[t, e, m, a, i, j, k, 1] = 10_000_000 + val + ww_array[t, e, m, a, i, j, k, 2] = 20_000_000 + val + + mcdc.simulation.weight_windows(ww_array, mesh=mesh, energy=energy, time=time, mu=mu, azimuthal=azimuthal) mcdc_container, data = preparation() return mcdc_container[0], data @@ -94,7 +106,7 @@ def make_ww_model_distinct(): # incorrect size ( {"mess_up_size": True}, - "Weight window array has shape (1, 3, 3, 4, 3), but expected (1, 3, 3, 3, 3)", + "Weight window array has shape (1, 1, 1, 1, 3, 3, 4, 3), but expected (1, 1, 1, 1, 3, 3, 3, 3)", ), # negative lower ( @@ -207,26 +219,36 @@ def test_query_weight_window(): dx, dy, dz = pitch / N, pitch / N, height / N # hardcode energy params + times = np.linspace(0.5, 3.5, 4) energies = np.linspace(0.5, 5.5, 6) + mus = np.linspace(-0.75, 0.75, 4) + azimuthals = np.linspace(-3.0 * np.pi / 4.0, 3.0 * np.pi / 4.0, 4) # loop over all bins, check query against expected ww - for ne, energy in enumerate(energies): - for ix in range(nx): - for iy in range(ny): - for iz in range(nz): - # put particle in center of current mesh bin - p[0]["x"] = xmin + dx * (ix + 0.5) - p[0]["y"] = ymin + dy * (iy + 0.5) - p[0]["z"] = zmin + dz * (iz + 0.5) - # assign energy to be in center of bins - p[0]["E"] = energy - - # query and predict - lower, target, upper = query_weight_window(p, simulation, data) - exp_lower = 1000 * ne + 100 * ix + 10 * iy + iz + 1 - exp_target = 10000 + exp_lower - exp_upper = 20000 + exp_lower - - assert lower == exp_lower - assert target == exp_target - assert upper == exp_upper + for it, time in enumerate(times): + for ie, energy in enumerate(energies): + for imu, mu in enumerate(mus): + for ia, azimuthal in enumerate(azimuthals): + for ix in range(nx): + for iy in range(ny): + for iz in range(nz): + # put particle in center of current mesh bin + p[0]["x"] = xmin + dx * (ix + 0.5) + p[0]["y"] = ymin + dy * (iy + 0.5) + p[0]["z"] = zmin + dz * (iz + 0.5) + # assign params to be in center of bins + p[0]["t"] = time + p[0]["E"] = energy + p[0]["ux"] = np.sqrt(1.0 - mu**2) * np.cos(azimuthal) + p[0]["uy"] = np.sqrt(1.0 - mu**2) * np.sin(azimuthal) + p[0]["uz"] = mu + + # query and predict + lower, target, upper = query_weight_window(p, simulation, data) + exp_lower = 1_000_000*it + 100_000*ie + 10_000*imu + 1_000*ia + 100 * ix + 10 * iy + iz + 1 + exp_target = 10_000_000 + exp_lower + exp_upper = 20_000_000 + exp_lower + + assert lower == exp_lower + assert target == exp_target + assert upper == exp_upper From 3fc83b3a4f6f13492e684fec4408b8a56c265caa Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 29 Apr 2026 17:59:34 -0700 Subject: [PATCH 06/15] back in black --- mcdc/code_factory/numba_objects_generator.py | 20 +++- mcdc/object_/technique.py | 18 +++- .../transport/technique/weight_windows.py | 98 ++++++++++++------- 3 files changed, 91 insertions(+), 45 deletions(-) diff --git a/mcdc/code_factory/numba_objects_generator.py b/mcdc/code_factory/numba_objects_generator.py index 28ab1b34..96e16669 100644 --- a/mcdc/code_factory/numba_objects_generator.py +++ b/mcdc/code_factory/numba_objects_generator.py @@ -1023,11 +1023,26 @@ def generate_mcdc_access(targets): ) elif len(shape) == 7: text_getter += _accessor_7d_element( - object_name, attribute_name, shape[1], shape[2], shape[3], shape[4], shape[5], shape[6] + object_name, + attribute_name, + shape[1], + shape[2], + shape[3], + shape[4], + shape[5], + shape[6], ) text_setter += _accessor_7d_element( - object_name, attribute_name, shape[1], shape[2], shape[3], shape[4], shape[5], shape[6], True + object_name, + attribute_name, + shape[1], + shape[2], + shape[3], + shape[4], + shape[5], + shape[6], + True, ) text_getter += _accessor_chunk(object_name, attribute_name) text_setter += _accessor_chunk(object_name, attribute_name, True) @@ -1221,6 +1236,7 @@ def _accessor_7d_element( text += f" return data[offset + index_1 * stride_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_2 * stride_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_3 * stride_4 * stride_5 * stride_6 * stride_7 + index_4 * stride_5 * stride_6 * stride_7 + index_5 * stride_6 * stride_7 + index_6 * stride_7 + index_7]\n\n\n" return text + # ====================================================================================== # Misc. # ====================================================================================== diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index e2df61ca..8a3745e9 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -100,9 +100,15 @@ class WeightWindows(ObjectSingleton): Nz: int # arrays of ww params - lower_weights: Annotated[NDArray[np.float64], ("Nt", "Ne", "Nmu", "Na", "Nx", "Ny", "Nz")] - target_weights: Annotated[NDArray[np.float64], ("Nt", "Ne", "Nmu", "Na", "Nx", "Ny", "Nz")] - upper_weights: Annotated[NDArray[np.float64], ("Nt", "Ne", "Nmu", "Na", "Nx", "Ny", "Nz")] + lower_weights: Annotated[ + NDArray[np.float64], ("Nt", "Ne", "Nmu", "Na", "Nx", "Ny", "Nz") + ] + target_weights: Annotated[ + NDArray[np.float64], ("Nt", "Ne", "Nmu", "Na", "Nx", "Ny", "Nz") + ] + upper_weights: Annotated[ + NDArray[np.float64], ("Nt", "Ne", "Nmu", "Na", "Nx", "Ny", "Nz") + ] def __init__(self): self.active = False @@ -122,7 +128,9 @@ def __init__(self): self.target_weights = np.array([1.0]).reshape(*shape) self.upper_weights = np.array([1.0]).reshape(*shape) - def __call__(self, weight_windows, mesh=None, energy=None, mu=None, azimuthal=None, time=None): + def __call__( + self, weight_windows, mesh=None, energy=None, mu=None, azimuthal=None, time=None + ): # fill in defaults if mesh is None: mesh = MeshUniform() @@ -151,7 +159,7 @@ def __call__(self, weight_windows, mesh=None, energy=None, mu=None, azimuthal=No f"{type(mesh).__name__} is not supported for weight windows" ) # validate energy and get size - self.__check_array(energy, "Energy") + self.__check_array(energy, "Energy") ne = energy.shape[0] - 1 # validate mu and get size self.__check_array(mu, "Mu") diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py index 3313b0c2..3a38b483 100644 --- a/test/unit/transport/technique/weight_windows.py +++ b/test/unit/transport/technique/weight_windows.py @@ -78,18 +78,29 @@ def make_ww_model_distinct(): # value at index is related to index, easy to predict during later test for t in range(Nt): - for e in range(Ne): - for m in range(Nmu): - for a in range(Na): - for i in range(N): - for j in range(N): - for k in range(N): - val = 1_000_000*t + 100_000*e + 10_000*m + 1_000*a + 100 * i + 10 * j + k + 1 - ww_array[t, e, m, a, i, j, k, 0] = val - ww_array[t, e, m, a, i, j, k, 1] = 10_000_000 + val - ww_array[t, e, m, a, i, j, k, 2] = 20_000_000 + val - - mcdc.simulation.weight_windows(ww_array, mesh=mesh, energy=energy, time=time, mu=mu, azimuthal=azimuthal) + for e in range(Ne): + for m in range(Nmu): + for a in range(Na): + for i in range(N): + for j in range(N): + for k in range(N): + val = ( + 1_000_000 * t + + 100_000 * e + + 10_000 * m + + 1_000 * a + + 100 * i + + 10 * j + + k + + 1 + ) + ww_array[t, e, m, a, i, j, k, 0] = val + ww_array[t, e, m, a, i, j, k, 1] = 10_000_000 + val + ww_array[t, e, m, a, i, j, k, 2] = 20_000_000 + val + + mcdc.simulation.weight_windows( + ww_array, mesh=mesh, energy=energy, time=time, mu=mu, azimuthal=azimuthal + ) mcdc_container, data = preparation() return mcdc_container[0], data @@ -226,29 +237,40 @@ def test_query_weight_window(): # loop over all bins, check query against expected ww for it, time in enumerate(times): - for ie, energy in enumerate(energies): - for imu, mu in enumerate(mus): - for ia, azimuthal in enumerate(azimuthals): - for ix in range(nx): - for iy in range(ny): - for iz in range(nz): - # put particle in center of current mesh bin - p[0]["x"] = xmin + dx * (ix + 0.5) - p[0]["y"] = ymin + dy * (iy + 0.5) - p[0]["z"] = zmin + dz * (iz + 0.5) - # assign params to be in center of bins - p[0]["t"] = time - p[0]["E"] = energy - p[0]["ux"] = np.sqrt(1.0 - mu**2) * np.cos(azimuthal) - p[0]["uy"] = np.sqrt(1.0 - mu**2) * np.sin(azimuthal) - p[0]["uz"] = mu - - # query and predict - lower, target, upper = query_weight_window(p, simulation, data) - exp_lower = 1_000_000*it + 100_000*ie + 10_000*imu + 1_000*ia + 100 * ix + 10 * iy + iz + 1 - exp_target = 10_000_000 + exp_lower - exp_upper = 20_000_000 + exp_lower - - assert lower == exp_lower - assert target == exp_target - assert upper == exp_upper + for ie, energy in enumerate(energies): + for imu, mu in enumerate(mus): + for ia, azimuthal in enumerate(azimuthals): + for ix in range(nx): + for iy in range(ny): + for iz in range(nz): + # put particle in center of current mesh bin + p[0]["x"] = xmin + dx * (ix + 0.5) + p[0]["y"] = ymin + dy * (iy + 0.5) + p[0]["z"] = zmin + dz * (iz + 0.5) + # assign params to be in center of bins + p[0]["t"] = time + p[0]["E"] = energy + p[0]["ux"] = np.sqrt(1.0 - mu**2) * np.cos(azimuthal) + p[0]["uy"] = np.sqrt(1.0 - mu**2) * np.sin(azimuthal) + p[0]["uz"] = mu + + # query and predict + lower, target, upper = query_weight_window( + p, simulation, data + ) + exp_lower = ( + 1_000_000 * it + + 100_000 * ie + + 10_000 * imu + + 1_000 * ia + + 100 * ix + + 10 * iy + + iz + + 1 + ) + exp_target = 10_000_000 + exp_lower + exp_upper = 20_000_000 + exp_lower + + assert lower == exp_lower + assert target == exp_target + assert upper == exp_upper From c677ccc1c604eecd14b64aa89b97bb4836de5278 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 29 Apr 2026 18:01:30 -0700 Subject: [PATCH 07/15] quick docs update for python api of ww --- docs/source/pythonapi/index.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/pythonapi/index.rst b/docs/source/pythonapi/index.rst index f3a8ed46..47385780 100644 --- a/docs/source/pythonapi/index.rst +++ b/docs/source/pythonapi/index.rst @@ -90,7 +90,7 @@ Techniques are enabled by calling methods on the ``mcdc.simulation`` singleton: - ``mcdc.simulation.global_weight_roulette(weight_threshold=0.0, weight_target=1.0)`` - ``mcdc.simulation.population_control(active=True)`` - ``mcdc.simulation.weighted_emission(active=True, weight_target=1.0)`` -- ``mcdc.simulation.weight_windows(weight_windows, mesh=None, energy=None)`` +- ``mcdc.simulation.weight_windows(weight_windows, mesh=None, energy=None, time=None, mu=None, azimuthal=None)`` Running ------- From bc24541ed35cbf34fd40fd56f7b70e1ca144d4b8 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 6 May 2026 20:03:26 -0700 Subject: [PATCH 08/15] added polar reference vector to ww object --- mcdc/numba_types.py | 3 +++ mcdc/object_/technique.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/mcdc/numba_types.py b/mcdc/numba_types.py index d4190c1f..9dbeac5c 100644 --- a/mcdc/numba_types.py +++ b/mcdc/numba_types.py @@ -593,6 +593,9 @@ ('Nx', int64), ('Ny', int64), ('Nz', int64), + ('px', float64), + ('py', float64), + ('pz', float64), ('lower_weights_offset', int64), ('lower_weights_length', int64), ('target_weights_offset', int64), diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index 8a3745e9..908f6bc8 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -99,6 +99,11 @@ class WeightWindows(ObjectSingleton): Ny: int Nz: int + # reference vector for calculating mu and azimuthal angle + px: float + py: float + pz: float + # arrays of ww params lower_weights: Annotated[ NDArray[np.float64], ("Nt", "Ne", "Nmu", "Na", "Nx", "Ny", "Nz") @@ -123,6 +128,9 @@ def __init__(self): self.mesh_ID = -1 # skirt around having to create a MeshBase instance self.Nx, self.Ny, self.Nz = 1, 1, 1 self.Nt = 1 + self.px = 0.0 + self.py = 0.0 + self.pz = 1.0 shape = (self.Nt, self.Ne, self.Nmu, self.Na, self.Nx, self.Ny, self.Nz) self.lower_weights = np.array([1.0]).reshape(*shape) self.target_weights = np.array([1.0]).reshape(*shape) @@ -208,6 +216,12 @@ def __call__( "Target weight can not be greater than the upper bound weight for any weight window" ) + def set_polar_reference(self, px, py, pz): + norm = np.linalg.norm([px, py, pz]) + self.px = px / norm + self.py = py / norm + self.pz = pz / norm + @staticmethod def __check_array(array: NDArray[np.float64], name: str): if not (np.diff(array) > 0).all(): From 950a7a1897f686191a84dcb9226f74d9ca7ced13 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 6 May 2026 20:04:02 -0700 Subject: [PATCH 09/15] added util helper for calculating mu and azi from reference --- mcdc/transport/technique.py | 10 +++++--- mcdc/transport/util.py | 51 +++++++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 4 deletions(-) diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index 05f8c2fc..bdcfa068 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -160,13 +160,15 @@ def get_ww_indices(particle_container, ww_obj, simulation, data): energy = particle["E"] ie = util.find_bin(energy, energy_bounds) - # get mu index + # get angular indices + mu, azimuthal = util.calculate_angles(particle_container, ww_obj["px"], ww_obj["py"], ww_obj["pz"]) + + # mu mu_bounds = ww_get.mu_bounds_all(ww_obj, data) - imu = util.find_bin(particle["uz"], mu_bounds) + imu = util.find_bin(mu, mu_bounds) - # get azimuthal index + # azimuthal azi_bounds = ww_get.azi_bounds_all(ww_obj, data) - azimuthal = np.arctan2(particle["uy"], particle["ux"]) ia = util.find_bin(azimuthal, azi_bounds) # get spatial index diff --git a/mcdc/transport/util.py b/mcdc/transport/util.py index 4627841d..a96a8b68 100644 --- a/mcdc/transport/util.py +++ b/mcdc/transport/util.py @@ -132,6 +132,57 @@ def log_interpolation(x, x1, x2, y1, y2): return math.exp(ly) +# = +# +# = + + +@njit +def calculate_angles(particle_container, px, py, pz): + particle = particle_container[0] + ux = particle["ux"] + uy = particle["uy"] + uz = particle["uz"] + + mu = ux*px + uy*py + uz*pz + + azimuthal = _calculate_azimuthal(ux, uy, uz, px, py, pz) + + return mu, azimuthal + + +def _calculate_azimuthal(ux, uy, uz, px, py, pz): + # get two orthonormal basis vectors u1, u2 perpendicular to p + # u1 done via gram-schmidt, u2 done via cross product + + # choose arbitrary guess v1, check to make sure its not accidentally parallel to p + if px < 0.9: + v1x, v1y, v1z = 1.0, 0.0, 0.0 + else: + v1x, v1y, v1z = 0.0, 1.0, 0.0 + + # u1 = (v1 - proj(v1, p)) / ||u1|| + v1dotp = v1x*px + v1y*py + v1z*pz + u1x = v1x - v1dotp * px + u1y = v1y - v1dotp * py + u1z = v1z - v1dotp * pz + u1norm = math.sqrt(u1x*u1x + u1y*u1y + u1z*u1z) + u1x /= u1norm + u1y /= u1norm + u1z /= u1norm + + # u2 = p x u1 + u2x = py*u1z - pz*u1y + u2y = pz*u1x - px*u1z + u2z = px*u1y - py*u1x + + # particle vector in terms of u1 and u2 (dot products) + rel_ux = ux*u1x + uy*u1y + uz*u1z + rel_uy = ux*u2x + uy*u2y + uz*u2z + + return math.atan2(rel_uy, rel_ux) + + # ====================================================================================== # Framework utilities # ====================================================================================== From d28f3358cc5c8fb52ee08b4c23b56178f99ad261 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 6 May 2026 20:12:06 -0700 Subject: [PATCH 10/15] added docstrings to added util functions --- mcdc/transport/util.py | 53 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 50 insertions(+), 3 deletions(-) diff --git a/mcdc/transport/util.py b/mcdc/transport/util.py index a96a8b68..8b0336ae 100644 --- a/mcdc/transport/util.py +++ b/mcdc/transport/util.py @@ -132,13 +132,34 @@ def log_interpolation(x, x1, x2, y1, y2): return math.exp(ly) -# = -# -# = +# ====================================================================================== +# Angle conversion utilities +# ====================================================================================== @njit def calculate_angles(particle_container, px, py, pz): + """ + Calculate particle mu and azimuthal angle from given reference vector + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + px : float + X-component of reference vector + py : float + Y-component of reference vector + pz : float + Z-component of reference vector + + Returns + ------- + mu : float + Mu of particle relative to reference vector + azimuthal : float + Azimuthal angle of particle relative to reference vector + """ particle = particle_container[0] ux = particle["ux"] uy = particle["uy"] @@ -152,6 +173,32 @@ def calculate_angles(particle_container, px, py, pz): def _calculate_azimuthal(ux, uy, uz, px, py, pz): + """ + Calculates the azimuthal angle of a particle relative to a reference vector. + This is done by finding two orthonormal basis vectors perpendicular to the reference vector. + The first orthonormal vector, u1, is found using the graham-schmidt procedure. + The second, u2, is found from the cross product of the reference vector and u1. + + Parameters + ---------- + ux : float + X-component of particle's direction + uy : float + Y-component of particle's direction + uz : float + Z-component of particle's direction + px : float + X-component of reference vector + py : float + Y-component of reference vector + pz : float + Z-component of reference vector + + Returns + ------- + azimuthal : float + Azimuthal angle of particle relative to reference vector + """ # get two orthonormal basis vectors u1, u2 perpendicular to p # u1 done via gram-schmidt, u2 done via cross product From dd78e013d6de70d6203555804a1642efe04df9d3 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 6 May 2026 20:27:35 -0700 Subject: [PATCH 11/15] add quick unit test for angle helper --- test/unit/transport/util/calculate_angles.py | 31 ++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 test/unit/transport/util/calculate_angles.py diff --git a/test/unit/transport/util/calculate_angles.py b/test/unit/transport/util/calculate_angles.py new file mode 100644 index 00000000..d3922237 --- /dev/null +++ b/test/unit/transport/util/calculate_angles.py @@ -0,0 +1,31 @@ +import numpy as np +import pytest + +from mcdc.transport.util import calculate_angles +from mcdc.numba_types import particle_data + + +@pytest.mark.parametrize( + "ux, uy, uz, expected_mu, expected_phi", + [ + (1.0, 0.0, 0.0, 0.0, 0.0), + (-1.0, 0.0, 0.0, 0.0, np.pi), + (0.0, 1.0, 0.0, 0.0, np.pi / 2.0), + (0.0, -1.0, 0.0, 0.0, -np.pi / 2.0), + (0.0, 0.0, 1.0, 1.0, 0.0), + (0.0, 0.0, -1.0, -1.0, 0.0) + ], +) +def test_calculate_angles(ux, uy, uz, expected_mu, expected_phi): + particle_container = np.zeros(1, particle_data) + particle = particle_container[0] + + particle["ux"] = ux + particle["uy"] = uy + particle["uz"] = uz + px, py, pz = 0.0, 0.0, 1.0 + + mu, phi = calculate_angles(particle_container, px, py, pz) + + assert np.isclose(mu, expected_mu) + assert np.isclose(phi, expected_phi) \ No newline at end of file From ac784560acddc07f382927aa40a55dac619efd57 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 6 May 2026 20:28:01 -0700 Subject: [PATCH 12/15] back in black --- mcdc/transport/technique.py | 4 +++- mcdc/transport/util.py | 20 ++++++++++---------- test/unit/transport/util/calculate_angles.py | 4 ++-- 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index bdcfa068..045acbaf 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -161,7 +161,9 @@ def get_ww_indices(particle_container, ww_obj, simulation, data): ie = util.find_bin(energy, energy_bounds) # get angular indices - mu, azimuthal = util.calculate_angles(particle_container, ww_obj["px"], ww_obj["py"], ww_obj["pz"]) + mu, azimuthal = util.calculate_angles( + particle_container, ww_obj["px"], ww_obj["py"], ww_obj["pz"] + ) # mu mu_bounds = ww_get.mu_bounds_all(ww_obj, data) diff --git a/mcdc/transport/util.py b/mcdc/transport/util.py index 8b0336ae..48058b87 100644 --- a/mcdc/transport/util.py +++ b/mcdc/transport/util.py @@ -165,7 +165,7 @@ def calculate_angles(particle_container, px, py, pz): uy = particle["uy"] uz = particle["uz"] - mu = ux*px + uy*py + uz*pz + mu = ux * px + uy * py + uz * pz azimuthal = _calculate_azimuthal(ux, uy, uz, px, py, pz) @@ -204,29 +204,29 @@ def _calculate_azimuthal(ux, uy, uz, px, py, pz): # choose arbitrary guess v1, check to make sure its not accidentally parallel to p if px < 0.9: - v1x, v1y, v1z = 1.0, 0.0, 0.0 + v1x, v1y, v1z = 1.0, 0.0, 0.0 else: v1x, v1y, v1z = 0.0, 1.0, 0.0 # u1 = (v1 - proj(v1, p)) / ||u1|| - v1dotp = v1x*px + v1y*py + v1z*pz + v1dotp = v1x * px + v1y * py + v1z * pz u1x = v1x - v1dotp * px u1y = v1y - v1dotp * py u1z = v1z - v1dotp * pz - u1norm = math.sqrt(u1x*u1x + u1y*u1y + u1z*u1z) + u1norm = math.sqrt(u1x * u1x + u1y * u1y + u1z * u1z) u1x /= u1norm u1y /= u1norm u1z /= u1norm # u2 = p x u1 - u2x = py*u1z - pz*u1y - u2y = pz*u1x - px*u1z - u2z = px*u1y - py*u1x + u2x = py * u1z - pz * u1y + u2y = pz * u1x - px * u1z + u2z = px * u1y - py * u1x # particle vector in terms of u1 and u2 (dot products) - rel_ux = ux*u1x + uy*u1y + uz*u1z - rel_uy = ux*u2x + uy*u2y + uz*u2z - + rel_ux = ux * u1x + uy * u1y + uz * u1z + rel_uy = ux * u2x + uy * u2y + uz * u2z + return math.atan2(rel_uy, rel_ux) diff --git a/test/unit/transport/util/calculate_angles.py b/test/unit/transport/util/calculate_angles.py index d3922237..087523c8 100644 --- a/test/unit/transport/util/calculate_angles.py +++ b/test/unit/transport/util/calculate_angles.py @@ -13,7 +13,7 @@ (0.0, 1.0, 0.0, 0.0, np.pi / 2.0), (0.0, -1.0, 0.0, 0.0, -np.pi / 2.0), (0.0, 0.0, 1.0, 1.0, 0.0), - (0.0, 0.0, -1.0, -1.0, 0.0) + (0.0, 0.0, -1.0, -1.0, 0.0), ], ) def test_calculate_angles(ux, uy, uz, expected_mu, expected_phi): @@ -28,4 +28,4 @@ def test_calculate_angles(ux, uy, uz, expected_mu, expected_phi): mu, phi = calculate_angles(particle_container, px, py, pz) assert np.isclose(mu, expected_mu) - assert np.isclose(phi, expected_phi) \ No newline at end of file + assert np.isclose(phi, expected_phi) From 9927c14144190aaaccd913550bf183b599e14c10 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 6 May 2026 20:51:35 -0700 Subject: [PATCH 13/15] update tally query of angular index --- mcdc/transport/tally/filter.py | 18 ++---------------- 1 file changed, 2 insertions(+), 16 deletions(-) diff --git a/mcdc/transport/tally/filter.py b/mcdc/transport/tally/filter.py index a743f785..8b06118b 100644 --- a/mcdc/transport/tally/filter.py +++ b/mcdc/transport/tally/filter.py @@ -12,7 +12,7 @@ COINCIDENCE_TOLERANCE_ENERGY, COINCIDENCE_TOLERANCE_TIME, ) -from mcdc.transport.util import find_bin_with_tolerance, find_bin_with_rules +from mcdc.transport.util import find_bin_with_tolerance, find_bin_with_rules, calculate_angles @njit @@ -33,26 +33,12 @@ def get_filter_indices(particle_container, tally, data, MG_mode): @njit def get_direction_index(particle_container, tally, data): - particle = particle_container[0] - - # Particle properties - ux = particle["ux"] - uy = particle["uy"] - uz = particle["uz"] - # Polar reference nx = tally["polar_reference"][0] ny = tally["polar_reference"][1] nz = tally["polar_reference"][2] - # TODO: Rotate direction based on the polar reference - if nz != 1.0: - pass - - mu = uz - azi = math.acos(ux / math.sqrt(ux * ux + uy * uy)) - if uy < 0.0: - azi *= -1 + mu, azi = calculate_angles(particle_container, nx, ny, nz) tolerance = COINCIDENCE_TOLERANCE_DIRECTION From 9d84091b06dc9fd885fc964af4bf6a644b8270f8 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 6 May 2026 21:13:34 -0700 Subject: [PATCH 14/15] update util to take a reference vector, and change ww to use vector --- mcdc/numba_types.py | 4 +--- mcdc/object_/technique.py | 14 ++++---------- mcdc/transport/tally/filter.py | 12 ++++++------ mcdc/transport/technique.py | 4 +--- mcdc/transport/util.py | 14 +++++++------- test/unit/transport/util/calculate_angles.py | 4 ++-- 6 files changed, 21 insertions(+), 31 deletions(-) diff --git a/mcdc/numba_types.py b/mcdc/numba_types.py index 9dbeac5c..4b31fa2c 100644 --- a/mcdc/numba_types.py +++ b/mcdc/numba_types.py @@ -593,9 +593,7 @@ ('Nx', int64), ('Ny', int64), ('Nz', int64), - ('px', float64), - ('py', float64), - ('pz', float64), + ('polar_reference', float64, (3,)), ('lower_weights_offset', int64), ('lower_weights_length', int64), ('target_weights_offset', int64), diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index 908f6bc8..b8affea0 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -100,9 +100,7 @@ class WeightWindows(ObjectSingleton): Nz: int # reference vector for calculating mu and azimuthal angle - px: float - py: float - pz: float + polar_reference: Annotated[NDArray[np.float64], (3,)] # arrays of ww params lower_weights: Annotated[ @@ -128,9 +126,7 @@ def __init__(self): self.mesh_ID = -1 # skirt around having to create a MeshBase instance self.Nx, self.Ny, self.Nz = 1, 1, 1 self.Nt = 1 - self.px = 0.0 - self.py = 0.0 - self.pz = 1.0 + self.polar_reference = np.array([0.0, 0.0, 1.0]) shape = (self.Nt, self.Ne, self.Nmu, self.Na, self.Nx, self.Ny, self.Nz) self.lower_weights = np.array([1.0]).reshape(*shape) self.target_weights = np.array([1.0]).reshape(*shape) @@ -217,10 +213,8 @@ def __call__( ) def set_polar_reference(self, px, py, pz): - norm = np.linalg.norm([px, py, pz]) - self.px = px / norm - self.py = py / norm - self.pz = pz / norm + polar_reference = np.array([px, py, pz]) + self.polar_reference = polar_reference / np.linalg.norm(polar_reference) @staticmethod def __check_array(array: NDArray[np.float64], name: str): diff --git a/mcdc/transport/tally/filter.py b/mcdc/transport/tally/filter.py index 8b06118b..2e1b4f1f 100644 --- a/mcdc/transport/tally/filter.py +++ b/mcdc/transport/tally/filter.py @@ -12,7 +12,11 @@ COINCIDENCE_TOLERANCE_ENERGY, COINCIDENCE_TOLERANCE_TIME, ) -from mcdc.transport.util import find_bin_with_tolerance, find_bin_with_rules, calculate_angles +from mcdc.transport.util import ( + find_bin_with_tolerance, + find_bin_with_rules, + calculate_angles, +) @njit @@ -34,11 +38,7 @@ def get_filter_indices(particle_container, tally, data, MG_mode): @njit def get_direction_index(particle_container, tally, data): # Polar reference - nx = tally["polar_reference"][0] - ny = tally["polar_reference"][1] - nz = tally["polar_reference"][2] - - mu, azi = calculate_angles(particle_container, nx, ny, nz) + mu, azi = calculate_angles(particle_container, tally["polar_reference"]) tolerance = COINCIDENCE_TOLERANCE_DIRECTION diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index 045acbaf..1d01e2f1 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -161,9 +161,7 @@ def get_ww_indices(particle_container, ww_obj, simulation, data): ie = util.find_bin(energy, energy_bounds) # get angular indices - mu, azimuthal = util.calculate_angles( - particle_container, ww_obj["px"], ww_obj["py"], ww_obj["pz"] - ) + mu, azimuthal = util.calculate_angles(particle_container, ww_obj["polar_reference"]) # mu mu_bounds = ww_get.mu_bounds_all(ww_obj, data) diff --git a/mcdc/transport/util.py b/mcdc/transport/util.py index 48058b87..fb50eec6 100644 --- a/mcdc/transport/util.py +++ b/mcdc/transport/util.py @@ -138,7 +138,7 @@ def log_interpolation(x, x1, x2, y1, y2): @njit -def calculate_angles(particle_container, px, py, pz): +def calculate_angles(particle_container, polar_reference): """ Calculate particle mu and azimuthal angle from given reference vector @@ -146,12 +146,8 @@ def calculate_angles(particle_container, px, py, pz): ---------- particle_container : ndarray Container holding the particle. - px : float - X-component of reference vector - py : float - Y-component of reference vector - pz : float - Z-component of reference vector + polar_reference : ndarray + 3D polar reference vector Returns ------- @@ -165,6 +161,10 @@ def calculate_angles(particle_container, px, py, pz): uy = particle["uy"] uz = particle["uz"] + px = polar_reference[0] + py = polar_reference[1] + pz = polar_reference[2] + mu = ux * px + uy * py + uz * pz azimuthal = _calculate_azimuthal(ux, uy, uz, px, py, pz) diff --git a/test/unit/transport/util/calculate_angles.py b/test/unit/transport/util/calculate_angles.py index 087523c8..0d86afd6 100644 --- a/test/unit/transport/util/calculate_angles.py +++ b/test/unit/transport/util/calculate_angles.py @@ -23,9 +23,9 @@ def test_calculate_angles(ux, uy, uz, expected_mu, expected_phi): particle["ux"] = ux particle["uy"] = uy particle["uz"] = uz - px, py, pz = 0.0, 0.0, 1.0 + polar_reference = np.array([0.0, 0.0, 1.0]) - mu, phi = calculate_angles(particle_container, px, py, pz) + mu, phi = calculate_angles(particle_container, polar_reference) assert np.isclose(mu, expected_mu) assert np.isclose(phi, expected_phi) From 8ca7a4187a988bdf6d00c54fc22fa493a6265df1 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Thu, 7 May 2026 00:20:39 -0700 Subject: [PATCH 15/15] add missing numba flag --- mcdc/transport/util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mcdc/transport/util.py b/mcdc/transport/util.py index fb50eec6..069ead06 100644 --- a/mcdc/transport/util.py +++ b/mcdc/transport/util.py @@ -171,7 +171,7 @@ def calculate_angles(particle_container, polar_reference): return mu, azimuthal - +@njit def _calculate_azimuthal(ux, uy, uz, px, py, pz): """ Calculates the azimuthal angle of a particle relative to a reference vector.