jax_sn.operators.transport_operator.implementations.transport_operator_warp module

Warp strategies and convenience constructors for the transport operator.

class WarpSequentialSweep

Bases: SingleAngleInversionStrategy

Single-angle inverse via the sequential Warp sweep kernel (no batching).

invert(data: TransportData, b)

check_compatible(data)

Asserts whether the sweep data is compatible with this inversion strategy (just typing information).

classmethod vmappable()

Whether an outer vmap over this inversion vectorises.

classmethod native_multi_angle()

The fused multi-angle lift of this kernel, if it has one. None (the default) means the only lift is a combinator over the single-angle sweep.

class WarpLevelSweep(maximum_number_of_levels: Optional[int] = None)

Bases: SingleAngleInversionStrategy

Single-angle inverse via the level-based-parallelism Warp sweep kernel. Requires sweep data providing the ProvidesLevels capability (the level setup for both sweep directions). The base check_compatible() enforces that capability; this strategy only adds the warp-data-layout check. The minimal concrete builder is LeveledSweepData.

maximum_number_of_levels: Optional[int] = None

Width the level arrays are padded to when this strategy builds its sweep data. None reads it off the levels themselves, which is a tracer whenever the domain the sweep is derived from is one – so pin it (to a width at least as large as the domain needs; too small is a clean equinox runtime error) whenever the sweep data is rebuilt under a trace, e.g. by a builder differentiating through a mesh deformation.

invert(data: TransportData, b)

classmethod required_capabilities()

Sweep-data capability interfaces this strategy requires (e.g. (ProvidesLevels,)). Empty (the default) means any SweepData is acceptable.

classmethod vmappable()

Whether an outer vmap over this inversion vectorises.

build_default_sweep_data(geometric_domain, angle)

The minimal sweep data this strategy needs, built for geometric_domain and angle (single or batched – dispatched on angle.ndim).

This builds rather than naming a type because a strategy may carry build-time knobs of its own (a level-padding width, say): those are properties of how this kernel sweeps, so they live as static fields on the strategy and are applied right here. Nothing upstream has to know they exist.

classmethod native_multi_angle()

The fused multi-angle lift of this kernel, if it has one. None (the default) means the only lift is a combinator over the single-angle sweep.

check_compatible(data)

Asserts whether the sweep data is compatible with this inversion strategy (just typing information).

class WarpSequentialAngles

Bases: MultiAngleInversionStrategy

Native multi-angle inverse: fused sequential Warp sweep.

single_angle_strategy()

The single-angle SingleAngleInversionStrategy this angle inversion corresponds to (used for composited strategy, error if does not exist)

invert(data: TransportData, b)

Invert over the angle axis: data is a batched TransportData (angle axis leading).

check_compatible(data)

Asserts whether the sweep data is compatible with this inversion strategy (just typing information).

class WarpLevelAngles(maximum_number_of_levels: Optional[int] = None)

Bases: MultiAngleInversionStrategy

Native multi-angle inverse: fused level-based-parallelism Warp sweep. Requires sweep data providing the ProvidesLevels capability.

maximum_number_of_levels: Optional[int] = None

Level-array padding width; see WarpLevelSweep.maximum_number_of_levels.

single_angle_strategy()

The single-angle SingleAngleInversionStrategy this angle inversion corresponds to (used for composited strategy, error if does not exist)

invert(data: TransportData, b)

Invert over the angle axis: data is a batched TransportData (angle axis leading).

required_capabilities()

Sweep-data capability interfaces this strategy requires (e.g. (ProvidesLevels,)). Empty (the default) means any SweepData is acceptable.

build_default_sweep_data(geometric_domain, angle)

The minimal sweep data this strategy needs, built for geometric_domain and angle (single or batched – dispatched on angle.ndim).

This builds rather than naming a type because a strategy may carry build-time knobs of its own (a level-padding width, say): those are properties of how this kernel sweeps, so they live as static fields on the strategy and are applied right here. Nothing upstream has to know they exist.

check_compatible(data)

Asserts whether the sweep data is compatible with this inversion strategy (just typing information).