Merge pull request #212 from alexandermote/better_tally

Decomposed mesh tallies when domain decomposition is active

mcdc/kernel.py

@@ -245,6 +245,7 @@ def dd_distribute_bank(mcdc, bank, dest_list):
     with objmode(send_delta="int64"):
         dest_count = len(dest_list)
         send_delta = 0
+
         for i, dest in enumerate(dest_list):
             size = get_bank_size(bank)
             ratio = int(size / dest_count)
@@ -2210,11 +2211,12 @@ def tally_reduce(data, mcdc):
     for i in range(N_bin):
         tally_bin[TALLY_SCORE][i] /= N_particle
 
-    # MPI Reduce
-    buff = np.zeros_like(tally_bin[TALLY_SCORE])
-    with objmode():
-        MPI.COMM_WORLD.Reduce(tally_bin[TALLY_SCORE], buff, MPI.SUM, 0)
-    tally_bin[TALLY_SCORE][:] = buff
+    if not mcdc["technique"]["domain_decomposition"]:
+        # MPI Reduce
+        buff = np.zeros_like(tally_bin[TALLY_SCORE])
+        with objmode():
+            MPI.COMM_WORLD.Reduce(tally_bin[TALLY_SCORE], buff, MPI.SUM, 0)
+        tally_bin[TALLY_SCORE][:] = buff
 
 
 @njit
@@ -2243,7 +2245,7 @@ def tally_closeout(data, mcdc):
     elif mcdc["setting"]["mode_eigenvalue"]:
         N_history = mcdc["setting"]["N_active"]
 
-    else:
+    elif not mcdc["technique"]["domain_decomposition"]:
         # MPI Reduce
         buff = np.zeros_like(tally[TALLY_SUM])
         buff_sq = np.zeros_like(tally[TALLY_SUM_SQ])

mcdc/main.py

@@ -314,6 +314,39 @@ def dd_prepare():
         input_deck.technique["dd_zn_neigh"] = []
 
 
+def dd_mesh_bounds(idx):
+    """
+    Defining mesh tally boundaries for domain decomposition.
+    Used in prepare() when domain decomposition is active.
+    """
+    # find DD mesh index of subdomain
+    d_idx = input_deck.technique["dd_idx"]  # subdomain index
+    d_Nx = input_deck.technique["dd_mesh"]["x"].size - 1
+    d_Ny = input_deck.technique["dd_mesh"]["y"].size - 1
+    d_Nz = input_deck.technique["dd_mesh"]["z"].size - 1
+    zmesh_idx = d_idx // (d_Nx * d_Ny)
+    ymesh_idx = (d_idx % (d_Nx * d_Ny)) // d_Nx
+    xmesh_idx = d_idx % d_Nx
+
+    # find spatial boundaries of subdomain
+    xn = input_deck.technique["dd_mesh"]["x"][xmesh_idx]
+    xp = input_deck.technique["dd_mesh"]["x"][xmesh_idx + 1]
+    yn = input_deck.technique["dd_mesh"]["y"][ymesh_idx]
+    yp = input_deck.technique["dd_mesh"]["y"][ymesh_idx + 1]
+    zn = input_deck.technique["dd_mesh"]["z"][zmesh_idx]
+    zp = input_deck.technique["dd_mesh"]["z"][zmesh_idx + 1]
+
+    # find boundary indices in tally mesh
+    mesh_xn = int(np.where(input_deck.mesh_tallies[idx].x == xn)[0])
+    mesh_xp = int(np.where(input_deck.mesh_tallies[idx].x == xp)[0]) + 1
+    mesh_yn = int(np.where(input_deck.mesh_tallies[idx].y == yn)[0])
+    mesh_yp = int(np.where(input_deck.mesh_tallies[idx].y == yp)[0]) + 1
+    mesh_zn = int(np.where(input_deck.mesh_tallies[idx].z == zn)[0])
+    mesh_zp = int(np.where(input_deck.mesh_tallies[idx].z == zp)[0]) + 1
+
+    return mesh_xn, mesh_xp, mesh_yn, mesh_yp, mesh_zn, mesh_zp
+
+
 def prepare():
     """
     Preparing the MC transport simulation:
@@ -627,11 +660,28 @@ def prepare():
         copy_field(mcdc["mesh_tallies"][i], input_deck.mesh_tallies[i], "N_bin")
 
         # Filters (variables with possible different sizes)
-        for name in ["x", "y", "z", "t", "mu", "azi", "g"]:
-            N = len(getattr(input_deck.mesh_tallies[i], name))
-            mcdc["mesh_tallies"][i]["filter"][name][:N] = getattr(
-                input_deck.mesh_tallies[i], name
-            )
+        if not input_deck.technique["domain_decomposition"]:
+            for name in ["x", "y", "z", "t", "mu", "azi", "g"]:
+                N = len(getattr(input_deck.mesh_tallies[i], name))
+                mcdc["mesh_tallies"][i]["filter"][name][:N] = getattr(
+                    input_deck.mesh_tallies[i], name
+                )
+
+        else:  # decomposed mesh filters
+            mxn, mxp, myn, myp, mzn, mzp = dd_mesh_bounds(i)
+
+            # Filters
+            new_x = input_deck.mesh_tallies[i].x[mxn:mxp]
+            new_y = input_deck.mesh_tallies[i].y[myn:myp]
+            new_z = input_deck.mesh_tallies[i].z[mzn:mzp]
+            mcdc["mesh_tallies"][i]["filter"]["x"] = new_x
+            mcdc["mesh_tallies"][i]["filter"]["y"] = new_y
+            mcdc["mesh_tallies"][i]["filter"]["z"] = new_z
+            for name in ["t", "mu", "azi", "g"]:
+                N = len(getattr(input_deck.mesh_tallies[i], name))
+                mcdc["mesh_tallies"][i]["filter"][name][:N] = getattr(
+                    input_deck.mesh_tallies[i], name
+                )
 
         # Set tally scores
         N_score = len(input_deck.mesh_tallies[i].scores)
@@ -662,6 +712,17 @@ def prepare():
         Nz = len(input_deck.mesh_tallies[i].z) - 1
         Nt = len(input_deck.mesh_tallies[i].t) - 1
 
+        # Decompose mesh tallies
+        if input_deck.technique["domain_decomposition"]:
+            Nmu = len(input_deck.mesh_tallies[i].mu) - 1
+            N_azi = len(input_deck.mesh_tallies[i].azi) - 1
+            Ng = len(input_deck.mesh_tallies[i].g) - 1
+            Nx = len(input_deck.mesh_tallies[i].x[mxn:mxp]) - 1
+            Ny = len(input_deck.mesh_tallies[i].y[myn:myp]) - 1
+            Nz = len(input_deck.mesh_tallies[i].z[mzn:mzp]) - 1
+            Nt = len(input_deck.mesh_tallies[i].t) - 1
+            mcdc["mesh_tallies"][i]["N_bin"] = Nx * Ny * Nz * Nt * Nmu * N_azi * Ng
+
         # Update N_bin
         mcdc["mesh_tallies"][i]["N_bin"] *= Ns * N_score
 
@@ -1144,6 +1205,20 @@ def dict_to_h5group(dict_, group):
 
 
 def generate_hdf5(data, mcdc):
+
+    # recombine tallies before output processing
+    if mcdc["technique"]["domain_decomposition"]:
+        tally = data[TALLY]
+        # create bin for recomposed tallies
+        d_Nx = input_deck.technique["dd_mesh"]["x"].size - 1
+        d_Ny = input_deck.technique["dd_mesh"]["y"].size - 1
+        d_Nz = input_deck.technique["dd_mesh"]["z"].size - 1
+        dd_tally = np.zeros((tally.shape[0], tally.shape[1] * d_Nx * d_Ny * d_Nz))
+        # gather tallies
+        MPI.COMM_WORLD.Gather(tally[0], dd_tally[0], root=0)
+        MPI.COMM_WORLD.Gather(tally[1], dd_tally[1], root=0)
+        MPI.COMM_WORLD.Gather(tally[2], dd_tally[2], root=0)
+
     if mcdc["mpi_master"]:
         if mcdc["setting"]["progress_bar"]:
             print_msg("")
@@ -1203,15 +1278,28 @@ def generate_hdf5(data, mcdc):
                 Nt = len(mesh["t"]) - 1
                 N_score = tally["N_score"]
 
+                if mcdc["technique"]["domain_decomposition"]:
+                    Nx *= input_deck.technique["dd_mesh"]["x"].size - 1
+                    Ny *= input_deck.technique["dd_mesh"]["y"].size - 1
+                    Nz *= input_deck.technique["dd_mesh"]["z"].size - 1
+
                 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"]
+                if mcdc["technique"]["domain_decomposition"]:
+                    # use recomposed N_bin
+                    N_bin = 1
+                    for elem in shape:
+                        N_bin *= elem
                 start = tally["stride"]["tally"]
                 tally_bin = data[TALLY][:, start : start + N_bin]
+                if mcdc["technique"]["domain_decomposition"]:
+                    # substitute recomposed tally
+                    tally_bin = dd_tally[:, start : start + N_bin]
                 tally_bin = tally_bin.reshape(shape)
 
                 # Roll tally so that score is in the front

mcdc/type_.py

@@ -674,6 +674,47 @@ def make_type_source(input_deck):
 # ==============================================================================
 
 
+def dd_meshtally(input_deck):
+    # find DD mesh index of subdomain
+    d_idx = input_deck.technique["dd_idx"]  # subdomain index
+    d_Nx = input_deck.technique["dd_mesh"]["x"].size - 1
+    d_Ny = input_deck.technique["dd_mesh"]["y"].size - 1
+    d_Nz = input_deck.technique["dd_mesh"]["z"].size - 1
+    zmesh_idx = d_idx // (d_Nx * d_Ny)
+    ymesh_idx = (d_idx % (d_Nx * d_Ny)) // d_Nx
+    xmesh_idx = d_idx % d_Nx
+
+    # find spatial boundaries of subdomain
+    xn = input_deck.technique["dd_mesh"]["x"][xmesh_idx]
+    xp = input_deck.technique["dd_mesh"]["x"][xmesh_idx + 1]
+    yn = input_deck.technique["dd_mesh"]["y"][ymesh_idx]
+    yp = input_deck.technique["dd_mesh"]["y"][ymesh_idx + 1]
+    zn = input_deck.technique["dd_mesh"]["z"][zmesh_idx]
+    zp = input_deck.technique["dd_mesh"]["z"][zmesh_idx + 1]
+
+    # Maximum numbers of mesh and filter grids and scores
+    Nx = 2
+    Ny = 2
+    Nz = 2
+    for card in input_deck.mesh_tallies:
+        # find boundary indices in tally mesh
+        mesh_xn = int(np.where(card.x == xn)[0])
+        mesh_xp = int(np.where(card.x == xp)[0]) + 1
+        mesh_yn = int(np.where(card.y == yn)[0])
+        mesh_yp = int(np.where(card.y == yp)[0]) + 1
+        mesh_zn = int(np.where(card.z == zn)[0])
+        mesh_zp = int(np.where(card.z == zp)[0]) + 1
+
+        # adjust Nmax numbers
+        new_x = card.x[mesh_xn:mesh_xp]
+        new_y = card.y[mesh_yn:mesh_yp]
+        new_z = card.z[mesh_zn:mesh_zp]
+        Nx = max(Nx, len(new_x))
+        Ny = max(Ny, len(new_y))
+        Nz = max(Nz, len(new_z))
+    return Nx, Ny, Nz
+
+
 def make_type_mesh_tally(input_deck):
     global mesh_tally
     struct = []
@@ -697,6 +738,10 @@ def make_type_mesh_tally(input_deck):
         Nmax_g = max(Nmax_g, len(card.g))
         Nmax_score = max(Nmax_score, len(card.scores))
 
+    # reduce tally sizes for subdomains
+    if input_deck.technique["domain_decomposition"]:
+        Nmax_x, Nmax_y, Nmax_z = dd_meshtally(input_deck)
+
     # Set the filter
     filter_ = [
         ("x", float64, (Nmax_x,)),