jax_sn.operators.scatter_operator.scatter_operator_base module

class ScatterData(energy_group_out: int, energy_group_in: int, cross_sections_indexed: CrossSectionsIndexed)

Bases: BaseOperatorData

energy_group_out: int

energy_group_in: int

cross_sections_indexed: CrossSectionsIndexed

class ScatterStrategy

Bases: SymmetricStrategy

How to apply the (symmetric-shaped) scatter contraction. Reads data.

class ScatterOperator(data: ScatterData, strategy: ScatterStrategy, moments_elements_basis_shape: ShapeDtypeStruct)

Bases: SymmetricDataStrategyLinearOperator

Maps the flux moments to scattering for a specific energy group in and a specific energy group out, using indexed cross sections.

Mathematical description

We have:

\[\mathbf{S}_{\tilde{g} g} = \overline{\mathbf{S}}_{\tilde{g} g} \otimes \mathbf{I}_{\mathcal{B}} : \mathcal{M} \otimes \mathcal{E}\otimes \mathcal{B} \to \mathcal{M} \otimes \mathcal{E}\otimes \mathcal{B}\]

with \(\tilde{g}\) the outgoing energy group, \(g\) the incoming energy group, \(\mathcal{F} = \mathcal{E}\otimes \mathcal{B}\) the space of elements and basis functions, and \(\mathbf{S}_{\tilde{g} g}\) the single-element scattering operator defined as:

\[\overline{\mathbf{S}}_{\tilde{g} g} [\psi] =\sigma_{\ell \tilde{g} g e} \psi_{g\ell e}\]

for an arbitrary \(\psi\in \mathcal{M} \otimes \mathcal{E}\) (we use \(\ell\) to denote \((l,m)\) moment index for brevity; for \(\sigma_{\ell}\) this just equals \(\sigma_{l}\) for all \(m\)).

Transposition

Since this is a element wise multiplication, the transpose is itself.

Technical details

We use a jax.ShapeDtypeStruct to represent the shape of the moments, elements and basis functions over which this operator is defined. This is used to define a consistent in and output structure. Since we want to have the flexibility to have a different number of moments than is present in the cross sections, we explicitly provide it here.

data: ScatterData

ScatterData containing the indexed cross sections for the scatter operator. [differentiable except for groups]

strategy: ScatterStrategy

How to apply the scatter contraction.

moments_elements_basis_shape: ShapeDtypeStruct

The shape of the moments, elements and basis functions over which this operator is batched. Identified with \(\mathcal{M}\otimes \mathcal{E}\otimes \mathcal{B}\) in the documentation above.

to_energy_group(energy_group_out: int, energy_group_in: int)

Return type:

ScatterOperator