jax_sn.operators.transport_operator.transport_data module¶

Data a transport operator hands to its forward / inversion strategies.

class TransportData(spatial_data: SpatialData, energy_group: int, angle: Array, sweep_data: SweepData)[source]¶

Data a transport operator hands to its forward / inversion strategies.

spatial_data: SpatialData¶: Geometry, FEM matrices, geometric factors and indexed cross sections. [differentiable]

energy_group: int¶: Energy group index. Selects the cross sections; see cross_sections.

angle: Array¶: Transport direction(s) \(\Omega\). Shape [n_dim], or [n_angles, n_dim] when batched over angles.

sweep_data: SweepData¶: Per-angle topological ordering. Ignored by forward strategies.

property cross_sections: CrossSectionDataSingleGroup¶: Cross sections mapped to full element array for this energy group. Shape [n_elements].

property forward_args¶: Positional args of a forward kernel (unpacked with *data.forward_args).

property inversion_args¶: Positional args of an inverse-sweep kernel (unpacked with *data.inversion_args).

property transposed: TransportData¶

the reversed direction. The ordering’s reverse is the partner of negating the angle. See here.

single_angle(idx)[source]¶

Slice the angle and the sweep ordering in lockstep down to angle idx.