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 handlers

  • slab_ordinates / sphere_ordinates: full sweep drivers

  • slab_known_source_sn / sphere_known_source_sn: known-source

  • slab_scatter_source_sn / sphere_scatter_source_sn: 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, P=None, P_weights=None)[source]

Dispatch to the appropriate spatial sweep for the geometry.

This function chooses the slab or sphere sweeping routine based on the geometry argument and returns the updated scalar flux array evaluated at cell centers.

Parameters:

  • flux_old (numpy.ndarray, shape (n_cells,)) – Previous iterate of the scalar flux.

  • xs_total (numpy.ndarray, shape (n_materials,)) – Total macroscopic cross sections per material for this group.

  • xs_scatter (numpy.ndarray) – Within-group scatter cross sections mapped to cells. - Isotropic: shape (n_cells,) — scalar xs_s0 per cell. - Anisotropic: shape (n_cells, L+1) — Legendre moments per cell.

  • off_scatter (numpy.ndarray, shape (n_cells,)) – Off-diagonal scatter source.

  • external (numpy.ndarray, shape (n_cells, n_angles or 1)) – External source for this group.

  • boundary (numpy.ndarray, shape (2, n_angles or 1)) – Boundary conditions for this group.

  • 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:

Updated scalar flux at cell centers (shape: n_cells,).

Return type:

numpy.ndarray

discrete1.spatial_sweep.known_source_sn(flux, xs_total, zero, source, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry, edges)[source]

Sweep with a fully prescribed source (no scatter iteration).

discrete1.spatial_sweep.scatter_source_sn(xs_total, scatter_source, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, geometry)[source]

Sweep with a prescribed scatter source (DJINN variant).

discrete1.spatial_sweep.slab_anisotropic_sn(flux_old, xs_total, xs_self_scatter, xs_scatter_w, off_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, edges, P, P_weights)[source]

Anisotropic slab-geometry discrete-ordinates sweep.

Parameters:

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

  • xs_total (numpy.ndarray, shape (n_materials,)) – Total cross sections per material.

  • xs_self_scatter (numpy.ndarray, shape (n_cells, L+1)) – Within-group Legendre scatter moments mapped to cells.

  • xs_scatter_w (numpy.ndarray, shape (n_cells, L+1)) – Precomputed (2l+1) * xs_self_scatter.

  • off_scatter (numpy.ndarray, shape (n_cells,)) – Off-diagonal scatter source.

  • external (numpy.ndarray, shape (n_cells, n_angles or 1)) – External source.

  • boundary (numpy.ndarray, shape (2, n_angles 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 and weights.

  • angle_w (numpy.ndarray, shape (angles,)) – Angular quadrature points and weights.

  • bc_x (sequence) – Boundary condition flags.

  • edges (int) – If 1, accumulate edge fluxes; else cell-centre fluxes.

  • P (numpy.ndarray, shape (L+1, angles)) – Precomputed P[l, n] = P_l(mu_n).

  • P_weights (numpy.ndarray, shape (L+1, angles)) – Precomputed w_n * P_l(mu_n).

Returns:

Converged scalar flux at cell centres.

Return type:

numpy.ndarray, shape (n_cells,)

discrete1.spatial_sweep.slab_backward_aniso(flux, xs_total, anisotropic, off_scatter, external, edge1, medium_map, delta_x, angle_x, angle_w, edges)[source]

Backward (right to left) anisotropic slab sweep for one ordinate.

Returns:

  • edge1 (float) – Outgoing edge flux at i = 0.

  • psi (numpy.ndarray, shape (n_cells,)) – Cell-centre angular flux for moment accumulation.

discrete1.spatial_sweep.slab_backward_iso(flux, flux_old, xs_total, xs_scatter, off_scatter, external, edge1, medium_map, delta_x, angle_x, angle_w, edges)[source]

Backward (right to left) isotropic slab sweep for one ordinate.

Returns:

Outgoing edge flux at i = 0.

Return type:

float

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, shape (n_cells,).

  • xs_total (numpy.ndarray) – Total cross sections by material, shape (n_materials,).

  • scatter_source (numpy.ndarray) – Known scatter source by cell, shape (n_cells).

  • external (numpy.ndarray) – External source by cell, shape (n_cells).

  • edge1 (float) – Right edge flux for current cell.

  • medium_map (numpy.ndarray) – Material indices by cell, shape (n_cells,).

  • delta_x (numpy.ndarray) – Cell widths, shape (n_cells,).

  • 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_aniso(flux, xs_total, anisotropic, off_scatter, external, edge1, medium_map, delta_x, angle_x, angle_w, edges)[source]

Forward (left to right) anisotropic slab sweep for one ordinate.

Returns:

  • edge1 (float) – Outgoing edge flux at i = I.

  • psi (numpy.ndarray, shape (n_cells,)) – Cell-centre angular flux for moment accumulation.

discrete1.spatial_sweep.slab_forward_iso(flux, flux_old, xs_total, xs_scatter, off_scatter, external, edge1, medium_map, delta_x, angle_x, angle_w, edges)[source]

Forward (left to right) isotropic slab sweep for one ordinate.

Returns:

Outgoing edge flux at i = I.

Return type:

float

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, shape (n_cells,)) – Scalar flux array to update.

  • xs_total (numpy.ndarray, shape (n_materials,)) – Total cross sections by material.

  • scatter_source (numpy.ndarray, shape (n_cells)) – Known scatter source by cell.

  • external (numpy.ndarray, shape (n_cells)) – External source by cell.

  • edge1 (float) – Right edge flux for current cell.

  • medium_map (numpy.ndarray, shape (n_cells,)) – Material indices by cell.

  • delta_x (numpy.ndarray, shape (n_cells,)) – 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_isotropic_sn(flux_old, xs_total, xs_scatter, off_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, edges)[source]

Isotropic 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 (numpy.ndarray, shape (n_materials,)) – Total cross sections per material.

  • xs_scatter (numpy.ndarray, shape (n_cells,)) – Within-group isotropic scatter cross section mapped to cells.

  • off_scatter (numpy.ndarray, shape (n_cells,)) – Off-diagonal scatter source.

  • external (numpy.ndarray, shape (n_cells, n_angles or 1)) – External source.

  • boundary (numpy.ndarray, shape (2, n_angles 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 and weights.

  • angle_w (numpy.ndarray, shape (angles,)) – Angular quadrature points and weights.

  • bc_x (sequence) – Boundary condition flags.

  • edges (int) – If 1, accumulate edge fluxes; else cell-center fluxes.

Returns:

Converged scalar flux at cell centers.

Return type:

numpy.ndarray, shape (n_cells,)

discrete1.spatial_sweep.slab_known_source_sn(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 flux input.

Parameters:

  • flux (numpy.ndarray) – Output array for scalar or angular fluxes (shape: n_cells, xdim).

  • xs_total (numpy.ndarray, shape (n_materials,)) – Total cross sections per material.

  • zero (numpy.ndarray, shape (n_cells,)) – Zero placeholder (passed to sweep kernels for unused arguments).

  • source (numpy.ndarray, shape (n_cells, n_angles or 1)) – Prescribed source for this group.

  • boundary (numpy.ndarray, shape (2, n_angles or 1)) – Boundary conditions.

  • medium_map – See slab_isotropic_sn().

  • delta_x – See slab_isotropic_sn().

  • angle_x – See slab_isotropic_sn().

  • angle_w – See slab_isotropic_sn().

  • bc_x – See slab_isotropic_sn().

  • edges – See slab_isotropic_sn().

Returns:

The flux array (same object as input) updated in-place.

Return type:

numpy.ndarray

discrete1.spatial_sweep.slab_scatter_source_sn(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 (numpy.ndarray, shape (n_materials,)) – Total cross sections indexed by material.

  • scatter_source (numpy.ndarray, shape (n_cells,)) – Precomputed scatter source.

  • external (numpy.ndarray, shape (n_cells, n_angles or 1)) – External source term.

  • boundary (numpy.ndarray, shape (2, n_angles or 1)) – Boundary conditions.

  • medium_map (numpy.ndarray, shape (n_cells,)) – Material indices.

  • delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.

  • angle_x (numpy.ndarray, shape (n_angles,)) – Angular ordinates.

  • angle_w (numpy.ndarray, shape (n_angles,)) – Quadrature weights.

  • bc_x (sequence) – Boundary condition flags.

Returns:

Scalar flux at cell centers (n_cells,).

Return type:

numpy.ndarray

discrete1.spatial_sweep.sphere_anisotropic_sn(flux_old, xs_total, xs_self_scatter, xs_scatter_w, off_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, edges, P, P_weights)[source]

Anisotropic sphere-geometry discrete-ordinates sweep.

Parameters match slab_anisotropic_sn() with sphere-specific half-angle bookkeeping. See that function for parameter descriptions.

Returns:

Converged scalar flux at cell centres.

Return type:

numpy.ndarray, shape (n_cells,)

discrete1.spatial_sweep.sphere_backward_aniso(flux, half_angle, xs_total, anisotropic, off_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, weight, tau, alpha_plus, alpha_minus, edges)[source]

Backward (edge to center) anisotropic sphere sweep for one ordinate.

Returns:

psi – Cell-centre angular flux for moment accumulation.

Return type:

numpy.ndarray, shape (n_cells,)

discrete1.spatial_sweep.sphere_backward_iso(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 flux in-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, shape (n_cells,)) – Scalar flux array to update.

  • half_angle (numpy.ndarray, shape (n_cells,)) – Half-angle flux values at cell edges.

  • xs_total (numpy.ndarray, shape (n_materials,)) – Total cross sections by material.

  • scatter_source (numpy.ndarray, shape (n_cells)) – Known scatter source by cell.

  • external (numpy.ndarray, shape (n_cells)) – External source by cell.

  • medium_map (numpy.ndarray, shape (n_cells,)) – Material indices by cell.

  • delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.

  • 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_aniso(flux, half_angle, xs_total, anisotropic, off_scatter, external, medium_map, delta_x, angle_x, angle_w, weight, tau, alpha_plus, alpha_minus, edges)[source]

Forward (center to edge) anisotropic sphere sweep for one ordinate.

Returns:

psi – Cell-centre angular flux for moment accumulation.

Return type:

numpy.ndarray, shape (n_cells,)

discrete1.spatial_sweep.sphere_forward_iso(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 flux in-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, shape (n_cells,)) – Scalar flux array to update.

  • half_angle (numpy.ndarray, shape (n_cells,)) – Half-angle flux values at cell edges.

  • xs_total (numpy.ndarray, shape (n_materials,)) – Total cross sections by material.

  • scatter_source (numpy.ndarray, shape (n_cells)) – Known scatter source by cell.

  • external (numpy.ndarray, shape (n_cells)) – External source by cell.

  • medium_map (numpy.ndarray, shape (n_cells,)) – Material indices by cell.

  • delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.

  • 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_isotropic_sn(flux_old, xs_total, xs_scatter, off_scatter, external, boundary, medium_map, delta_x, angle_x, angle_w, bc_x, edges)[source]

Isotropic 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:
Returns:

Converged scalar flux at cell centers (n_cells,).

Return type:

numpy.ndarray

discrete1.spatial_sweep.sphere_known_source_sn(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, 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

discrete1.spatial_sweep.sphere_scatter_source_sn(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:

  • flux (numpy.ndarray, shape (n_cells,)) – Scalar flux array to update.

  • xs_total (numpy.ndarray, shape (n_materials,)) – Total cross sections by material.

  • scatter_source (numpy.ndarray, shape (n_cells)) – Known scatter source by cell.

  • external (numpy.ndarray, shape (n_cells)) – External source by cell.

  • edge1 (float) – Right edge flux for current cell.

  • medium_map (numpy.ndarray, shape (n_cells,)) – Material indices by cell.

  • delta_x (numpy.ndarray, shape (n_cells,)) – Cell widths.

  • angle_x (float) – Current angular ordinate.

  • angle_w (float) – Quadrature weight.

  • bc_x (sequence) – Boundary condition flags.

Returns:

Scalar flux at cell centers.

Return type:

numpy.ndarray, shape (n_cells)