Automated time dependant weight windows

examples/fixed_source/kobayashi3-TD-ww/input.py

@@ -0,0 +1,85 @@
+import numpy as np
+import mcdc
+
+# =============================================================================
+# Set model
+# =============================================================================
+# Based on Kobayashi dog-leg benchmark problem
+# (PNE 2001, https://doi.org/10.1016/S0149-1970(01)00007-5)
+
+# Set materials
+m = mcdc.material(capture=np.array([0.05]), scatter=np.array([[0.05]]))
+m_void = mcdc.material(capture=np.array([5e-5]), scatter=np.array([[5e-5]]))
+
+# Set surfaces
+sx1 = mcdc.surface("plane-x", x=0.0, bc="reflective")
+sx2 = mcdc.surface("plane-x", x=10.0)
+sx3 = mcdc.surface("plane-x", x=30.0)
+sx4 = mcdc.surface("plane-x", x=40.0)
+sx5 = mcdc.surface("plane-x", x=60.0, bc="vacuum")
+sy1 = mcdc.surface("plane-y", y=0.0, bc="reflective")
+sy2 = mcdc.surface("plane-y", y=10.0)
+sy3 = mcdc.surface("plane-y", y=50.0)
+sy4 = mcdc.surface("plane-y", y=60.0)
+sy5 = mcdc.surface("plane-y", y=100.0, bc="vacuum")
+sz1 = mcdc.surface("plane-z", z=0.0, bc="reflective")
+sz2 = mcdc.surface("plane-z", z=10.0)
+sz3 = mcdc.surface("plane-z", z=30.0)
+sz4 = mcdc.surface("plane-z", z=40.0)
+sz5 = mcdc.surface("plane-z", z=60.0, bc="vacuum")
+
+# Set cells
+# Source
+source_cell = mcdc.cell(+sx1 & -sx2 & +sy1 & -sy2 & +sz1 & -sz2, m)
+# Voids
+channel_1 = +sx1 & -sx2 & +sy2 & -sy3 & +sz1 & -sz2
+channel_2 = +sx1 & -sx3 & +sy3 & -sy4 & +sz1 & -sz2
+channel_3 = +sx3 & -sx4 & +sy3 & -sy4 & +sz1 & -sz3
+channel_4 = +sx3 & -sx4 & +sy3 & -sy5 & +sz3 & -sz4
+void_channel = channel_1 | channel_2 | channel_3 | channel_4
+void_cell = mcdc.cell(void_channel, m_void)
+# Shield
+box = +sx1 & -sx5 & +sy1 & -sy5 & +sz1 & -sz5
+shield_cell = mcdc.cell(box & ~void_channel, m)
+
+# =============================================================================
+# Set source
+# =============================================================================
+# The source pulses in t=[0,5]
+
+mcdc.source(
+    x=[0.0, 10.0], y=[0.0, 10.0], z=[0.0, 10.0], time=[0.0, 50.0], isotropic=True
+)
+
+# =============================================================================
+# Set tally, setting, and run mcdc
+# =============================================================================
+
+# Tally: z-integrated flux (X-Y section view)
+mcdc.tally.mesh_tally(
+    scores=["flux"],
+    x=np.linspace(0.0, 60.0, 31),
+    y=np.linspace(0.0, 100.0, 51),
+    t=np.linspace(0.0, 200.0, 21),
+)
+
+mcdc.tally.cell_tally(source_cell, scores=["flux"])
+mcdc.tally.cell_tally(void_cell, scores=["flux"])
+mcdc.tally.cell_tally(shield_cell, scores=["flux"])
+
+mcdc.time_census(t=np.linspace(0.0, 200.0, 21, endpoint=True)[1:])
+mcdc.setting(census_bank_buff=1e2, active_bank_buff=1e2)
+
+mcdc.weight_window(
+    x=np.linspace(0.0, 60.0, 31),
+    y=np.linspace(0.0, 100.0, 51),
+    t=np.linspace(0.0, 200.0, 21),
+    width=2.5,
+    epsilon=2e-2,
+)
+
+# Setting
+mcdc.setting(N_particle=1e4)
+
+# Run
+mcdc.run()

examples/fixed_source/kobayashi3-TD-ww/process.py

@@ -0,0 +1,44 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.colors import LogNorm
+import h5py
+import matplotlib.animation as animation
+
+
+# =============================================================================
+# Plot results
+# =============================================================================
+
+# Results
+with h5py.File("output.h5", "r") as f:
+    tallies = f["tallies/mesh_tally_0"]
+    flux = tallies["flux"]
+    grid = tallies["grid"]
+    x = grid["x"][:]
+    x_mid = 0.5 * (x[:-1] + x[1:])
+    y = grid["y"][:]
+    y_mid = 0.5 * (y[:-1] + y[1:])
+    t = grid["t"][:]
+    t_mid = 0.5 * (t[:-1] + t[1:])
+    X, Y = np.meshgrid(y, x)
+
+    phi = flux["mean"][:]
+    phi_sd = flux["sdev"][:]
+
+
+fig, ax = plt.subplots()
+cax = ax.pcolormesh(X, Y, phi[0], vmin=phi[0].min(), vmax=phi[0].max())
+text = ax.text(0.02, 1.02, "", transform=ax.transAxes)
+ax.set_aspect("equal", "box")
+ax.set_xlabel("$y$ [cm]")
+ax.set_ylabel("$x$ [cm]")
+
+
+def animate(i):
+    cax.set_array(phi[i])
+    cax.set_clim(phi[i].min(), phi[i].max())
+    text.set_text(r"$t \in [%.1f,%.1f]$ s" % (t[i], t[i + 1]))
+
+
+anim = animation.FuncAnimation(fig, animate, interval=10, frames=len(t) - 1)
+plt.show()

examples/fixed_source/kobayashi3-TD/input.py

@@ -68,7 +68,7 @@
 mcdc.tally.cell_tally(shield_cell, scores=["flux"])
 
 # Setting
-mcdc.setting(N_particle=80)
+mcdc.setting(N_particle=1e4)
 
 # Run
 mcdc.run()

examples/fixed_source/kobayashi3-TD/process.py

@@ -25,11 +25,6 @@
     phi = flux["mean"][:]
     phi_sd = flux["sdev"][:]
 
-    for i in range(len(f["input_deck"]["cell_tallies"])):
-        flux_score = f[f"tallies/cell_tally_{i}/flux"]
-        print(
-            f'cell {i+1} mean = {flux_score["mean"][()]}, sdev = {flux_score["sdev"][()]}'
-        )
 
 fig, ax = plt.subplots()
 cax = ax.pcolormesh(X, Y, phi[0], vmin=phi[0].min(), vmax=phi[0].max())

mcdc/global_.py

@@ -81,6 +81,8 @@ def reset(self):
             "ww": np.ones([1, 1, 1, 1]),
             "ww_width": 2.5,
             "ww_mesh": make_card_mesh(),
+            "ww_auto": False,
+            "ww_epsilon": 1e-4,
             "domain_decomposition": False,
             "dd_idx": 0,
             "dd_mesh": make_card_mesh(),

mcdc/input_.py

@@ -1224,37 +1224,29 @@ def branchless_collision():
     card["weighted_emission"] = False
 
 
-def time_census(t):
+def time_census(time_grid):
     """
     Set time-census boundaries.
 
     Parameters
     ----------
-    t : array_like[float]
+    time_grid : array_like[float]
         The time-census boundaries.
 
     Returns
     -------
         None (in-place card alterations).
     """
 
-    # Remove census beyond the final tally time grid point
-    while True:
-        if t[-1] >= global_.input_deck.tally["mesh"]["t"][-1]:
-            t = t[:-1]
-        else:
-            break
-
-    # Add the default, final census-at-infinity
-    t = np.append(t, INF)
-
     # Set the time census parameters
     card = global_.input_deck.setting
-    card["census_time"] = t
-    card["N_census"] = len(t)
+    card["census_time"] = time_grid
+    card["N_census"] = len(time_grid)
 
 
-def weight_window(x=None, y=None, z=None, t=None, window=None, width=None):
+def weight_window(
+    x=None, y=None, z=None, t=None, window=None, width=None, epsilon=1e-4
+):
     """
     Activate weight window variance reduction technique.
 
@@ -1296,6 +1288,19 @@ def weight_window(x=None, y=None, z=None, t=None, window=None, width=None):
         card["ww_mesh"]["t"] = t
 
     # Set window
+    if window is None:
+        card["ww_auto"] = True
+        window_ax = []
+        if t is not None:
+            window_ax.append(len(t) - 1)
+        if x is not None:
+            window_ax.append(len(x) - 1)
+        if y is not None:
+            window_ax.append(len(y) - 1)
+        if z is not None:
+            window_ax.append(len(z) - 1)
+        window = np.ones(window_ax)
+
     ax_expand = []
     if t is None:
         ax_expand.append(0)
@@ -1309,7 +1314,7 @@ def weight_window(x=None, y=None, z=None, t=None, window=None, width=None):
     for ax in ax_expand:
         window = np.expand_dims(window, axis=ax)
     card["ww"] = window
-
+    card["ww_epsilon"] = epsilon
     return card
 
 

mcdc/kernel.py

@@ -3164,6 +3164,118 @@ def weight_window(P_arr, prog):
             P["w"] = w_target
 
 
+# =============================================================================
+# Weight widow
+# =============================================================================
+
+
+def get_flux(idx, data, mcdc):
+    # Determine the correct tally list based on tally_type
+    tallies = mcdc["mesh_tallies"]
+
+    # Loop over the tallies to find the requested score
+    for ID, tally in enumerate(tallies):
+        if mcdc["technique"]["iQMC"]:
+            break
+
+        mesh = tally["filter"]
+
+        # Shape based on technique and tally type
+        N_sensitivity = 0
+        Ns = 1 + N_sensitivity
+
+        Nmu = len(mesh["mu"]) - 1
+        N_azi = len(mesh["azi"]) - 1
+        Ng = len(mesh["g"]) - 1
+        Nt = len(mesh["t"]) - 1
+
+        Nx = len(mesh["x"]) - 1
+        Ny = len(mesh["y"]) - 1
+        Nz = len(mesh["z"]) - 1
+
+        N_score = tally["N_score"]
+
+        if not mcdc["technique"]["uq"]:
+            shape = (3, Ns, Nmu, N_azi, Ng, Nt, Nx, Ny, Nz, N_score)
+        else:
+            shape = (5, Ns, Nmu, N_azi, Ng, Nt, Nx, Ny, Nz, N_score)
+
+        # Reshape tally
+        N_bin = tally["N_bin"]
+        start = tally["stride"]["tally"]
+        tally_bin = np.ascontiguousarray(data[TALLY][:, start : start + N_bin])
+        tally_bin = tally_bin.reshape(shape)
+
+        # Roll tally so that score is in the front
+        tally_bin = np.rollaxis(tally_bin, 9, 0)
+
+        # Iterate over scores to find the requested score
+        for i in range(N_score):
+            score_type = tally["scores"][i]
+            if score_type == SCORE_FLUX:
+                score_tally_bin = np.squeeze(tally_bin[i])
+                mean = score_tally_bin[TALLY_SUM]
+                sdev = score_tally_bin[TALLY_SUM_SQ]
+                old_flux = mean[idx][:]
+                if Nx == 1 and old_flux.ndim < 3:
+                    old_flux = np.expand_dims(old_flux, axis=0)
+                if Ny == 1 and old_flux.ndim < 3:
+                    old_flux = np.expand_dims(old_flux, axis=1)
+                if Nz == 1 and old_flux.ndim < 3:
+                    old_flux = np.expand_dims(old_flux, axis=2)
+                return old_flux
+
+
+@njit
+def update_weight_window(data, mcdc):
+    # Get current global census index
+    idx_census = mcdc["idx_census"]
+
+    # Get the weight window
+    window_centers = mcdc["technique"]["ww"]
+
+    # Get the flux from the previous time census
+    with objmode(flux_old="float64[:,:,:]"):
+        flux_old = get_flux(idx_census - 1, data, mcdc)
+
+    # Next, we ensure that we don't have a zero old flux
+
+    # If the old flux is all-zero, assign the small value epsilon
+    eps = mcdc["technique"]["ww_epsilon"]
+    if np.max(flux_old) == 0:
+        flux_old += eps
+        # Assign as the current weight window
+        window_centers[idx_census] = flux_old
+        return
+
+    # Get minimum, non-zero flux
+    flux_min = INF
+    Nx, Ny, Nz = flux_old.shape
+    for ix in range(Nx):
+        for iy in range(Ny):
+            for iz in range(Nz):
+                value = flux_old[ix, iy, iz]
+                if 0.0 < value and value < flux_min:
+                    flux_min = value
+
+    # Replace zeros with the half of the minimum
+    flux_min *= 0.5
+    for ix in range(Nx):
+        for iy in range(Ny):
+            for iz in range(Nz):
+                if flux_old[ix, iy, iz] == 0.0:
+                    flux_old[ix, iy, iz] = flux_min
+
+    # Normalize the old flux
+    flux_old /= np.max(flux_old)
+
+    # Ensure that the flux is not smaller than the epsilon
+    flux_old = (1 - eps) * flux_old + eps
+
+    # Assign as the current weight window
+    window_centers[idx_census] = flux_old
+
+
 # =============================================================================
 # Weight Roulette
 # =============================================================================

mcdc/loop.py

@@ -61,7 +61,6 @@ def teardown_gpu(mcdc):
 
 @njit
 def loop_fixed_source(data_arr, mcdc_arr):
-
     # Ensure `data` and `mcdc` exist for the lifetime of the program
     # by intentionally leaking their memory
     adapt.leak(data_arr)
@@ -71,6 +70,7 @@ def loop_fixed_source(data_arr, mcdc_arr):
 
     # Loop over batches
     for idx_batch in range(mcdc["setting"]["N_batch"]):
+        # Update global batch index and generate batch seed
         mcdc["idx_batch"] = idx_batch
         seed_batch = kernel.split_seed(idx_batch, mcdc["setting"]["rng_seed"])
 
@@ -84,9 +84,14 @@ def loop_fixed_source(data_arr, mcdc_arr):
 
         # Loop over time censuses
         for idx_census in range(mcdc["setting"]["N_census"]):
+            # Update global census index and generate census seed
             mcdc["idx_census"] = idx_census
             seed_census = kernel.split_seed(seed_batch, SEED_SPLIT_CENSUS)
 
+            # Update automated weight window
+            if mcdc["technique"]["ww_auto"] and idx_census > 0:
+                kernel.update_weight_window(data, mcdc)
+
             # Loop over source particles
             seed_source = kernel.split_seed(seed_census, SEED_SPLIT_SOURCE)
             loop_source(seed_source, data, mcdc)