jax_sn.operators.transport_operator.implementations.multi_angle_inversions module

Composed multi-angle strategies, which lift a single-angle strategy over the angle axis. Both axes (forward and inverse sweep) get the same two lifts: vmap and lax.map.

map_over_angles(data: TransportData, leading, fn)

vmap closure over angle & sweep data. Closure because we don’t know a priori other arguments (i.e. batching.)

class VmapAnglesForward(inner: SingleAngleForwardStrategy)

Bases: ComposedMultiAngleForwardStrategy

Lift the single-angle forward by vmap-ing it over the angle axis.

forward(data: TransportData, x)

class ScanAnglesForward(inner: SingleAngleForwardStrategy)

Bases: ComposedMultiAngleForwardStrategy

Lift the single-angle forward by lax.map-ing it over the angle axis.

forward(data: TransportData, x)

class VmapAnglesInversion(inner: SingleAngleInversionStrategy)

Bases: ComposedMultiAngleInversionStrategy

Lift the single-angle inverse sweep by vmap-ing it over the angle axis.

invert(data: TransportData, b)

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

class ScanAnglesInversion(inner: SingleAngleInversionStrategy)

Bases: ComposedMultiAngleInversionStrategy

Lift the single-angle inverse sweep by lax.map-ing it over the angle axis.

invert(data: TransportData, b)

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