discrete1.tools module

Utility routines and performance-critical kernels for neutron transport.

This module provides computational tools for multigroup neutron transport calculations, including source term calculations, scattering operations, fission source handling, and hybrid machine learning / physics methods. Functions are optimized using Numba JIT compilation where appropriate for performance-critical operations.

The module is organized into several functional categories:

Multigroup Transport Functions

  • scatter_prod: Compute scattering source terms

  • fission_prod: Compute fission source terms for eigenvalue problems

  • reaction_rates: Calculate per-cell reaction rates

Hybrid ML/Physics Functions

  • scatter_prod_predict: Scatter source with ML model predictions

  • fission_prod_predict: Fission source with ML model predictions

Geometry and Boundary Utilities

  • reflector_corrector: Manage reflective boundary conditions

Dynamic Mode Decomposition

  • dmd: Accelerate convergence using DMD extrapolation

Spatially Independent (0D) Variants

Functions prefixed with _0d provide zero-dimensional versions of transport operations for homogeneous problems.

Internal Kernels

Functions prefixed with _ are Numba-decorated internal kernels that implement heavy numerical work. These are used by higher-level solvers and are generally not part of the public API.

Notes

Many functions use Numba’s JIT compilation with nopython=True and cache=True for optimal performance. Type signatures are explicitly specified for compiled functions to ensure type stability.

The hybrid ML/physics functions allow material-specific predictions using trained machine learning models while falling back to traditional physics calculations when models are unavailable.

Examples

Compute scattering source for a multigroup problem:

>>> source = np.zeros((n_cells, n_groups))
>>> scatter_prod(flux, xs_scatter, source, medium_map)

Use hybrid ML approach for fission source:

>>> fission_prod_predict(flux, xs_fission, source, medium_map, keff, models)
discrete1.tools.dmd(flux_old, y_minus, y_plus, K)[source]

Estimate flux via Dynamic Mode Decomposition (DMD) extrapolation.

Given snapshot differences (y_minus, y_plus) collected over K-1 timesteps/iterations, compute a low-rank DMD model and use it to extrapolate the flux from the last known iterate.

Parameters:

  • flux_old (numpy.ndarray) – Last known flux (n_cells, n_groups).

  • y_minus (numpy.ndarray) – Snapshot difference array with shape (n_cells, n_groups, K-1) containing previous (backward) differences.

  • y_plus (numpy.ndarray) – Snapshot difference array with shape (n_cells, n_groups, K-1) containing forward differences.

  • K (int) – Number of snapshots (including the base) used for DMD.

Returns:

Extrapolated flux array with shape (n_cells, n_groups).

Return type:

numpy.ndarray

discrete1.tools.fission_prod(flux, xs_fission, source, medium_map, keff)[source]

Calculate fission source term for criticality problems.

Computes the fission source for each spatial cell and energy group by integrating the product of flux and fission cross sections, scaled by the effective multiplication factor (keff). This is used in eigenvalue (criticality) calculations.

Parameters:

  • flux (numpy.ndarray) – Scalar flux array of shape (cells_x, groups). Contains the neutron flux at each spatial cell and energy group.

  • xs_fission (numpy.ndarray) – Fission production cross section array of shape (materials, groups, groups). xs_fission[m, og, ig] is the cross section for neutrons in group ig causing fission that produces neutrons in group og for material m.

  • source (numpy.ndarray) – Fission source array of shape (cells_x, 1, groups) to be updated. Zeroed out and overwritten with computed fission source.

  • medium_map (numpy.ndarray) – Material index map of shape (cells_x,). Maps each spatial cell to its material index.

  • keff (float) – Effective multiplication factor. Used to normalize the fission source.

Returns:

Updated fission source array of shape (cells_x, 1, groups).

Return type:

numpy.ndarray

Notes

This function is JIT-compiled with Numba for performance. The fission source for cell ii and output group og is:

\[\begin{split}S_{f,ii,og} = \\frac{1}{k_{eff}} \\sum_{ig} \\phi_{ii,ig} \\sigma_{f,mat,og,ig}\end{split}\]

where mat = medium_map[ii].

discrete1.tools.fission_prod_predict(flux, xs_fission, source, medium_map, keff, models, label=None)[source]

Calculate fission source using hybrid ML/physics approach.

Computes the fission source using a combination of machine learning models and traditional physics calculations. For each material, if a trained ML model is available, it predicts the fission source; otherwise, the standard physics-based calculation is used.

Parameters:

  • flux (numpy.ndarray) – Scalar flux array of shape (cells_x, groups). Contains the neutron flux at each spatial cell and energy group.

  • xs_fission (numpy.ndarray) – Fission production cross section array of shape (materials, groups, groups). Used for scaling predictions and fallback calculations.

  • source (numpy.ndarray) – Fission source array of shape (cells_x, 1, groups) to be updated. Zeroed out and overwritten with computed fission source.

  • medium_map (numpy.ndarray) – Material index map of shape (cells_x,). Maps each spatial cell to its material index. Also used to index into the models list.

  • keff (float) – Effective multiplication factor. Used to normalize the fission source.

  • models (list) – List of trained ML models (or integers as placeholders). models[nn] is used for material nn. If models[nn] is an integer (typically 0), falls back to physics-based calculation.

  • label (numpy.ndarray, optional) – Parameter/label array for parametric ML predictions. Passed to model.predict() when making predictions. Default is None.

Returns:

The source array is modified in-place.

Return type:

None

Notes

The ML models predict the energy distribution of fission neutrons, which is then scaled to conserve the total fission rate. The scaling factor is computed as:

\[\begin{split}scale_i = \\sum_g \\phi_{i,g} \\sigma_{f,mat,g,0}\end{split}\]

The predicted source is then normalized by keff and rescaled to match the physics-based total fission rate.

discrete1.tools.reaction_rates(flux, xs_matrix, medium_map)[source]

Compute per-cell reaction rates from flux and cross sections.

Parameters:

  • flux (numpy.ndarray) – Scalar or angular-integrated flux with shape (n_cells, n_groups).

  • xs_matrix (numpy.ndarray) – Cross-section matrix indexed by material (n_materials, n_groups, …) where the last axes align with group dimensions.

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

Returns:

Reaction rate per cell and per group with shape (n_cells, n_groups).

Return type:

numpy.ndarray

discrete1.tools.reflector_corrector(reflector, angle_x, edge, nn, bc_x)[source]

Apply reflector boundary correction for a paired angle.

When reflective boundary conditions are active, the outgoing edge flux for one discrete ordinate is used to set the corresponding reflected ordinate’s incoming reflector buffer. This helper writes the edge value into the reflector array at the symmetric angle index when reflection flags are set in bc_x.

Parameters:

  • reflector (numpy.ndarray) – Array storing reflected edge values for each angle (n_angles,).

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

  • edge (float) – Computed outgoing edge flux for the current angle.

  • nn (int) – Current angle index.

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

discrete1.tools.scatter_prod(flux, xs_scatter, source, medium_map)[source]

Calculate scatter source term for multigroup transport.

Computes the scattering source for each spatial cell and energy group by integrating the product of flux and scattering cross sections. This represents neutrons scattering from all energy groups into each output energy group.

Parameters:

  • flux (numpy.ndarray) – Scalar flux array of shape (cells_x, groups). Contains the neutron flux at each spatial cell and energy group.

  • xs_scatter (numpy.ndarray) – Scattering cross section array of shape (materials, groups, groups). xs_scatter[m, og, ig] is the cross section for scattering from group ig to group og for material m.

  • source (numpy.ndarray) – Scatter source array of shape (cells_x, groups) to be updated. Overwritten with computed scatter source.

  • medium_map (numpy.ndarray) – Material index map of shape (cells_x,). Maps each spatial cell to its material index.

Returns:

Updated scatter source array of shape (cells_x, groups).

Return type:

numpy.ndarray

Notes

This function is JIT-compiled with Numba for performance. The scatter source for cell ii and output group og is:

\[\begin{split}S_{s,ii,og} = \\sum_{ig} \\phi_{ii,ig} \\sigma_{s,mat,og,ig}\end{split}\]

where mat = medium_map[ii].

discrete1.tools.scatter_prod_predict(flux, xs_scatter, source, medium_map, models, label=None)[source]

Calculate scatter source using hybrid ML/physics approach.

Computes the scatter source using a combination of machine learning models and traditional physics calculations. For each material, if a trained ML model is available, it predicts the scatter source; otherwise, the standard physics-based calculation is used.

Parameters:

  • flux (numpy.ndarray) – Scalar flux array of shape (cells_x, groups). Contains the neutron flux at each spatial cell and energy group.

  • xs_scatter (numpy.ndarray) – Scattering cross section array of shape (materials, groups, groups). xs_scatter[m, og, ig] is the cross section for scattering from group ig to group og for material m.

  • source (numpy.ndarray) – Scatter source array of shape (cells_x, groups) to be updated. Zeroed out and overwritten with computed scatter source.

  • medium_map (numpy.ndarray) – Material index map of shape (cells_x,). Maps each spatial cell to its material index. Also used to index into the models list.

  • models (list) – List of trained ML models (or integers as placeholders). models[nn] is used for material nn. If models[nn] is an integer (typically 0), falls back to physics-based calculation.

  • label (numpy.ndarray, optional) – Parameter/label array for parametric ML predictions. Passed to model.predict() when making predictions. Default is None.

Returns:

The source array is modified in-place.

Return type:

None

Notes

The ML models predict the energy distribution of scattered neutrons, which is then scaled to conserve the total scatter rate. The scaling factor is computed as:

\[\begin{split}scale_i = \\sum_g \\phi_{i,g} \\sigma_{s,mat,g,0}\end{split}\]

The predicted source is rescaled to match the physics-based total scatter rate:

\[\begin{split}S_{s,predicted} = S_{ML} \\cdot \\frac{scale}{\\sum_g S_{ML,g}}\end{split}\]