jax_sn.operators.single_group_operator.single_group_operators_base module

class SingleGroupData(operator_data: OperatorData, energy_group: int)

Bases: BaseOperatorData

operator_data: OperatorData

energy_group: int

class SingleGroupStrategyBase

Bases: TransposableStrategy, RequiresNonDifferentiableData

Base class for single-group strategies, which are linear operators that act on a single energy group.

Like the multigroup strategies, it states its requirement on the whole non-differentiable tier (see RequiresNonDifferentiableData) rather than on the sweep data alone.

abstractmethod forward(single_group_data: SingleGroupData, x)

abstractmethod transpose(single_group_data: SingleGroupData, x)

class SingleGroupStrategy

Bases: SingleGroupStrategyBase

Apply the full \(\mathbf{I} - \mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\mathbf{D}\mathbf{S}_{gg}\) operator, either forward or backward.

Abstract chain and source operators: a SingleGroupStrategy can have a very logical chain and source operator (e.g. if composed of sub-operators), which it can then yield (or just raise an error if this does not exist).

class SingleGroupChainStrategy

Bases: SingleGroupStrategyBase

Apply the chain \(\mathbf{C}_g = \mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\mathbf{D}\) operator, either forward or backward.

class SingleGroupSourceStrategy

Bases: SingleGroupStrategyBase

Apply the source \(\mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\) operator, either forward or backward.

single_group_moment_structure(data: SingleGroupData)

Return type:

ShapeDtypeStruct

single_group_angle_structure(data: SingleGroupData)

Return type:

ShapeDtypeStruct

class SingleGroupOperatorBase(data: SingleGroupData, strategy: SingleGroupStrategyBase, transpose_bool: bool = False)

Bases: TransposeBoolDataStrategyLinearOperator

Base operator acting on a single group.

data: SingleGroupData

strategy: SingleGroupStrategyBase

transpose_bool: bool = False

to_energy_group(energy_group)

Return type:

SingleGroupOperatorBase

class SingleGroupOperator(data: SingleGroupData, strategy: SingleGroupStrategy, transpose_bool: bool = False)

Bases: SingleGroupOperatorBase

Single energy group linear operator.

This implements the total, single-group linear operator:

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

where \(\mathbf{T}\) is the (multi-angle) transport operator (see TransportOperatorMultiAngle), \(\mathbf{H}\) the batched (over angles) mass operator (see MassOperator), \(\mathbf{M}\) the batched (over element and basis function) moment operator (see MomentOperator), \(\mathbf{S}_{gg}\) the inscattering operator (see ScatterOperator).

Transposition

The transpose follows from transposing the individual operators:

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

where the individual transposes are defined in their respective classes or directly in the fused operator.

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.

strategy: SingleGroupStrategy

class SingleGroupChainOperator(data: SingleGroupData, strategy: SingleGroupChainStrategy, transpose_bool: bool = False)

Bases: SingleGroupOperatorBase

The chain operator \(\mathbf{C}_g = \mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\mathbf{D}\), built from operator_data and the energy group.

strategy: SingleGroupChainStrategy

class SingleGroupSourceOperator(data: SingleGroupData, strategy: SingleGroupSourceStrategy, transpose_bool: bool = False)

Bases: SingleGroupOperatorBase

The source operator \(\mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\), built from operator_data and the energy group.

strategy: SingleGroupSourceStrategy