jax_sn.operators.transport_operator.sweep_data_implementations module

Concrete sweep orderings: the precomputed topological-sort structure required to invert (sweep) the transport operator. The abstract SweepData base lives one level up (referenced by the shared operator-data bundle); this module holds the concrete plain / level-parallel sweep data.

Level-parallel data exists to allow for different capabilities: a normal sweep does not require data on the levels. Since this can be particularly large (on the same order as the base data), it is not given by default.

class SweepData[source]

Bases: BaseOperatorData

Abstract base for the data required to invert (sweep) the transport operator.

levels: AbstractVar[Array]

Levels of the elements in the topological sort. Each level can be computed in parallel.

levels_argsort: AbstractVar[Array]

The indices that sort the elements by their levels.

abstract property reverse: SweepData

Reverses the topological sort (the partner of negating the angle, done on the operator). This is also correct for batched data.

Returns:

reversed_sweep_data – The reversed sweep data, of the same concrete type as self.

Return type:

SweepData

abstractmethod reconstruct(geometric_domain: GeometricDomain, angle: Array, n_iterations_per_loop: int = 100)[source]

Recompute this sweep data for a new (deformed) geometric domain and the given angle – so a parametrisation can rebuild it after moving the vertices. Type-preserving: each leaf returns its own concrete type and rebuilds its own extra data (e.g. a leveled sweep rebuilds its level setup with the same padding).

Parameters:

  • geometric_domain (GeometricDomain) – The deformed domain to recompute the sort for.

  • angle (Array) – The angles the sort is for.

Returns:

sweep_data – The recomputed sweep data for geometric_domain, of the same concrete type as self.

Return type:

SweepData

class BasicSweepData(levels: Array, levels_argsort: Array)[source]

Bases: SweepData

The plain sweep data for a geometric domain: just the topological sort itself, with no precomputed level-parallelism data. Used by the vmap/batched/pallas sweeps. Also the home of the factories that build a plain sort from a domain/quadrature set.

levels: Array

Levels of the elements in the topological sort. Each level can be computed in parallel.

levels_argsort: Array

The indices that sort the elements by their levels.

classmethod from_domain_angle(geometric_domain: GeometricDomain, angle: Array, n_iterations_per_loop: int = 100)[source]

Computes a (plain) topological sort for the given domain and angle(s). The angle is used to compute the order but is not stored (it lives in the angular setup). Accepts either a single angle or a batch of angles (dispatched on angle.ndim), so this is the single generic entry point a creation helper can call without knowing which case it has.

Can throw if there are cycles using equinox.

Parameters:

  • geometric_domain (GeometricDomain) – The geometric domain for which to compute the topological sort.

  • angle (Array) – The angle(s) for which to compute the topological sort.

  • n_iterations_per_loop (int) – The number of iterations per loop for the level computation. The default is 100.

Returns:

sweep_data – The computed sweep data, batched over angles iff angle was.

Return type:

BasicSweepData

classmethod from_domain_quadrature_set(geometric_domain: GeometricDomain, quadrature_set: QuadratureSetBase, n_iterations_per_loop: int = 100)[source]

Computes (plain) topological sorts for all angles in the given quadrature set and geometric domain.

Parameters:

  • geometric_domain (GeometricDomain) – The geometric domain for which to compute the topological sorts.

  • quadrature_set (QuadratureSetBase) – The quadrature set containing the angles for which to compute the topological sorts.

  • n_iterations_per_loop (int) – The number of iterations per loop for the level computation. The default is 100.

Returns:

sweep_data – The computed sweep data for all angles in the quadrature set. Batched over angles.

Return type:

BasicSweepData

property reverse: BasicSweepData

Reverses the topological sort (the partner of negating the angle, done on the operator). This is also correct for batched data.

Returns:

reversed_sweep_data – The reversed sweep data, of the same concrete type as self.

Return type:

SweepData

reconstruct(geometric_domain: GeometricDomain, angle: Array, n_iterations_per_loop: int = 100)[source]

Recompute this sweep data for a new (deformed) geometric domain and the given angle – so a parametrisation can rebuild it after moving the vertices. Type-preserving: each leaf returns its own concrete type and rebuilds its own extra data (e.g. a leveled sweep rebuilds its level setup with the same padding).

Parameters:

  • geometric_domain (GeometricDomain) – The deformed domain to recompute the sort for.

  • angle (Array) – The angles the sort is for.

Returns:

sweep_data – The recomputed sweep data for geometric_domain, of the same concrete type as self.

Return type:

BasicSweepData

class LevelBasedParallelism(first_occurences: Array, ends: Array, sizes: Array)[source]

Bases: Module

Class providing utilities for level-based parallelism in the context of transport sweeps. Stores the first occurrence, end, and size of each level for each angle, based on the sweep data provided by the transport operator.

Since JAX does not allow dynamic shapes, the class method from_sweep_data optionally is given a maximum number of levels, which is used to pad the resulting arrays to a fixed shape. If this is not sufficiently high, an equinox error is given.

first_occurences: Array

First occurrence of each level, per angle. Shape […, max_levels]

ends: Array

End of each level (i.e. start of next level), per angle. Shape […, max_levels]

sizes: Array

Number of elements in each level, per angle. Shape […, max_levels]

classmethod from_sweep_data(sweep_data: SweepData, maximum_number_of_levels: Optional[int] = None)[source]

Create a LevelBasedParallelism instance from a SweepData, which contains the levels and their argsort for each angle. Optionally, a maximum number of levels can be provided to pad the resulting arrays to a fixed shape. This has the advantage that array sizes are known at compile time. If not sufficiently high enough, it will raise an equinox error.

Parameters:

  • sweep_data (SweepData) – The sweep data containing the levels and their argsort for each angle.

  • maximum_number_of_levels (Optional[int]) – The maximum number of levels to pad to. If None, it will be set to the maximum number of levels present in the sweep data. Default is None.

Returns:

An instance of LevelBasedParallelism containing the first occurrence, end, and size of each level for each angle, based on the sweep data provided by the transport operator.

Return type:

LevelBasedParallelism

property n_levels

Number of distinct levels

class ProvidesLevels[source]

Bases: Module

Mixin for sweep data the provides levels forward and backward.

level_fwd: AbstractVar[LevelBasedParallelism]
level_rev: AbstractVar[LevelBasedParallelism]
class LeveledSweepData(levels: Array, levels_argsort: Array, level_fwd: LevelBasedParallelism, level_rev: LevelBasedParallelism)[source]

Bases: SweepData, ProvidesLevels

Sweep data that also provides Levels for LevelBasedParallelism.

levels: Array

Levels of the elements in the topological sort. Each level can be computed in parallel.

levels_argsort: Array

The indices that sort the elements by their levels.

level_fwd: LevelBasedParallelism

Level-based-parallelism setup for the (forward) sweep along the angle.

level_rev: LevelBasedParallelism

Level-based-parallelism setup for the reversed sweep (along the negated angle).

classmethod from_basic_sweep_data(sweep_data: SweepData, maximum_number_of_levels: Optional[int] = None)[source]

Build leveled sweep data from a plain one, precomputing the level setup for both the forward and the reversed sweep.

Return type:

LeveledSweepData

classmethod from_domain_quadrature_set(geometric_domain: GeometricDomain, quadrature_set: QuadratureSetBase, maximum_number_of_levels: Optional[int] = None, n_iterations_per_loop: int = 100)[source]

Build leveled sweep data for all angles in a quadrature set directly from a domain (the level computation runs once), padding the level arrays to maximum_number_of_levels.

Return type:

LeveledSweepData

classmethod from_domain_angle(geometric_domain: GeometricDomain, angle: Array, maximum_number_of_levels: Optional[int] = None, n_iterations_per_loop: int = 100)[source]

Build leveled sweep data directly from a domain and angle(s) (single or batched, dispatched on angle.ndim – see BasicSweepData.from_domain_angle()), padding the level arrays to maximum_number_of_levels.

Return type:

LeveledSweepData

property reverse: LeveledSweepData

Reverses the topological sort (the partner of negating the angle, done on the operator). This is also correct for batched data.

Returns:

reversed_sweep_data – The reversed sweep data, of the same concrete type as self.

Return type:

SweepData

reconstruct(geometric_domain: GeometricDomain, angle: Array, n_iterations_per_loop: int = 100)[source]

Recompute this sweep data for a new (deformed) geometric domain and the given angle – so a parametrisation can rebuild it after moving the vertices. Type-preserving: each leaf returns its own concrete type and rebuilds its own extra data (e.g. a leveled sweep rebuilds its level setup with the same padding).

Parameters:

  • geometric_domain (GeometricDomain) – The deformed domain to recompute the sort for.

  • angle (Array) – The angles the sort is for.

Returns:

sweep_data – The recomputed sweep data for geometric_domain, of the same concrete type as self.

Return type:

LeveledSweepData