jax_sn.operators.transport_operator.transport_operator_base module

class SingleAngleForwardStrategy

Bases: Strategy

abstractmethod forward(data: TransportData, x)

classmethod native_multi_angle()

The fused multi-angle lift of this kernel, if it has one. None (the default) means the only lift is a combinator over the single-angle forward.

Return type:

Type[MultiAngleForwardStrategy]

class MultiAngleForwardStrategy

Bases: Strategy

How to apply the forward operator over the angle axis. Consumes a batched TransportData (angle axis leading). Mirrors MultiAngleInversionStrategy.

abstractmethod forward(data: TransportData, x)

abstractmethod single_angle_strategy()

The single-angle SingleAngleForwardStrategy this angle lift corresponds to.

Return type:

SingleAngleForwardStrategy

class AbstractInversionStrategy

Bases: Module

How to sweep-invert the operator, and the sweep-data precondition that inversion declares. Shared by the single- and multi-angle hierarchies, which differ only in whether data carries an angle axis – so this is what a transport operator types its inversion strategy against.

abstractmethod invert(data: TransportData, b)

required_capabilities()

Sweep-data capability interfaces this strategy requires (e.g. (ProvidesLevels,)). Empty (the default) means any SweepData is acceptable.

Return type:

Tuple[type, ...]

build_default_sweep_data(geometric_domain: GeometricDomain, angle: Array)

The minimal sweep data this strategy needs, built for geometric_domain and angle (single or batched – dispatched on angle.ndim).

This builds rather than naming a type because a strategy may carry build-time knobs of its own (a level-padding width, say): those are properties of how this kernel sweeps, so they live as static fields on the strategy and are applied right here. Nothing upstream has to know they exist.

Return type:

SweepData

accepts(sweep_data: SweepData)

Accepts any sweep data that provides every capability in required_capabilities().

Return type:

bool

check_compatible(data: TransportData)

Asserts whether the sweep data is compatible with this inversion strategy (just typing information).

class SingleAngleInversionStrategy

Bases: AbstractInversionStrategy

How to sweep-invert the single-angle operator.

abstractmethod classmethod vmappable()

Whether an outer vmap over this inversion vectorises.

Return type:

bool

classmethod native_multi_angle()

The fused multi-angle lift of this kernel, if it has one. None (the default) means the only lift is a combinator over the single-angle sweep.

Return type:

Type[MultiAngleInversionStrategy]

class MultiAngleInversionStrategy

Bases: AbstractInversionStrategy

How to sweep-invert over the angle axis. Consumes a batched TransportData (angle axis leading).

abstractmethod invert(data: TransportData, b)

Invert over the angle axis: data is a batched TransportData (angle axis leading).

abstractmethod single_angle_strategy()

The single-angle SingleAngleInversionStrategy this angle inversion corresponds to (used for composited strategy, error if does not exist)

Return type:

SingleAngleInversionStrategy

class AbstractTransportOperator

Bases: DataTransposeDataStrategyLinearOperator, DirectlyInvertibleOperator

Shared for single-angle and multi-angle transport operators (and even for future multi-group transport operators).

data: AbstractVar[TransportData]

Transport Data

strategy: AbstractVar[Strategy]

How to apply the forward operator.

inversion_strategy: AbstractVar[AbstractInversionStrategy]

How to invert the operator

to_energy_group(energy_group: int)

Derive a new transport operator for a different energy group, keeping the same strategies.

Return type:

AbstractTransportOperator

abstract property n_angles: int

Number of angles this operator covers.

class TransportOperatorSingleAngle(data: TransportData, strategy: SingleAngleForwardStrategy, inversion_strategy: SingleAngleInversionStrategy)

Bases: AbstractTransportOperator

Base class for transport operators. Since the implementation of the transport operator can vary (e.g., vmap-based, batched, Pallas, etc.), this is an abstract base class that defines the interface for transport operators.

This operator can be constructed on its own or through a MultiAngleTransportSpec and a OperatorData. The spec defines the strategies, the operator_data provides the data, and the energy group is passed to MultiAngleTransportSpec.build().

Mathematical description

The transport operator maps angular fluxes to transported angular fluxes for a specific energy group \(g\) and angle \(\Omega_k\):

\[\mathbf{T}_{gk} : \mathcal{F} \to \mathcal{F}\]

where \(\mathcal{F} = \mathcal{E}\otimes \mathcal{B}\) is the space of elements and basis functions.

Its definition is:

\[\mathbf{T}_{gk}[\psi_{gk, ej}] = \sum_j \psi_j \int \Omega_k \cdot \nabla (H_j) H_i + \sigma^t_{ge} H_j H_i dV\]

This is then implemented using discontinuous Galerkin elements, where the surface terms are handled through upwinding based on the angle direction:

\[\begin{split}\mathbf{T}_{gk}[\psi_{gk, ej}] &= \sum_j \psi_{ej} \left[ \int_{e} - \Omega_k \cdot \nabla (H_i) H_j + \sigma^t_{ge} H_i H_j dV + \sum_{\gamma^+} |\Omega_k \cdot n_{\gamma}| \int_{\gamma} H_i H_j dA \right] \\ & - \sum_{\gamma^-} |\Omega_k \cdot n_{\gamma}| \sum_{p} \psi_{gk, e(\gamma),p } \int_{\gamma} H_i H_p^{(e(\gamma))}dA\end{split}\]

where \(\gamma\) denotes a face of element \(e\), with \(n_{\gamma}\) the outward normal of that face, and \(e(\gamma)\) the neighboring element on the other side of face \(\gamma\). \(\gamma^+\) and \(\gamma^-\) denote the faces where the dot product with the angle direction is positive (outflow) and negative (inflow) respectively.

The required geometric factors and finite elements are provided through the SpatialData object.

Inversion

The inverse can be computed through a so-called sweep. Since a topological ordering can be computed for the elements based on the angle direction, the transport equation can be solved element by element in that order. For each element, one then computes:

\[\begin{split}G_{e, ij} &= \int_{e} - \Omega_k \cdot \nabla (H_i) H_j + \sigma^t_{ge} H_i H_j dV + \sum_{\gamma^+} |\Omega_k \cdot n_{\gamma}| \int_{\gamma} H_i H_j dA \\ G^{-1}_{e, ij} &= \left( G_{e, ij} \right)^{-1} \\ D_{e \gamma, ij} &= |\Omega_k \cdot n_{\gamma}| \int_{\gamma} H_i H_j^{(e(\gamma))}dA \\\end{split}\]

These can be computed in parallel since they do not depend on the angular fluxes, or even stored a priori (although this is expensive in memory).

The ‘sweep’ then consists of iterating over the elements in topological order, and for each element collapsing (\(b_{ej}\) is the right-hand side for element \(e\), \(a_{ej}\) the unknown angular fluxes for element e):

\[a_{ei} = \sum_j G_{e,ij}^{-1} \left( b_{ej} + \sum_{\gamma^-} \sum_{p} D_{e \gamma,jp} a_{e(\gamma),p} \right)\]

This is the only truly sequential part of the transport operator inversion.

This inversion can be implemented in various ways, e.g., using vmap over angles, batched over angles, using Pallas, etc. These are implemented in subclasses.

Transposition

The transpose of the transport operator is obtained by simply reversing the angle (including the topological sort information). This is implemented in transpose(). See here for a detailed proof of this fact.

The transpose of the inverse is obtained by the inverse of the transpose, i.e., by performing a sweep in reverse topological order with reversed angles.

data: TransportData

Transport Data

strategy: SingleAngleForwardStrategy

How to apply the forward operator.

inversion_strategy: SingleAngleInversionStrategy

How to invert the operator

property n_angles: int

Number of angles in the transport operator. Here is just one.

class TransportOperatorMultiAngle(data: TransportData, strategy: MultiAngleForwardStrategy, inversion_strategy: MultiAngleInversionStrategy)

Bases: AbstractTransportOperator

Transport operator for multiple angles.

Mathematical description

The transport operator for multiple angles maps angular fluxes to transported angular fluxes for a specific energy group \(g\) and multiple angles \(\Omega_k\):

\[\begin{split}&\mathbf{T}_{g} : \mathcal{D} \otimes \mathcal{F} \to \mathcal{D} \otimes \mathcal{F} \\ &\mathbf{T}_g = \bigoplus_{k} \mathbf{T}_{gk}\end{split}\]

where \(\mathcal{F} = \mathcal{E}\otimes \mathcal{B}\) is the space of elements and basis functions, and \(\mathcal{D}\) is the space of discrete angles. The \(\mathbf{T}_{gk}\) are the single-angle transport operators defined in TransportOperatorSingleAngle.

Therefore, we have for an angular flux \(\psi_g \in \mathcal{D} \otimes \mathcal{F}\) with components \(\psi_{gk} \in \mathcal{F}\) for each angle \(k\):

\[(\mathbf{T}_{g}[\psi_{g}])_k = T_{gk}[\psi_{gk}]\]

Inversion

The inverse can be computed by inverting each single-angle transport operator independently due to the block-diagonal structure:

\[\mathbf{T}_g^{-1} = \bigoplus_{k} \mathbf{T}_{gk}^{-1}\]

Transpose

The transpose of the multiple-angle transport operator is obtained by transposing each single-angle transport operator, as in TransportOperatorSingleAngle.

Structure

Since this is by definition a block-diagonal operator with the single-angle transport operators on the diagonal, this operator does not couple angles and can be computed independently for each angle (and thus in parallel)

data: TransportData

Transport Data

strategy: MultiAngleForwardStrategy

How to apply the forward operator.

inversion_strategy: MultiAngleInversionStrategy

How to invert the operator

property n_angles: int

Number of angles in the transport operator.

to_single_angle(idx)

Derive the genuine single-angle transport operator for angle idx, slicing the angle and the sweep ordering in lockstep and using the natural single-angle inversion named by inversion_strategy.

Return type:

TransportOperatorSingleAngle