From 529067f95841a9410e15d9b2a8dc906ed9f24ac2 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 1 Apr 2026 18:44:43 -0700 Subject: [PATCH 01/18] add split method to techniques --- mcdc/transport/simulation.py | 3 ++- mcdc/transport/technique.py | 17 +++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/mcdc/transport/simulation.py b/mcdc/transport/simulation.py index ba0be763..796e32be 100644 --- a/mcdc/transport/simulation.py +++ b/mcdc/transport/simulation.py @@ -305,8 +305,9 @@ def step_particle(particle_container, mcdc, data): if particle["event"] & EVENT_TIME_BOUNDARY: particle["alive"] = False - # Weight roulette + # Weight splitting / rouletting if particle["alive"]: + technique.weight_split(particle_container, mcdc) technique.weight_roulette(particle_container, mcdc) diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index f8b15c05..340e73c1 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -10,6 +10,23 @@ import mcdc.transport.particle_bank as particle_bank_module import mcdc.transport.rng as rng +# ====================================================================================== +# Weight Splitting +# ====================================================================================== + + +@njit +def weight_split(particle_container, mcdc): + particle = particle_container[0] + weight = particle["w"] + threshold = mcdc["weight_split"]["weight_threshold"] + if weight > threshold: + num_split = math.ceil(weight / threshold) + particle["w"] = weight / num_split + for _ in range(num_split - 1): + particle_bank_module.bank_active_particle(particle_container, mcdc) + + # ====================================================================================== # Weight Roulette # ====================================================================================== From 023cbbec0071e45f708622d53a84496a48a62f57 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 1 Apr 2026 18:45:34 -0700 Subject: [PATCH 02/18] add split method to simulation --- mcdc/object_/simulation.py | 3 +++ mcdc/object_/technique.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/mcdc/object_/simulation.py b/mcdc/object_/simulation.py index f7a1156e..3acd888a 100644 --- a/mcdc/object_/simulation.py +++ b/mcdc/object_/simulation.py @@ -4,6 +4,7 @@ from mcdc.object_.technique import ( ImplicitCapture, PopulationControl, + WeightSplit, WeightRoulette, WeightedEmission, ) @@ -77,6 +78,7 @@ class Simulation(ObjectSingleton): # Techniques implicit_capture: ImplicitCapture weighted_emission: WeightedEmission + weight_split: WeightSplit weight_roulette: WeightRoulette population_control: PopulationControl @@ -159,6 +161,7 @@ def __init__(self): # Techniques self.implicit_capture = ImplicitCapture() self.weighted_emission = WeightedEmission() + self.weight_split = WeightSplit() self.weight_roulette = WeightRoulette() self.population_control = PopulationControl() diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index 1c3bacc1..0a7a051b 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -39,6 +39,24 @@ def __call__(self, active: bool = True, weight_target: float = 1.0): self.weight_target = weight_target +# ====================================================================================== +# Weight split +# ====================================================================================== + + +class WeightSplit(ObjectSingleton): + # Annotations for Numba mode + label: str = "weight_split" + + weight_threshold: float + + def __init__(self): + self.weight_threshold = 1.0 + + def __call__(self, weight_threshold: float = 1.0): + self.weight_threshold = weight_threshold + + # ====================================================================================== # Weight roulette # ====================================================================================== From 1ce5146a194c204c6aeffa5ee92cd9a8c4a26c22 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 1 Apr 2026 21:52:11 -0700 Subject: [PATCH 03/18] remove initial split, add weight window mockup --- mcdc/object_/technique.py | 54 ++++++++++++++++++++++++------------- mcdc/transport/technique.py | 53 +++++++++++++++++++++++++++--------- 2 files changed, 75 insertions(+), 32 deletions(-) diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index 0a7a051b..f8848863 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -1,4 +1,7 @@ -from mcdc.object_.base import ObjectSingleton +from dataclasses import dataclass +from numpy.typing import NDArray +from mcdc.object_.base import ObjectBase, ObjectSingleton +from mcdc.object_.mesh import MeshBase from mcdc.print_ import print_error # ====================================================================================== @@ -39,24 +42,6 @@ def __call__(self, active: bool = True, weight_target: float = 1.0): self.weight_target = weight_target -# ====================================================================================== -# Weight split -# ====================================================================================== - - -class WeightSplit(ObjectSingleton): - # Annotations for Numba mode - label: str = "weight_split" - - weight_threshold: float - - def __init__(self): - self.weight_threshold = 1.0 - - def __call__(self, weight_threshold: float = 1.0): - self.weight_threshold = weight_threshold - - # ====================================================================================== # Weight roulette # ====================================================================================== @@ -82,6 +67,37 @@ def __call__(self, weight_threshold: float = 0.0, weight_target: float = 1.0): self.weight_target = weight_target +# ====================================================================================== +# Weight Windows +# ====================================================================================== + + +@dataclass +class WeightWindow(ObjectBase): + label: str = "weight_window" + upper_bound: float = 1.0 + lower_bound: float = 0.0 + target_weight: float = 1.0 + + +class WeightWindows(ObjectSingleton): + label: str = "weight_windows" + + mesh: MeshBase + # 3d array of datatype WeightWindow + weight_windows: NDArray[tuple[int, int, int], WeightWindow] + + def __call__(self, mesh, weight_windows): + ww_shape = weight_windows.shape + mesh_shape = (mesh.Nx, mesh.Ny, mesh.Nz) + if ww_shape != mesh_shape: + print_error( + f"Weight window array has shape {ww_shape}, but expected {mesh_shape}" + ) + self.mesh = mesh + self.weight_windows = weight_windows + + # ====================================================================================== # Population control # ====================================================================================== diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index 340e73c1..26891590 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -6,37 +6,64 @@ #### 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 + # ====================================================================================== -# Weight Splitting +# Weight Roulette # ====================================================================================== @njit -def weight_split(particle_container, mcdc): +def weight_roulette(particle_container, mcdc): particle = particle_container[0] - weight = particle["w"] - threshold = mcdc["weight_split"]["weight_threshold"] - if weight > threshold: - num_split = math.ceil(weight / threshold) - particle["w"] = weight / num_split - for _ in range(num_split - 1): - particle_bank_module.bank_active_particle(particle_container, mcdc) + if particle["w"] < mcdc["weight_roulette"]["weight_threshold"]: + w_target = mcdc["weight_roulette"]["weight_target"] + survival_probability = particle["w"] / w_target + if rng.lcg(particle_container) < survival_probability: + particle["w"] = w_target + else: + particle["alive"] = False # ====================================================================================== -# Weight Roulette +# Weight Windows # ====================================================================================== @njit -def weight_roulette(particle_container, mcdc): +def weight_windows(particle_container, mcdc, data): + weight_window = query_weight_window(particle_container, mcdc, data) + split_from_weight_window(particle_container, weight_window["upper_bound"], mcdc) + roulette_from_weight_window(particle_container, weight_window["lower_bound"], weight_window["target_weight"]) + + +@njit +def query_weight_window(particle_container, mcdc, data): + mesh = mcdc["weight_windows"]["mesh"] + ww_array = mcdc["weight_windows"]["weight_windows"] + idx, idy, idz = get_mesh_indices(particle_container, mesh, mcdc, data) + return ww_array[idx, idy, idz] + + +@njit +def split_from_weight_window(particle_container, threshold_weight, mcdc): particle = particle_container[0] - if particle["w"] < mcdc["weight_roulette"]["weight_threshold"]: - w_target = mcdc["weight_roulette"]["weight_target"] + weight = particle["w"] + if weight > threshold_weight: + num_split = math.ceil(weight / threshold_weight) + particle["w"] = weight / num_split + for _ in range(num_split - 1): + particle_bank_module.bank_active_particle(particle_container, mcdc) + + +@njit +def roulette_from_weight_window(particle_container, w_threshold, w_target): + particle = particle_container[0] + if particle["w"] < w_threshold: survival_probability = particle["w"] / w_target if rng.lcg(particle_container) < survival_probability: particle["w"] = w_target From 3d3e239bd9e93494ceeca4ae892c3d36b5143e9f Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Wed, 1 Apr 2026 22:00:02 -0700 Subject: [PATCH 04/18] add notion of active for weight windows, fully removed initial split mockup --- mcdc/object_/simulation.py | 6 +++--- mcdc/object_/technique.py | 5 +++++ mcdc/transport/simulation.py | 5 +++-- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/mcdc/object_/simulation.py b/mcdc/object_/simulation.py index 3acd888a..6b75fc43 100644 --- a/mcdc/object_/simulation.py +++ b/mcdc/object_/simulation.py @@ -4,8 +4,8 @@ from mcdc.object_.technique import ( ImplicitCapture, PopulationControl, - WeightSplit, WeightRoulette, + WeightWindows, WeightedEmission, ) @@ -78,8 +78,8 @@ class Simulation(ObjectSingleton): # Techniques implicit_capture: ImplicitCapture weighted_emission: WeightedEmission - weight_split: WeightSplit weight_roulette: WeightRoulette + weight_windows: WeightWindows population_control: PopulationControl # Particle banks @@ -161,8 +161,8 @@ def __init__(self): # Techniques self.implicit_capture = ImplicitCapture() self.weighted_emission = WeightedEmission() - self.weight_split = WeightSplit() 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 f8848863..2ed660e4 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -83,10 +83,14 @@ class WeightWindow(ObjectBase): class WeightWindows(ObjectSingleton): label: str = "weight_windows" + active: bool mesh: MeshBase # 3d array of datatype WeightWindow weight_windows: NDArray[tuple[int, int, int], WeightWindow] + def __init__(self): + self.active = False + def __call__(self, mesh, weight_windows): ww_shape = weight_windows.shape mesh_shape = (mesh.Nx, mesh.Ny, mesh.Nz) @@ -94,6 +98,7 @@ def __call__(self, mesh, weight_windows): print_error( f"Weight window array has shape {ww_shape}, but expected {mesh_shape}" ) + self.active = True self.mesh = mesh self.weight_windows = weight_windows diff --git a/mcdc/transport/simulation.py b/mcdc/transport/simulation.py index 796e32be..cafd5a00 100644 --- a/mcdc/transport/simulation.py +++ b/mcdc/transport/simulation.py @@ -306,8 +306,9 @@ def step_particle(particle_container, mcdc, data): particle["alive"] = False # Weight splitting / rouletting - if particle["alive"]: - technique.weight_split(particle_container, mcdc) + if particle["alive"] & mcdc.weight_windows["active"]: + technique.weight_windows(particle_container, mcdc, data) + elif particle["alive"]: technique.weight_roulette(particle_container, mcdc) From 520986de371b5122b58984dd090c83e745e0300d Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 18:22:12 -0700 Subject: [PATCH 05/18] switch to 1darray from nested, fix mesh get --- mcdc/object_/technique.py | 41 ++++++++++++++++++++++++------------ mcdc/transport/simulation.py | 2 +- mcdc/transport/technique.py | 15 +++++++------ 3 files changed, 36 insertions(+), 22 deletions(-) diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index 2ed660e4..535dbb1f 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -1,5 +1,6 @@ from dataclasses import dataclass from numpy.typing import NDArray +from numpy import float64 from mcdc.object_.base import ObjectBase, ObjectSingleton from mcdc.object_.mesh import MeshBase from mcdc.print_ import print_error @@ -72,35 +73,47 @@ def __call__(self, weight_threshold: float = 0.0, weight_target: float = 1.0): # ====================================================================================== -@dataclass -class WeightWindow(ObjectBase): - label: str = "weight_window" - upper_bound: float = 1.0 - lower_bound: float = 0.0 - target_weight: float = 1.0 - - class WeightWindows(ObjectSingleton): label: str = "weight_windows" active: bool mesh: MeshBase - # 3d array of datatype WeightWindow - weight_windows: NDArray[tuple[int, int, int], WeightWindow] + Nx: int + Ny: int + Nz: int + + # flattened array of ww + weight_windows: 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], mesh.y.shape[0], mesh.z.shape[0] + case _: + print_error( + f"{type(mesh).__name__} is not supported for weight windows" + ) + + mesh_shape = (nx, ny, nz) ww_shape = weight_windows.shape - mesh_shape = (mesh.Nx, mesh.Ny, mesh.Nz) - if ww_shape != mesh_shape: + expected_shape = (*mesh_shape, 3) + if ww_shape != expected_shape: print_error( - f"Weight window array has shape {ww_shape}, but expected {mesh_shape}" + f"Weight window array has shape {ww_shape}, but expected {expected_shape}" ) + self.active = True self.mesh = mesh - self.weight_windows = weight_windows + self.Nx = nx + self.Ny = ny + self.Nz = nz + self.weight_windows = weight_windows.reshape(-1) # ====================================================================================== diff --git a/mcdc/transport/simulation.py b/mcdc/transport/simulation.py index cafd5a00..89b18fef 100644 --- a/mcdc/transport/simulation.py +++ b/mcdc/transport/simulation.py @@ -306,7 +306,7 @@ def step_particle(particle_container, mcdc, data): particle["alive"] = False # Weight splitting / rouletting - if particle["alive"] & mcdc.weight_windows["active"]: + if particle["alive"] & mcdc["weight_windows"]["active"]: technique.weight_windows(particle_container, mcdc, data) elif particle["alive"]: technique.weight_roulette(particle_container, mcdc) diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index 26891590..56a9b49e 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -4,7 +4,7 @@ from numba import njit #### - +from mcdc.mcdc_get.weight_windows import weight_windows_all, weight_windows_chunk import mcdc.numba_types as type_ from mcdc.transport.mesh import get_indices as get_mesh_indices import mcdc.transport.particle as particle_module @@ -36,17 +36,18 @@ def weight_roulette(particle_container, mcdc): @njit def weight_windows(particle_container, mcdc, data): - weight_window = query_weight_window(particle_container, mcdc, data) - split_from_weight_window(particle_container, weight_window["upper_bound"], mcdc) - roulette_from_weight_window(particle_container, weight_window["lower_bound"], weight_window["target_weight"]) + [lower, target, upper] = query_weight_window(particle_container, mcdc, data) + split_from_weight_window(particle_container, upper, mcdc) + roulette_from_weight_window(particle_container, lower, target) @njit def query_weight_window(particle_container, mcdc, data): - mesh = mcdc["weight_windows"]["mesh"] - ww_array = mcdc["weight_windows"]["weight_windows"] + ww_obj = mcdc["weight_windows"] + mesh = mcdc["meshes"][ww_obj["mesh_ID"]] idx, idy, idz = get_mesh_indices(particle_container, mesh, mcdc, data) - return ww_array[idx, idy, idz] + start = 3*(idx * mesh["Nx"] + idy * mesh["Ny"] + idz * mesh["Nz"]) + return weight_windows_chunk(start, 3, ww_obj, data) @njit From f9d6940cc6a4680aaf09f82c02bfe6b8a5ceb959 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 18:22:30 -0700 Subject: [PATCH 06/18] generate mcdc_get and mcdc_set --- mcdc/mcdc_get/__init__.py | 2 ++ mcdc/mcdc_get/weight_windows.py | 32 ++++++++++++++++++++++++++++++++ mcdc/mcdc_set/__init__.py | 2 ++ mcdc/mcdc_set/weight_windows.py | 32 ++++++++++++++++++++++++++++++++ 4 files changed, 68 insertions(+) create mode 100644 mcdc/mcdc_get/weight_windows.py create mode 100644 mcdc/mcdc_set/weight_windows.py diff --git a/mcdc/mcdc_get/__init__.py b/mcdc/mcdc_get/__init__.py index 0f51357b..8138c476 100644 --- a/mcdc/mcdc_get/__init__.py +++ b/mcdc/mcdc_get/__init__.py @@ -78,6 +78,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..143d55fc --- /dev/null +++ b/mcdc/mcdc_get/weight_windows.py @@ -0,0 +1,32 @@ +# The following is automatically generated by code_factory.py + +from numba import njit + + +@njit +def weight_windows(index, weight_windows, data): + offset = weight_windows["weight_windows_offset"] + return data[offset + index] + + +@njit +def weight_windows_all(weight_windows, data): + start = weight_windows["weight_windows_offset"] + size = weight_windows["weight_windows_length"] + end = start + size + return data[start:end] + + +@njit +def weight_windows_last(weight_windows, data): + start = weight_windows["weight_windows_offset"] + size = weight_windows["weight_windows_length"] + end = start + size + return data[end - 1] + + +@njit +def weight_windows_chunk(start, length, weight_windows, data): + start += weight_windows["weight_windows_offset"] + end = start + length + return data[start:end] diff --git a/mcdc/mcdc_set/__init__.py b/mcdc/mcdc_set/__init__.py index 840de390..1d2486ed 100644 --- a/mcdc/mcdc_set/__init__.py +++ b/mcdc/mcdc_set/__init__.py @@ -78,6 +78,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..9f599f3c --- /dev/null +++ b/mcdc/mcdc_set/weight_windows.py @@ -0,0 +1,32 @@ +# The following is automatically generated by code_factory.py + +from numba import njit + + +@njit +def weight_windows(index, weight_windows, data, value): + offset = weight_windows["weight_windows_offset"] + data[offset + index] = value + + +@njit +def weight_windows_all(weight_windows, data, value): + start = weight_windows["weight_windows_offset"] + size = weight_windows["weight_windows_length"] + end = start + size + data[start:end] = value + + +@njit +def weight_windows_last(weight_windows, data, value): + start = weight_windows["weight_windows_offset"] + size = weight_windows["weight_windows_length"] + end = start + size + data[end - 1] = value + + +@njit +def weight_windows_chunk(start, length, weight_windows, data, value): + start += weight_windows["weight_windows_offset"] + end = start + length + data[start:end] = value From 30357378ccc400b75dfa1045520784798c59f87c Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 18:46:03 -0700 Subject: [PATCH 07/18] switch to tracking different arrays of lower, target, upper --- mcdc/mcdc_get/weight_windows.py | 78 ++++++++++++++++++++++++++++----- mcdc/mcdc_set/weight_windows.py | 78 ++++++++++++++++++++++++++++----- mcdc/object_/technique.py | 16 ++++--- mcdc/transport/technique.py | 9 ++-- 4 files changed, 151 insertions(+), 30 deletions(-) diff --git a/mcdc/mcdc_get/weight_windows.py b/mcdc/mcdc_get/weight_windows.py index 143d55fc..00252e54 100644 --- a/mcdc/mcdc_get/weight_windows.py +++ b/mcdc/mcdc_get/weight_windows.py @@ -4,29 +4,87 @@ @njit -def weight_windows(index, weight_windows, data): - offset = weight_windows["weight_windows_offset"] +def lower_weights(index, weight_windows, data): + offset = weight_windows["lower_weights_offset"] return data[offset + index] @njit -def weight_windows_all(weight_windows, data): - start = weight_windows["weight_windows_offset"] - size = weight_windows["weight_windows_length"] +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 weight_windows_last(weight_windows, data): - start = weight_windows["weight_windows_offset"] - size = weight_windows["weight_windows_length"] +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 weight_windows_chunk(start, length, weight_windows, data): - start += weight_windows["weight_windows_offset"] +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/weight_windows.py b/mcdc/mcdc_set/weight_windows.py index 9f599f3c..95613786 100644 --- a/mcdc/mcdc_set/weight_windows.py +++ b/mcdc/mcdc_set/weight_windows.py @@ -4,29 +4,87 @@ @njit -def weight_windows(index, weight_windows, data, value): - offset = weight_windows["weight_windows_offset"] +def lower_weights(index, weight_windows, data, value): + offset = weight_windows["lower_weights_offset"] data[offset + index] = value @njit -def weight_windows_all(weight_windows, data, value): - start = weight_windows["weight_windows_offset"] - size = weight_windows["weight_windows_length"] +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 weight_windows_last(weight_windows, data, value): - start = weight_windows["weight_windows_offset"] - size = weight_windows["weight_windows_length"] +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 weight_windows_chunk(start, length, weight_windows, data, value): - start += weight_windows["weight_windows_offset"] +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_/technique.py b/mcdc/object_/technique.py index 535dbb1f..5ebec3d9 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -82,8 +82,10 @@ class WeightWindows(ObjectSingleton): Ny: int Nz: int - # flattened array of ww - weight_windows: NDArray[float64] + # flattened arrays of ww params + lower_weights: NDArray[float64] + target_weights: NDArray[float64] + upper_weights: NDArray[float64] def __init__(self): self.active = False @@ -94,7 +96,7 @@ def __call__(self, mesh, weight_windows): case "uniform_mesh": nx, ny, nz = mesh.Nx, mesh.Ny, mesh.Nz case "structured_mesh": - nx, ny, nz = mesh.x.shape[0], mesh.y.shape[0], mesh.z.shape[0] + 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" @@ -110,10 +112,10 @@ def __call__(self, mesh, weight_windows): self.active = True self.mesh = mesh - self.Nx = nx - self.Ny = ny - self.Nz = nz - self.weight_windows = weight_windows.reshape(-1) + 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) # ====================================================================================== diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index 56a9b49e..b44a6af8 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -4,7 +4,7 @@ from numba import njit #### -from mcdc.mcdc_get.weight_windows import weight_windows_all, weight_windows_chunk +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 @@ -46,8 +46,11 @@ def query_weight_window(particle_container, mcdc, data): ww_obj = mcdc["weight_windows"] mesh = mcdc["meshes"][ww_obj["mesh_ID"]] idx, idy, idz = get_mesh_indices(particle_container, mesh, mcdc, data) - start = 3*(idx * mesh["Nx"] + idy * mesh["Ny"] + idz * mesh["Nz"]) - return weight_windows_chunk(start, 3, ww_obj, data) + index = ((idx * ww_obj["Ny"]) + idy) * ww_obj["Nz"] + idz + 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 From 5707fe187852022951e6fe089989b68a33e40999 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 18:57:35 -0700 Subject: [PATCH 08/18] add checks for weight window validity --- mcdc/object_/technique.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index 5ebec3d9..137363d0 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -1,6 +1,7 @@ from dataclasses import dataclass from numpy.typing import NDArray from numpy import float64 +import numpy as np from mcdc.object_.base import ObjectBase, ObjectSingleton from mcdc.object_.mesh import MeshBase from mcdc.print_ import print_error @@ -102,6 +103,7 @@ def __call__(self, mesh, weight_windows): 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) @@ -117,6 +119,21 @@ def __call__(self, mesh, weight_windows): 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( + f"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( + f"Target weight can not be greater than the upper bound weight for any weight window" + ) + + # ====================================================================================== # Population control From e9a6afa13384c3f6450eabca3a903ac24d7441d7 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 19:49:42 -0700 Subject: [PATCH 09/18] add test for instantiation --- mcdc/object_/technique.py | 4 +- .../transport/technique/weight_windows.py | 84 +++++++++++++++++++ 2 files changed, 86 insertions(+), 2 deletions(-) create mode 100644 test/unit/transport/technique/weight_windows.py diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index 137363d0..4017f3b7 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -126,11 +126,11 @@ def __call__(self, mesh, weight_windows): ) if (self.lower_weights > self.target_weights).any(): print_error( - f"Lower bound weight can not be greater than the target weight for any weight window" + "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( - f"Target weight can not be greater than the upper bound weight for any weight window" + "Target weight can not be greater than the upper bound weight for any weight window" ) diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py new file mode 100644 index 00000000..fa39a0d9 --- /dev/null +++ b/test/unit/transport/technique/weight_windows.py @@ -0,0 +1,84 @@ +import mcdc +import numpy as np +import pytest +import os + +from mcdc.main import preparation + +def make_ww_model(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): + # Geometry + # surfaces + cylinder = mcdc.Surface.CylinderZ(radius=0.5) + pitch = 2.0 + height = 10.0 + 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") + z0 = mcdc.Surface.PlaneZ(z=0.0, boundary_condition="reflective") + z1 = mcdc.Surface.PlaneZ(z=height, boundary_condition="reflective") + + # cells + mcdc.Cell(-cylinder) + mcdc.Cell(+x0 & -x1 & +y0 & -y1 & +z0 & -z1 & +cylinder) + + # 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 + + # weight windows + N = 3 + mesh = mcdc.MeshUniform( + "mesh", + x=(-pitch/2, pitch/2, N), + y=(-pitch/2, pitch/2, N), + z=(0.0, height, N) + ) + if mess_up_size: + shape = (N, N, 4, 3) + else: + shape = (N, N, N, 3) + ww_array = np.ones(shape) + ww_array[:, :, :, 0] = lower + ww_array[:, :, :, 1] = target + ww_array[:, :, :, 2] = upper + mcdc.simulation.weight_windows(mesh, ww_array) + + return preparation() + +@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(**kwargs) + + out = capsys.readouterr().out + assert expected_msg in out \ No newline at end of file From 28d9f2cfe219ee09790b81f3258ce31044b41d1c Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 20:22:54 -0700 Subject: [PATCH 10/18] add test for ww roulette method --- .../transport/technique/weight_windows.py | 41 ++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py index fa39a0d9..20da77ca 100644 --- a/test/unit/transport/technique/weight_windows.py +++ b/test/unit/transport/technique/weight_windows.py @@ -4,6 +4,21 @@ 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 +) +from mcdc.constant import ( + TINY +) + +# =========================================================================== # +# Helper method for creating a dummy model +# =========================================================================== # + def make_ww_model(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): # Geometry @@ -51,6 +66,12 @@ def make_ww_model(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): return preparation() + +# =========================================================================== # +# Test error throwing in object creation +# =========================================================================== # + + @pytest.mark.parametrize( "kwargs, expected_msg", [ @@ -81,4 +102,22 @@ def test_error_throw(capsys, kwargs, expected_msg): make_ww_model(**kwargs) out = capsys.readouterr().out - assert expected_msg in out \ No newline at end of file + 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 i 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"]) From e832317492534c207b274c889ede07f41ec13e7b Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 20:50:11 -0700 Subject: [PATCH 11/18] Add split test --- .../transport/technique/weight_windows.py | 33 +++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py index 20da77ca..63933d7c 100644 --- a/test/unit/transport/technique/weight_windows.py +++ b/test/unit/transport/technique/weight_windows.py @@ -9,7 +9,8 @@ roulette_from_weight_window, split_from_weight_window, query_weight_window, - weight_windows + weight_windows, + particle_bank_module ) from mcdc.constant import ( TINY @@ -64,7 +65,8 @@ def make_ww_model(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): ww_array[:, :, :, 2] = upper mcdc.simulation.weight_windows(mesh, ww_array) - return preparation() + mcdc_container, data = preparation() + return mcdc_container[0], data # =========================================================================== # @@ -121,3 +123,30 @@ def test_roullete_from_weight_window(): 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() + + # 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 + assert particle_bank_module.get_bank_size(bank) == init_bank_size + num_new + for i in range(2): + pnew = bank["particles"][init_bank_size + i] + assert pnew["w"] == p1["w"] + From eb037e9392332accfac49e0095e996e02e13ba09 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 21:37:39 -0700 Subject: [PATCH 12/18] Add query of weight windows test --- .../transport/technique/weight_windows.py | 159 ++++++++++++------ 1 file changed, 111 insertions(+), 48 deletions(-) diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py index 63933d7c..7974d691 100644 --- a/test/unit/transport/technique/weight_windows.py +++ b/test/unit/transport/technique/weight_windows.py @@ -12,6 +12,7 @@ weight_windows, particle_bank_module ) +from mcdc.transport.mesh.interface import get_indices from mcdc.constant import ( TINY ) @@ -20,53 +21,82 @@ # Helper method for creating a dummy model # =========================================================================== # +def make_base_model(): -def make_ww_model(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): - # Geometry - # surfaces - cylinder = mcdc.Surface.CylinderZ(radius=0.5) - pitch = 2.0 - height = 10.0 - 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") - z0 = mcdc.Surface.PlaneZ(z=0.0, boundary_condition="reflective") - z1 = mcdc.Surface.PlaneZ(z=height, boundary_condition="reflective") - - # cells - mcdc.Cell(-cylinder) - mcdc.Cell(+x0 & -x1 & +y0 & -y1 & +z0 & -z1 & +cylinder) - - # 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 - - # weight windows - N = 3 - mesh = mcdc.MeshUniform( - "mesh", - x=(-pitch/2, pitch/2, N), - y=(-pitch/2, pitch/2, N), - z=(0.0, height, N) - ) - if mess_up_size: - shape = (N, N, 4, 3) - else: - shape = (N, N, N, 3) - ww_array = np.ones(shape) - 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 + # Geometry + cylinder = mcdc.Surface.CylinderZ(radius=0.5) + pitch = 2.0 + height = 10.0 + + 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") + z0 = mcdc.Surface.PlaneZ(z=0.0, boundary_condition="reflective") + z1 = mcdc.Surface.PlaneZ(z=height, boundary_condition="reflective") + + mcdc.Cell(-cylinder) + mcdc.Cell(+x0 & -x1 & +y0 & -y1 & +z0 & -z1 & +cylinder) + + # Source + mcdc.Source(position=[0.0, 0.0, 0.0], isotropic=True, time=0.0, energy=14.1e6) + + # Settings + mcdc.settings.N_particle = 20 + mcdc.settings.N_batch = 2 + mcdc.settings.time_boundary = 1.0 + mcdc.settings.active_bank_buffer = 1000 + + # Mesh + 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_base_model() + + if mess_up_size: + ww_array = np.ones((N, N, 4, 3)) + else: + 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_base_model() + + ww_array = np.empty((N, N, N, 3)) + + 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 # =========================================================================== # @@ -101,7 +131,7 @@ def make_ww_model(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): ) def test_error_throw(capsys, kwargs, expected_msg): with pytest.raises(SystemExit): - make_ww_model(**kwargs) + make_ww_model_params(**kwargs) out = capsys.readouterr().out assert expected_msg in out @@ -130,7 +160,7 @@ def test_split_from_weight_window(): init_weight = 2.0 + TINY particles[0]["w"] = init_weight threshold = 1.0 - mcdc_obj, data = make_ww_model() + mcdc_obj, data = make_ww_model_distinct() # get bank and init size bank = mcdc_obj["bank_active"] @@ -150,3 +180,36 @@ def test_split_from_weight_window(): pnew = bank["particles"][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() + ww_obj = mcdc_obj["weight_windows"] + + mesh = mcdc_obj["meshes"][ww_obj["mesh_ID"]] + # hardcode mesh params + nx, ny, nz = 3, 3, 3 + xmin, ymin, zmin = -1.0, -1.0, 0.0 + dx, dy, dz = 2/3, 2/3, 10/3 + + + # loop over all bins, check query against expected ww + for ix in range(nx): + for iy in range(ny): + for iz in range(nz): + p[0]["x"] = xmin + dx * (ix + 0.5) + p[0]["y"] = ymin + dy * (iy + 0.5) + p[0]["z"] = zmin + dz * (iz + 0.5) + + 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 + From 8c2e62661fb5b6acff7992703d8ddb5e3723f845 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 21:42:36 -0700 Subject: [PATCH 13/18] Clean up tests a bit, removed needless setup --- .../transport/technique/weight_windows.py | 41 ++++--------------- 1 file changed, 9 insertions(+), 32 deletions(-) diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py index 7974d691..0687c26e 100644 --- a/test/unit/transport/technique/weight_windows.py +++ b/test/unit/transport/technique/weight_windows.py @@ -21,34 +21,12 @@ # Helper method for creating a dummy model # =========================================================================== # -def make_base_model(): - - # Geometry - cylinder = mcdc.Surface.CylinderZ(radius=0.5) +def make_mesh(): + # constants pitch = 2.0 height = 10.0 - - 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") - z0 = mcdc.Surface.PlaneZ(z=0.0, boundary_condition="reflective") - z1 = mcdc.Surface.PlaneZ(z=height, boundary_condition="reflective") - - mcdc.Cell(-cylinder) - mcdc.Cell(+x0 & -x1 & +y0 & -y1 & +z0 & -z1 & +cylinder) - - # Source - mcdc.Source(position=[0.0, 0.0, 0.0], isotropic=True, time=0.0, energy=14.1e6) - - # Settings - mcdc.settings.N_particle = 20 - mcdc.settings.N_batch = 2 - mcdc.settings.time_boundary = 1.0 - mcdc.settings.active_bank_buffer = 1000 - - # Mesh N = 3 + mesh = mcdc.MeshUniform( "mesh", x=(-pitch/2, pitch/N, N), @@ -62,7 +40,7 @@ def make_base_model(): def make_ww_model_params(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): import mcdc - mesh, N = make_base_model() + mesh, N = make_mesh() if mess_up_size: ww_array = np.ones((N, N, 4, 3)) @@ -81,7 +59,7 @@ def make_ww_model_params(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): def make_ww_model_distinct(): import mcdc - mesh, N = make_base_model() + mesh, N = make_mesh() ww_array = np.empty((N, N, N, 3)) @@ -185,13 +163,12 @@ def test_query_weight_window(): p = np.zeros(1, type_.particle_data) mcdc_obj, data = make_ww_model_distinct() - ww_obj = mcdc_obj["weight_windows"] - mesh = mcdc_obj["meshes"][ww_obj["mesh_ID"]] # hardcode mesh params - nx, ny, nz = 3, 3, 3 - xmin, ymin, zmin = -1.0, -1.0, 0.0 - dx, dy, dz = 2/3, 2/3, 10/3 + 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 From 4933873c499cc5d7f8e0e9b94dd55f96a00e550b Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 21:49:32 -0700 Subject: [PATCH 14/18] integrate the changed names into ww tests --- test/unit/transport/technique/weight_windows.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py index 0687c26e..32d5c071 100644 --- a/test/unit/transport/technique/weight_windows.py +++ b/test/unit/transport/technique/weight_windows.py @@ -155,7 +155,7 @@ def test_split_from_weight_window(): assert p1["w"] == init_weight / num_split assert particle_bank_module.get_bank_size(bank) == init_bank_size + num_new for i in range(2): - pnew = bank["particles"][init_bank_size + i] + pnew = bank["particle_data"][init_bank_size + i] assert pnew["w"] == p1["w"] From 6e8cb5b4bfafab9b3c566f14da66624329afdabb Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 21:51:23 -0700 Subject: [PATCH 15/18] black format --- mcdc/object_/technique.py | 9 +- mcdc/transport/technique.py | 3 +- .../transport/technique/weight_windows.py | 147 +++++++++--------- 3 files changed, 78 insertions(+), 81 deletions(-) diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index 4017f3b7..5509774e 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -97,12 +97,16 @@ def __call__(self, mesh, weight_windows): 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 + 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 @@ -134,7 +138,6 @@ def __call__(self, mesh, weight_windows): ) - # ====================================================================================== # Population control # ====================================================================================== diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index b14d25b0..5ca54c1a 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -12,7 +12,6 @@ import mcdc.transport.rng as rng import mcdc.transport.util as util - # ====================================================================================== # Weight Roulette # ====================================================================================== @@ -50,7 +49,7 @@ def query_weight_window(particle_container, mcdc, data): index = ((idx * ww_obj["Ny"]) + idy) * ww_obj["Nz"] + idz 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) + upper = ww_get.upper_weights(index, ww_obj, data) return lower, target, upper diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py index 32d5c071..b90b6a25 100644 --- a/test/unit/transport/technique/weight_windows.py +++ b/test/unit/transport/technique/weight_windows.py @@ -6,21 +6,20 @@ 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 + 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 -) +from mcdc.constant import TINY # =========================================================================== # # Helper method for creating a dummy model # =========================================================================== # + def make_mesh(): # constants pitch = 2.0 @@ -29,9 +28,9 @@ def make_mesh(): mesh = mcdc.MeshUniform( "mesh", - x=(-pitch/2, pitch/N, N), - y=(-pitch/2, pitch/N, N), - z=(0.0, height/N, N) + x=(-pitch / 2, pitch / N, N), + y=(-pitch / 2, pitch / N, N), + z=(0.0, height / N, N), ) return mesh, N @@ -66,7 +65,7 @@ def make_ww_model_distinct(): for i in range(N): for j in range(N): for k in range(N): - val = 100*i + 10*j + k + 1 + 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 @@ -108,85 +107,81 @@ def make_ww_model_distinct(): ], ) def test_error_throw(capsys, kwargs, expected_msg): - with pytest.raises(SystemExit): - make_ww_model_params(**kwargs) + with pytest.raises(SystemExit): + make_ww_model_params(**kwargs) - out = capsys.readouterr().out - assert expected_msg in out + out = capsys.readouterr().out + assert expected_msg in out # =========================================================================== # -# Tests for helper methods +# Tests for helper methods # =========================================================================== # def test_roullete_from_weight_window(): - # because of rng, want to loop over to hit both branches - for i 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"]) + # because of rng, want to loop over to hit both branches + for i 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() + 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) + # 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) + # split + split_from_weight_window(particles, threshold, mcdc_obj) - p1 = particles[0] - num_split = np.ceil(init_weight / threshold) - num_new = num_split - 1 + 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 - 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"] + # check weight of original particle + assert p1["w"] == init_weight / num_split + 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): - p[0]["x"] = xmin + dx * (ix + 0.5) - p[0]["y"] = ymin + dy * (iy + 0.5) - p[0]["z"] = zmin + dz * (iz + 0.5) - - 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 - + 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): + p[0]["x"] = xmin + dx * (ix + 0.5) + p[0]["y"] = ymin + dy * (iy + 0.5) + p[0]["z"] = zmin + dz * (iz + 0.5) + + 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 From a57e72e6d70ed74dc06c4d51c1ef8a46c700b813 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 22:15:05 -0700 Subject: [PATCH 16/18] add oneliner on how to add weight windows to a simulation --- docs/source/pythonapi/index.rst | 1 + 1 file changed, 1 insertion(+) 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 From 4bc9c774fbeb43aad0b0cb5e1992b0f6b308caf0 Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 22:25:29 -0700 Subject: [PATCH 17/18] add docstrings and some comments for clarity --- mcdc/transport/technique.py | 83 +++++++++++++++++++++++++++++++++---- 1 file changed, 74 insertions(+), 9 deletions(-) diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index 5ca54c1a..a9c72fac 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -35,18 +35,55 @@ def weight_roulette(particle_container, simulation): @njit -def weight_windows(particle_container, mcdc, data): - [lower, target, upper] = query_weight_window(particle_container, mcdc, data) - split_from_weight_window(particle_container, upper, mcdc) +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, mcdc, data): - ww_obj = mcdc["weight_windows"] - mesh = mcdc["meshes"][ww_obj["mesh_ID"]] - idx, idy, idz = get_mesh_indices(particle_container, mesh, mcdc, data) +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) @@ -54,21 +91,49 @@ def query_weight_window(particle_container, mcdc, data): @njit -def split_from_weight_window(particle_container, threshold_weight, mcdc): +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): - particle_bank_module.bank_active_particle(particle_container, mcdc) + # 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: From 33438fe261bf1350e42529574821fee009b023ca Mon Sep 17 00:00:00 2001 From: Nathan Glaser Date: Sun, 12 Apr 2026 23:33:56 -0700 Subject: [PATCH 18/18] clean up imports, a few more comments for clarity --- mcdc/object_/technique.py | 4 +--- mcdc/transport/technique.py | 1 + test/unit/transport/technique/weight_windows.py | 7 ++++++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/mcdc/object_/technique.py b/mcdc/object_/technique.py index 5509774e..f18fcafd 100644 --- a/mcdc/object_/technique.py +++ b/mcdc/object_/technique.py @@ -1,8 +1,6 @@ -from dataclasses import dataclass from numpy.typing import NDArray from numpy import float64 -import numpy as np -from mcdc.object_.base import ObjectBase, ObjectSingleton +from mcdc.object_.base import ObjectSingleton from mcdc.object_.mesh import MeshBase from mcdc.print_ import print_error diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index a9c72fac..cb586e69 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -4,6 +4,7 @@ from numba import njit #### + 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 diff --git a/test/unit/transport/technique/weight_windows.py b/test/unit/transport/technique/weight_windows.py index b90b6a25..7a12eb5d 100644 --- a/test/unit/transport/technique/weight_windows.py +++ b/test/unit/transport/technique/weight_windows.py @@ -44,6 +44,7 @@ def make_ww_model_params(lower=0.1, target=1.0, upper=1.0, mess_up_size=False): 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 @@ -62,6 +63,7 @@ def make_ww_model_distinct(): 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): @@ -121,7 +123,7 @@ def test_error_throw(capsys, kwargs, expected_msg): def test_roullete_from_weight_window(): # because of rng, want to loop over to hit both branches - for i in range(10): + for _ in range(10): particles = np.zeros(1, type_.particle_data) particles[0]["w"] = 0.1 target = 0.2 @@ -152,6 +154,7 @@ def test_split_from_weight_window(): # 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] @@ -173,10 +176,12 @@ def test_query_weight_window(): 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