jax_sn.operators.multi_group_operator

class MultiGroupOperatorBase(data: OperatorData, strategy: MultiGroupStrategyBase, transpose_bool: bool = False)

Bases: TransposeBoolDataStrategyLinearOperator

Base operator acting on the full energy-group axis: shared operand structure and data accessors.

data: OperatorData

strategy: MultiGroupStrategyBase

transpose_bool: bool = False

property n_energy_groups: int

Number of energy groups, derived from operator_data’s cross sections.

property solution_domain

Reconstruct the SolutionDomain this operator was built from (e.g. for interpolation).

class MultiGroupOperator(data: OperatorData, strategy: MultiGroupStrategy, transpose_bool: bool = False)

Bases: MultiGroupOperatorBase

Multi energy group linear operator implementation base class.

Mathematical Description

This operator maps as:

\[\mathbf{A} : \mathcal{G} \otimes \mathcal{M} \otimes \mathcal{F} \rightarrow \mathcal{G}\otimes \mathcal{M} \otimes \mathcal{F}\]

The structure of the operator is

\[\begin{split}\mathbf{A} = \begin{bmatrix} \mathbf{A}_0 & B_{01} & B_{02} & \dots & B_{0G} \\ B_{10} & \mathbf{A}_1 & B_{12} & \dots & B_{1G} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ B_{G0} & B_{G1} & B_{G2} & \dots & \mathbf{A}_G \end{bmatrix}\end{split}\]

where \(\mathbf{A}_g\) is the single group operator for energy group g, and \(B_{ij}\) is the scattering operator from group j to group i.

These operators are defined as:

\[\mathbf{A}_g = \mathbf{I} - \mathbf{M}\mathbf{T}_g^{-1}\mathbf{H} \mathbf{D} \mathbf{S}_{gg}\]

\[\mathbf{B}_{ij} = - \mathbf{M}\mathbf{T}_i^{-1}\mathbf{H} \mathbf{D}\mathbf{S}_{ij}\]

where \(\mathbf{M}\) is the batched moment operator (see MomentOperator), \(\mathbf{T}_g\) is the transport operator for group g (see TransportOperatorMultiAngle), \(\mathbf{H}\) is the mass operator (see MassOperator), \(\mathbf{D}\) is the moment-to-discrete operator (see DiscreteOperator), and \(\mathbf{S}_{ij}\) is the scattering operator from group j to group i (see ScatterOperator). These operators are all defined for a single energy group, as in their original definitions.

These are combined to form the multi-group operator as shown above. However, the naive implementation would require \(G^2\) transport sweeps to apply the operator (one for each block in the matrix). This is clearly very inefficient. We can exploit the forward and transpose structure instead to limit this to only \(G\) sweeps per application of the operator.

Forward Structure

Note that for a specific row (energy group), the operator acts as a single group operator plus a scattering source from all other energy groups. Writing out explicitly the action on a vector \(\psi\) for energy group g, we have:

\[\mathbf{A}[\psi]_g = \left[\mathbf{I} - \mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\mathbf{D}\mathbf{S}_{gg} \right] \psi_{g} - \sum_{g' \neq g} \mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\mathbf{D}\mathbf{S}_{gg'} \psi_{g'}\]

which is equivalent to:

\[\mathbf{A}[\psi]_g = \left[\mathbf{I} - \mathbf{M} \mathbf{T}_g^{-1} \mathbf{H}\mathbf{D}\left(\mathbf{S}_{gg}[\psi_g] - \sum_{g'\neq g} \mathbf{S}_{gg'}[\psi_{g'}]\right)\right] \equiv \left[\mathbf{I} - \mathbf{M} \mathbf{T}_g^{-1} \mathbf{H} \mathbf{D}\mathbf{U}_{g}[\psi]\right]\]

This latter operator:

\[\mathbf{U}_g[\psi] = \sum_{\tilde{g}} \mathbf{S}_{g\tilde{g}}[\psi_{\tilde{g}}]\]

is implemented efficiently in the GatherScatterOperator class.

This procedure minimizes the number of sweeps to just one per group in the forward direction.

Transposed Structure

The transposition structure is not as straightforward as the forward structure. We have

\[\begin{split}\mathbf{A}^T = \begin{bmatrix} \mathbf{A}_0^T & B_{10}^T & B_{20}^T & \dots & B_{G0}^T \\ B_{01}^T & \mathbf{A}_1^T & B_{21}^T & \dots & B_{G1}^T \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ B_{0G}^T & B_{1G}^T & B_{2G}^T & \dots & \mathbf{A}_G^T \end{bmatrix}\end{split}\]

For a single row, we have

\[\mathbf{A}^T[\psi]_g = \left[\mathbf{I} - \mathbf{S}_{gg}^T \mathbf{D}^T \mathbf{H}^T (\mathbf{T}_g^{-1})^T \mathbf{M}^T \right] \psi_{g} - \sum_{g' \neq g} \mathbf{S}_{g'g}^T \mathbf{D}^T\mathbf{H}^T (\mathbf{T}_{g'}^{-1})^T \mathbf{M}^T \psi_{g'}\]

Now, we note that we cannot take together the \(\mathbf{T}_{g'}^{-1}\) terms since this is a different energy group for each term in the sum.

So, instead we precompute, for all energy groups \(g\):

\[\zeta_g = \mathbf{D}^T \mathbf{H}^T (\mathbf{T}_{g}^{-1})^T \mathbf{M}^T \psi_g\]

and then combine these with the transposed indexed scattering operators to form the final result (note that the scatter operator transposition is itself, since it is element-wise multiplication):

\[\mathbf{A}^T[\psi]_g = \psi_g - \sum_{g'} \mathbf{S}_{g'g}^T[\zeta_{g'}] \equiv \psi_g - \sum_{g'}\mathbf{S}_{g'g}[\zeta_{g'}] \equiv \psi_g - \mathbf{V}_g[g]\]

with

\[\mathbf{V}_g[\zeta] = \sum_{\tilde{g}} \mathbf{S}_{\tilde{g} g}[\zeta_{\tilde{g}}]\]

again implemented efficiently in the GatherScatterOperator class.

This again minimizes the number of sweeps to just one per group in the transposed direction, although some extra storage is required for the intermediate \(\zeta_g\) vector

Solution

Since this is a iteratively defined linear operator, the standard way to solve systems involving this operator is to use an iterative solver (e.g., GMRES, CG, etc.) with the mv() method defined here. This can be accomplished using e.g. lineax solvers. This allows for the propagation of the derivatives through the solver as well. However, in some cases, e.g. downscattering only problems (such as often found in fusion applications), the multi-group operator is lower triangular. In these cases, it is much more efficient to solve the system directly using a block forward substitution approach, solving each energy group sequentially using the single-group operator for that group and the already computed higher energy group solutions to form the scattering source. This is implemented in the function solve_lower_triangular().

strategy: MultiGroupStrategy

class MultiGroupChainOperator(data: OperatorData, strategy: MultiGroupChainStrategy, transpose_bool: bool = False)

Bases: MultiGroupOperatorBase

The chain \(\mathbf{C} = \mathbf{M}\mathbf{T}^{-1}\mathbf{H}\mathbf{D}\) applied independently to every group (no gather-scatter, so groups don’t couple). Same operand structure as MultiGroupOperator; used to precondition a per-group right-hand side.

Wrapping the strategy as an operator is what buys transpose() and lineax composition for free. It is a view: it shares data with the multigroup operator it accompanies, so build it at the call site – a stored copy would not follow an eqx.tree_at on the original’s data.

strategy: MultiGroupChainStrategy

class MultiGroupStrategyBase

Bases: TransposableStrategy, RequiresNonDifferentiableData

Apply a multigroup operator, base class for both full operator and chain.

States its requirement on the whole non-differentiable tier (see RequiresNonDifferentiableData), not on the sweep data alone: it is handed an OperatorData, so what it needs of that tier is its own business and may grow past the sweep ordering.

abstractmethod forward(operator_data: OperatorData, x)

abstractmethod transpose(operator_data: OperatorData, x)

class MultiGroupStrategy

Bases: MultiGroupStrategyBase

Multigroup strategy for $mathbf{A}= mathbf{I} - mathbf{M}mathbf{T}^{-1}mathbf{H}mathbf{D}mathbf{S}$, with $mathbf{S}$ the full scattering operator.

class MultiGroupChainGatherScanStrategy(single_group_chain_strategy: SingleGroupChainStrategy, single_group_gather_scatter_strategy: SingleGroupGatherScatterStrategy, checkpoint: bool = False)

Bases: _ScanLift, MultiGroupChainGatherStrategy

Lift the per-group body by a serialising scan. Optionally checkpoint for reduced memory usage

class MultiGroupChainGatherVmapStrategy(single_group_chain_strategy: SingleGroupChainStrategy, single_group_gather_scatter_strategy: SingleGroupGatherScatterStrategy)

Bases: _VmapLift, MultiGroupChainGatherStrategy

Lift by batching the entire vmap. Naturally most parallelizable but significantly increases memory (use on small problems only)

class MultiGroupChainStrategy

Bases: MultiGroupStrategyBase

Multigroup strategy for chain $mathbf{M}mathbf{T}^{-1}mathbf{H}mathbf{D}$

class ComposedMultiGroupChainStrategy(single_group_chain_strategy: SingleGroupChainStrategy)

Bases: MultiGroupChainStrategy

Applies a SingleGroupChainStrategy independently over every group. Used to precondition a per-group right-hand side with the same group-lift (scan/vmap, checkpoint) as the operator it accompanies; see from_composed().

single_group_chain_strategy: SingleGroupChainStrategy

forward(operator_data: OperatorData, x)

transpose(operator_data: OperatorData, x)

classmethod from_composed(composed: MultiGroupChainGatherStrategy)

The chain alone, lifted the same way (scan/vmap, checkpoint) as composed.

Return type:

ComposedMultiGroupChainStrategy

accepts(non_differentiable_data: NonDifferentiableData)

Whether it can accept this non-differentiable data.

Return type:

bool

build_default_non_differentiable_data(geometric_domain, angles)

Require each strategy to provide a way to construct this non-differentiable data from the geometric domain and angles.

Return type:

NonDifferentiableData

class MultiGroupChainScanStrategy(single_group_chain_strategy: SingleGroupChainStrategy, checkpoint: bool = False)

Bases: _ScanLift, ComposedMultiGroupChainStrategy

Lift the chain by a serialising scan. Optionally checkpoint for reduced memory usage

class MultiGroupChainVmapStrategy(single_group_chain_strategy: SingleGroupChainStrategy)

Bases: _VmapLift, ComposedMultiGroupChainStrategy

Lift the chain by batching the entire vmap.

Submodules

• jax_sn.operators.multi_group_operator.composed_multi_group_strategies module
  • MultiGroupChainGatherStrategy
  • MultiGroupChainGatherScanStrategy
  • MultiGroupChainGatherVmapStrategy
  • ComposedMultiGroupChainStrategy
  • MultiGroupChainScanStrategy
  • MultiGroupChainVmapStrategy
• jax_sn.operators.multi_group_operator.multi_group_operator_base module
  • MultiGroupStrategyBase
  • MultiGroupStrategy
  • MultiGroupChainStrategy
  • MultiGroupOperatorBase
  • MultiGroupOperator
  • MultiGroupChainOperator