diff --git a/docs/source/pythonapi/index.rst b/docs/source/pythonapi/index.rst index 3b108838..c16a2fa5 100644 --- a/docs/source/pythonapi/index.rst +++ b/docs/source/pythonapi/index.rst @@ -89,6 +89,7 @@ Techniques are enabled by calling methods on the ``mcdc.simulation`` singleton: - ``mcdc.simulation.implicit_capture(active=True)`` - ``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(mesh, weight_windows)`` - ``mcdc.simulation.population_control(active=True)`` Running diff --git a/mcdc/mcdc_get/__init__.py b/mcdc/mcdc_get/__init__.py index d3b53338..b5a5d9a6 100644 --- a/mcdc/mcdc_get/__init__.py +++ b/mcdc/mcdc_get/__init__.py @@ -92,6 +92,8 @@ 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/weight_windows.py b/mcdc/mcdc_get/weight_windows.py new file mode 100644 index 00000000..00252e54 --- /dev/null +++ b/mcdc/mcdc_get/weight_windows.py @@ -0,0 +1,90 @@ +# The following is automatically generated by code_factory.py + +from numba import njit + + +@njit +def lower_weights(index, weight_windows, data): + offset = weight_windows["lower_weights_offset"] + return data[offset + index] + + +@njit +def lower_weights_all(weight_windows, data): + start = weight_windows["lower_weights_offset"] + size = weight_windows["lower_weights_length"] + end = start + size + return data[start:end] + + +@njit +def lower_weights_last(weight_windows, data): + start = weight_windows["lower_weights_offset"] + size = weight_windows["lower_weights_length"] + end = start + size + return data[end - 1] + + +@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, weight_windows, data): + offset = weight_windows["target_weights_offset"] + return data[offset + index] + + +@njit +def target_weights_all(weight_windows, data): + start = weight_windows["target_weights_offset"] + size = weight_windows["target_weights_length"] + end = start + size + return data[start:end] + + +@njit +def target_weights_last(weight_windows, data): + start = weight_windows["target_weights_offset"] + size = weight_windows["target_weights_length"] + end = start + size + return data[end - 1] + + +@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, weight_windows, data): + offset = weight_windows["upper_weights_offset"] + return data[offset + index] + + +@njit +def upper_weights_all(weight_windows, data): + start = weight_windows["upper_weights_offset"] + size = weight_windows["upper_weights_length"] + end = start + size + return data[start:end] + + +@njit +def upper_weights_last(weight_windows, data): + start = weight_windows["upper_weights_offset"] + size = weight_windows["upper_weights_length"] + end = start + size + return data[end - 1] + + +@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 38ce0ddd..cfaf2cdb 100644 --- a/mcdc/mcdc_set/__init__.py +++ b/mcdc/mcdc_set/__init__.py @@ -92,6 +92,8 @@ 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/weight_windows.py b/mcdc/mcdc_set/weight_windows.py new file mode 100644 index 00000000..95613786 --- /dev/null +++ b/mcdc/mcdc_set/weight_windows.py @@ -0,0 +1,90 @@ +# The following is automatically generated by code_factory.py + +from numba import njit + + +@njit +def lower_weights(index, weight_windows, data, value): + offset = weight_windows["lower_weights_offset"] + data[offset + index] = value + + +@njit +def lower_weights_all(weight_windows, data, value): + start = weight_windows["lower_weights_offset"] + size = weight_windows["lower_weights_length"] + end = start + size + data[start:end] = value + + +@njit +def lower_weights_last(weight_windows, data, value): + start = weight_windows["lower_weights_offset"] + size = weight_windows["lower_weights_length"] + end = start + size + data[end - 1] = 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, weight_windows, data, value): + offset = weight_windows["target_weights_offset"] + data[offset + index] = value + + +@njit +def target_weights_all(weight_windows, data, value): + start = weight_windows["target_weights_offset"] + size = weight_windows["target_weights_length"] + end = start + size + data[start:end] = value + + +@njit +def target_weights_last(weight_windows, data, value): + start = weight_windows["target_weights_offset"] + size = weight_windows["target_weights_length"] + end = start + size + data[end - 1] = 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, weight_windows, data, value): + offset = weight_windows["upper_weights_offset"] + data[offset + index] = value + + +@njit +def upper_weights_all(weight_windows, data, value): + start = weight_windows["upper_weights_offset"] + size = weight_windows["upper_weights_length"] + end = start + size + data[start:end] = value + + +@njit +def upper_weights_last(weight_windows, data, value): + start = weight_windows["upper_weights_offset"] + size = weight_windows["upper_weights_length"] + end = start + size + data[end - 1] = 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/object_/simulation.py b/mcdc/object_/simulation.py index fa5b5f9b..57eba631 100644 --- a/mcdc/object_/simulation.py +++ b/mcdc/object_/simulation.py @@ -5,6 +5,7 @@ ImplicitCapture, PopulationControl, WeightRoulette, + WeightWindows, WeightedEmission, ) @@ -82,6 +83,7 @@ class Simulation(ObjectSingleton): implicit_capture: ImplicitCapture weighted_emission: WeightedEmission weight_roulette: WeightRoulette + weight_windows: WeightWindows population_control: PopulationControl # Particle banks @@ -167,6 +169,7 @@ def __init__(self): self.implicit_capture = ImplicitCapture() 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 1c3bacc1..f18fcafd 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -1,4 +1,7 @@ +from numpy.typing import NDArray +from numpy import float64 from mcdc.object_.base import ObjectSingleton +from mcdc.object_.mesh import MeshBase from mcdc.print_ import print_error # ====================================================================================== @@ -64,6 +67,75 @@ 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 + mesh: MeshBase + Nx: int + Ny: int + Nz: int + + # flattened arrays of ww params + lower_weights: NDArray[float64] + target_weights: NDArray[float64] + upper_weights: NDArray[float64] + + def __init__(self): + self.active = False + + def __call__(self, mesh, weight_windows): + # 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" + ) + + # check correct shape + mesh_shape = (nx, ny, nz) + ww_shape = weight_windows.shape + expected_shape = (*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.mesh = mesh + self.Nx, self.Ny, self.Nz = mesh_shape + self.lower_weights = weight_windows[..., 0].reshape(-1) + self.target_weights = weight_windows[..., 1].reshape(-1) + self.upper_weights = weight_windows[..., 2].reshape(-1) + + # 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/simulation.py b/mcdc/transport/simulation.py index 6f2e3f06..74d3849d 100644 --- a/mcdc/transport/simulation.py +++ b/mcdc/transport/simulation.py @@ -346,8 +346,10 @@ def step_particle(particle_container, program, data): if particle["event"] & EVENT_TIME_BOUNDARY: particle["alive"] = False - # Weight roulette - if particle["alive"]: + # Weight splitting / rouletting + if particle["alive"] & simulation["weight_windows"]["active"]: + technique.weight_windows(particle_container, simulation, data) + elif particle["alive"]: technique.weight_roulette(particle_container, simulation) diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index fe629af3..cb586e69 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -5,7 +5,9 @@ #### +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.particle as particle_module import mcdc.transport.particle_bank as particle_bank_module import mcdc.transport.rng as rng @@ -28,6 +30,117 @@ def weight_roulette(particle_container, simulation): particle["alive"] = False +# ====================================================================================== +# Weight Windows +# ====================================================================================== + + +@njit +def weight_windows(particle_container, simulation, data): + """ + Apply weight window splitting and rouletting to a 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. + """ + [lower, target, upper] = query_weight_window(particle_container, simulation, data) + split_from_weight_window(particle_container, upper, simulation) + roulette_from_weight_window(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"] + mesh = simulation["meshes"][ww_obj["mesh_ID"]] + # find spatial indices from mesh + idx, idy, idz = get_mesh_indices(particle_container, mesh, simulation, data) + # get index of ww in ww arrays + index = ((idx * ww_obj["Ny"]) + idy) * ww_obj["Nz"] + idz + # grab the actual ww parameters + lower = ww_get.lower_weights(index, ww_obj, data) + target = ww_get.target_weights(index, ww_obj, data) + upper = ww_get.upper_weights(index, ww_obj, data) + return lower, target, upper + + +@njit +def split_from_weight_window(particle_container, threshold_weight, simulation): + """ + 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. + simulation : object + Simulation state 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, simulation) + + +@njit +def roulette_from_weight_window(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"] < 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: + particle["alive"] = False + + # ====================================================================================== # Population Control # ====================================================================================== diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py new file mode 100644 index 00000000..7a12eb5d --- /dev/null +++ b/test/unit/transport/technique/weight_windows.py @@ -0,0 +1,192 @@ +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_window, + 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 + +# =========================================================================== # +# 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() + + if mess_up_size: + ww_array = np.ones((N, N, 4, 3)) + else: + # global assign for simplicity + ww_array = np.ones((N, N, N, 3)) + ww_array[..., 0] = lower + ww_array[..., 1] = target + ww_array[..., 2] = upper + + mcdc.simulation.weight_windows(mesh, ww_array) + + mcdc_container, data = preparation() + return mcdc_container[0], data + + +def make_ww_model_distinct(): + import mcdc + + mesh, N = make_mesh() + + ww_array = np.empty((N, N, N, 3)) + + # value at index is related to index, easy to predict during later test + for i in range(N): + for j in range(N): + for k in range(N): + val = 100 * i + 10 * j + k + 1 + ww_array[i, j, k, 0] = val + ww_array[i, j, k, 1] = 1000 + val + ww_array[i, j, k, 2] = 2000 + val + + mcdc.simulation.weight_windows(mesh, ww_array) + + 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 (3, 3, 4, 3), but expected (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_roullete_from_weight_window(): + # 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_window(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 + mcdc_obj, data = make_ww_model_distinct() + + # get bank and init size + bank = mcdc_obj["bank_active"] + init_bank_size = particle_bank_module.get_bank_size(bank) + + # split + split_from_weight_window(particles, threshold, mcdc_obj) + + 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) + + mcdc_obj, data = make_ww_model_distinct() + + # 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 + + # loop over all bins, check query against expected ww + 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) + + # query and predict + lower, target, upper = query_weight_window(p, mcdc_obj, data) + exp_lower = 100 * ix + 10 * iy + iz + 1 + exp_target = 1000 + exp_lower + exp_upper = 2000 + exp_lower + + assert lower == exp_lower + assert target == exp_target + assert upper == exp_upper