jax_sn.operators.moment_operator

class MomentOperator(data: MomentData, elements_basis_shape: ShapeDtypeStruct, strategy: FluxMomentStrategy, transpose_bool: bool = False)[source]

Bases: FluxMomentOperator

Moment operator that maps angular fluxes to moments.

Mathematical description

Implements the following operation (the tensor dot indicates that it operates on the leading dimension, with no coupling in the trailing dimensions):

\[M = \overline{\mathbf{M}} \otimes \mathbf{I}_F: \mathcal{D}\otimes \mathcal{F} \rightarrow \mathcal{M}\otimes \mathcal{F}\]

with

\[\overline{\mathbf{M}}[\{a_k\}] = \sum_{k=1}^{N_{angles}} a_k w_k Y_{lm}(\Omega_k)\]

where \(\mathcal{D}\) is the discrete angular space of size \(N_{angles}\), \(\mathcal{M}\) is the moment space of size \((L+1)^2\) (for moment order L), \(F\) is the FEM basis function space of size \(N_{elements} \times N_{basis}\), \(a_k\) is the angular flux for angle \(k\), \(w_k\) is the weight for angle \(k\) and \(Y_{lm}(\Omega_k)\) is the spherical harmonic of order \((l,m)\) evaluated at direction \(\Omega_k\). For the latter definition, see jax_sn.quadrature_set.QuadratureSet.wk_ylm_real.

Transposition

Its transposition is defined as (transpose again does not act on the \(F\) space):

\[M^T = \overline{\mathbf{M}}^T \otimes \mathbf{I}_F: \mathcal{M}\otimes \mathcal{F} \rightarrow \mathcal{D}\otimes \mathcal{F}\]

with

\[\overline{\mathbf{M}}^T[\{ b_{lm} \}] = \sum_{lm} b_{lm} w_k Y_{lm}(\Omega_k)\]

This is not the same as the discrete-to-moment operator, which is defined without the weights. Mathematically, the transpose of this operator is w.r.t to the standard inner product, while the discrete-to-moment operator is defined w.r.t. a weighted inner product (with weights being the quadrature weights).

data: MomentData

MomentData containing the \(w_k Y_{lm}(\Omega_k)\) weights for the moment operator [differentiable]

elements_basis_shape: ShapeDtypeStruct

The shape of the elements and basis functions over which this operator is batched. Identified with \(F\) in the documentation above.

strategy: FluxMomentStrategy

How to apply the contraction. See fluxmoment_strategies.

transpose_bool: bool = False

Whether this operator applies the transposed action (moments to angles) instead of the forward action (angles to moments).

class DiscreteOperator(data: DiscreteData, elements_basis_shape: ShapeDtypeStruct, strategy: FluxMomentStrategy, transpose_bool: bool = False)[source]

Bases: FluxMomentOperator

Discrete operator that maps moments to angular fluxes. Approximate inverse of jax_sn.operators.moment_operator.MomentOperator.

Mathematical description

Implements:

\[D = \overline{\mathbf{D}} \otimes \mathbf{I}_F: \mathcal{M}\otimes \mathcal{F} \rightarrow \mathcal{D}\otimes \mathcal{F}\]

with

\[\overline{\mathbf{D}}[\{ b_{lm} \}] = \sum_{lm} b_{lm} Y_{lm}(\Omega_k)\]

where \(\mathcal{D}\) is the space of discrete angles, \(\mathcal{M}\) the space of moments, and \(\mathcal{F}\) the space of elements and basis functions.

Transposition

The transposed action (selected via transpose_bool) implements:

\[D^T = \overline{\mathbf{D}}^T \otimes \mathbf{I}_F: \mathcal{D}\otimes \mathcal{F} \rightarrow \mathcal{M}\otimes \mathcal{F}\]

with

\[\overline{\mathbf{D}}^T[s_k] = \sum_{k} s_k Y_{lm}(\Omega_k)\]

This is not the same as the moment operator, since the weights \(w_k\) are not included here.

data: DiscreteData

Matrix of [n_angles, n_moments] containing \(Y_{l,m}(\Omega_k)\) values for each angle \(k\) and moment \((l,m)\) (collapsed into a single index).

elements_basis_shape: ShapeDtypeStruct

The shape of the elements and basis functions over which this operator is batched. Identified with \(F\) in the documentation above.

strategy: FluxMomentStrategy

How to apply the contraction. See fluxmoment_strategies.

transpose_bool: bool = False

Whether this operator applies the transposed action (angles to moments) instead of the forward action (moments to angles).

class FluxMomentOperator[source]

Bases: TransposeBoolDataStrategyLinearOperator

Shared base for the moment and discrete operators: a harmonics contraction on the leading axis, batched over elements_basis_shape. The direction is carried by data.

data: AbstractVar[DiscreteMomentData]
elements_basis_shape: AbstractVar[ShapeDtypeStruct]
strategy: AbstractVar[FluxMomentStrategy]
transpose_bool: AbstractVar[bool]
property n_moments

Number of moments (= \((L+1)^2\) for moment order L)

property n_angles

Number of discrete angles

class DiscreteMomentData[source]

Bases: Module

A real-spherical-harmonics matrix of shape [n_angles, n_moments], plus the axis its forward action consumes. The matrix layout is the same for the moment and discrete operators; only in_axis distinguishes them, which is what lets both share one strategy hierarchy.

in_axis: AbstractClassVar[int]

0 = angles, 1 = moments.

Type:

Axis of matrix the forward action contracts over

abstract property matrix: Array

The harmonics matrix [n_angles, n_moments].

abstractmethod classmethod from_quadrature_set(quadrature_set: QuadratureSetBase, l_max: int)[source]

Computes wk_yklm or yklm from the quadrature set. Prefer to initialise directly.

property n_angles

Number of discrete angles.

property n_moments

Number of moments (= \((L+1)^2\) for moment order L).

class MomentData(wk_yklm: Array)[source]

Bases: DiscreteMomentData

Weighted harmonics \(w_k Y_{lm}(\Omega_k)\). Applied angles \(\to\) moments.

wk_yklm: Array
in_axis: ClassVar[int] = 0

0 = angles, 1 = moments.

Type:

Axis of matrix the forward action contracts over

property matrix

The harmonics matrix [n_angles, n_moments].

classmethod from_quadrature_set(quadrature_set: QuadratureSetBase, l_max: int)[source]

Computes wk_yklm or yklm from the quadrature set. Prefer to initialise directly.

class DiscreteData(yklm: Array)[source]

Bases: DiscreteMomentData

Unweighted harmonics \(Y_{lm}(\Omega_k)\). Applied moments \(\to\) angles.

yklm: Array
in_axis: ClassVar[int] = 1

0 = angles, 1 = moments.

Type:

Axis of matrix the forward action contracts over

property matrix

The harmonics matrix [n_angles, n_moments].

classmethod from_quadrature_set(quadrature_set: QuadratureSetBase, l_max: int)[source]

Computes wk_yklm or yklm from the quadrature set. Prefer to initialise directly.

class FluxMomentStrategy[source]

Bases: TransposableStrategy

abstractmethod contract(data: DiscreteMomentData, x, axis: int)[source]

Contract x over axis of data.matrix (0 = angles, 1 = moments).

forward(data: DiscreteMomentData, x)[source]
transpose(data: DiscreteMomentData, x)[source]
class EinsumFluxMoment[source]

Bases: FluxMomentStrategy

Batched-matrix (einsum) implementation. Batches over the trailing dimensions.

contract(data: DiscreteMomentData, x, axis: int)[source]

Contract x over axis of data.matrix (0 = angles, 1 = moments).

class StreamingFluxMoment[source]

Bases: FluxMomentStrategy

Streaming (scan-over-the-contracted-axis) implementation. Accumulates the contributions one slice at a time, trading compute for a smaller peak memory footprint than the batched einsum.

contract(data: DiscreteMomentData, x, axis: int)[source]

Contract x over axis of data.matrix (0 = angles, 1 = moments).

Submodules