diff --git a/docs/source/pythonapi/index.rst b/docs/source/pythonapi/index.rst index 3b108838a..86a687a62 100644 --- a/docs/source/pythonapi/index.rst +++ b/docs/source/pythonapi/index.rst @@ -87,8 +87,10 @@ Defining techniques Techniques are enabled by calling methods on the ``mcdc.simulation`` singleton: - ``mcdc.simulation.implicit_capture(active=True)`` +- ``mcdc.simulation.forced_collisions(cells=[], weight_thresholds=[], weight_targets=[])`` - ``mcdc.simulation.weighted_emission(active=True, weight_target=1.0)`` - ``mcdc.simulation.weight_roulette(weight_threshold=0.0, weight_target=1.0)`` +- ``mcdc.simulation.weight_windows(weight_windows, mesh=None, energy=None)`` - ``mcdc.simulation.population_control(active=True)`` Running diff --git a/mcdc/code_factory/numba_objects_generator.py b/mcdc/code_factory/numba_objects_generator.py index d07da87ad..c466b92d6 100644 --- a/mcdc/code_factory/numba_objects_generator.py +++ b/mcdc/code_factory/numba_objects_generator.py @@ -1013,6 +1013,14 @@ def generate_mcdc_access(targets): object_name, attribute_name, shape[1], shape[2], True ) + elif len(shape) == 4: + text_getter += _accessor_4d_element( + object_name, attribute_name, shape[1], shape[2], shape[3] + ) + + text_setter += _accessor_4d_element( + object_name, attribute_name, shape[1], shape[2], shape[3], True + ) text_getter += _accessor_chunk(object_name, attribute_name) text_setter += _accessor_chunk(object_name, attribute_name, True) @@ -1157,6 +1165,25 @@ def _accessor_3d_element(object_name, attribute_name, stride_2, stride_3, setter return text +def _accessor_4d_element( + object_name, attribute_name, stride_2, stride_3, stride_4, setter=False +): + text = f"@njit\n" + if setter: + text += f"def {attribute_name}(index_1, index_2, index_3, index_4, {object_name}, data, value):\n" + else: + text += f"def {attribute_name}(index_1, index_2, index_3, index_4, {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' + if setter: + text += f" data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4] = value\n\n\n" + else: + text += f" return data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4]\n\n\n" + return text + + # ====================================================================================== # Misc. # ====================================================================================== diff --git a/mcdc/mcdc_get/__init__.py b/mcdc/mcdc_get/__init__.py index d3b533388..7e486388c 100644 --- a/mcdc/mcdc_get/__init__.py +++ b/mcdc/mcdc_get/__init__.py @@ -86,12 +86,16 @@ import mcdc.mcdc_get.settings as settings +import mcdc.mcdc_get.forced_collisions as forced_collisions + import mcdc.mcdc_get.implicit_capture as implicit_capture import mcdc.mcdc_get.population_control as population_control import mcdc.mcdc_get.weight_roulette as weight_roulette +import mcdc.mcdc_get.weight_windows as weight_windows + import mcdc.mcdc_get.weighted_emission as weighted_emission import mcdc.mcdc_get.source as source diff --git a/mcdc/mcdc_get/forced_collisions.py b/mcdc/mcdc_get/forced_collisions.py new file mode 100644 index 000000000..24e9b1bbd --- /dev/null +++ b/mcdc/mcdc_get/forced_collisions.py @@ -0,0 +1,90 @@ +# The following is automatically generated by code_factory.py + +from numba import njit + + +@njit +def cell_IDs(index, forced_collisions, data): + offset = forced_collisions["cell_IDs_offset"] + return data[offset + index] + + +@njit +def cell_IDs_all(forced_collisions, data): + start = forced_collisions["cell_IDs_offset"] + size = forced_collisions["cell_IDs_length"] + end = start + size + return data[start:end] + + +@njit +def cell_IDs_last(forced_collisions, data): + start = forced_collisions["cell_IDs_offset"] + size = forced_collisions["cell_IDs_length"] + end = start + size + return data[end - 1] + + +@njit +def cell_IDs_chunk(start, length, forced_collisions, data): + start += forced_collisions["cell_IDs_offset"] + end = start + length + return data[start:end] + + +@njit +def threshold_weights(index, forced_collisions, data): + offset = forced_collisions["threshold_weights_offset"] + return data[offset + index] + + +@njit +def threshold_weights_all(forced_collisions, data): + start = forced_collisions["threshold_weights_offset"] + size = forced_collisions["threshold_weights_length"] + end = start + size + return data[start:end] + + +@njit +def threshold_weights_last(forced_collisions, data): + start = forced_collisions["threshold_weights_offset"] + size = forced_collisions["threshold_weights_length"] + end = start + size + return data[end - 1] + + +@njit +def threshold_weights_chunk(start, length, forced_collisions, data): + start += forced_collisions["threshold_weights_offset"] + end = start + length + return data[start:end] + + +@njit +def target_weights(index, forced_collisions, data): + offset = forced_collisions["target_weights_offset"] + return data[offset + index] + + +@njit +def target_weights_all(forced_collisions, data): + start = forced_collisions["target_weights_offset"] + size = forced_collisions["target_weights_length"] + end = start + size + return data[start:end] + + +@njit +def target_weights_last(forced_collisions, data): + start = forced_collisions["target_weights_offset"] + size = forced_collisions["target_weights_length"] + end = start + size + return data[end - 1] + + +@njit +def target_weights_chunk(start, length, forced_collisions, data): + start += forced_collisions["target_weights_offset"] + end = start + length + return data[start:end] diff --git a/mcdc/mcdc_get/weight_windows.py b/mcdc/mcdc_get/weight_windows.py new file mode 100644 index 000000000..b1b7cbece --- /dev/null +++ b/mcdc/mcdc_get/weight_windows.py @@ -0,0 +1,80 @@ +# The following is automatically generated by code_factory.py + +from numba import njit + + +@njit +def energy_bounds(index, weight_windows, data): + offset = weight_windows["energy_bounds_offset"] + return data[offset + index] + + +@njit +def energy_bounds_all(weight_windows, data): + start = weight_windows["energy_bounds_offset"] + size = weight_windows["energy_bounds_length"] + end = start + size + return data[start:end] + + +@njit +def energy_bounds_last(weight_windows, data): + start = weight_windows["energy_bounds_offset"] + size = weight_windows["energy_bounds_length"] + end = start + size + return data[end - 1] + + +@njit +def energy_bounds_chunk(start, length, weight_windows, data): + start += weight_windows["energy_bounds_offset"] + end = start + length + return data[start:end] + + +@njit +def lower_weights(index_1, index_2, index_3, index_4, 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] + + +@njit +def lower_weights_chunk(start, length, weight_windows, data): + start += weight_windows["lower_weights_offset"] + end = start + length + return data[start:end] + + +@njit +def target_weights(index_1, index_2, index_3, index_4, 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] + + +@njit +def target_weights_chunk(start, length, weight_windows, data): + start += weight_windows["target_weights_offset"] + end = start + length + return data[start:end] + + +@njit +def upper_weights(index_1, index_2, index_3, index_4, 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] + + +@njit +def upper_weights_chunk(start, length, weight_windows, data): + start += weight_windows["upper_weights_offset"] + end = start + length + return data[start:end] diff --git a/mcdc/mcdc_set/__init__.py b/mcdc/mcdc_set/__init__.py index 38ce0dddd..2699dbb34 100644 --- a/mcdc/mcdc_set/__init__.py +++ b/mcdc/mcdc_set/__init__.py @@ -86,12 +86,16 @@ import mcdc.mcdc_set.settings as settings +import mcdc.mcdc_set.forced_collisions as forced_collisions + import mcdc.mcdc_set.implicit_capture as implicit_capture import mcdc.mcdc_set.population_control as population_control import mcdc.mcdc_set.weight_roulette as weight_roulette +import mcdc.mcdc_set.weight_windows as weight_windows + import mcdc.mcdc_set.weighted_emission as weighted_emission import mcdc.mcdc_set.source as source diff --git a/mcdc/mcdc_set/forced_collisions.py b/mcdc/mcdc_set/forced_collisions.py new file mode 100644 index 000000000..eda92e70e --- /dev/null +++ b/mcdc/mcdc_set/forced_collisions.py @@ -0,0 +1,90 @@ +# The following is automatically generated by code_factory.py + +from numba import njit + + +@njit +def cell_IDs(index, forced_collisions, data, value): + offset = forced_collisions["cell_IDs_offset"] + data[offset + index] = value + + +@njit +def cell_IDs_all(forced_collisions, data, value): + start = forced_collisions["cell_IDs_offset"] + size = forced_collisions["cell_IDs_length"] + end = start + size + data[start:end] = value + + +@njit +def cell_IDs_last(forced_collisions, data, value): + start = forced_collisions["cell_IDs_offset"] + size = forced_collisions["cell_IDs_length"] + end = start + size + data[end - 1] = value + + +@njit +def cell_IDs_chunk(start, length, forced_collisions, data, value): + start += forced_collisions["cell_IDs_offset"] + end = start + length + data[start:end] = value + + +@njit +def threshold_weights(index, forced_collisions, data, value): + offset = forced_collisions["threshold_weights_offset"] + data[offset + index] = value + + +@njit +def threshold_weights_all(forced_collisions, data, value): + start = forced_collisions["threshold_weights_offset"] + size = forced_collisions["threshold_weights_length"] + end = start + size + data[start:end] = value + + +@njit +def threshold_weights_last(forced_collisions, data, value): + start = forced_collisions["threshold_weights_offset"] + size = forced_collisions["threshold_weights_length"] + end = start + size + data[end - 1] = value + + +@njit +def threshold_weights_chunk(start, length, forced_collisions, data, value): + start += forced_collisions["threshold_weights_offset"] + end = start + length + data[start:end] = value + + +@njit +def target_weights(index, forced_collisions, data, value): + offset = forced_collisions["target_weights_offset"] + data[offset + index] = value + + +@njit +def target_weights_all(forced_collisions, data, value): + start = forced_collisions["target_weights_offset"] + size = forced_collisions["target_weights_length"] + end = start + size + data[start:end] = value + + +@njit +def target_weights_last(forced_collisions, data, value): + start = forced_collisions["target_weights_offset"] + size = forced_collisions["target_weights_length"] + end = start + size + data[end - 1] = value + + +@njit +def target_weights_chunk(start, length, forced_collisions, data, value): + start += forced_collisions["target_weights_offset"] + end = start + length + data[start:end] = value diff --git a/mcdc/mcdc_set/weight_windows.py b/mcdc/mcdc_set/weight_windows.py new file mode 100644 index 000000000..e6a2ab8a2 --- /dev/null +++ b/mcdc/mcdc_set/weight_windows.py @@ -0,0 +1,80 @@ +# The following is automatically generated by code_factory.py + +from numba import njit + + +@njit +def energy_bounds(index, weight_windows, data, value): + offset = weight_windows["energy_bounds_offset"] + data[offset + index] = value + + +@njit +def energy_bounds_all(weight_windows, data, value): + start = weight_windows["energy_bounds_offset"] + size = weight_windows["energy_bounds_length"] + end = start + size + data[start:end] = value + + +@njit +def energy_bounds_last(weight_windows, data, value): + start = weight_windows["energy_bounds_offset"] + size = weight_windows["energy_bounds_length"] + end = start + size + data[end - 1] = value + + +@njit +def energy_bounds_chunk(start, length, weight_windows, data, value): + start += weight_windows["energy_bounds_offset"] + end = start + length + data[start:end] = value + + +@njit +def lower_weights(index_1, index_2, index_3, index_4, 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 + + +@njit +def lower_weights_chunk(start, length, weight_windows, data, value): + start += weight_windows["lower_weights_offset"] + end = start + length + data[start:end] = value + + +@njit +def target_weights(index_1, index_2, index_3, index_4, 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 + + +@njit +def target_weights_chunk(start, length, weight_windows, data, value): + start += weight_windows["target_weights_offset"] + end = start + length + data[start:end] = value + + +@njit +def upper_weights(index_1, index_2, index_3, index_4, 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 + + +@njit +def upper_weights_chunk(start, length, weight_windows, data, value): + start += weight_windows["upper_weights_offset"] + end = start + length + data[start:end] = value diff --git a/mcdc/numba_types.py b/mcdc/numba_types.py index 8ca7ba4f4..46c5c4d01 100644 --- a/mcdc/numba_types.py +++ b/mcdc/numba_types.py @@ -561,6 +561,16 @@ ('gpu_storage', int64), ]) +forced_collisions = into_dtype([ + ('active', bool), + ('cell_IDs_offset', int64), + ('cell_IDs_length', int64), + ('threshold_weights_offset', int64), + ('threshold_weights_length', int64), + ('target_weights_offset', int64), + ('target_weights_length', int64), +]) + implicit_capture = into_dtype([ ('active', bool), ]) @@ -574,6 +584,23 @@ ('weight_target', float64), ]) +weight_windows = into_dtype([ + ('active', bool), + ('energy_bounds_offset', int64), + ('energy_bounds_length', int64), + ('Ne', int64), + ('mesh_ID', int64), + ('Nx', int64), + ('Ny', int64), + ('Nz', int64), + ('lower_weights_offset', int64), + ('lower_weights_length', int64), + ('target_weights_offset', int64), + ('target_weights_length', int64), + ('upper_weights_offset', int64), + ('upper_weights_length', int64), +]) + weighted_emission = into_dtype([ ('active', bool), ('weight_target', float64), @@ -813,8 +840,10 @@ def set_simulation(N: dict): ('N_tracklength_tally', int64), ('settings', settings), ('implicit_capture', implicit_capture), + ('forced_collisions', forced_collisions), ('weighted_emission', weighted_emission), ('weight_roulette', weight_roulette), + ('weight_windows', weight_windows), ('population_control', population_control), ('gpu_meta', gpu_meta), ('bank_future', bank_future), diff --git a/mcdc/object_/simulation.py b/mcdc/object_/simulation.py index fa5b5f9bd..89344f67c 100644 --- a/mcdc/object_/simulation.py +++ b/mcdc/object_/simulation.py @@ -3,8 +3,10 @@ from mcdc.object_.technique import ( ImplicitCapture, + ForcedCollisions, PopulationControl, WeightRoulette, + WeightWindows, WeightedEmission, ) @@ -80,8 +82,10 @@ class Simulation(ObjectSingleton): # Techniques implicit_capture: ImplicitCapture + forced_collisions: ForcedCollisions weighted_emission: WeightedEmission weight_roulette: WeightRoulette + weight_windows: WeightWindows population_control: PopulationControl # Particle banks @@ -165,8 +169,10 @@ def __init__(self): # Techniques self.implicit_capture = ImplicitCapture() + self.forced_collisions = ForcedCollisions() self.weighted_emission = WeightedEmission() self.weight_roulette = WeightRoulette() + self.weight_windows = WeightWindows() self.population_control = PopulationControl() # ============================================================================== diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index 1c3bacc1b..af719f6a2 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -1,4 +1,12 @@ +from typing import TYPE_CHECKING, Annotated +from numpy.typing import NDArray +import numpy as np +from mcdc.constant import FILL_MATERIAL, INF from mcdc.object_.base import ObjectSingleton + +if TYPE_CHECKING: + from mcdc.object_.cell import Cell +from mcdc.object_.mesh import MeshBase, MeshUniform from mcdc.print_ import print_error # ====================================================================================== @@ -18,6 +26,47 @@ def __call__(self, active: bool = True): self.active = active +# ====================================================================================== +# ForcedCollisions +# ====================================================================================== + + +class ForcedCollisions(ObjectSingleton): + # Annotations for Numba mode + label: str = "forced_collisions" + active: bool + + cell_IDs: list[np.int64] + threshold_weights: list[float] + target_weights: list[float] + + def __init__(self): + self.active = False + self.cell_IDs = [] + self.threshold_weights = [] + self.target_weights = [] + + def __call__(self, cells, threshold_weights=None, target_weights=None): + self.active = True + if threshold_weights is None: + threshold_weights = [0.5] * len(cells) + if target_weights is None: + target_weights = [1.0] * len(cells) + if len(cells) != len(threshold_weights) or len(cells) != len(target_weights): + print_error( + f"Expected cells, threshold_weights, and target_weights to be the same size, but got {len(cells)}, {len(threshold_weights)}, and {len(target_weights)} instead" + ) + + for cell in cells: + if cell.fill_type != FILL_MATERIAL: + print_error( + f"Invalid cell fill on cell: \n{cell}\nForced collision technique is only valid on cells with material fill" + ) + self.cell_IDs.append(cell.ID) + self.threshold_weights = threshold_weights + self.target_weights = target_weights + + # ====================================================================================== # Weighted emission # ====================================================================================== @@ -64,6 +113,108 @@ def __call__(self, weight_threshold: float = 0.0, weight_target: float = 1.0): self.weight_target = weight_target +# ====================================================================================== +# Weight Windows +# ====================================================================================== + + +class WeightWindows(ObjectSingleton): + label: str = "weight_windows" + + active: bool + + # energy + energy_bounds: NDArray[np.float64] + Ne: int + # space + mesh: MeshBase + Nx: int + Ny: int + 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")] + + def __init__(self): + self.active = False + self.energy_bounds = np.array([0.0, 1.0]) + 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.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): + # 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]) + + # get mesh size + match mesh.label: + case "uniform_mesh": + nx, ny, nz = mesh.Nx, mesh.Ny, mesh.Nz + case "structured_mesh": + nx, ny, nz = ( + mesh.x.shape[0] - 1, + mesh.y.shape[0] - 1, + mesh.z.shape[0] - 1, + ) + case _: + 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" + ) + ne = energy.shape[0] - 1 + + # check correct shape + mesh_shape = (nx, ny, nz) + 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.energy_bounds = energy + self.Ne = ne + self.mesh = mesh + self.Nx, self.Ny, self.Nz = mesh_shape + shape = (self.Ne, *mesh_shape) + self.lower_weights = weight_windows[..., 0] + self.target_weights = weight_windows[..., 1] + self.upper_weights = weight_windows[..., 2] + + # check weight windows are valid + if (self.lower_weights <= 0.0).any(): + print_error( + "Lower bound weights must be strictly positive to avoid invalid roulette behavior" + ) + if (self.lower_weights > self.target_weights).any(): + print_error( + "Lower bound weight can not be greater than the target weight for any weight window" + ) + if (self.target_weights > self.upper_weights).any(): + print_error( + "Target weight can not be greater than the upper bound weight for any weight window" + ) + + # ====================================================================================== # Population control # ====================================================================================== diff --git a/mcdc/transport/particle.py b/mcdc/transport/particle.py index e524e1c21..edeb0bd67 100644 --- a/mcdc/transport/particle.py +++ b/mcdc/transport/particle.py @@ -49,3 +49,24 @@ def copy_as_child(child_particle_container, parent_particle_container): # Evolve parent seed rng.lcg(parent_particle_container) + + +@njit +def copy_run_state(target_particle_container, source_particle_container): + """ + Helper for copying runtime particle state into new particle. + + Parameters + ---------- + target_particle_container : ndarray + Container holding the particle to copy data into. + source_particle_container : ndarray + Container holding the particle to copy data from. + """ + target_particle = target_particle_container[0] + source_particle = source_particle_container[0] + target_particle["alive"] = source_particle["alive"] + target_particle["material_ID"] = source_particle["material_ID"] + target_particle["cell_ID"] = source_particle["cell_ID"] + target_particle["surface_ID"] = source_particle["surface_ID"] + target_particle["event"] = source_particle["event"] diff --git a/mcdc/transport/physics/interface.py b/mcdc/transport/physics/interface.py index ef3cef349..f1e2c8a6a 100644 --- a/mcdc/transport/physics/interface.py +++ b/mcdc/transport/physics/interface.py @@ -30,6 +30,35 @@ def particle_speed(particle_container, simulation, data): # ====================================================================================== +@njit +def total_xs(particle_container, simulation, data): + """ + Convenience helper for getting specifically the total cross section. + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + simulation : object + Simulation object. + data : object + Simulation data for array access. + + Returns + ------- + float + Total macroscopic cross section. + """ + particle = particle_container[0] + if particle["particle_type"] == PARTICLE_NEUTRON: + total = NEUTRON_REACTION_TOTAL + return neutron.macro_xs(total, particle_container, simulation, data) + elif particle["particle_type"] == PARTICLE_ELECTRON: + total = ELECTRON_REACTION_TOTAL + return electron.macro_xs(total, particle_container, simulation, data) + return 0.0 + + @njit def macro_xs(reaction_type, particle_container, simulation, data): particle = particle_container[0] @@ -57,14 +86,8 @@ def neutron_production_xs(reaction_type, particle_container, simulation, data): @njit def collision_distance(particle_container, simulation, data): - particle = particle_container[0] - # Get total cross-section - SigmaT = 0.0 - if particle["particle_type"] == PARTICLE_NEUTRON: - SigmaT = macro_xs(NEUTRON_REACTION_TOTAL, particle_container, simulation, data) - elif particle["particle_type"] == PARTICLE_ELECTRON: - SigmaT = macro_xs(ELECTRON_REACTION_TOTAL, particle_container, simulation, data) + SigmaT = total_xs(particle_container, simulation, data) # Vacuum material? if SigmaT == 0.0: @@ -76,6 +99,41 @@ def collision_distance(particle_container, simulation, data): return distance +@njit +def forced_collision_distance(particle_container, surface_distance, simulation, data): + """ + Method for finding the distance for a forced collision particle to travel. + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + surface_distance: + The distance to the next surface along the particles direction. + simulation : object + Simulation object. + data : object + Simulation data for array access. + + Returns + ------- + distance : float + Distance for particle to travel. + """ + # Get total cross-section + SigmaT = total_xs(particle_container, simulation, data) + + # Vacuum material? + if SigmaT == 0.0: + return INF + + # Sample collision distance + xi = rng.lcg(particle_container) + + distance = -math.log(1 - xi * (1 - math.exp(-surface_distance * SigmaT))) / SigmaT + return distance + + @njit def collision(particle_container, collision_data_container, program, data): particle = particle_container[0] diff --git a/mcdc/transport/simulation.py b/mcdc/transport/simulation.py index 6f2e3f06f..cd72f7f44 100644 --- a/mcdc/transport/simulation.py +++ b/mcdc/transport/simulation.py @@ -346,9 +346,18 @@ def step_particle(particle_container, program, data): if particle["event"] & EVENT_TIME_BOUNDARY: particle["alive"] = False - # Weight roulette + # Weight splitting / rouletting + # Apply techniques if particle["alive"]: - technique.weight_roulette(particle_container, simulation) + # Weight windows + if simulation["weight_windows"]["active"]: + technique.weight_windows(particle_container, program, data) + + elif simulation["forced_collisions"]["active"]: + technique.forced_collision_roulette(particle_container, program, data) + # Weight roulette + else: + technique.weight_roulette(particle_container, simulation) @njit @@ -401,7 +410,12 @@ def move_to_event(particle_container, simulation, data): ) # Distance to next collision - d_collision = physics.collision_distance(particle_container, simulation, data) + if technique.in_forced_collision_cell(particle_container, simulation, data): + d_collision = technique.forced_collisions( + particle_container, d_boundary, simulation, data + ) + else: + d_collision = physics.collision_distance(particle_container, simulation, data) # ================================================================================== # Determine event(s) @@ -435,40 +449,10 @@ def move_to_event(particle_container, simulation, data): # ================================================================================== # Move particle # ================================================================================== - # Score tracklength tallies - if simulation["cycle_active"]: - # Cell tallies - cell = simulation["cells"][particle["cell_ID"]] - for i in range(cell["N_tally"]): - tally_base_ID = int(mcdc_get.cell.tally_IDs(i, cell, data)) - tally_base = simulation["tallies"][tally_base_ID] - - # Skip non-tracklength tallies - if tally_base["child_type"] != TALLY_TRACKLENGTH: - continue - - tally = simulation["tracklength_tallies"][tally_base["child_ID"]] - tally_module.score.tracklength_tally( - particle_container, distance, tally, simulation, data - ) - - # Other tracklength tallies - for i in range(simulation["N_tracklength_tally"]): - tally = simulation["tracklength_tallies"][i] - - # Skip cell tallies - if tally["spatial_filter_type"] == SPATIAL_FILTER_CELL: - continue - - tally_module.score.tracklength_tally( - particle_container, distance, tally, simulation, data - ) - - if settings["neutron_eigenvalue_mode"]: - tally_module.score.eigenvalue_tally( - particle_container, distance, simulation, data - ) + tally_module.score.score_tracklength_tallies( + particle_container, distance, simulation, data + ) # Move particle particle_module.move(particle_container, distance, simulation, data) diff --git a/mcdc/transport/tally/score.py b/mcdc/transport/tally/score.py index 924c03386..6c3633d51 100644 --- a/mcdc/transport/tally/score.py +++ b/mcdc/transport/tally/score.py @@ -26,6 +26,8 @@ SCORE_NET_CURRENT, SCORE_ENERGY_DEPOSITION, SPATIAL_FILTER_MESH, + SPATIAL_FILTER_CELL, + TALLY_TRACKLENGTH, ) from mcdc.transport.geometry.surface import get_normal_component from mcdc.transport.tally.filter import get_filter_indices @@ -140,6 +142,52 @@ def collision_tally( # ====================================================================================== +@njit +def score_tracklength_tallies(particle_container, distance, simulation, data): + """ + Helper for scoring traveled distance on all track length tallies in the simulation. + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + distance: + The distance traveled to score. + simulation : object + Simulation object. + data : object + Simulation data for array access. + """ + particle = particle_container[0] + + if simulation["cycle_active"]: + # Cell tallies + cell = simulation["cells"][particle["cell_ID"]] + for i in range(cell["N_tally"]): + tally_base_ID = int(mcdc_get.cell.tally_IDs(i, cell, data)) + tally_base = simulation["tallies"][tally_base_ID] + + # Skip non-tracklength tallies + if tally_base["child_type"] != TALLY_TRACKLENGTH: + continue + + tally = simulation["tracklength_tallies"][tally_base["child_ID"]] + tracklength_tally(particle_container, distance, tally, simulation, data) + + # Other tracklength tallies + for i in range(simulation["N_tracklength_tally"]): + tally = simulation["tracklength_tallies"][i] + + # Skip cell tallies + if tally["spatial_filter_type"] == SPATIAL_FILTER_CELL: + continue + + tracklength_tally(particle_container, distance, tally, simulation, data) + + if simulation["settings"]["neutron_eigenvalue_mode"]: + eigenvalue_tally(particle_container, distance, simulation, data) + + @njit def tracklength_tally(particle_container, distance, tally, simulation, data): particle = particle_container[0] diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index fe629af33..a82d0c796 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -2,14 +2,221 @@ import math from numba import njit +import numba #### +import mcdc.mcdc_get.weight_windows as ww_get import mcdc.numba_types as type_ +from mcdc.transport.mesh import get_indices as get_mesh_indices +import mcdc.transport.geometry as geometry_module import mcdc.transport.particle as particle_module import mcdc.transport.particle_bank as particle_bank_module +import mcdc.transport.tally as tally_module import mcdc.transport.rng as rng +from mcdc.transport.physics import interface as physics import mcdc.transport.util as util +import mcdc.mcdc_get as mcdc_get +from mcdc.print_ import print_error + +# ====================================================================================== +# Forced Collisions +# ====================================================================================== + + +@njit +def forced_collisions(particle_container, surface_distance, program, data): + """ + Applies the method of forced collisions, splitting the source particle into a collided and transmitted particle. This method returns the distance for the collided particle to travel, letting the `simulation.move_to_event` handle the actual transport for the collided particle. + + Parameters + ---------- + particle_container : ndarray + Container holding the original particle to copy over all data from. + surface_distance : float + The distance to the surface the transmitted particle will be moved to. + program : object + Program object containing simulation state with forced collision data. + data : object + Simulation data for array access. + + Returns + ------- + collision_distance : float + Distance for the collided component to travel. + """ + simulation = util.access_simulation(program) + + # find weight multiplier + SigmaT = physics.total_xs(particle_container, simulation, data) + weight_multiplier = math.exp(-surface_distance * SigmaT) + + # transmitted particle + bank_transmitted_particle( + particle_container, weight_multiplier, surface_distance, program, data + ) + + # alias input particle as collided particle + collided_container = particle_container + # update collided particle + collided = collided_container[0] + collided["w"] *= 1 - weight_multiplier + + # return distance to forced collision, let simulation handle the rest (tallies) + collision_distance = physics.forced_collision_distance( + collided_container, surface_distance, simulation, data + ) + return collision_distance + + +@njit +def bank_transmitted_particle( + particle_container, weight_multiplier, surface_distance, program, data +): + """ + Helper for creating and banking the transmitted component. If the + transmitted particle leaves the simulation through a boundary surface, the + particle is not banked. Additionally, the particle is scored over all + tracklength tallies via the helper in `tally.score`. + + Parameters + ---------- + particle_container : ndarray + Container holding the original particle to copy over all data from. + weight_multiplier : float + The multiplier to adjust the particle weight by. + surface_distance : float + The distance to the surface the transmitted particle will be moved to. + program : object + Program object containing simulation state with forced collision data. + data : object + Simulation data for array access. + """ + simulation = util.access_simulation(program) + + # create child copy of collided particle history + transmitted_container = util.local_array(1, type_.particle) + particle_module.copy_as_child(transmitted_container, particle_container) + particle_module.copy_run_state(transmitted_container, particle_container) + + # assign weight + transmitted = transmitted_container[0] + transmitted["w"] *= weight_multiplier + + # score tracklength tallies + tally_module.score.score_tracklength_tallies( + transmitted_container, surface_distance, simulation, data + ) + + # update position and perform surface crossing + particle_module.move(transmitted_container, surface_distance, simulation, data) + geometry_module.surface_crossing(transmitted_container, simulation, data) + + # particle could leave through BC, so check if alive before banking + if transmitted["alive"]: + particle_bank_module.bank_active_particle(transmitted_container, program) + + +@njit +def forced_collision_roulette(particle_container, program, data): + """ + Roulette procedure for forced collision. Particle is only rouletted if in a cell marked for forced collisions + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + program : object + Program object containing simulation state with forced collision data. + data : object + Simulation data for array access. + """ + simulation = util.access_simulation(program) + fc_object = simulation["forced_collisions"] + + # check if particle is in a cell with forced collisions + if not in_forced_collision_cell(particle_container, simulation, data): + return + + # get index into arrays + index = get_forced_collision_cell_index(particle_container, fc_object, data) + if index < 0: + return + + # get weights + threshold = mcdc_get.forced_collisions.threshold_weights(index, fc_object, data) + target = mcdc_get.forced_collisions.target_weights(index, fc_object, data) + + # roulette + roulette_from_weight_bounds(particle_container, threshold, target) + + +@njit +def get_forced_collision_cell_index(particle_container, fc_object, data): + """ + Helper for getting the getter index for weight roulette parameters + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + fc_object : object + Forced collision object for use in mcdc_get methods. + data : object + Simulation data for array access. + + Returns + ------- + index : int + The flattened index for getting weight roulette parameters. + """ + particle = particle_container[0] + + # grab all cell ids + cell_ids = mcdc_get.forced_collisions.cell_IDs_all(fc_object, data) + + # find the cell index + for index in range(len(cell_ids)): + if cell_ids[index] == particle["cell_ID"]: + return index + + # should never hit this, but just to be safe + return -1 + + +@njit +def in_forced_collision_cell(particle_container, simulation, data): + """ + Check if particle is in a cell marked for forced collision + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + program : object + Program object containing simulation state with forced collision data. + data : object + Simulation data for array access. + + Returns + ------- + bool + True if particle in cell marked for forced collision. + """ + fc_object = simulation["forced_collisions"] + + # not active, dont need to query cells + if not fc_object["active"]: + return False + + # active, need to check if in active cell + cell_ids = mcdc_get.forced_collisions.cell_IDs_all(fc_object, data) + if particle_container[0]["cell_ID"] in cell_ids: + return True + + # not in active cell + return False + # ====================================================================================== # Weight Roulette @@ -18,10 +225,151 @@ @njit def weight_roulette(particle_container, simulation): + threshold = simulation["weight_roulette"]["weight_threshold"] + target = simulation["weight_roulette"]["weight_target"] + roulette_from_weight_bounds(particle_container, threshold, target) + + +# ====================================================================================== +# Weight Windows +# ====================================================================================== + + +@njit +def weight_windows(particle_container, program, data): + """ + Apply weight window splitting and rouletting to a particle. + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + program : object + Program object containing simulation state with weight window and mesh data. + data : object + Simulation data for array access. + """ + simulation = util.access_simulation(program) + [lower, target, upper] = query_weight_window(particle_container, simulation, data) + split_from_weight_window(particle_container, upper, simulation) + roulette_from_weight_bounds(particle_container, lower, target) + + +@njit +def query_weight_window(particle_container, simulation, data): + """ + Query weight window bounds for the particle. + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + simulation : object + Simulation state containing weight window and mesh data. + data : object + Simulation data for array access. + + Returns + ------- + lower : float + Lower weight bound. + target : float + Target weight. + upper : float + Upper weight bound. + """ + # grab objects + ww_obj = simulation["weight_windows"] + indices = get_ww_indices(particle_container, ww_obj, simulation, data) + # grab the actual ww parameters + lower = ww_get.lower_weights(*indices, ww_obj, data) + target = ww_get.target_weights(*indices, ww_obj, data) + upper = ww_get.upper_weights(*indices, ww_obj, data) + return lower, target, upper + + +@njit +def get_ww_indices(particle_container, ww_obj, simulation, data): + """ + Get flattened weight window index from particle information + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + weight_window_object : object + The weight window object containing index information + simulation : object + Simulation state containing weight window and mesh data. + data : object + Simulation data for array access. + + Returns + ------- + indices: Tuple[int] + the flattened index in the weight window array + """ + particle = particle_container[0] + + # get energy index + energy_bounds = ww_get.energy_bounds_all(ww_obj, data) + if simulation["settings"]["neutron_multigroup_mode"]: + energy = particle["g"] + else: + energy = particle["E"] + ie = util.find_bin(energy, energy_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) + + +@njit +def split_from_weight_window(particle_container, threshold_weight, program): + """ + Split a particle if its weight exceeds the threshold. + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + threshold_weight : float + Upper weight bound triggering splitting. + program : object + Program object used for banking split particles. + """ + particle = particle_container[0] + weight = particle["w"] + if weight > threshold_weight: + # determine how many to split into + num_split = math.ceil(weight / threshold_weight) + # distribute weight + particle["w"] = weight / num_split + for _ in range(num_split - 1): + # bank split particles into the active bank + particle_bank_module.bank_active_particle(particle_container, program) + + +@njit +def roulette_from_weight_bounds(particle_container, w_threshold, w_target): + """ + Russian roulette particle if weight is below threshold. + + Parameters + ---------- + particle_container : ndarray + Container holding the particle. + w_threshold : float + Lower weight bound triggering roulette. + w_target : float + Target weight assigned upon survival. + """ particle = particle_container[0] - if particle["w"] < simulation["weight_roulette"]["weight_threshold"]: - w_target = simulation["weight_roulette"]["weight_target"] + if particle["w"] < w_threshold: survival_probability = particle["w"] / w_target + # sample random number to determine survival if rng.lcg(particle_container) < survival_probability: particle["w"] = w_target else: diff --git a/test/regression/basic_weight_windows/answer.h5 b/test/regression/basic_weight_windows/answer.h5 new file mode 100644 index 000000000..32f66b1dd Binary files /dev/null and b/test/regression/basic_weight_windows/answer.h5 differ diff --git a/test/regression/basic_weight_windows/input.py b/test/regression/basic_weight_windows/input.py new file mode 100644 index 000000000..38d3e9805 --- /dev/null +++ b/test/regression/basic_weight_windows/input.py @@ -0,0 +1,63 @@ +import mcdc +from mcdc.constant import INF +import numpy as np +import os + +os.environ["MCDC_LIB"] = "../mcdc-regression_test_data/" + +# Uses Pincell model as the basis + +# Material +fuel = mcdc.Material( + nuclide_composition={ + "U235": 0.0001654509603995036, + "U238": 0.022801089905717036, + "O16": 0.04593308173223308, + } +) +moderator = mcdc.Material( + nuclide_composition={ + "H1": 0.05129627050184732, + "O16": 0.024622209840886707, + "B10": 4.103701640147785e-05, + } +) + +# Geometry +cylinder = mcdc.Surface.CylinderZ(radius=0.45720) +pitch = 1.25984 +x0 = mcdc.Surface.PlaneX(x=-pitch / 2, boundary_condition="reflective") +x1 = mcdc.Surface.PlaneX(x=pitch / 2, boundary_condition="reflective") +y0 = mcdc.Surface.PlaneY(y=-pitch / 2, boundary_condition="reflective") +y1 = mcdc.Surface.PlaneY(y=pitch / 2, boundary_condition="reflective") +# +mcdc.Cell(-cylinder, fill=fuel) +mcdc.Cell(+x0 & -x1 & +y0 & -y1 & +cylinder, fill=moderator) + +# Source +mcdc.Source(position=[0.0, 0.0, 0.0], isotropic=True, time=0.0, energy=14.1e6) + +# Setting +mcdc.settings.N_particle = 20 +mcdc.settings.N_batch = 2 +mcdc.settings.time_boundary = 1.0 +mcdc.settings.active_bank_buffer = 1000 + +# Mesh +Nx, Ny = 20, 20 +x0, y0 = -pitch / 2, -pitch / 2 +dx, dy = pitch / Nx, pitch / Ny +mesh = mcdc.MeshUniform(x=(x0, dx, Nx), y=(y0, dy, Ny)) + +# Tally +mcdc.Tally(mesh=mesh, scores=["flux"]) + +# Weight windows +ww_array = np.ones((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 +ww_array[..., 2] = 0.9 # forces splitting on all particles born with w=1.0 +mcdc.simulation.weight_windows(ww_array, mesh=mesh) + +mcdc.run() diff --git a/test/regression/slab_beam_forced_collision/answer.h5 b/test/regression/slab_beam_forced_collision/answer.h5 new file mode 100644 index 000000000..6d43b155f Binary files /dev/null and b/test/regression/slab_beam_forced_collision/answer.h5 differ diff --git a/test/regression/slab_beam_forced_collision/input.py b/test/regression/slab_beam_forced_collision/input.py new file mode 100644 index 000000000..7b391604e --- /dev/null +++ b/test/regression/slab_beam_forced_collision/input.py @@ -0,0 +1,56 @@ +import numpy as np +import os +import mcdc + +os.environ["MCDC_LIB"] = "../mcdc-regression_test_data/" + +# ====================================================================================== +# Set model +# ====================================================================================== +# Finite homogeneous pure-absorbing slab + +# Set materials +# Set materials +generate_material = lambda atomdensity: mcdc.Material( + nuclide_composition={"H1": atomdensity} +) +m1 = generate_material(0.0001) + +# Set surfaces +s1 = mcdc.Surface.PlaneX(x=0.0, boundary_condition="vacuum") +s2 = mcdc.Surface.PlaneX(x=1.0, boundary_condition="vacuum") + +# Set cells +low_xs_cell = mcdc.Cell(region=+s1 & -s2, fill=m1) + +# ====================================================================================== +# Set source +# ====================================================================================== +# Isotropic beam from left-end + +mcdc.Source(position=(0.0, 0.0, 0.0), direction=(1.0, 0.0, 0.0)) + +# ====================================================================================== +# Set tallies, settings, and run MC/DC +# ====================================================================================== + +# Tallies +# energy deposition for actual VR against analog +mesh = mcdc.MeshUniform() +mcdc.Tally( + mesh=mesh, + scores=["energy_deposition"], +) +# flux to make sure tracklength is unbiased +mcdc.Tally(cell=low_xs_cell, scores=["flux"]) +# net-current to make sure surface is unbiased +mcdc.Tally(surface=s2, scores=["net-current"]) + +# Settings +mcdc.settings.N_particle = 5000 +mcdc.settings.N_batch = 2 + +mcdc.simulation.forced_collisions(cells=[low_xs_cell]) + +# Run +mcdc.run() diff --git a/test/unit/transport/technique/forced_collisions.py b/test/unit/transport/technique/forced_collisions.py new file mode 100644 index 000000000..498b3c83f --- /dev/null +++ b/test/unit/transport/technique/forced_collisions.py @@ -0,0 +1,89 @@ +import numpy as np +import os +import pytest +from numba import njit +import mcdc +from mcdc.main import preparation +from mcdc.transport import technique +import mcdc.numba_types as types_ +from mcdc.transport.util import access_simulation +from mcdc.transport import geometry +from mcdc.mcdc_get import forced_collisions as get_fc + +os.environ["MCDC_LIB"] = "../../../regression/mcdc-regression_test_data/" + + +# ============================================================================= +# Model base fixture +# ============================================================================= + + +@pytest.fixture +def pin_cell_model(): + # Material + fuel = mcdc.Material(nuclide_composition={"U235": 1.0}) + moderator = mcdc.Material(nuclide_composition={"H1": 1.0}) + + # Geometry + cylinder = mcdc.Surface.CylinderZ(radius=0.5) + pitch = 1.5 + x0 = mcdc.Surface.PlaneX(x=-pitch / 2, boundary_condition="reflective") + x1 = mcdc.Surface.PlaneX(x=pitch / 2, boundary_condition="reflective") + y0 = mcdc.Surface.PlaneY(y=-pitch / 2, boundary_condition="reflective") + y1 = mcdc.Surface.PlaneY(y=pitch / 2, boundary_condition="reflective") + # + fuel_cell = mcdc.Cell(-cylinder, fill=fuel) + mod_cell = mcdc.Cell(+x0 & -x1 & +y0 & -y1 & +cylinder, fill=moderator) + + yield fuel_cell, mod_cell + + +# ============================================================================= +# Error throwing in object creation +# ============================================================================= + + +@pytest.mark.parametrize( + "cells_builder, thresholds, targets, expected_msg", + [ + ( + lambda fuel_cell: [fuel_cell], + [0.5, 0.5], + [1.0], + "Expected cells, threshold_weights, and target_weights to be the same size", + ), + ( + lambda fuel_cell: [fuel_cell], + [0.5], + [1.0, 1.0], + "Expected cells, threshold_weights, and target_weights to be the same size", + ), + ( + lambda fuel_cell: [mcdc.Cell(fill=mcdc.Universe(cells=[fuel_cell]))], + None, + None, + "Invalid cell fill on cell", + ), + ], +) +def test_forced_collisions_error_throw( + pin_cell_model, capsys, cells_builder, thresholds, targets, expected_msg +): + fuel_cell, mod_cell = pin_cell_model + + cells = cells_builder(fuel_cell) + + with pytest.raises(SystemExit): + mcdc.simulation.forced_collisions( + cells, + threshold_weights=thresholds, + target_weights=targets, + ) + + captured = capsys.readouterr() + assert expected_msg in captured.out + + +# ============================================================================= +# Method tests +# ============================================================================= diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py new file mode 100644 index 000000000..e29709cd2 --- /dev/null +++ b/test/unit/transport/technique/weight_windows.py @@ -0,0 +1,205 @@ +import mcdc +import numpy as np +import pytest +import os + +from mcdc.main import preparation +import mcdc.numba_types as type_ +from mcdc.transport.technique import ( + roulette_from_weight_bounds, + split_from_weight_window, + query_weight_window, + weight_windows, + particle_bank_module, +) +from mcdc.transport.mesh.interface import get_indices +from mcdc.constant import TINY +import mcdc.transport.util as util + +# =========================================================================== # +# Helper method for creating a dummy model +# =========================================================================== # + + +def make_mesh(): + # constants + pitch = 2.0 + height = 10.0 + N = 3 + + mesh = mcdc.MeshUniform( + "mesh", + x=(-pitch / 2, pitch / N, N), + y=(-pitch / 2, pitch / N, N), + z=(0.0, height / N, N), + ) + + return mesh, N + + +def make_ww_model_params(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): + import mcdc + + mesh, N = make_mesh() + Ne = 1 + + if mess_up_size: + ww_array = np.ones((Ne, N, N, 4, 3)) + else: + # global assign for simplicity + ww_array = np.ones((Ne, N, N, N, 3)) + ww_array[..., 0] = lower + ww_array[..., 1] = target + ww_array[..., 2] = upper + + mcdc.simulation.weight_windows(ww_array, mesh=mesh) + + mcdc_container, data = preparation() + return mcdc_container[0], data + + +def make_ww_model_distinct(): + import mcdc + + mesh, N = make_mesh() + energy = np.linspace(0.0, 6.0, 7) + Ne = 6 + + ww_array = np.empty((Ne, 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) + + mcdc_container, data = preparation() + return mcdc_container[0], data + + +# =========================================================================== # +# Test error throwing in object creation +# =========================================================================== # + + +@pytest.mark.parametrize( + "kwargs, expected_msg", + [ + # incorrect size + ( + {"mess_up_size": True}, + "Weight window array has shape (1, 3, 3, 4, 3), but expected (1, 3, 3, 3, 3)", + ), + # negative lower + ( + {"lower": -1.0}, + "Lower bound weights must be strictly positive", + ), + # lower > target + ( + {"lower": 1.0, "target": 0.5}, + "Lower bound weight can not be greater than the target weight", + ), + # target > upper + ( + {"target": 1.5}, + "Target weight can not be greater than the upper bound weight", + ), + ], +) +def test_error_throw(capsys, kwargs, expected_msg): + with pytest.raises(SystemExit): + make_ww_model_params(**kwargs) + + out = capsys.readouterr().out + assert expected_msg in out + + +# =========================================================================== # +# Tests for helper methods +# =========================================================================== # + + +def test_roulette_from_weight_bounds(): + # because of rng, want to loop over to hit both branches + for _ in range(10): + particles = np.zeros(1, type_.particle_data) + particles[0]["w"] = 0.1 + target = 0.2 + threshold = 0.1 + TINY + + roulette_from_weight_bounds(particles, threshold, target) + p = particles[0] + assert p["w"] == target or not p["alive"] + + +def test_split_from_weight_window(): + particles = np.zeros(1, type_.particle_data) + init_weight = 2.0 + TINY + particles[0]["w"] = init_weight + threshold = 1.0 + program, data = make_ww_model_distinct() + simulation = util.access_simulation(program) + + # get bank and init size + bank = simulation["bank_active"] + init_bank_size = particle_bank_module.get_bank_size(bank) + + # split + split_from_weight_window(particles, threshold, program) + + p1 = particles[0] + num_split = np.ceil(init_weight / threshold) + num_new = num_split - 1 + + # check weight of original particle + assert p1["w"] == init_weight / num_split + # check particles were created and check their weights + assert particle_bank_module.get_bank_size(bank) == init_bank_size + num_new + for i in range(2): + pnew = bank["particle_data"][init_bank_size + i] + assert pnew["w"] == p1["w"] + + +def test_query_weight_window(): + p = np.zeros(1, type_.particle_data) + + program, data = make_ww_model_distinct() + simulation = util.access_simulation(program) + simulation["settings"]["neutron_multigroup_mode"] = False + # hardcode mesh params + pitch, height, N = 2.0, 10.0, 3 + nx, ny, nz = N, N, N + xmin, ymin, zmin = -pitch / 2, -pitch / 2, 0.0 + dx, dy, dz = pitch / N, pitch / N, height / N + + # hardcode energy params + energies = np.linspace(0.5, 5.5, 6) + + # 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