Mesh-based weight windows

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

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

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]

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

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

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()
 
         # ==============================================================================

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
 # ======================================================================================

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)