discrete1.spatial_sweep module
Spatial sweep routines for slab and spherical geometries.
This module implements discrete-ordinates spatial sweeps used by the solvers. It provides high-level driver functions that perform slab or sphere sweeps for steady-state and time-dependent problems, plus specialized variants for known external sources or scatter-only operators used by DJINN integrations.
The performance-critical kernels are decorated with numba.jit and
operate on NumPy arrays. Public API (typical callers):
discrete_ordinates: generic dispatcher to slab/sphere handlersslab_ordinates/sphere_ordinates: full sweep driversslab_ordinates_source/sphere_ordinates_source: known-sourceslab_ordinates_scatter/sphere_ordinates_scatter: DJINN scatter
Docstrings on the individual functions describe parameter shapes and return values expected by the rest of the package.
- discrete1.spatial_sweep.angle_coef_corrector(alpha_minus, angle_x, angle_w, nn, angles)[source]
Calculate angular differencing coefficient for spherical geometry.
Computes the angular differencing coefficient used in spherical geometry sweeps for proper treatment of angular derivatives.
- Parameters:
alpha_minus (
float) – Previous angular coefficient.angle_x (
float) – Current angular ordinate.angle_w (
float) – Quadrature weight for current ordinate.nn (
int) – Current angle index.angles (
int) – Total number of angles.
- Returns:
Updated angular coefficient. Returns 0 for last angle.
- Return type:
float
- discrete1.spatial_sweep.discrete_ordinates(flux_old, xs_total, xs_scatter, off_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry)[source]
Dispatch to the appropriate spatial sweep for the geometry.
This function chooses the slab or sphere sweeping routine based on the
geometryargument and returns the updated scalar flux array evaluated at cell centers.- Parameters:
flux_old (
numpy.ndarray) – Previous iterate of the scalar flux (shape: n_cells,).xs_total (
array_like) – Total macroscopic cross sections per material.xs_scatter (
array_like) – Scatter cross section (per material) used in source evaluation.off_scatter (
numpy.ndarray) – Off-diagonal scatter correction/source term (shape: n_cells,).external (
numpy.ndarray) – External source array (shape varies: spatial x angles x groups).boundary (
numpy.ndarray) – Boundary condition array (2, n_angles, …).medium_map (
array_like) – Integer material index per spatial cell (length n_cells).delta_x (
array_like) – Cell widths (length n_cells).angle_x (
array_like) – Angular ordinates (length n_angles).angle_w (
array_like) – Quadrature weights (length n_angles).bc_x (
sequence) – Boundary condition flags for left/right boundaries.geometry (
int) – Geometry selector: 1 -> slab, 2 -> sphere.
- Returns:
Updated scalar flux at cell centers (shape: n_cells,).
- Return type:
numpy.ndarray
- discrete1.spatial_sweep.slab_backward(flux, flux_old, xs_total, xs_scatter, off_scatter, external, edge1, medium_map, delta_x, angle_x, angle_w, edges)[source]
Perform the backward sweep for slab geometry for one ordinate.
Marches from the right boundary toward the left and updates the provided
fluxarray with center or edge contributions for the given angular ordinate. This kernel is jitted with numba and mutates its inputs in-place for performance.
- discrete1.spatial_sweep.slab_backward_scatter(flux, xs_total, scatter_source, external, edge1, medium_map, delta_x, angle_x, angle_w)[source]
Backward sweep kernel for slab geometry with known scatter source.
Performs a backward (right-to-left) sweep through slab cells using diamond difference with prescribed scatter source. Used by DJINN accelerated variants.
- Parameters:
flux (
numpy.ndarray) – Scalar flux array to update (n_cells).xs_total (
array_like) – Total cross sections by material.scatter_source (
array_like) – Known scatter source by cell.external (
array_like) – External source by cell.edge1 (
float) – Right edge flux for current cell.medium_map (
array_like) – Material indices by cell.delta_x (
array_like) – Cell widths.angle_x (
float) – Current angular ordinate.angle_w (
float) – Quadrature weight.
- Returns:
Outgoing edge flux for sweep continuation.
- Return type:
float
- discrete1.spatial_sweep.slab_forward(flux, flux_old, xs_total, xs_scatter, off_scatter, external, edge1, medium_map, delta_x, angle_x, angle_w, edges)[source]
Perform the forward sweep for slab geometry for one ordinate.
This kernel marches from the left boundary toward the right and computes either edge or center flux contributions for a single angular ordinate.
Notes
This function is JIT-compiled with numba for performance and operates in-place on the
fluxarray.
- discrete1.spatial_sweep.slab_forward_scatter(flux, xs_total, scatter_source, external, edge1, medium_map, delta_x, angle_x, angle_w)[source]
Forward sweep kernel for slab geometry with known scatter source.
Performs a forward (left-to-right) sweep through slab cells using diamond difference with prescribed scatter source. Used by DJINN accelerated variants.
- Parameters:
flux (
numpy.ndarray) – Scalar flux array to update (n_cells).xs_total (
array_like) – Total cross sections by material.scatter_source (
array_like) – Known scatter source by cell.external (
array_like) – External source by cell.edge1 (
float) – Left edge flux for current cell.medium_map (
array_like) – Material indices by cell.delta_x (
array_like) – Cell widths.angle_x (
float) – Current angular ordinate.angle_w (
float) – Quadrature weight.
- Returns:
Outgoing edge flux for sweep continuation.
- Return type:
float
- discrete1.spatial_sweep.slab_ordinates(flux_old, xs_total, xs_scatter, off_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, edges)[source]
Perform a slab-geometry discrete-ordinates spatial sweep.
The solver iterates over angular ordinates and performs forward and backward diamond-difference passes until the scalar flux converges (or a maximum iteration count is reached). This routine returns the converged cell-centered scalar flux.
- Parameters:
flux_old (
numpy.ndarray) – Initial flux guess (n_cells,).xs_total (
array_like) – Total cross sections indexed by material.xs_scatter (
array_like) – Scatter cross sections indexed by material.off_scatter (
numpy.ndarray) – Precomputed off-scatter source (n_cells,).external (
numpy.ndarray) – External source (n_x, n_angles or 1, 1).boundary (
numpy.ndarray) – Boundary conditions (2, n_angles, …).medium_map (
array_like) – Material index per cell (n_cells,).delta_x (
array_like) – Cell widths (n_cells,).angle_x (
array_like) – Angular ordinates and weights.angle_w (
array_like) – Angular ordinates and weights.bc_x (
sequence) – Boundary condition flags for left/right.edges (
int) – If 1, compute edge flux contributions; otherwise compute center flux.
- Returns:
Converged scalar flux at cell centers (n_cells,).
- Return type:
numpy.ndarray
- discrete1.spatial_sweep.slab_ordinates_scatter(xs_total, scatter_source, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x)[source]
Specialized slab sweep for known scatter source (DJINN variant).
This variant is optimized for DJINN-accelerated transport where the scatter source is provided externally rather than computed from flux moments. The function performs a single pass without iteration.
- Parameters:
xs_total (
array_like) – Total cross sections indexed by material.scatter_source (
numpy.ndarray) – Precomputed scatter source (n_cells,).external (
numpy.ndarray) – External source term (n_cells, n_angles or 1).boundary (
numpy.ndarray) – Boundary conditions (2, n_angles or 1).medium_map (
array_like) – Material indices (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.
- Returns:
Scalar flux at cell centers (n_cells,).
- Return type:
numpy.ndarray
- discrete1.spatial_sweep.slab_ordinates_source(flux, xs_total, zero, source, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, edges)[source]
Slab sweep driver for a known (prescribed) external source.
This variant is used when the external/source term is provided and the scatter contribution is not iterated. The function supports scalar or angular flux shapes depending on the
fluxinput.- Parameters:
flux (
numpy.ndarray) – Output array for scalar or angular fluxes (shape: n_cells x xdim).xs_total (
array_like) – Total cross sections indexed by material.zero (
numpy.ndarray) – Zero-array placeholder used for kernels that expect an array.source (
numpy.ndarray) – Prescribed source (spatial x angular-or-1).boundary – See
slab_ordinates()for descriptions.medium_map – See
slab_ordinates()for descriptions.delta_x – See
slab_ordinates()for descriptions.angle_x – See
slab_ordinates()for descriptions.angle_w – See
slab_ordinates()for descriptions.bc_x – See
slab_ordinates()for descriptions.edges – See
slab_ordinates()for descriptions.
- Returns:
The flux array (same object as input) updated in-place.
- Return type:
numpy.ndarray
- discrete1.spatial_sweep.sphere_backward(flux, flux_old, half_angle, xs_total, xs_scatter, off_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, weight, tau, alpha_plus, alpha_minus, edges)[source]
Backward sweep kernel for spherical geometry for a single ordinate.
Marches from the outer radius toward the center, updating
fluxin-place. The kernel uses geometry-specific surface area and volume factors to compute center contributions.
- discrete1.spatial_sweep.sphere_backward_scatter(flux, half_angle, xs_total, scatter_source, external, boundary, medium_map, delta_x, angle_x, angle_w, weight, tau, alpha_plus, alpha_minus)[source]
Backward sweep kernel for spherical geometry with known scatter source.
Performs a backward (edge-to-center) sweep through spherical shells using diamond difference with prescribed scatter source and proper spherical geometry factors. Used by DJINN accelerated variants.
- Parameters:
flux (
numpy.ndarray) – Scalar flux array to update (n_cells).half_angle (
array_like) – Half-angle flux values at cell edges.xs_total (
array_like) – Total cross sections by material.scatter_source (
array_like) – Known scatter source by cell.external (
array_like) – External source by cell.boundary (
float) – Outer boundary condition.medium_map (
array_like) – Material indices by cell.delta_x (
array_like) – Shell thicknesses.angle_x (
float) – Current angular ordinate.angle_w (
float) – Quadrature weight.weight (
float) – Angular weight factor.tau (
float) – Weighted diamond difference factor.alpha_plus (
float) – Forward angular coefficient.alpha_minus (
float) – Backward angular coefficient.
- discrete1.spatial_sweep.sphere_forward(flux, flux_old, half_angle, xs_total, xs_scatter, off_scatter, external, medium_map, delta_x, angle_x, angle_w, weight, tau, alpha_plus, alpha_minus, edges)[source]
Forward sweep kernel for spherical geometry for a single ordinate.
Updates
fluxin-place by marching from center toward the outer radius using the provided half-angle and differencing coefficients. This kernel is performance-critical and compiled with numba.
- discrete1.spatial_sweep.sphere_forward_scatter(flux, half_angle, xs_total, scatter_source, external, medium_map, delta_x, angle_x, angle_w, weight, tau, alpha_plus, alpha_minus)[source]
Forward sweep kernel for spherical geometry with known scatter source.
Performs a forward (center-to-edge) sweep through spherical shells using diamond difference with prescribed scatter source and proper spherical geometry factors. Used by DJINN accelerated variants.
- Parameters:
flux (
numpy.ndarray) – Scalar flux array to update (n_cells).half_angle (
array_like) – Half-angle flux values at cell edges.xs_total (
array_like) – Total cross sections by material.scatter_source (
array_like) – Known scatter source by cell.external (
array_like) – External source by cell.medium_map (
array_like) – Material indices by cell.delta_x (
array_like) – Shell thicknesses.angle_x (
float) – Current angular ordinate.angle_w (
float) – Quadrature weight.weight (
float) – Angular weight factor.tau (
float) – Weighted diamond difference factor.alpha_plus (
float) – Forward angular coefficient.alpha_minus (
float) – Backward angular coefficient.
- discrete1.spatial_sweep.sphere_ordinates(flux_old, xs_total, xs_scatter, off_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, edges)[source]
Perform a spherical-geometry discrete-ordinates spatial sweep.
The spherical sweep uses half-angle bookkeeping and angular differencing coefficients appropriate for radial geometry. The implementation iterates over ordinates and updates a converged scalar flux similar to the slab driver.
- Parameters:
flux_old – See
slab_ordinates()for parameter shapes and meanings.xs_total – See
slab_ordinates()for parameter shapes and meanings.xs_scatter – See
slab_ordinates()for parameter shapes and meanings.off_scatter – See
slab_ordinates()for parameter shapes and meanings.external – See
slab_ordinates()for parameter shapes and meanings.boundary – See
slab_ordinates()for parameter shapes and meanings.medium_map – See
slab_ordinates()for parameter shapes and meanings.delta_x – See
slab_ordinates()for parameter shapes and meanings.angle_x – See
slab_ordinates()for parameter shapes and meanings.angle_w – See
slab_ordinates()for parameter shapes and meanings.bc_x – See
slab_ordinates()for parameter shapes and meanings.edges – See
slab_ordinates()for parameter shapes and meanings.
- Returns:
Converged scalar flux at cell centers (n_cells,).
- Return type:
numpy.ndarray
- discrete1.spatial_sweep.sphere_ordinates_scatter(xs_total, scatter_source, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x)[source]
Specialized sphere sweep for known scatter source (DJINN variant).
This variant is optimized for DJINN-accelerated transport where the scatter source is provided externally rather than computed from flux moments. The function performs a single pass without iteration.
- Parameters:
xs_total (
array_like) – Total cross sections indexed by material.scatter_source (
numpy.ndarray) – Precomputed scatter source (n_cells,).external (
numpy.ndarray) – External source term (n_cells, n_angles or 1).boundary (
numpy.ndarray) – Boundary conditions (2, n_angles or 1).medium_map (
array_like) – Material indices (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.
- Returns:
Scalar flux at cell centers (n_cells,).
- Return type:
numpy.ndarray
- discrete1.spatial_sweep.sphere_ordinates_source(flux, xs_total, zero, source, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, edges)[source]
Sphere sweep driver for prescribed external source.
This variant implements the spherical sweep for cases with a known external source without iterating on scatter. Supports both scalar and angular flux output shapes based on input array dimensions.
- Parameters:
flux (
numpy.ndarray) – Output array for scalar/angular fluxes (n_cells x xdim).xs_total (
array_like) – Total cross sections indexed by material.zero (
numpy.ndarray) – Zero-array placeholder for kernel compatibility.source (
numpy.ndarray) – External source (n_cells, n_angles or 1).boundary (
numpy.ndarray) – Boundary conditions (2, n_angles or 1).medium_map (
array_like) – Material indices (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.edges (
numpy.ndarray) – Edge fluxes at cell boundaries.
- Returns:
Flux array (same as input) updated in-place.
- Return type:
numpy.ndarray