jax_sn.operators.scatter_operator

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

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

ScatterOperator

class ScatterData(energy_group_out: int, energy_group_in: int, cross_sections_indexed: CrossSectionsIndexed)[source]

Bases: BaseOperatorData

energy_group_out: int
energy_group_in: int
cross_sections_indexed: CrossSectionsIndexed
class ScatterStrategy[source]

Bases: SymmetricStrategy

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

class IndexedScatter[source]

Bases: ScatterStrategy

forward(data: ScatterData, x)[source]
class GatherScatterData(cross_sections_indexed: CrossSectionsIndexed)[source]

Bases: BaseOperatorData

Data for the gather-scatter operator: the indexed cross sections it contracts against.

cross_sections_indexed: CrossSectionsIndexed
class SingleGroupGatherScatterStrategy[source]

Bases: Module

Gathers the scatter source in a single energy group, either from incoming groups (normal scattering) or from outgoing groups (transpose/adjoint scattering).

\[\sum_{g'} \sigma_{\ell g g',e}\, x_{g'\ell e b} \quad\text{or}\quad \sum_{g'} \sigma_{\ell g' g,e}\, x_{g'\ell e b }\]

Should be batched over trailing axes (e.g. elements, basis functions). Note this is not a Transpose: it doesn’t satisfy the adjoint identity in this single-group form! This is why this is a separate object, not a strategy.

This can be used either to define the total GatherScatter operation (by lifting over the group axis) or to define a single-group operator (by fixing the group index). This just applies the single-group version.

abstractmethod gather_group(data: GatherScatterData, x: Array, g: int, sum_over_outgoing: bool)[source]
Return type:

Array

incoming_gather(data: GatherScatterData, x: Array, g: int)[source]
Return type:

Array

outgoing_gather(data: GatherScatterData, x: Array, g: int)[source]
Return type:

Array

class GatherScatterStrategy[source]

Bases: TransposableStrategy

The total gather scatter operator:

\[\mathbf{U} : [G,M,E,B] \to [G,M,E,B], \quad (\mathbf{U} x)_{g\ell e b} = \sum_{g'} \sigma_{\ell g g',e}\, x_{g'\ell e b}\]

This is transposable:

\[\mathbf{U}^T : [G,M,E,B] \to [G,M,E,B], \quad (\mathbf{U}^T x)_{g\ell e b} = \sum_{g'} \sigma_{\ell g' g,e}\, x_{g' \ell e b}\]

and is therefore a TransposableStrategy.

class GatherScatterOperator(data: GatherScatterData, strategy: GatherScatterStrategy, group_moments_elements_basis_shape: ShapeDtypeStruct, transpose_bool: bool = False)[source]

Bases: TransposeBoolDataStrategyLinearOperator

The all-groups scatter operator \(\mathbf{U}\) (or \(\mathbf{U}^T\)) as a lineax operator on [n_groups, n_moments, n_elements, n_basis].

data: GatherScatterData
strategy: GatherScatterStrategy
group_moments_elements_basis_shape: ShapeDtypeStruct

Full operand structure [n_groups, n_moments, n_elements, n_basis] (cross sections fix only groups + elements).

transpose_bool: bool = False
class ComposedGatherScatterStrategy(single_group_strategy: SingleGroupGatherScatterStrategy)[source]

Bases: GatherScatterStrategy

The full scatter operator built from a per-group row + something that maps it to the full operator

single_group_strategy: SingleGroupGatherScatterStrategy
abstractmethod forward(data: GatherScatterData, x: Array)[source]
Return type:

Array

abstractmethod transpose(data: GatherScatterData, x: Array)[source]
Return type:

Array

class BroadcastedSingleGroupGather[source]

Bases: SingleGroupGatherScatterStrategy

Broadcasted (vmap-over-elements) single-group gather kernel.

gather_group(data: GatherScatterData, x: Array, g: int, sum_over_outgoing: bool)[source]
Return type:

Array

class VmapGatherScatterStrategy(single_group_strategy: SingleGroupGatherScatterStrategy)[source]

Bases: ComposedGatherScatterStrategy

Full \(\mathbf{U}\) / \(\mathbf{U}^T\) by vmapping the single-group block over the group axis.

forward(data: GatherScatterData, x: Array)[source]
Return type:

Array

transpose(data: GatherScatterData, x: Array)[source]
Return type:

Array

Submodules