jax_sn.operators.sweep_data module

Non-differentiable sweep-execution data: the abstract interface the shared operator-data bundle

class SweepData

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)

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