jax_sn.operators.operator_data module¶

Common data structures for all operators, containing differentiable and non-differentiable data.

class AngularData(quadrature_set: QuadratureSetBase, wk_ylm: Array, ylm: Array)[source]¶

Bases: Module

Operator-defining angular data: the quadrature set and its precomputed harmonics.

quadrature_set: QuadratureSetBase¶: Angular quadrature set (angles + weights).

wk_ylm: Array¶: Weighted harmonics \(w_k Y_{l,m}(\Omega_k)\) – the moment operator’s matrix. [n_angles, n_moments]

ylm: Array¶: Unweighted harmonics \(Y_{l,m}(\Omega_k)\) – the discrete operator’s matrix. [n_angles, n_moments]

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

Precompute the harmonics up to moment order l_max for quadrature_set.

Return type:: AngularData

class DiscretisationData(spatial: SpatialData, angular: AngularData)[source]¶

Bases: Module

Operator-defining differentiable data.

spatial: SpatialData¶

geometry, FEM matrices, geometric factors, cross sections.

Type:: Spatial + material data

angular: AngularData¶

quadrature + harmonics.

Type:: Angular data

class NonDifferentiableData(sweep_data: SweepData)[source]¶

Bases: Module

Non-differentiable data (setup data) that is not part of the operator’s definition and as such does not have to take part in the VJP/JVP of the operator.

sweep_data: SweepData¶: The sweep execution plan (per-angle topological ordering).

class RequiresNonDifferentiableData[source]¶

Bases: Module

Mixin for strategies that require non-differentiable data (transport / single-group/ multi-group).

abstractmethod accepts(non_differentiable_data: NonDifferentiableData)[source]¶

Whether it can accept this non-differentiable data.

Return type:: bool

check_compatible(non_differentiable_data: NonDifferentiableData)[source]¶: The asserting twin of accepts(), for data assembled by hand.

abstractmethod build_default_non_differentiable_data(geometric_domain: GeometricDomain, angles: Array)[source]¶

Require each strategy to provide a way to construct this non-differentiable data from the geometric domain and angles.

Return type:: NonDifferentiableData

class OperatorData(discretisation_data: DiscretisationData, non_differentiable_data: NonDifferentiableData)[source]¶

Bases: Module

Data for neutron transport operators, split into differentiable and non-differentiable data.

discretisation_data: DiscretisationData¶: Operator-defining, differentiable data.

non_differentiable_data: NonDifferentiableData¶: Discrete, non-differentiable execution data

property spatial_data: SpatialData¶

property angular_data: AngularData¶

property sweep_data: SweepData¶

property domain¶

property solution_domain: SolutionDomain¶

property elements_basis_shape¶

property n_moments¶

property n_angles¶

property n_basis¶

property n_energy_groups¶