Redesign and track numba_types; exclude black-formatting for scripts generated by code_factory

.github/workflows/black_lint.yml

@@ -16,7 +16,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.13"
+          python-version: "3.14"
           cache: pip
 
       - name: Install Black

.github/workflows/docker.yml

@@ -112,20 +112,20 @@ jobs:
           - Numba-MPI
         include:
           - mode: Python-Serial
-            run_cmd: "python run.py --include_native_physics"
+            run_cmd: "python run.py"
             timeout: 10
           - mode: Python-MPI
             run_cmd: |
-              python run.py --include_native_physics --mpiexec=4
-              python run.py --include_native_physics --mpiexec=16 --name=slab_reed_dd_3d
+              python run.py --mpiexec=4
+              python run.py --mpiexec=16 --name=slab_reed_dd_3d
             timeout: 20
           - mode: Numba-Serial
-            run_cmd: "python run.py --include_native_physics --mode numba"
+            run_cmd: "python run.py --mode numba"
             timeout: 120
           - mode: Numba-MPI
             run_cmd: |
-              python run.py --include_native_physics --mode numba --mpiexec=4
-              python run.py --include_native_physics --mode numba --mpiexec=16 --name=slab_reed_dd_3d
+              python run.py --mode numba --mpiexec=4
+              python run.py --mode numba --mpiexec=16 --name=slab_reed_dd_3d
             timeout: 120
 
     steps:

.github/workflows/regression_test.yml

@@ -63,24 +63,24 @@ jobs:
         if: matrix.mode == 'Python-Serial'
         working-directory: test/regression
         run: |
-          python run.py --include_native_physics
+          python run.py
 
       - name: Run regression tests (${{ matrix.mode }})
         if: matrix.mode == 'Python-MPI'
         working-directory: test/regression
         run: |
-          python run.py --include_native_physics --mpiexec=4
-          python run.py --include_native_physics --mpiexec=16 --name=slab_reed_dd_3d
+          python run.py --mpiexec=4
+          python run.py --mpiexec=16 --name=slab_reed_dd_3d
 
       - name: Run regression tests (${{ matrix.mode }})
         if: matrix.mode == 'Numba-Serial'
         working-directory: test/regression
         run: |
-          python run.py --include_native_physics --mode numba
+          python run.py --mode numba
 
       - name: Run regression tests (${{ matrix.mode }})
         if: matrix.mode == 'Numba-MPI'
         working-directory: test/regression
         run: |
-          python run.py --include_native_physics --mode numba --mpiexec=4
-          python run.py --include_native_physics --mode numba --mpiexec=16 --name=slab_reed_dd_3d
+          python run.py --mode numba --mpiexec=4
+          python run.py --mode numba --mpiexec=16 --name=slab_reed_dd_3d

mcdc/code_factory/numba_objects_generator.py

@@ -323,42 +323,78 @@ def generate_numba_objects(simulation):
             )
 
             for label in structures.keys():
+                # Skip special types
+                if label in ["gpu_meta"] + bank_names + ["simulation"]:
+                    continue
+
                 text += f"{label} = into_dtype([\n"
                 structure = structures[label]
 
-                # GPU meta override
-                if label == "gpu_meta":
-                    for item in structure:
-                        if type(item[1]) != np.dtypes.VoidDType:
-                            if isinstance(item[1], str):
-                                dtype = f"'{item[1]}'"
-                            elif item[1].__name__ == "uint64":
-                                dtype = "uintp"
-                            else:
-                                dtype = item[1].__name__
-                            text += f"    ('{item[0]}', {dtype}),\n"
-                    text += "])\n\n"
-                    continue
+                for item in structure:
+                    text += decode_structure_item(item)
+                text += "])\n\n"
+
+            # GPU meta
+            text += "gpu_meta = into_dtype([\n"
+            for item in structures["gpu_meta"]:
+                if item[0].endswith("pointer"):
+                    text += f"    ('{item[0]}', uintp),\n"
+                else:
+                    text += decode_structure_item(item)
+            text += "])\n\n"
 
+            # Particle banks
+            for label in bank_names:
+                structure = structures[label]
+
+                text += f"{label} = None\n"
+                text += f"def set_{label}(N: dict):\n"
+                text += f"    global {label}\n"
+                text += f"    {label} = into_dtype([\n"
                 for item in structure:
-                    if type(item[1]) != np.dtypes.VoidDType:
-                        if isinstance(item[1], str):
-                            dtype = f"'{item[1]}'"
-                        else:
-                            dtype = item[1].__name__
-                        if len(item) == 3:
-                            text += f"    ('{item[0]}', {dtype}, {item[2]}),\n"
-                        else:
-                            text += f"    ('{item[0]}', {dtype}),\n"
+                    if item[0] == "particle_data":
+                        text += (
+                            f"        ('{item[0]}', {item[0]}, (N['{item[0]}'],)),\n"
+                        )
                     else:
-                        if len(item) == 3:
-                            text += f"    ('{item[0]}', {plural_to_singular(item[0])}, {item[2]}),\n"
-                        else:
-                            text += f"    ('{item[0]}', {item[0]}),\n"
-                text += "])\n\n"
+                        text += decode_structure_item(item, "    ")
+                text += "    ])\n\n"
+
+            # Simulation
+            text += f"simulation = None\n"
+            text += f"def set_simulation(N: dict):\n"
+            text += f"    global simulation\n"
+            text += f"    simulation = into_dtype([\n"
+            for item in structures["simulation"]:
+                if type(item[1]) == np.dtypes.VoidDType and len(item) == 3:
+                    singular_field = plural_to_singular(item[0])
+                    text += f"        ('{item[0]}', {singular_field}, (N['{singular_field}'])),\n"
+                else:
+                    text += decode_structure_item(item, "    ")
+            text += "    ])\n\n"
 
             f.write(text)
 
+    # ==================================================================================
+    # Set numba_types.py
+    # ==================================================================================
+
+    import mcdc.numba_types as type_
+
+    # Particle banks
+    type_.set_bank_active({"particle_data": simulation.bank_active.size[0]})
+    type_.set_bank_census({"particle_data": simulation.bank_census.size[0]})
+    type_.set_bank_source({"particle_data": simulation.bank_source.size[0]})
+    type_.set_bank_future({"particle_data": simulation.bank_future.size[0]})
+
+    # Simulation
+    N = {}
+    for item in structures["simulation"]:
+        if type(item[1]) == np.dtypes.VoidDType and len(item) == 3:
+            singular_field = plural_to_singular(item[0])
+            N[singular_field] = item[2]
+    type_.set_simulation(N)
+
     # ==================================================================================
     # GPU preparation: Adapt transport functions, forward declare, and build program
     # ==================================================================================
@@ -1206,3 +1242,24 @@ def singular_to_plural(word: str) -> str:
         parts[-1] = w + "s"
 
     return "_".join(parts)
+
+
+def decode_structure_item(item, prefix=""):
+    """
+    A structure item is a list describing a member field in the structure.
+    The list contains [field name, field type, size].
+    """
+    if type(item[1]) != np.dtypes.VoidDType:
+        if isinstance(item[1], str):
+            dtype = f"'{item[1]}'"
+        else:
+            dtype = item[1].__name__
+        if len(item) == 3:
+            return f"{prefix}    ('{item[0]}', {dtype}, {item[2]}),\n"
+        else:
+            return f"{prefix}    ('{item[0]}', {dtype}),\n"
+    else:
+        if len(item) == 3:
+            return f"{prefix}    ('{item[0]}', {plural_to_singular(item[0])}, {item[2]}),\n"
+        else:
+            return f"{prefix}    ('{item[0]}', {item[0]}),\n"

mcdc/mcdc_get/neutron_inelastic_scattering_reaction.py

@@ -26,9 +26,7 @@ def spectrum_probability_grid_last(neutron_inelastic_scattering_reaction, data):
 
 
 @njit
-def spectrum_probability_grid_chunk(
-    start, length, neutron_inelastic_scattering_reaction, data
-):
+def spectrum_probability_grid_chunk(start, length, neutron_inelastic_scattering_reaction, data):
     start += neutron_inelastic_scattering_reaction["spectrum_probability_grid_offset"]
     end = start + length
     return data[start:end]
@@ -51,9 +49,7 @@ def spectrum_probability(index_1, index_2, neutron_inelastic_scattering_reaction
 
 
 @njit
-def spectrum_probability_chunk(
-    start, length, neutron_inelastic_scattering_reaction, data
-):
+def spectrum_probability_chunk(start, length, neutron_inelastic_scattering_reaction, data):
     start += neutron_inelastic_scattering_reaction["spectrum_probability_offset"]
     end = start + length
     return data[start:end]
@@ -82,9 +78,7 @@ def energy_spectrum_IDs_last(neutron_inelastic_scattering_reaction, data):
 
 
 @njit
-def energy_spectrum_IDs_chunk(
-    start, length, neutron_inelastic_scattering_reaction, data
-):
+def energy_spectrum_IDs_chunk(start, length, neutron_inelastic_scattering_reaction, data):
     start += neutron_inelastic_scattering_reaction["energy_spectrum_IDs_offset"]
     end = start + length
     return data[start:end]

mcdc/mcdc_set/electron_ionization_reaction.py

@@ -55,9 +55,7 @@ def subshell_product_IDs_last(electron_ionization_reaction, data, value):
 
 
 @njit
-def subshell_product_IDs_chunk(
-    start, length, electron_ionization_reaction, data, value
-):
+def subshell_product_IDs_chunk(start, length, electron_ionization_reaction, data, value):
     start += electron_ionization_reaction["subshell_product_IDs_offset"]
     end = start + length
     data[start:end] = value

mcdc/mcdc_set/element.py

@@ -316,9 +316,7 @@ def electron_ionization_subshell_binding_energy_last(element, data, value):
 
 
 @njit
-def electron_ionization_subshell_binding_energy_chunk(
-    start, length, element, data, value
-):
+def electron_ionization_subshell_binding_energy_chunk(start, length, element, data, value):
     start += element["electron_ionization_subshell_binding_energy_offset"]
     end = start + length
     data[start:end] = value

mcdc/mcdc_set/neutron_inelastic_scattering_reaction.py

@@ -4,9 +4,7 @@
 
 
 @njit
-def spectrum_probability_grid(
-    index, neutron_inelastic_scattering_reaction, data, value
-):
+def spectrum_probability_grid(index, neutron_inelastic_scattering_reaction, data, value):
     offset = neutron_inelastic_scattering_reaction["spectrum_probability_grid_offset"]
     data[offset + index] = value
 
@@ -28,18 +26,14 @@ def spectrum_probability_grid_last(neutron_inelastic_scattering_reaction, data,
 
 
 @njit
-def spectrum_probability_grid_chunk(
-    start, length, neutron_inelastic_scattering_reaction, data, value
-):
+def spectrum_probability_grid_chunk(start, length, neutron_inelastic_scattering_reaction, data, value):
     start += neutron_inelastic_scattering_reaction["spectrum_probability_grid_offset"]
     end = start + length
     data[start:end] = value
 
 
 @njit
-def spectrum_probability_vector(
-    index_1, neutron_inelastic_scattering_reaction, data, value
-):
+def spectrum_probability_vector(index_1, neutron_inelastic_scattering_reaction, data, value):
     offset = neutron_inelastic_scattering_reaction["spectrum_probability_offset"]
     stride = neutron_inelastic_scattering_reaction["N_spectrum"]
     start = offset + index_1 * stride
@@ -48,18 +42,14 @@ def spectrum_probability_vector(
 
 
 @njit
-def spectrum_probability(
-    index_1, index_2, neutron_inelastic_scattering_reaction, data, value
-):
+def spectrum_probability(index_1, index_2, neutron_inelastic_scattering_reaction, data, value):
     offset = neutron_inelastic_scattering_reaction["spectrum_probability_offset"]
     stride = neutron_inelastic_scattering_reaction["N_spectrum"]
     data[offset + index_1 * stride + index_2] = value
 
 
 @njit
-def spectrum_probability_chunk(
-    start, length, neutron_inelastic_scattering_reaction, data, value
-):
+def spectrum_probability_chunk(start, length, neutron_inelastic_scattering_reaction, data, value):
     start += neutron_inelastic_scattering_reaction["spectrum_probability_offset"]
     end = start + length
     data[start:end] = value
@@ -88,9 +78,7 @@ def energy_spectrum_IDs_last(neutron_inelastic_scattering_reaction, data, value)
 
 
 @njit
-def energy_spectrum_IDs_chunk(
-    start, length, neutron_inelastic_scattering_reaction, data, value
-):
+def energy_spectrum_IDs_chunk(start, length, neutron_inelastic_scattering_reaction, data, value):
     start += neutron_inelastic_scattering_reaction["energy_spectrum_IDs_offset"]
     end = start + length
     data[start:end] = value

mcdc/mcdc_set/nuclide.py

@@ -258,9 +258,7 @@ def neutron_inelastic_scattering_reaction_IDs_last(nuclide, data, value):
 
 
 @njit
-def neutron_inelastic_scattering_reaction_IDs_chunk(
-    start, length, nuclide, data, value
-):
+def neutron_inelastic_scattering_reaction_IDs_chunk(start, length, nuclide, data, value):
     start += nuclide["neutron_inelastic_scattering_reaction_IDs_offset"]
     end = start + length
     data[start:end] = value

mcdc/mcdc_set/tabulated_energy_angle_distribution.py

@@ -171,9 +171,7 @@ def cosine_offset__last(tabulated_energy_angle_distribution, data, value):
 
 
 @njit
-def cosine_offset__chunk(
-    start, length, tabulated_energy_angle_distribution, data, value
-):
+def cosine_offset__chunk(start, length, tabulated_energy_angle_distribution, data, value):
     start += tabulated_energy_angle_distribution["cosine_offset__offset"]
     end = start + length
     data[start:end] = value