Merge pull request #222 from ilhamv/refactor_csg

Refactor CSG tracking with Reverse Polish Notation

mcdc/adapt.py

@@ -416,6 +416,16 @@ def local_j_array():
     return cuda.local.array(1, type_.j_array)[0]
 
 
+@for_cpu()
+def local_RPN_array():
+    return np.zeros(1, dtype=type_.RPN_array)[0]
+
+
+@for_gpu()
+def local_RPN_array():
+    return cuda.local.array(1, type_.RPN_array)[0]
+
+
 @for_cpu()
 def local_particle():
     return np.zeros(1, dtype=type_.particle)[0]

mcdc/card.py

@@ -1,6 +1,13 @@
 import numpy as np
+import sympy
 
-from mcdc.constant import INF, SHIFT
+from mcdc.constant import (
+    BOOL_AND,
+    BOOL_OR,
+    BOOL_NOT,
+    INF,
+    SHIFT,
+)
 
 # Get the global variable container
 import mcdc.global_ as global_
@@ -16,7 +23,10 @@ def __str__(self):
         for name in [
             a
             for a in dir(self)
-            if not a.startswith("__") and not callable(getattr(self, a)) and a != "tag"
+            if not a.startswith("__")
+            and not callable(getattr(self, a))
+            and a != "tag"
+            and not a.startswith("_")
         ]:
             text += "  %s : %s\n" % (name, str(getattr(self, name)))
         return text
@@ -88,8 +98,8 @@ def __init__(self, type_):
         # Set card data
         self.ID = None
         self.type = type_
-        self.A = -1
-        self.B = -1
+        self.A = None
+        self.B = None
 
     def __and__(self, other):
         region = RegionCard("intersection")
@@ -117,6 +127,21 @@ def __invert__(self):
         global_.input_deck.regions.append(region)
         return region
 
+    def __str__(self):
+        if self.type == "halfspace":
+            if self.B > 0:
+                return "+s%i" % self.A
+            else:
+                return "-s%i" % self.A
+        elif self.type == "intersection":
+            return "r%i & r%i" % (self.A, self.B)
+        elif self.type == "union":
+            return "r%i | r%i" % (self.A, self.B)
+        elif self.type == "complement":
+            return "~r%i" % (self.A)
+        elif self.type == "all":
+            return "all"
+
 
 class SurfaceCard(InputCard):
     def __init__(self):
@@ -146,23 +171,33 @@ def __init__(self):
         self.sensitivity_ID = 0
         self.dsm_Np = 1.0
 
-    def __pos__(self):
+    def _create_halfspace(self, positive):
         region = RegionCard("halfspace")
         region.A = self.ID
-        region.B = 1
+        if positive:
+            region.B = 1
+        else:
+            region.B = -1
+
+        # Check if an identical halfspace region already existed
+        for idx, existing_region in enumerate(global_.input_deck.regions):
+            if (
+                existing_region.type == "halfspace"
+                and region.A == existing_region.A
+                and region.B == existing_region.B
+            ):
+                return global_.input_deck.regions[idx]
+
         # Set ID and push to deck
         region.ID = len(global_.input_deck.regions)
         global_.input_deck.regions.append(region)
         return region
 
+    def __pos__(self):
+        return self._create_halfspace(True)
+
     def __neg__(self):
-        region = RegionCard("halfspace")
-        region.A = self.ID
-        region.B = 0
-        # Set ID and push to deck
-        region.ID = len(global_.input_deck.regions)
-        global_.input_deck.regions.append(region)
-        return region
+        return self._create_halfspace(False)
 
 
 class CellCard(InputCard):
@@ -171,12 +206,87 @@ def __init__(self):
 
         # Set card data
         self.ID = None
-        self.region_ID = -1
+        self.region_ID = None
+        self.region = "all"
         self.fill_type = "material"
-        self.fill_ID = -1
+        self.fill_ID = None
         self.translation = np.array([0.0, 0.0, 0.0])
-        self.N_surface = 0
-        self.surface_ID = np.zeros(0, dtype=int)
+        self.surface_IDs = np.zeros(0, dtype=int)
+        self._region_RPN = []  # Reverse Polish Notation
+
+    def set_region_RPN(self):
+        # Make alias and reset
+        rpn = self._region_RPN
+        rpn.clear()
+
+        # Build RPN based on the assigned region
+        region = global_.input_deck.regions[self.region_ID]
+        stack = [region]
+        while len(stack) > 0:
+            token = stack.pop()
+            if isinstance(token, RegionCard):
+                if token.type == "halfspace":
+                    rpn.append(token.ID)
+                elif token.type == "intersection":
+                    region_A = global_.input_deck.regions[token.A]
+                    region_B = global_.input_deck.regions[token.B]
+                    stack += ["&", region_A, region_B]
+                elif token.type == "union":
+                    region_A = global_.input_deck.regions[token.A]
+                    region_B = global_.input_deck.regions[token.B]
+                    stack += ["|", region_A, region_B]
+                elif token.type == "complement":
+                    region = global_.input_deck.regions[token.A]
+                    stack += ["~", region]
+            else:
+                if token == "&":
+                    rpn.append(BOOL_AND)
+                elif token == "|":
+                    rpn.append(BOOL_OR)
+                elif token == "~":
+                    rpn.append(BOOL_NOT)
+                else:
+                    print_error("Something is wrong with cell RPN creation.")
+
+    def set_region(self):
+        stack = []
+
+        for token in self._region_RPN:
+            if token >= 0:
+                stack.append(token)
+            else:
+                if token == BOOL_AND or token == BOOL_OR:
+                    item_1 = stack.pop()
+                    if isinstance(item_1, int):
+                        item_1 = sympy.symbols(str(global_.input_deck.regions[item_1]))
+
+                    item_2 = stack.pop()
+                    if isinstance(item_2, int):
+                        item_2 = sympy.symbols(str(global_.input_deck.regions[item_2]))
+
+                    if token == BOOL_AND:
+                        stack.append(item_1 & item_2)
+                    else:
+                        stack.append(item_1 | item_2)
+
+                elif token == BOOL_NOT:
+                    item = stack.pop()
+                    if isinstance(item, int):
+                        item = sympy.symbols(str(global_.input_deck.regions[item]))
+                    stack.append(~item)
+
+        self.region = sympy.logic.boolalg.simplify_logic(stack[0])
+
+    def set_surface_IDs(self):
+        surface_IDs = []
+
+        for token in self._region_RPN:
+            if token >= 0:
+                ID = global_.input_deck.regions[token].A
+                if not ID in surface_IDs:
+                    surface_IDs.append(ID)
+
+        self.surface_IDs = np.sort(np.array(surface_IDs))
 
 
 class UniverseCard(InputCard):

mcdc/constant.py

@@ -20,7 +20,12 @@
 REGION_COMPLEMENT = 3
 REGION_ALL = 4
 
-# UNIVERSE
+# Boolean operator
+BOOL_AND = -1
+BOOL_OR = -2
+BOOL_NOT = -3
+
+# Universe
 UNIVERSE_ROOT = 0
 
 # Events

mcdc/input_.py

@@ -732,6 +732,11 @@ def cell(region=None, fill=None, translation=(0.0, 0.0, 0.0)):
     # Assign region
     card.region_ID = region.ID
 
+    # Set region Reverse Polish Notation and region description
+    if region.type != "all":
+        card.set_region_RPN()
+        card.set_region()
+
     # Assign fill type and ID
     if fill.tag == "Material":
         card.fill_type = "material"
@@ -745,39 +750,14 @@ def cell(region=None, fill=None, translation=(0.0, 0.0, 0.0)):
     card.translation[:] = translation
 
     # Get all surface IDs
-    surface_IDs = []
-
-    region = global_.input_deck.regions[card.region_ID]
-    get_all_surface_IDs(region, surface_IDs)
-    surface_IDs = list(set(surface_IDs))
-
-    card.N_surface = len(surface_IDs)
-    card.surface_IDs = np.array(surface_IDs)
+    card.set_surface_IDs()
 
     # Add to deck
     global_.input_deck.cells.append(card)
 
     return card
 
 
-def get_all_surface_IDs(region, surface_IDs):
-    if region.type == "halfspace":
-        surface = global_.input_deck.surfaces[region.A]
-        surface_IDs.append(surface.ID)
-
-    elif region.type in ["intersection", "union"]:
-        region1 = global_.input_deck.regions[region.A]
-        region2 = global_.input_deck.regions[region.B]
-
-        get_all_surface_IDs(region1, surface_IDs)
-        get_all_surface_IDs(region2, surface_IDs)
-
-    elif region.type in ["complement"]:
-        region1 = global_.input_deck.regions[region.A]
-
-        get_all_surface_IDs(region1, surface_IDs)
-
-
 def universe(cells, root=False):
     """
     Define a list of cells as a universe.

mcdc/kernel.py

@@ -1708,8 +1708,40 @@ def split_particle(P):
 
 @njit
 def cell_check(P, cell, trans, mcdc):
-    region = mcdc["regions"][cell["region_ID"]]
-    return region_check(P, region, trans, mcdc)
+    # Access RPN data
+    idx = cell["region_data_idx"]
+    N_token = mcdc["cell_region_data"][idx]
+
+    # Create local value array
+    value_struct = adapt.local_RPN_array()
+    value = value_struct["values"]
+    N_value = 0
+
+    # March forward through RPN tokens
+    idx += 1
+    idx_end = idx + N_token
+    while idx < idx_end:
+        token = mcdc["cell_region_data"][idx]
+
+        if token >= 0:
+            surface = mcdc["surfaces"][token]
+            value[N_value] = surface_evaluate(P, surface, trans) > 0.0
+            N_value += 1
+
+        elif token == BOOL_NOT:
+            value[N_value - 1] = not value[N_value - 1]
+
+        elif token == BOOL_AND:
+            value[N_value - 2] = value[N_value - 2] & value[N_value - 1]
+            N_value -= 1
+
+        elif token == BOOL_OR:
+            value[N_value - 2] = value[N_value - 2] | value[N_value - 1]
+            N_value -= 1
+
+        idx += 1
+
+    return value[0]
 
 
 @njit
@@ -1864,7 +1896,6 @@ def surface_shift(P, surface, trans, mcdc):
         dot = ux * nx + uy * ny + uz * nz + J1 / v
     else:
         dot = ux * nx + uy * ny + uz * nz
-
     if dot > 0.0:
         P["x"] += shift_x
         P["y"] += shift_y
@@ -2531,7 +2562,6 @@ def move_to_event(P, mcdc):
     # Also set particle material and speed
     d_boundary, event = distance_to_boundary(P, mcdc)
 
-    # print(P["material_ID"])
     # Distance to tally mesh
     d_mesh = INF
     if mcdc["cycle_active"]:
@@ -2693,13 +2723,22 @@ def distance_to_nearest_surface(P, cell, trans, mcdc):
     surface_ID = -1
     surface_move = False
 
-    for i in range(cell["N_surface"]):
-        surface = mcdc["surfaces"][cell["surface_IDs"][i]]
+    # Access cell surface data
+    idx = cell["surface_data_idx"]
+    N_surface = mcdc["cell_surface_data"][idx]
+
+    # Iterate over all surfaces
+    idx += 1
+    idx_end = idx + N_surface
+    while idx < idx_end:
+        candidate_surface_ID = mcdc["cell_surface_data"][idx]
+        surface = mcdc["surfaces"][candidate_surface_ID]
         d, sm = surface_distance(P, surface, trans, mcdc)
         if d < distance:
             distance = d
             surface_ID = surface["ID"]
             surface_move = sm
+        idx += 1
     return distance, surface_ID, surface_move