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)

Perform Dynamic Mode Decomposition (DMD) accelerated iteration.

Runs R+K source-iteration steps collecting snapshot differences then performs a DMD extrapolation to estimate the converged flux.

Parameters:

• flux_old (numpy.ndarray) – Initial scalar flux guess with shape (n_cells, n_groups).

• xs_total – See source_iteration() for descriptions.

• xs_scatter – See source_iteration() for descriptions.

• external – See source_iteration() for descriptions.

• boundary – See source_iteration() for descriptions.

• medium_map – See source_iteration() for descriptions.

• delta_x – See source_iteration() for descriptions.

• angle_x – See source_iteration() for descriptions.

• angle_w – See source_iteration() for descriptions.

• bc_x – See source_iteration() for descriptions.

• geometry – See source_iteration() for descriptions.

• R (int) – Number of iterations before collecting DMD snapshots.

• K (int) – Number of DMD snapshots to collect and use for extrapolation.

Returns:

Estimated flux after DMD extrapolation (n_cells, n_groups).

Return type:

numpy.ndarray

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

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 spatial_sweep._known_source kernel invoked per group.

Parameters:

• xs_total (array_like) – Total cross sections (per material or per cell) indexed by group.

• source (numpy.ndarray) – Prescribed source with shape (n_cells, n_angles, n_groups).

• boundary – See source_iteration() for descriptions.

• medium_map – See source_iteration() for descriptions.

• delta_x – See source_iteration() for descriptions.

• angle_x – See source_iteration() for descriptions.

• angle_w – See source_iteration() for descriptions.

• bc_x – See source_iteration() for descriptions.

• geometry – See source_iteration() for descriptions.

• edges (int) – Number of edge cells included in the angular flux array.

Returns:

Angular flux array with shape (n_cells + edges, n_angles, n_groups).

Return type:

numpy.ndarray

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

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

Wrapper that solves for scalar flux using a prescribed source and the internal _known_source kernel. Returns a (n_cells+edges, n_groups) scalar flux array suitable for post-processing.

Parameters:

• xs_total – See known_source_angular() for descriptions.

• source – See known_source_angular() for descriptions.

• boundary – See known_source_angular() for descriptions.

• medium_map – See known_source_angular() for descriptions.

• delta_x – See known_source_angular() for descriptions.

• angle_x – See known_source_angular() for descriptions.

• angle_w – See known_source_angular() for descriptions.

• bc_x – See known_source_angular() for descriptions.

• geometry – See known_source_angular() for descriptions.

• edges – See known_source_angular() for descriptions.

Returns:

Scalar flux array with shape (n_cells + edges, n_groups).

Return type:

numpy.ndarray

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 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 provided discrete-ordinates known-scatter kernel. This is intended to accelerate or replace explicit scattering computations.

Parameters:

• flux_old (numpy.ndarray) – Initial scalar flux guess with shape (n_cells, n_groups).

• xs_total – See source_iteration() for descriptions.

• xs_scatter – See source_iteration() for descriptions.

• external – See source_iteration() for descriptions.

• boundary – See source_iteration() for descriptions.

• medium_map – See source_iteration() for descriptions.

• delta_x – See source_iteration() for descriptions.

• angle_x – See source_iteration() for descriptions.

• angle_w – See source_iteration() for descriptions.

• bc_x – See source_iteration() for descriptions.

• geometry – See source_iteration() for descriptions.

• scatter_models (list, optional) – Sequence of loaded DJINN models (or placeholders) indexed by material id; models produce group-wise scatter predictions.

• scatter_labels (numpy.ndarray, optional) – Optional label array used as additional model inputs (n_cells,).

Returns:

Converged scalar flux (n_cells, n_groups).

Return type:

numpy.ndarray

discrete1.multi_group.source_iteration(flux_old, xs_total, xs_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry)

Multigroup source iteration solver.

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) – Initial scalar flux guess with shape (n_cells, n_groups).

• xs_total (array_like) – Total cross sections (n_materials or n_cells, n_groups).

• xs_scatter (array_like) – Scatter cross sections (indexed by material, GxG).

• external (numpy.ndarray) – External source array (n_cells, n_angles, n_groups) or with singleton dimensions where appropriate.

• boundary (numpy.ndarray) – Boundary conditions array (2, n_angles, n_groups) or with singleton dimensions where appropriate.

• medium_map (array_like) – Material index per spatial cell (n_cells,).

• delta_x (array_like) – Cell widths (n_cells,).

• angle_x (array_like) – Angular ordinates (n_angles,).

• angle_w (array_like) – Quadrature weights (n_angles,).

• bc_x (sequence) – Boundary condition flags for left/right boundaries.

• geometry (int) – Geometry selector (1=slab, 2=sphere).

Returns:

Converged scalar flux array with shape (n_cells, n_groups).

Return type:

numpy.ndarray

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 iteration that collects intermediate flux snapshots.

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

Parameters:

• flux_old (numpy.ndarray) – Initial scalar flux guess with shape (n_cells, n_groups).

• xs_total – See source_iteration() for descriptions.

• xs_scatter – See source_iteration() for descriptions.

• external – See source_iteration() for descriptions.

• boundary – See source_iteration() for descriptions.

• medium_map – See source_iteration() for descriptions.

• delta_x – See source_iteration() for descriptions.

• angle_x – See source_iteration() for descriptions.

• angle_w – See source_iteration() for descriptions.

• bc_x – See source_iteration() for descriptions.

• geometry – See source_iteration() for descriptions.

• iteration (int) – Iteration index used to name the saved snapshot file.

• filepath (str) – Directory/filepath prefix where snapshots are saved.

Returns:

Final converged flux (n_cells, n_groups).

Return type:

numpy.ndarray

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)

Source iteration for problems with variable coarse/fine groups.

This variant performs multigroup source iteration when a coarse energy grid is constructed from a fine grid. It computes coarse group cross-sections and off-scatter terms appropriately and then calls the discrete-ordinates solver per coarse group.

Parameters:

• flux_old (numpy.ndarray) – Initial scalar flux guess on the fine grid (n_cells, n_groups_fine).

• xs_total (array_like) – Total cross sections on the fine grid.

• star_coef_c (array_like) – Additional coarse-group correction terms added to xs_total_c.

• xs_scatter (array_like) – Scatter cross sections on the fine grid.

• external – See source_iteration() for descriptions.

• boundary – See source_iteration() for descriptions.

• medium_map – See source_iteration() for descriptions.

• delta_x – See source_iteration() for descriptions.

• angle_x – See source_iteration() for descriptions.

• angle_w – See source_iteration() for descriptions.

• bc_x – See source_iteration() for descriptions.

• delta_coarse (array_like) – Coarse-group widths (n_groups_coarse,).

• delta_fine (array_like) – Fine-group widths (n_groups_fine,).

• edges_gidx_c (array_like) – Indices mapping coarse groups to fine-group indices.

• geometry (int) – Geometry selector (1=slab, 2=sphere).

Returns:

Converged scalar flux on the fine grid (n_cells, n_groups_fine).

Return type:

numpy.ndarray