jax_sn.solvers.base_solvers module¶

Direct block-substitution solver for lower-block-triangular (downscattering-only) multigroup systems.

Design¶

The system is solved directly by block forward substitution over energy groups, one single-group solve per group. Two things are supplied separately and play distinct roles:

operator (the MultiGroupOperator passed to lx.linear_solve) is the differentiation surface only. lineax differentiates linear_solve implicitly through operator.mv / operator.transpose() – the solver’s compute below is not a differentiation surface – so the operator may be an arbitrary (even fused) differentiable kernel.
the strategies carried by the solver (diagonal_strategy, chain_strategy, gather) plus the inner single_group_solver define the value of the direct solve. They are decoupled from operator.strategy on purpose: the diagonal inverse A_g^{-1} may be a non-differentiable / fused single-group solver kernel, expressed simply as the single_group_solver (a fused direct inverse is an lx.AbstractLinearSolver).

The solver holds the three single-group operators it genuinely needs, each at its abstract type:

diagonal_strategy – the full block A_g = I - M T_g^{-1} H D S_gg (diagonal solve),
chain_strategy – the chain C_g = M T_g^{-1} H D alone (off-diagonal scatter source),
gather – the per-group scatter gather U_g / V_g.

The full block and the chain overlap in the composed case (they share the same sub-strategies), but that redundancy is resolved once, in MultiGroupLowerTriangularBlockSolver.from_composed(), not by reconstructing one from the other inside the solver.

Consistency is the caller’s responsibility: the blocks the substitution builds must be the same linear operator operator represents, else the (implicit) gradient describes a different system than the computed value.

class SolverDiagnostics(print_sweep: bool, print_single_group: bool)[source]¶

Bases: Module

print_sweep: bool¶

print_single_group: bool¶

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

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))[source]¶

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])[source]¶

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])[source]¶

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])[source]¶

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()[source]¶

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])[source]¶

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]]