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 ------- diff --git a/mcdc/code_factory/numba_objects_generator.py b/mcdc/code_factory/numba_objects_generator.py index c466b92d..96e16669 100644 --- a/mcdc/code_factory/numba_objects_generator.py +++ b/mcdc/code_factory/numba_objects_generator.py @@ -1021,6 +1021,29 @@ 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 +1207,36 @@ 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 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), diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index e2adfb80..8a3745e9 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,49 @@ 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 +158,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 +208,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 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 diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py index 6ad4eb44..3a38b483 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,40 @@ 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 +117,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 +230,47 @@ 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