jax_sn.solvers.right_hand_side module

class RightHandSide(structure: ShapeDtypeStruct)

Bases: Module

Class that can override the default right-hand-side behaviour for specific applications.

Example is a fixed source problem with only neutrons on the first energy group.

structure: ShapeDtypeStruct

ShapeDtypeStruct for static shape/dtype checking. [n_groups, n_moments, n_elements, n_basis]

property shape: tuple

class IsotropicSourceSingleGroup(structure: ShapeDtypeStruct, energy_group: int, isotropic_value: Array)

Bases: RightHandSide

energy_group: int

isotropic_value: Array

classmethod from_shape(shape: ShapeDtypeStruct, isotropic_value: Array, energy_group: int = 0)

Isotropic source living in a single energy group; all other groups are zero.

classmethod from_operator(operator: MultiGroupOperator, isotropic_value: Array, energy_group: int = 0)

As from_shape(), taking the operand structure from operator.out_structure().

class TotalSource(structure: ShapeDtypeStruct, total_source: Array)

Bases: RightHandSide

total_source: Array

classmethod from_operator_and_isotropic_source(operator: MultiGroupOperator, total_isotropic_source: Array)

Isotropic total source (all groups), zero in the higher moments.

solve_multigroup_operator_rhs(operator: MultiGroupOperator, rhs: RightHandSide, solver: AbstractLinearSolver, chain_strategy: MultiGroupChainStrategy)

Solve \(\mathbf{A}\psi = \mathbf{C}[b]\) for the multi-group operator operator.

Parameters:

• operator (MultiGroupOperator) – The multi-group operator; its mv defines the differentiation surface.

• rhs (RightHandSide) – The physical right-hand side (or a bare array of matching structure).

• solver (AbstractLinearSolver) – The lineax solver applied to the preconditioned system.

• chain_strategy (MultiGroupChainStrategy) – Per-group chain, lifted over groups, used to precondition the source. Build via ComposedMultiGroupChainStrategy.from_composed(operator.strategy) to match the operator’s own lift.

Return type:

lx.LinearSolveOutput