Merge pull request CEMeNT-PSAAP#227 from CEMeNT-PSAAP/better_tally

Merge the new tally structure

mcdc/__init__.py

@@ -8,7 +8,6 @@
     universe,
     lattice,
     source,
-    tally,
     setting,
     eigenmode,
     implicit_capture,
@@ -26,6 +25,7 @@
     make_particle_bank,
     save_particle_bank,
 )
+import mcdc.tally
 from mcdc.main import run, prepare
 
 __version__ = importlib.metadata.version("mcdc")

mcdc/card.py

@@ -6,7 +6,7 @@
     BOOL_OR,
     BOOL_NOT,
     INF,
-    TINY,
+    PI,
 )
 
 # Get the global variable container
@@ -33,11 +33,17 @@ def __str__(self):
 
 
 class NuclideCard(InputCard):
-    def __init__(self, G=1, J=0):
+    def __init__(self, G=1, J=0, name=None):
         InputCard.__init__(self, "Nuclide")
 
+        # Continuous energy?
+        if name is not None:
+            G = 0
+            J = 0
+
         # Set card data
         self.ID = None
+        self.name = name
         self.G = G
         self.J = J
         self.fissionable = False
@@ -161,6 +167,8 @@ def __init__(self):
         self.nx = 0.0
         self.ny = 0.0
         self.nz = 0.0
+        self.N_tally = 0
+        self.tally_IDs = []
 
     def _create_halfspace(self, positive):
         region = RegionCard("halfspace")
@@ -331,6 +339,47 @@ def __init__(self):
         self.white_y = 0.0
         self.white_z = 0.0
         self.group = np.array([1.0])
-        self.energy = np.array([[14e6, 14e6], [1.0, 1.0]])
+        self.energy = np.array([[1e6 - 1.0, 1e6 + 1.0], [1.0, 1.0]])
         self.time = np.array([0.0, 0.0])
         self.prob = 1.0
+
+
+# ======================================================================================
+# Tally cards
+# ======================================================================================
+
+
+class TallyCard(InputCard):
+    def __init__(self, type_):
+        InputCard.__init__(self, type_)
+
+        # Set card data
+        self.ID = None
+        self.scores = []
+        self.N_bin = 0
+
+        # Filters
+        self.t = np.array([-INF, INF])
+        self.mu = np.array([-1.0, 1.0])
+        self.azi = np.array([-PI, PI])
+        self.g = np.array([-INF, INF])
+
+
+class MeshTallyCard(TallyCard):
+    def __init__(self):
+        TallyCard.__init__(self, "Mesh 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
+
+
+class SurfaceTallyCard(TallyCard):
+    def __init__(self, surface_ID):
+        TallyCard.__init__(self, "Surface tally")
+
+        # Set card data
+        self.surface_ID = surface_ID
+        self.N_bin = 1

mcdc/constant.py

@@ -3,6 +3,22 @@
 import numba as nb
 
 
+# Data index
+TALLY = 0
+
+# Tally bins
+TALLY_SCORE = 0
+TALLY_SUM = 1
+TALLY_SUM_SQ = 2
+TALLY_UQ_BATCH = 3
+TALLY_UQ_BATCH_VAR = 4
+
+# Tally scores
+SCORE_FLUX = 0
+SCORE_TOTAL = 1
+SCORE_FISSION = 2
+SCORE_NET_CURRENT = 3
+
 # Boundary condition
 BC_NONE = 0
 BC_VACUUM = 1

mcdc/global_.py

@@ -33,19 +33,8 @@ def reset(self):
         self.universes = [None]  # Placeholder for the root universe
         self.lattices = []
         self.sources = []
-
-        self.tally = {
-            "tag": "Tally",
-            "tracklength": True,
-            "flux": False,
-            "density": False,
-            "fission": False,
-            "total": False,
-            "current": False,
-            "eddington": False,
-            "exit": False,
-            "mesh": make_card_mesh(),
-        }
+        self.mesh_tallies = []
+        self.surface_tallies = []
 
         self.setting = {
             "tag": "Setting",

mcdc/input_.py

@@ -10,8 +10,6 @@
 import numpy as np
 import scipy as sp
 
-import mcdc.type_ as type_
-
 from mcdc.card import (
     NuclideCard,
     MaterialCard,
@@ -39,6 +37,7 @@
     TINY,
 )
 from mcdc.print_ import print_error
+import mcdc.type_ as type_
 
 
 def nuclide(
@@ -902,6 +901,16 @@ def source(**kw):
             G = global_.input_deck.materials[0].G
             group = np.ones(G)
             card.group = group / np.sum(group)
+        # Default for CE
+        if global_.input_deck.setting["mode_CE"]:
+            # Normalize pdf
+            card.energy[1, :] = card.energy[1, :] / np.trapz(
+                card.energy[1, :], x=card.energy[0, :]
+            )
+            # Make cdf
+            card.energy[1, :] = sp.integrate.cumulative_trapezoid(
+                card.energy[1], x=card.energy[0], initial=0.0
+            )
 
     # Set time
     if time is not None:
@@ -917,82 +926,6 @@ def source(**kw):
     return card
 
 
-def tally(
-    scores,
-    x=np.array([-INF, INF]),
-    y=np.array([-INF, INF]),
-    z=np.array([-INF, INF]),
-    t=np.array([-INF, INF]),
-    mu=np.array([-1.0, 1.0]),
-    azi=np.array([-PI, PI]),
-    g=np.array([-INF, INF]),
-    E=np.array([0.0, INF]),
-):
-    """
-    Create a tally card to collect MC solutions.
-
-    Parameters
-    ----------
-    scores : list of str
-        List of tally types (default ["tracklength"]).
-    x : array_like[float], optional
-        x-coordinates that demarcate tally bins (default numpy.ndarray([-INF, INF])).
-    y : array_like[float], optional
-        y-coordinates that demarcate tally bins (default numpy.ndarray([-INF, INF])).
-    z : array_like[float], optional
-        z-coordinates that demarcate tally bins (default numpy.ndarray([-INF, INF])).
-    t : array_like[float], optional
-        Times that demarcate tally bins (default numpy.ndarray([-INF, INF])).
-    mu : array_like[float], optional
-        Angles that demarcate axial angular tally bins (default numpy.ndarray([-1.0, 1.0])).
-    azi : array_like[float], optional
-        Angles that demarcate azimuthal angular tally bins (default numpy.ndarray([-1.0, 1.0])).
-    g : array_like[float], optional
-        Energies that demarcate energy tally bins (default numpy.ndarray([-INF, INF])).
-    E : array_like[float], optional
-        Continuous energy functionality, (default numpy.ndarray([0.0, INF])).
-
-    Returns
-    -------
-    dictionary
-        A tally card.
-    """
-
-    # Get tally card
-    card = global_.input_deck.tally
-
-    # Set mesh
-    card["mesh"]["x"] = x
-    card["mesh"]["y"] = y
-    card["mesh"]["z"] = z
-    card["mesh"]["t"] = t
-    card["mesh"]["mu"] = mu
-    card["mesh"]["azi"] = azi
-
-    # Set energy group grid
-    if type(g) == type("string") and g == "all":
-        G = global_.input_deck.materials[0].G
-        card["mesh"]["g"] = np.linspace(0, G, G + 1) - 0.5
-    else:
-        card["mesh"]["g"] = g
-    if global_.input_deck.setting["mode_CE"]:
-        card["mesh"]["g"] = E
-
-    # Set score flags
-    for s in scores:
-        found = False
-        for score_name in type_.score_list:
-            if s.replace("-", "_") == score_name:
-                card["tracklength"] = True
-                card[score_name] = True
-                found = True
-                break
-        if not found:
-            print_error("Unknown tally score %s" % s)
-
-    return card
-
-
 # ==============================================================================
 # Setting
 # ==============================================================================
@@ -1203,10 +1136,6 @@ def eigenmode(
         else:
             print_error("Unknown gyration radius type")
 
-    # Update tally card
-    card = global_.input_deck.tally
-    card["tracklength"] = True
-
 
 # ==============================================================================
 # Technique

mcdc/iqmc/iqmc_kernel.py

@@ -455,6 +455,24 @@ def iqmc_continuous_weight_reduction(P, distance, mcdc):
     P["w"] = P["iqmc"]["w"].sum()
 
 
+# =============================================================================
+# Surface crossing
+# =============================================================================
+
+
+@toggle("iQMC")
+def iqmc_surface_crossing(P, prog):
+    mcdc = adapt.device(prog)
+
+    # Implement BC
+    surface = mcdc["surfaces"][P["surface_ID"]]
+    geometry.surface_bc(P, surface)
+
+    # Need to check new cell later?
+    if P["alive"] and not surface["BC"] == BC_REFLECTIVE:
+        P["cell_ID"] = -1
+
+
 # =============================================================================
 # iQMC Source Operations
 # =============================================================================

mcdc/iqmc/iqmc_loop.py

@@ -372,7 +372,7 @@ def iqmc_step_particle(P, prog):
 
     # Surface crossing
     if event & EVENT_SURFACE_CROSSING:
-        kernel.surface_crossing(P, prog)
+        iqmc_kernel.iqmc_surface_crossing(P, prog)
         if event & EVENT_DOMAIN_CROSSING:
             if not (
                 mcdc["surfaces"][P["surface_ID"]]["BC"] == BC_REFLECTIVE