jax_sn.solvers

class MultiGroupLowerTriangularBlockSolver(single_group_solver: AbstractLinearSolver, diagonal_strategy: SingleGroupStrategy, chain_strategy: SingleGroupChainStrategy, gather: SingleGroupGatherScatterStrategy, diagnostics: SolverDiagnostics = SolverDiagnostics(print_sweep=False, print_single_group=False))

Bases: AbstractLinearSolver[MultiGroupOperator]

Block forward-substitution solver for a lower-block-triangular (downscattering-only) multigroup operator. Both directions are implemented, as implicit differentiation requires the transpose solve: the primal is forward substitution over ascending groups; its transpose is an upper-triangular back substitution over descending groups, carrying the per-group intermediate zeta_g = C_g^T x_g.

single_group_solver: AbstractLinearSolver

diagonal_strategy: SingleGroupStrategy

chain_strategy: SingleGroupChainStrategy

gather: SingleGroupGatherScatterStrategy

diagnostics: SolverDiagnostics = SolverDiagnostics(print_sweep=False, print_single_group=False)

classmethod from_composed(composed_strategy: MultiGroupChainGatherStrategy, single_group_solver: AbstractLinearSolver, scatter_strategy: ScatterStrategy = IndexedScatter(), diagnostics: SolverDiagnostics = SolverDiagnostics(print_sweep=False, print_single_group=False))

Build the solver from a MultiGroupChainGatherStrategy with a ComposedSingleGroupChainStrategy. The diagonal block is assembled from the chain’s sub-strategies plus scatter_strategy for the within-group scatter \(\mathbf{S}_{gg}\) – the chain strategy itself has no scatter operator (\(C_g = \mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\mathbf{D}\) doesn’t need one), so it isn’t available to read off the chain.

Return type:

Self

init(operator: MultiGroupOperator, options: Dict[str, Any])

Do any initial computation on just the operator.

For example, an LU solver would compute the LU decomposition of the operator (and this does not require knowing the vector yet).

It is common to need to solve the linear system Ax=b multiple times in succession, with the same operator A and multiple vectors b. This method improves efficiency by making it possible to re-use the computation performed on just the operator.

!!! Example

`python operator = lx.MatrixLinearOperator(...) vector1 = ... vector2 = ... solver = lx.LU() state = solver.init(operator, options={}) solution1 = lx.linear_solve(operator, vector1, solver, state=state) solution2 = lx.linear_solve(operator, vector2, solver, state=state) `

Arguments:

• operator: a linear operator.

• options: a dictionary of any extra options that the solver may wish to

  accept.

Returns:

A PyTree of arbitrary Python objects.

Return type:

MultiGroupOperator

transpose(state: MultiGroupOperator, options: Dict[str, Any])

Transposes the result of [lineax.AbstractLinearSolver.init][].

That is, it should be the case that `python state_transpose, _ = solver.transpose(solver.init(operator, options), options) state_transpose2 = solver.init(operator.T, options) ` must be identical to each other.

It is relatively common (in particular when differentiating through a linear solve) to need to solve both Ax = b and A^T x = b. This method makes it possible to avoid computing both solver.init(operator) and solver.init(operator.T) if one can be cheaply computed from the other.

Arguments:

• state: as returned from solver.init.

• options: any extra options that were passed to solve.init.

Returns:

A 2-tuple of:

• The state of the transposed operator.

• The options for the transposed operator.

Return type:

Tuple[MultiGroupOperator, Dict[str, Any]]

conj(state: MultiGroupOperator, options: Dict[str, Any])

Conjugate the result of [lineax.AbstractLinearSolver.init][].

That is, it should be the case that `python state_conj, _ = solver.conj(solver.init(operator, options), options) state_conj2 = solver.init(conj(operator), options) ` must be identical to each other.

Arguments:

• state: as returned from solver.init.

• options: any extra options that were passed to solve.init.

Returns:

A 2-tuple of:

• The state of the conjugated operator.

• The options for the conjugated operator.

Return type:

Tuple[MultiGroupOperator, Dict[str, Any]]

assume_full_rank()

Does this solver assume that all operators are full rank?

When False, a more expensive backward pass is needed to account for the extra generality. In a custom linear solver, it is always safe to return False.

Arguments:

Nothing.

Returns:

Either True or False.

Return type:

bool

compute(state: MultiGroupOperator, rhs: Array, options: Dict[str, Any])

Solves a linear system.

Arguments:

• state: as returned from [lineax.AbstractLinearSolver.init][].

• vector: the vector to solve against.

• options: a dictionary of any extra options that the solver may wish to

  accept. For example, [lineax.CG][] accepts a preconditioner option.

Returns:

A 3-tuple of:

• The solution to the linear system.

• An integer indicating the success or failure of the solve. This is an integer

  which may be converted to a human-readable error message via lx.RESULTS[…].

• A dictionary of an extra statistics about the solve, e.g. the number of steps

  taken.

Return type:

Tuple[Array, RESULTS, Dict[str, Array]]

class SolverDiagnostics(print_sweep: bool, print_single_group: bool)

Bases: Module

print_sweep: bool

print_single_group: bool

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

Submodules

• jax_sn.solvers.base_solvers module
  • Design
  • SolverDiagnostics
  • MultiGroupLowerTriangularBlockSolver
• jax_sn.solvers.right_hand_side module
  • RightHandSide
  • IsotropicSourceSingleGroup
  • TotalSource
  • solve_multigroup_operator_rhs()