Merge pull request #266 from ethan-lame/dev

Compressed Sensing Implementation

examples/fixed_source/cooper_combo/input.py

@@ -56,12 +56,12 @@
 
 mcdc.tally.mesh_tally(
     scores=["flux"],
-    x=np.linspace(0.0, 4.0, 40),
-    y=np.linspace(0.0, 4.0, 40),
+    x=np.linspace(0.0, 4.0, 41),
+    y=np.linspace(0.0, 4.0, 41),
 )
 
 # Setting
-mcdc.setting(N_particle=50)
+mcdc.setting(N_particle=1e2)
 mcdc.implicit_capture()
 
 # Run

examples/fixed_source/kobayashi3-TD/input.py

@@ -60,15 +60,26 @@
     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),
+    # t=np.linspace(0.0, 200.0, 21),
+    # g=np.array([-0.5, 3.5, 6.5]) # fast (0, 1, 2, 3) and thermal (4, 5, 6) groups
 )
 
 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.tally.cs_tally(
+    N_cs_bins=[150],
+    cs_bin_size=[8.0, 8.0],
+    x=np.linspace(0.0, 60.0, 31),
+    y=np.linspace(0.0, 100.0, 51),
+    scores=["flux"],
+)
+
+
 # Setting
-mcdc.setting(N_particle=80)
+mcdc.setting(N_particle=1e2)
 
 # Run
 mcdc.run()

examples/fixed_source/kobayashi3-TD/process.py

@@ -11,6 +11,10 @@
 
 # Results
 with h5py.File("output.h5", "r") as f:
+    cs_recon = f["tallies/cs_tally_0/flux/reconstruction"][:]
+    plt.imshow(cs_recon)
+    plt.show()
+
     tallies = f["tallies/mesh_tally_0"]
     flux = tallies["flux"]
     grid = tallies["grid"]
@@ -30,20 +34,3 @@
         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())
-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/sphere_in_cube/input.py

@@ -7,7 +7,7 @@
 # Homogeneous pure-fission sphere inside a pure-scattering cube
 
 # Set materials
-pure_f = mcdc.material(fission=np.array([1.0]), nu_p=np.array([1.2]))
+pure_f = mcdc.material(fission=np.array([1.0]), nu_p=np.array([1.1]))
 pure_s = mcdc.material(scatter=np.array([[1.0]]))
 
 # Set surfaces
@@ -40,17 +40,24 @@
 # =============================================================================
 mcdc.tally.mesh_tally(
     scores=["fission"],
-    x=np.linspace(0.0, 4.0, 2),
-    y=np.linspace(0.0, 4.0, 2),
-    z=np.linspace(0.0, 4.0, 2),
+    x=np.linspace(0.0, 4.0, 41),
+    y=np.linspace(0.0, 4.0, 41),
+    # z=np.linspace(0.0, 4.0, 41),
     # t=np.linspace(0.0, 200.0, 2),
 )
 
-
 mcdc.tally.cell_tally(sphere_cell, scores=["fission"])
 
+mcdc.tally.cs_tally(
+    N_cs_bins=[150],
+    cs_bin_size=np.array([3.0, 3.0]),
+    x=np.linspace(0.0, 4.0, 41),
+    y=np.linspace(0.0, 4.0, 41),
+    scores=["fission"],
+)
+
 # Setting
-mcdc.setting(N_particle=100)
+mcdc.setting(N_particle=1e3)
 mcdc.implicit_capture()
 
 # Run

examples/fixed_source/sphere_in_cube/process.py

@@ -1,23 +1,62 @@
 import h5py
+import numpy as np
+import matplotlib.pyplot as plt
+import scipy.fft as spfft
+import cvxpy as cp
+
+# Note: there are some lines in main.py with np.save(...) that I added
+# for ease of post-processing, like getting the center points used and
+# the sampling matrix S. None are required for the input file and this
+# script to run, but may be useful for debugging purposes
 
-# Load results
 with h5py.File("output.h5", "r") as f:
-    # print(f["tallies"].keys())
-    print(f["input_deck"]["cell_tallies"].keys())
+    S = f["tallies"]["cs_tally_0"]["S"][:]
+    recon = f["tallies"]["cs_tally_0"]["fission"]["reconstruction"]
+    plt.imshow(recon)
+    plt.title("Reconstruction, $\lambda$ = 0.5")  # assuming l in main.py remains at 0.5
+    plt.colorbar()
+    plt.show()
+
+    cs_results = f["tallies"]["cs_tally_0"]["fission"]["mean"][:]
+
+    mesh_results = f["tallies"]["mesh_tally_0"]["fission"]["mean"][:]
+    plt.imshow(mesh_results)
+    plt.title("mesh results")
+    plt.colorbar()
+    plt.show()
+
+Nx = 40
+Ny = 40
+N_fine_cells = Nx * Ny
+
+# Can use this for post-processing
+# mesh_b = S @ mesh_results.flatten()
+# b = mesh_b
+
+# Use this for analyzing the in-situ results
+cs_b = cs_results
+b = cs_b
 
-    for i in range(len(f["input_deck"]["cell_tallies"])):
-        fission_score = f[f"tallies/cell_tally_{i}/fission"]
+# Constructing T and A
+idct_basis_x = spfft.idct(np.identity(Nx), axis=0)
+idct_basis_y = spfft.idct(np.identity(Ny), axis=0)
 
-        print(
-            f'for sphere {i+1}, mean = {fission_score["mean"][()]}, sdev = {fission_score["sdev"][()]}'
-        )
+T_inv = np.kron(idct_basis_y, idct_basis_x)
+A = S @ T_inv
 
-        # print(fission_score["mean"][()])
-        # print(fission_score["sdev"][()])
+# Basis pursuit denoising solver - change l to get different results
+vx = cp.Variable(N_fine_cells)
+l = 10
+objective = cp.Minimize(0.5 * cp.norm(A @ vx - b, 2) + l * cp.norm(vx, 1))
+prob = cp.Problem(objective)
+result = prob.solve(verbose=False)
+sparse_solution = np.array(vx.value).squeeze()
 
-        # print(f"fission_score mean = {fission_score["mean"][()]}")
-        # print(f"fission_score mean = {fission_score["sdev"][()]}")
+# Obtaining the reconstruction
+recon = T_inv @ sparse_solution
+recon_reshaped = recon.reshape(Ny, Nx)
 
-    # cell = f["tallies/cell_tally_0/fission"]
-    # print(f'sphere1 mean = {cell["mean"][()]}')
-    # print(f'sphere2 sdev = {cell["sdev"][()]}')
+plt.imshow(recon_reshaped)
+plt.title(f"Reconstruction, $\lambda$ = {l}")
+plt.colorbar()
+plt.show()

mcdc/card.py

@@ -414,3 +414,16 @@ def __init__(self, cell_ID):
         # Set card data
         self.cell_ID = cell_ID
         self.N_bin = 1
+
+
+class CSTallyCard(TallyCard):
+    def __init__(self):
+        TallyCard.__init__(self, "CS tally")
+
+        # Set card data
+        self.x = np.array([-INF, INF])
+        self.y = np.array([-INF, INF])
+        self.z = np.array([-INF, INF])
+        self.N_bin = 1
+        self.N_cs_bins = 1
+        self.cs_bin_size = np.array([1.0, 1.0])

mcdc/global_.py

@@ -36,6 +36,7 @@ def reset(self):
         self.mesh_tallies = []
         self.surface_tallies = []
         self.cell_tallies = []
+        self.cs_tallies = []
 
         self.setting = {
             "tag": "Setting",

mcdc/kernel.py

@@ -2007,7 +2007,142 @@ def score_cell_tally(P_arr, distance, tally, data, mcdc):
             SigmaF = get_MacroXS(XS_FISSION, material, P_arr, mcdc)
             score = flux * SigmaF
 
-        tally_bin[TALLY_SCORE, cell_idx] += score
+        adapt.global_add(tally_bin, (TALLY_SCORE, cell_idx + i), round(score))
+        # tally_bin[TALLY_SCORE, cell_idx + i] += score
+
+
+@njit
+def score_cs_tally(P_arr, distance, tally, data, mcdc):
+    # Each time that this function is called, EVERY cs bin needs to be checked to see if the particle is in it.
+    # The particle needs to score into all the bins that it is within
+    P = P_arr[0]
+    tally_bin = data[TALLY]
+    material = mcdc["materials"][P["material_ID"]]
+    N_cs_bins = tally["filter"]["N_cs_bins"]
+
+    cs_bin_size = tally["filter"]["cs_bin_size"]
+    cs_centers = tally["filter"]["cs_centers"]
+
+    stride = tally["stride"]
+    bin_idx = stride["tally"]
+
+    # Particle 4D direction
+    ux = P["ux"]
+    uy = P["uy"]
+    uz = P["uz"]
+    ut = 1.0 / physics.get_speed(P_arr, mcdc)
+
+    # Particle initial and final coordinate
+    x = P["x"]
+    y = P["y"]
+    z = P["z"]
+    t = P["t"]
+    x_final = x + ux * distance
+    y_final = y + uy * distance
+    z_final = z + uz * distance
+    t_final = t + ut * distance
+
+    # Check each coarse bin
+    for j in range(N_cs_bins):
+        center = np.array([cs_centers[0][j], cs_centers[1][j]])
+        start = np.array([x, y])
+        end = np.array([x_final, y_final])
+
+        distance_inside = calculate_distance_in_coarse_bin(
+            start, end, distance, center, cs_bin_size
+        )
+
+        # Last bin covers the whole problem
+        if j == N_cs_bins - 1:
+            cs_bin_size_full_problem = np.array([INF, INF], dtype=np.float64)
+            distance_inside = calculate_distance_in_coarse_bin(
+                start, end, distance, center, cs_bin_size_full_problem
+            )
+
+        distance_in_bin = np.minimum(distance, distance_inside)  # this line is good
+
+        # Calculate flux and other scores
+        flux = distance_in_bin * P["w"]
+        for i in range(tally["N_score"]):
+            score_type = tally["scores"][i]
+            if score_type == SCORE_FLUX:
+                score = flux
+            elif score_type == SCORE_DENSITY:
+                score = flux / physics.get_speed(P_arr, mcdc)
+            elif score_type == SCORE_TOTAL:
+                SigmaT = get_MacroXS(XS_TOTAL, material, P_arr, mcdc)
+                score = flux * SigmaT
+            elif score_type == SCORE_FISSION:
+                SigmaF = get_MacroXS(XS_FISSION, material, P_arr, mcdc)
+                score = flux * SigmaF
+
+            adapt.global_add(
+                tally_bin,
+                (TALLY_SCORE, bin_idx + j * tally["N_score"] + i),
+                round(score),
+            )
+
+
+@njit
+def cs_clip(p, q, t0, t1):
+    if p < 0:
+        t = q / p
+        if t > t1:
+            return False, t0, t1
+        if t > t0:
+            t0 = t
+    elif p > 0:
+        t = q / p
+        if t < t0:
+            return False, t0, t1
+        if t < t1:
+            t1 = t
+    elif q < 0:
+        return False, t0, t1
+    return True, t0, t1
+
+
+@njit
+def cs_tracklength_in_box(start, end, x_min, x_max, y_min, y_max):
+    # Uses Liang-Barsky algorithm for finding tracklength in box
+    t0, t1 = 0.0, 1.0
+    dx = end[0] - start[0]
+    dy = end[1] - start[1]
+
+    # Perform clipping for each boundary
+    result, t0, t1 = cs_clip(-dx, start[0] - x_min, t0, t1)
+    if not result:
+        return 0.0
+    result, t0, t1 = cs_clip(dx, x_max - start[0], t0, t1)
+    if not result:
+        return 0.0
+    result, t0, t1 = cs_clip(-dy, start[1] - y_min, t0, t1)
+    if not result:
+        return 0.0
+    result, t0, t1 = cs_clip(dy, y_max - start[1], t0, t1)
+    if not result:
+        return 0.0
+
+    # Update start and end points based on clipping results
+    if t1 < 1:
+        end = start + t1 * np.array([dx, dy])
+    if t0 > 0:
+        start = start + t0 * np.array([dx, dy])
+
+    return np.linalg.norm(end - start)
+
+
+@njit
+def calculate_distance_in_coarse_bin(start, end, distance, center, cs_bin_size):
+    # Edges of the coarse bin
+    x_min = center[0] - cs_bin_size[0] / 2
+    x_max = center[0] + cs_bin_size[0] / 2
+    y_min = center[1] - cs_bin_size[1] / 2
+    y_max = center[1] + cs_bin_size[1] / 2
+
+    distance_inside = cs_tracklength_in_box(start, end, x_min, x_max, y_min, y_max)
+
+    return distance_inside
 
 
 @njit
@@ -2030,20 +2165,12 @@ def tally_reduce(data, mcdc):
 
 @njit
 def tally_accumulate(data, mcdc):
-
-    # print(f'fixed_source_data = {data[TALLY][TALLY_SCORE][-1]}')
-
     tally_bin = data[TALLY]
-    # print(data[0][0, 30000], data[0][0, 30001])
     N_bin = tally_bin.shape[1]
 
-    # print(f'N_bin = {N_bin}')
-
     for i in range(N_bin):
         # Accumulate score and square of score into sum and sum_sq
         score = tally_bin[TALLY_SCORE, i]
-        # if (score != 0 and i > 29999):
-        #     print(f'score = {score}, i = {i}')
         tally_bin[TALLY_SUM, i] += score
         tally_bin[TALLY_SUM_SQ, i] += score * score
 
@@ -2412,6 +2539,11 @@ def move_to_event(P_arr, data, mcdc):
             ID = cell["tally_IDs"][i]
             tally = mcdc["cell_tallies"][ID]
             score_cell_tally(P_arr, distance, tally, data, mcdc)
+
+        # CS tallies
+        for tally in mcdc["cs_tallies"]:
+            score_cs_tally(P_arr, distance, tally, data, mcdc)
+
     if mcdc["setting"]["mode_eigenvalue"]:
         eigenvalue_tally(P_arr, distance, mcdc)