discrete1.multi_group module

Multigroup solvers and DJINN integration utilities.

This module provides multigroup-level driver routines that orchestrate outer iterations over energy groups and call the spatial sweep implementations (discrete ordinates) for each group. It includes standard source-iteration, variable/coarsened-group iteration, Dynamic Mode Decomposition (DMD) accelerated routines, and helpers for generating training data or using DJINN-predicted scatter/fission sources.

Public routines include: - source_iteration: standard multigroup source iteration - variable_source_iteration: coarse/fine-group-aware iteration - dynamic_mode_decomp: DMD-accelerated solver - known_source_angular / known_source_scalar: wrappers for known-source problems - ml_source_iteration: use DJINN models to predict scatter sources

The multigroup drivers expect arrays shaped as used by the rest of the package (cell-major ordering, angular and group axes as described in function docstrings).

discrete1.multi_group.dynamic_mode_decomp(flux_old, xs_total, xs_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry, R, K, P=None, P_weights=None)[source]

DMD-accelerated multigroup source iteration.

Runs R warm-up sweeps followed by K snapshot-collection sweeps, then performs a Dynamic Mode Decomposition (DMD) extrapolation to estimate the converged flux from the collected difference snapshots. Dispatches to isotropic or anisotropic inner sweeps based on xs_scatter.ndim.

Parameters:

flux_old (numpy.ndarray, shape (n_cells, n_groups)) – Initial scalar flux guess.
xs_total (numpy.ndarray, shape (n_materials, n_groups)) – Total cross sections.
xs_scatter (numpy.ndarray) –
Scatter cross sections.
- Isotropic: shape (n_materials, n_groups, n_groups).
- Anisotropic: shape (n_materials, n_groups, n_groups, L+1).
external (numpy.ndarray, shape (n_cells, n_angles or 1, n_groups or 1)) – External source array.
boundary (numpy.ndarray, shape (2, n_angles or 1, n_groups or 1)) – Boundary conditions.
medium_map (numpy.ndarray, shape (n_cells,)) – Material index per cell.
delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.
angle_x (numpy.ndarray, shape (n_angles,)) – Angular quadrature points.
angle_w (numpy.ndarray, shape (n_angles,)) – Angular quadrature weights.
bc_x (sequence of int) – Boundary condition flags [left, right]. 0 = vacuum, 1 = reflect.
geometry (int) – Geometry type. 1 for slab, 2 for sphere.
R (int) – Number of warm-up iterations performed before snapshot collection begins.
K (int) – Number of DMD snapshot iterations used for extrapolation.
P (numpy.ndarray, shape (L+1, n_angles), optional) – Precomputed Legendre polynomials. Computed automatically for anisotropic problems when not provided.
P_weights (numpy.ndarray, shape (L+1, n_angles), optional) – Precomputed w_n * P_l(mu_n). Computed automatically when not provided.

Returns:

Estimated converged flux after DMD extrapolation.

Return type:

numpy.ndarray, shape (n_cells, n_groups)

discrete1.multi_group.known_source_angular(xs_total, source, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry, edges)[source]

Compute angular flux for a prescribed multigroup source.

Solves the transport equation for a supplied angular source term and returns the angular flux field (including any edge terms indicated by edges). This is a convenience wrapper around the known_source_sn() kernel invoked per group.

Parameters:

xs_total (numpy.ndarray, shape (n_materials, n_groups)) – Total cross sections.
source (numpy.ndarray, shape (n_cells, n_angles, n_groups)) – Prescribed angular source.
boundary (numpy.ndarray, shape (2, n_angles or 1, n_groups or 1)) – Boundary conditions.
medium_map (numpy.ndarray, shape (n_cells,)) – Material index per cell.
delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.
angle_x (numpy.ndarray, shape (n_angles,)) – Angular quadrature points.
angle_w (numpy.ndarray, shape (n_angles,)) – Angular quadrature weights.
bc_x (sequence of int) – Boundary condition flags [left, right]. 0 = vacuum, 1 = reflect.
geometry (int) – Geometry type. 1 for slab, 2 for sphere.
edges (int) – If 1, include an extra cell for the right-edge flux; 0 otherwise.

Returns:

Angular flux.

Return type:

numpy.ndarray, shape (n_cells + edges, n_angles, n_groups)

discrete1.multi_group.known_source_scalar(xs_total, source, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry, edges)[source]

Compute scalar (group-integrated) flux for a prescribed source.

Solves for the scalar (angle-integrated) flux using a prescribed source and the known_source_sn() kernel. Returns a scalar flux array suitable for post-processing.

Parameters:

xs_total (numpy.ndarray, shape (n_materials, n_groups)) – Total cross sections.
source (numpy.ndarray, shape (n_cells, n_angles, n_groups)) – Prescribed angular source.
boundary (numpy.ndarray, shape (2, n_angles or 1, n_groups or 1)) – Boundary conditions.
medium_map (numpy.ndarray, shape (n_cells,)) – Material index per cell.
delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.
angle_x (numpy.ndarray, shape (n_angles,)) – Angular quadrature points.
angle_w (numpy.ndarray, shape (n_angles,)) – Angular quadrature weights.
bc_x (sequence of int) – Boundary condition flags [left, right]. 0 = vacuum, 1 = reflect.
geometry (int) – Geometry type. 1 for slab, 2 for sphere.
edges (int) – If 1, include an extra cell for the right-edge flux; 0 otherwise.

Returns:

Scalar flux.

Return type:

numpy.ndarray, shape (n_cells + edges, n_groups)

discrete1.multi_group.ml_source_iteration(flux_old, xs_total, xs_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry, scatter_models=[], scatter_labels=None)[source]

Source iteration using DJINN-predicted scatter source.

Uses pre-trained DJINN models to predict the scatter source each outer iteration and then solves each group with the scatter_source_sn() kernel. Falls back to the standard physics-based scatter calculation for materials without a trained model.

Parameters:

flux_old (numpy.ndarray, shape (n_cells, n_groups)) – Initial scalar flux guess.
xs_total (numpy.ndarray, shape (n_materials, n_groups)) – Total cross sections.
xs_scatter (numpy.ndarray, shape (n_materials, n_groups, n_groups)) – Isotropic scatter cross sections used as fallback for materials without a trained model.
external (numpy.ndarray, shape (n_cells, n_angles or 1, n_groups or 1)) – External source array.
boundary (numpy.ndarray, shape (2, n_angles or 1, n_groups or 1)) – Boundary conditions.
medium_map (numpy.ndarray, shape (n_cells,)) – Material index per cell.
delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.
angle_x (numpy.ndarray, shape (n_angles,)) – Angular quadrature points.
angle_w (numpy.ndarray, shape (n_angles,)) – Angular quadrature weights.
bc_x (sequence of int) – Boundary condition flags [left, right]. 0 = vacuum, 1 = reflect.
geometry (int) – Geometry type. 1 for slab, 2 for sphere.
scatter_models (list) – Trained DJINN models indexed by material id. An integer placeholder (typically 0) triggers the physics fallback.
scatter_labels (numpy.ndarray, optional) – Label array passed to model.predict() for parametric models. Default is None.

Returns:

Converged scalar flux.

Return type:

numpy.ndarray, shape (n_cells, n_groups)

discrete1.multi_group.source_iteration(flux_old, xs_total, xs_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry, P=None, P_weights=None)[source]

Multigroup source iteration solver - dispatches to iso or aniso variant.

Performs outer iterations over energy groups using inner discrete ordinates sweeps for each group. Off-diagonal scatter (up- and down-scatter) contributions are handled via an “off_scatter” accumulator that uses the previous and current flux iterates.

Parameters:

flux_old (numpy.ndarray, shape (n_cells, n_groups)) – Initial scalar flux guess.
xs_total (numpy.ndarray, shape (n_materials, n_groups)) – Total cross sections.
xs_scatter (numpy.ndarray) – Scatter cross sections. - Isotropic: shape (n_materials, n_groups, n_groups). - Anisotropic: shape (n_materials, n_groups, n_groups, L+1).
external (numpy.ndarray, shape (n_cells, n_angles or 1, n_groups or 1)) – External source array.
boundary (numpy.ndarray, shape (2, n_angles or 1, n_groups or 1)) – Boundary conditions.
medium_map (numpy.ndarray, shape (n_cells,)) – Material index per cell.
delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.
angle_x (numpy.ndarray, shape (angles,)) – Angular quadrature points.
angle_w (numpy.ndarray, shape (angles,)) – Angular quadrature weights.
bc_x (sequence) – Boundary condition flags for left/right boundaries.
geometry (int) – 1 = slab, 2 = sphere.
P (numpy.ndarray, shape (L+1, angles), optional) – Precomputed Legendre polynomials. Required for anisotropic.
P_weights (numpy.ndarray, shape (L+1, angles), optional) – Precomputed w_n * P_l(mu_n). Required for anisotropic.

Returns:

Converged scalar flux.

Return type:

numpy.ndarray, shape (n_cells, n_groups)

discrete1.multi_group.source_iteration_aniso(flux_old, xs_total, xs_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry, P, P_weights)[source]

Multigroup source iteration for anisotropic scattering.

Identical workflow to source_iteration_iso() but each group sweep uses the full Legendre expansion of the scatter kernel. The within-group anisotropic correction is applied inside the spatial sweep; only the L=0 moment drives the off-diagonal scatter term.

Parameters:

flux_old (numpy.ndarray, shape (n_cells, n_groups)) – Initial scalar flux guess.
xs_total (numpy.ndarray, shape (n_materials, n_groups)) – Total cross sections.
xs_scatter (numpy.ndarray, shape (n_materials, n_groups, n_groups, L+1)) – Legendre moments of scatter cross sections. The last axis indexes the moment order (0 to L).
external (numpy.ndarray, shape (n_cells, n_angles or 1, n_groups or 1)) – External source array.
boundary (numpy.ndarray, shape (2, n_angles or 1, n_groups or 1)) – Boundary conditions.
medium_map (numpy.ndarray, shape (n_cells,)) – Material index per cell.
delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.
angle_x (numpy.ndarray, shape (n_angles,)) – Angular quadrature points.
angle_w (numpy.ndarray, shape (n_angles,)) – Angular quadrature weights.
bc_x (sequence of int) – Boundary condition flags [left, right]. 0 = vacuum, 1 = reflect.
geometry (int) – Geometry type. 1 for slab, 2 for sphere.
P (numpy.ndarray, shape (L+1, n_angles)) – Precomputed Legendre polynomials P[l, n] = P_l(mu_n).
P_weights (numpy.ndarray, shape (L+1, n_angles)) – Precomputed w_n * P_l(mu_n) used for flux moment accumulation.

Returns:

Converged scalar flux.

Return type:

numpy.ndarray, shape (n_cells, n_groups)

discrete1.multi_group.source_iteration_collect(flux_old, xs_total, xs_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry, iteration, filepath)[source]

Source iteration that collects intermediate flux snapshots.

Runs standard isotropic multigroup source iteration while recording per-outer-iteration flux snapshots into an array which is saved to disk. Useful for generating training data for DJINN scatter models.

Parameters:

flux_old (numpy.ndarray, shape (n_cells, n_groups)) – Initial scalar flux guess.
xs_total (numpy.ndarray, shape (n_materials, n_groups)) – Total cross sections.
xs_scatter (numpy.ndarray, shape (n_materials, n_groups, n_groups)) – Isotropic scatter cross sections.
external (numpy.ndarray, shape (n_cells, n_angles or 1, n_groups or 1)) – External source array.
boundary (numpy.ndarray, shape (2, n_angles or 1, n_groups or 1)) – Boundary conditions.
medium_map (numpy.ndarray, shape (n_cells,)) – Material index per cell.
delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.
angle_x (numpy.ndarray, shape (n_angles,)) – Angular quadrature points.
angle_w (numpy.ndarray, shape (n_angles,)) – Angular quadrature weights.
bc_x (sequence of int) – Boundary condition flags [left, right]. 0 = vacuum, 1 = reflect.
geometry (int) – Geometry type. 1 for slab, 2 for sphere.
iteration (int) – Iteration index appended to the saved snapshot filename.
filepath (str) – Directory or filepath prefix for the saved snapshot file.

Returns:

Final converged flux.

Return type:

numpy.ndarray, shape (n_cells, n_groups)

discrete1.multi_group.source_iteration_iso(flux_old, xs_total, xs_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry)[source]

Multigroup source iteration for isotropic scattering.

Performs outer iterations over energy groups using inner discrete ordinates sweeps for each group. Convergence is measured with the relative L2 norm of the flux change across all cells and groups.

Parameters:

flux_old (numpy.ndarray, shape (n_cells, n_groups)) – Initial scalar flux guess.
xs_total (numpy.ndarray, shape (n_materials, n_groups)) – Total cross sections.
xs_scatter (numpy.ndarray, shape (n_materials, n_groups, n_groups)) – Isotropic scatter cross sections. xs_scatter[m, og, ig] is the cross section for scattering from group ig to group og for material m.
external (numpy.ndarray, shape (n_cells, n_angles or 1, n_groups or 1)) – External source array.
boundary (numpy.ndarray, shape (2, n_angles or 1, n_groups or 1)) – Boundary conditions.
medium_map (numpy.ndarray, shape (n_cells,)) – Material index per cell.
delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.
angle_x (numpy.ndarray, shape (n_angles,)) – Angular quadrature points.
angle_w (numpy.ndarray, shape (n_angles,)) – Angular quadrature weights.
bc_x (sequence of int) – Boundary condition flags [left, right]. 0 = vacuum, 1 = reflect.
geometry (int) – Geometry type. 1 for slab, 2 for sphere.

Returns:

Converged scalar flux.

Return type:

numpy.ndarray, shape (n_cells, n_groups)

discrete1.multi_group.variable_source_iteration(flux_old, xs_total, star_coef_c, xs_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, delta_coarse, delta_fine, edges_gidx_c, geometry, P=None, P_weights=None)[source]

Source iteration for problems with variable coarse/fine groups.

Dispatches to the isotropic or anisotropic variant based on xs_scatter.ndim. Coarse-group cross sections and off-scatter terms are collapsed from the fine-group data before each group sweep, allowing efficient treatment of problems where a coarse energy grid is derived from a fine one.

Parameters:

flux_old (numpy.ndarray, shape (n_cells, n_groups_coarse)) – Initial scalar flux guess on the coarse group structure.
xs_total (numpy.ndarray, shape (n_materials, n_groups_fine)) – Fine-group total cross sections.
star_coef_c (numpy.ndarray, shape (n_groups_coarse,)) – Coarse-group streaming/removal coefficient added to the collapsed total cross section.
xs_scatter (numpy.ndarray) –
Fine-group scatter cross sections.
- Isotropic: shape (n_materials, n_groups_fine, n_groups_fine).
- Anisotropic: shape (n_materials, n_groups_fine, n_groups_fine, L+1).
external (numpy.ndarray, shape (n_cells, n_angles or 1, n_groups_coarse or 1)) – External source array.
boundary (numpy.ndarray, shape (2, n_angles or 1, n_groups_coarse or 1)) – Boundary conditions.
medium_map (numpy.ndarray, shape (n_cells,)) – Material index per cell.
delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.
angle_x (numpy.ndarray, shape (n_angles,)) – Angular quadrature points.
angle_w (numpy.ndarray, shape (n_angles,)) – Angular quadrature weights.
bc_x (sequence of int) – Boundary condition flags [left, right]. 0 = vacuum, 1 = reflect.
delta_coarse (numpy.ndarray, shape (n_groups_coarse,)) – Coarse-group energy widths.
delta_fine (numpy.ndarray, shape (n_groups_fine,)) – Fine-group energy widths.
edges_gidx_c (numpy.ndarray, shape (n_groups_coarse + 1,)) – Fine-group index boundaries for each coarse group.
geometry (int) – Geometry type. 1 for slab, 2 for sphere.
P (numpy.ndarray, shape (L+1, n_angles), optional) – Precomputed Legendre polynomials. Computed automatically for anisotropic problems when not provided.
P_weights (numpy.ndarray, shape (L+1, n_angles), optional) – Precomputed w_n * P_l(mu_n). Computed automatically when not provided.

Returns:

Converged scalar flux on the coarse group structure.

Return type:

numpy.ndarray, shape (n_cells, n_groups_coarse)