jax_sn.operators.scatter_operator.gather_scatter_operator_base module

class GatherScatterData(cross_sections_indexed: CrossSectionsIndexed)

Bases: BaseOperatorData

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

cross_sections_indexed: CrossSectionsIndexed

class SingleGroupGatherScatterStrategy

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)

Return type:

Array

incoming_gather(data: GatherScatterData, x: Array, g: int)

Return type:

Array

outgoing_gather(data: GatherScatterData, x: Array, g: int)

Return type:

Array

class GatherScatterStrategy

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)

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)

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)

Return type:

Array

abstractmethod transpose(data: GatherScatterData, x: Array)

Return type:

Array