Numba mode works for capture xsec uq

input.py

@@ -41,7 +41,7 @@
 
 # Setting
 seed = random.randint(1,1000)
-mcdc.setting(N_particle=1E4, N_batch=2, rng_seed=seed)
+mcdc.setting(N_particle=1E3, N_batch=1E4, rng_seed=seed)
 
 mcdc.uq(material=m1, distribution="uniform", capture=np.array([0.7]))
 mcdc.uq(material=m2, distribution="uniform", capture=np.array([0.12]))

mcdc/card.py

@@ -297,4 +297,6 @@ def make_card_uq():
         "delta": 0.0,
         "distribution": 'd',
         "rng_seed": 0,
+        "group" : False,
+        "group_group" : False,
     }

mcdc/kernel.py

@@ -154,11 +154,12 @@ def rng_from_seed(seed):
 
 
 @njit
-def rng_array(seed, d):
-    xi = np.zeros(d)
-    for i in range(d):
+def rng_array(seed, shape, size):
+    xi = np.zeros(size)
+    for i in range(size):
         xi_seed = split_seed(i, seed)
         xi[i] = rng_from_seed(xi_seed)
+    xi.reshape(shape)
     return xi
 
 
@@ -3454,23 +3455,36 @@ def binary_search(val, grid):
 @njit
 def uq_resample(parameter):
     # Currently only uniform distribution
-    d = len(parameter["mean"])
-    xi = rng_array(parameter["rng_seed"], d)
+    shape = parameter["mean"].shape
+    size = parameter["mean"].size
+    xi = rng_array(parameter["rng_seed"], shape, size)
+
+    return parameter["mean"] + (2*xi - 1)*parameter["delta"]
 
-    if parameter["distribution"] == "uniform":
-        return parameter["mean"] + (2*xi - 1)*parameter["delta"]
+
+@njit
+def uq_reset_material(param, mcdc):
+    material = mcdc["materials"][param["ID"]]
+    # Assumes just capture xsec for now
+    capture = uq_resample(param)
+    for key in literal_unroll(("capture",)):
+        material[key] = capture
+    for key in literal_unroll(("total",)):
+        # No scatter or fission for now
+        material[key] = capture
 
 
 @njit
 def uq_reset(mcdc, seed):
-    # Loop over number of uq parameters
-    parameters = mcdc["technique"]["uq_"]["parameters"]
-
-    # Numba work-around
-    for idx_param in literal_unroll(mcdc["technique"]["uq_"]["names"]):
-        param = parameters[idx_param]
-#        param["rng_seed"] = split_seed(idx_param, seed)
-#        mcdc[param["tag"]][param["ID"]][param["key"]] = uq_resample(param)
+    # Loop over uq parameters based on shape
+    parameters = mcdc["technique"]["uq_"]
+    g_list = ("capture", "scatter", "fission", "nu_s", "nu_p", "chi_d", "speed", "decay")
+    gg_list = ("chi_p")
+    for idx_param in range(parameters["N_group"]):
+        param = parameters["group_idx"][idx_param]
+        param["rng_seed"] = split_seed(idx_param, seed)
+        if param["tag"] == "materials":
+            uq_reset_material(param, mcdc)
 
 
 @njit

mcdc/main.py

@@ -411,15 +411,25 @@ def prepare():
     # Variance Deconvolution - UQ
     # =========================================================================
     mcdc["technique"]["uq"] = True
-    mcdc["technique"]["uq_"]["N_params"] = len(input_deck.uq_parameters)
+    N_uq = len(input_deck.uq_parameters)
     for name in type_.uq_tally.names:
         if name != "score":
             mcdc["technique"]["uq_tally"][name] = input_deck.tally[name]
 
-    for i in range(len(input_deck.uq_parameters)):
-        mcdc["technique"]["uq_"]["names"][i] = str(i)
-        for name in type_.uq["parameters"][0].names:
-            mcdc["technique"]["uq_"]["parameters"][str(i)][name] = input_deck.uq_parameters[i][name]
+    g = 0
+    gg = 0
+    for parameter in input_deck.uq_parameters:
+        if parameter["group"]:
+            for name in type_.type_group.names:
+                mcdc["technique"]["uq_"]["group_idx"][g][name] = input_deck.uq_parameters[g][name]
+            g += 1
+        elif parameter["group_group"]:
+            for name in param_names:
+                mcdc["technique"]["uq_"]["groupgroup_idx"][gg][name] = input_deck.uq_parameters[gg][name]
+            gg += 1
+    mcdc["technique"]["uq_"]["N_group"] = g
+    mcdc["technique"]["uq_"]["N_group_group"] = gg
+
 
     # =========================================================================
     # MPI

mcdc/type_.py

@@ -631,7 +631,7 @@ def make_type_uq_score(shape):
 
 
 def make_type_uq(uq_deck):
-    global uq
+    global uq, type_group
 
     def make_type_parameter(shape):
         return np.dtype(
@@ -646,19 +646,37 @@ def make_type_parameter(shape):
             ]
         )
 
-    struct = [("N_params", int64)]
     N_uq = len(uq_deck)
-    parameters_struct = []
+
+    # Materials and nuclides have two shapes, (G,) and (G, G)
+    for i in range(N_uq):
+        if uq_deck[i]["tag"] in ("materials", "nuclides"):
+            G = len(uq_deck[i]["mean"])
+            break
+        else:
+            G = 0
+    type_group = make_type_parameter((G,))
+    type_group_group = make_type_parameter((G, G))
+
+    N_group = 0
+    N_group_group = 0
     for i in range(N_uq):
         shape = np.shape(uq_deck[i]["mean"])
-        parameters_struct += [(str(i), make_type_parameter(shape))]
-    parameters = np.dtype(parameters_struct)
-    struct += [("parameters", parameters)]
-    struct += [("names", str_, (N_uq,))]
+        if shape == (G,):
+            N_group += 1
+            uq_deck[i]["group"] = True
+        elif shape == (G,G):
+            N_group_group += 1
+            uq_deck[i]["group_group"] = True
+    struct = [("N_group", int64),
+              ("N_group_group", int64),
+              ("group_idx", type_group, (N_group,)),
+              ("groupgroup_idx", type_group_group, (N_group_group,)),
+              ]
     uq = np.dtype(struct)
 
 
-param_names = ["tag", "ID", "key", "mean", "delta", "dist"]
+param_names = ["tag", "ID", "key", "mean", "delta", "distribution", "rng_seed"]
 
 
 # ==============================================================================

read_output.py

@@ -0,0 +1,18 @@
+import h5py
+import numpy as np
+
+f = h5py.File('output.h5','r')
+Nb = 1E4
+mean = f['tally']['exit']['mean'][:][1]
+sdev = f['tally']['exit']['sdev'][:][1]
+varp = f['tally']['exit']['uq_var'][:][1]
+
+exact = 5.504/1000
+
+print('Var param: ' + str(varp))
+total = (sdev**2)*Nb
+print('Var total: ' +  str(total))
+print('param error: ' + str(abs(exact-varp)))
+print('total error: ' + str(abs(exact-total)))
+
+f.close()

submission_script.py

@@ -0,0 +1,57 @@
+import numpy as np
+import mcdc
+
+
+def submit_mcdc(seed, output):
+    # =============================================================================
+    # Set model
+    # =============================================================================
+    # Three slab layers with different purely-absorbing materials
+
+    # Set materials
+    m1 = mcdc.material(capture=np.array([0.90]))
+    m2 = mcdc.material(capture=np.array([0.15]))
+    m3 = mcdc.material(capture=np.array([0.60]))
+
+    # Set surfaces
+    s1 = mcdc.surface("plane-x", x=-1.0, bc="vacuum")
+    s2 = mcdc.surface("plane-x", x=2.0)
+    s3 = mcdc.surface("plane-x", x=5.0)
+    s4 = mcdc.surface("plane-x", x=6.0, bc="vacuum")
+
+    # Set cells
+    mcdc.cell([+s1, -s2], m1)
+    mcdc.cell([+s2, -s3], m2)
+    mcdc.cell([+s3, -s4], m3)
+
+    # =============================================================================
+    # Set source
+    # =============================================================================
+    # Incident beam source
+    mcdc.source(point=[0.0, 0.0, 0.0], direction=[1.0, 0.0, 0.0])
+
+    # =============================================================================
+    # Set tally, setting, and run mcdc
+    # =============================================================================
+
+    # Tally: cell-average fluxes
+    mcdc.tally(
+        scores=["exit"],
+        x=np.linspace(0.0, 6.0, 2)
+    )
+
+    # Setting
+    mcdc.setting(N_particle=1E3, N_batch=1E3, rng_seed=seed, output_name=output)
+
+    mcdc.uq(material=m1, distribution="uniform", capture=np.array([0.7]))
+    mcdc.uq(material=m2, distribution="uniform", capture=np.array([0.12]))
+    mcdc.uq(material=m3, distribution="uniform", capture=np.array([0.5]))
+
+    # Run
+    mcdc.run()
+
+
+for i in range(10):
+    seed = np.random.randint(10000000)
+    output = 'output' + str(i)
+    submit_mcdc(seed, output)