diff --git a/docs/source/pythonapi/index.rst b/docs/source/pythonapi/index.rst index f3a8ed46f..473857808 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 c466b92d6..96e166697 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 b1b7cbece..f15b61d26 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 e6a2ab8a2..5a77acea2 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 056ce1778..4b31fa2c3 100644 --- a/mcdc/numba_types.py +++ b/mcdc/numba_types.py @@ -577,13 +577,23 @@ 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), ('Nz', int64), + ('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 e2adfb803..b8affea0b 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,39 +82,71 @@ 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 Ny: int Nz: int + # reference vector for calculating mu and azimuthal angle + polar_reference: Annotated[NDArray[np.float64], (3,)] + # 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) + 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) 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 +162,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 +212,19 @@ def __call__(self, weight_windows, mesh=None, energy=None): "Target weight can not be greater than the upper bound weight for any weight window" ) + def set_polar_reference(self, px, py, pz): + 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): + 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/tally/filter.py b/mcdc/transport/tally/filter.py index a743f785f..2e1b4f1f4 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 +from mcdc.transport.util import ( + find_bin_with_tolerance, + find_bin_with_rules, + calculate_angles, +) @njit @@ -33,26 +37,8 @@ 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, tally["polar_reference"]) tolerance = COINCIDENCE_TOLERANCE_DIRECTION diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index 5de637998..1d01e2f12 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,22 @@ def get_ww_indices(particle_container, ww_obj, simulation, data): energy = particle["E"] ie = util.find_bin(energy, energy_bounds) + # get angular indices + mu, azimuthal = util.calculate_angles(particle_container, ww_obj["polar_reference"]) + + # mu + mu_bounds = ww_get.mu_bounds_all(ww_obj, data) + imu = util.find_bin(mu, mu_bounds) + + # azimuthal + azi_bounds = ww_get.azi_bounds_all(ww_obj, data) + 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/mcdc/transport/util.py b/mcdc/transport/util.py index 4627841d0..069ead06c 100644 --- a/mcdc/transport/util.py +++ b/mcdc/transport/util.py @@ -132,6 +132,104 @@ def log_interpolation(x, x1, x2, y1, y2): return math.exp(ly) +# ====================================================================================== +# Angle conversion utilities +# ====================================================================================== + + +@njit +def calculate_angles(particle_container, polar_reference): + """ + Calculate particle mu and azimuthal angle from given reference vector + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + polar_reference : ndarray + 3D polar 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"] + 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) + + 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. + 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 + + # 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 # ====================================================================================== diff --git a/test/regression/basic_weight_windows/input.py b/test/regression/basic_weight_windows/input.py index 38d3e9805..84c4735cf 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 6ad4eb44a..3a38b483e 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 diff --git a/test/unit/transport/util/calculate_angles.py b/test/unit/transport/util/calculate_angles.py new file mode 100644 index 000000000..0d86afd6d --- /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 + polar_reference = np.array([0.0, 0.0, 1.0]) + + mu, phi = calculate_angles(particle_container, polar_reference) + + assert np.isclose(mu, expected_mu) + assert np.isclose(phi, expected_phi)