jax_sn.operators.single_group_operator

This sub-module defines operators on a single group.

This includes several sub-operators which are used in other places (i.e. lifting to multigroup / rhs):

  1. The total single group operator (used for lower tri solves):

    I - M T_g^{-1} H D S_{gg}

  2. The chain operator (used for gathering the scattering source):

    M T_g^{-1} H D

  3. The source operator (used for right-hand-side construction):

    M T_g^{-1} H

Each operator carries a strategy of its matching type (full / chain / source); the composed strategies factor the product into the mass / discrete / moment / transport sub-operators.

class SingleGroupData(operator_data: OperatorData, energy_group: int)[source]

Bases: BaseOperatorData

operator_data: OperatorData
energy_group: int
class SingleGroupStrategyBase[source]

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)[source]
abstractmethod transpose(single_group_data: SingleGroupData, x)[source]
class SingleGroupStrategy[source]

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

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

Bases: SingleGroupStrategyBase

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

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

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

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

Bases: SingleGroupOperatorBase

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

strategy: SingleGroupSourceStrategy
class ComposedSingleGroupStrategy(transport_strategy: MultiAngleForwardStrategy, transport_inversion_strategy: MultiAngleInversionStrategy, moment_strategy: FluxMomentStrategy, mass_strategy: MassStrategy, discrete_strategy: FluxMomentStrategy, scatter_strategy: ScatterStrategy)[source]

Bases: ComposedSingleGroupTransportBase, SingleGroupStrategy

The full \(\mathbf{I} - \mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\mathbf{D}\mathbf{S}_{gg}\) operator.

discrete_strategy: FluxMomentStrategy
scatter_strategy: ScatterStrategy
all_operators(single_group_data: SingleGroupData)[source]

Construct the entire chain of operators. This is a free construction in any case (after JIT): just assigning a type.

Return type:

Tuple[TransportOperatorMultiAngle, MomentOperator, DiscreteOperator, MassOperator, ScatterOperator]

forward(data: SingleGroupData, x)[source]
transpose(data: SingleGroupData, x)[source]
class ComposedSingleGroupChainStrategy(transport_strategy: MultiAngleForwardStrategy, transport_inversion_strategy: MultiAngleInversionStrategy, moment_strategy: FluxMomentStrategy, mass_strategy: MassStrategy, discrete_strategy: FluxMomentStrategy)[source]

Bases: ComposedSingleGroupTransportBase, SingleGroupChainStrategy

The chain operator \(\mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\mathbf{D}\) – no scatter operator needed.

discrete_strategy: FluxMomentStrategy
all_operators(single_group_data: SingleGroupData)[source]
Return type:

Tuple[TransportOperatorMultiAngle, MomentOperator, DiscreteOperator, MassOperator]

forward(data: SingleGroupData, x)[source]
transpose(data: SingleGroupData, x)[source]
class ComposedSingleGroupSourceStrategy(transport_strategy: MultiAngleForwardStrategy, transport_inversion_strategy: MultiAngleInversionStrategy, moment_strategy: FluxMomentStrategy, mass_strategy: MassStrategy)[source]

Bases: ComposedSingleGroupTransportBase, SingleGroupSourceStrategy

The source operator \(\mathbf{M}\mathbf{T}_g^{-1}\mathbf{H}\) – no discrete or scatter operator needed.

all_operators(single_group_data: SingleGroupData)[source]
Return type:

Tuple[TransportOperatorMultiAngle, MomentOperator, MassOperator]

forward(data: SingleGroupData, x)[source]
transpose(data: SingleGroupData, x)[source]
single_group_moment_structure(data: SingleGroupData)[source]
Return type:

ShapeDtypeStruct

single_group_angle_structure(data: SingleGroupData)[source]
Return type:

ShapeDtypeStruct

Submodules