jax_sn.operators.transport_operator

Transport operator implementations; both single angle and multi-angle.

transport_operator_base defines the single-angle and multi-angle lineax operators, strategies, and data.

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 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

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 TransportData(spatial_data: SpatialData, energy_group: int, angle: Array, sweep_data: SweepData)[source]

Bases: TransposableOperatorData

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.

Type:

Data of the transposed operator

single_angle(idx)[source]

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

Return type:

TransportData

class AbstractTransportOperator[source]

Bases: DataTransposeDataStrategyLinearOperator, DirectlyInvertibleOperator

Shared for single-angle and multi-angle transport operators (and even for future multi-group transport operators).

data: AbstractVar[TransportData]

Transport Data

strategy: AbstractVar[Strategy]

How to apply the forward operator.

inversion_strategy: AbstractVar[AbstractInversionStrategy]

How to invert the operator

to_energy_group(energy_group: int)[source]

Derive a new transport operator for a different energy group, keeping the same strategies.

Return type:

AbstractTransportOperator

abstract property n_angles: int

Number of angles this operator covers.

class TransportOperatorSingleAngle(data: TransportData, strategy: SingleAngleForwardStrategy, inversion_strategy: SingleAngleInversionStrategy)[source]

Bases: AbstractTransportOperator

Base class for transport operators. Since the implementation of the transport operator can vary (e.g., vmap-based, batched, Pallas, etc.), this is an abstract base class that defines the interface for transport operators.

This operator can be constructed on its own or through a MultiAngleTransportSpec and a OperatorData. The spec defines the strategies, the operator_data provides the data, and the energy group is passed to MultiAngleTransportSpec.build().

Mathematical description

The transport operator maps angular fluxes to transported angular fluxes for a specific energy group \(g\) and angle \(\Omega_k\):

\[\mathbf{T}_{gk} : \mathcal{F} \to \mathcal{F}\]

where \(\mathcal{F} = \mathcal{E}\otimes \mathcal{B}\) is the space of elements and basis functions.

Its definition is:

\[\mathbf{T}_{gk}[\psi_{gk, ej}] = \sum_j \psi_j \int \Omega_k \cdot \nabla (H_j) H_i + \sigma^t_{ge} H_j H_i dV\]

This is then implemented using discontinuous Galerkin elements, where the surface terms are handled through upwinding based on the angle direction:

\[\begin{split}\mathbf{T}_{gk}[\psi_{gk, ej}] &= \sum_j \psi_{ej} \left[ \int_{e} - \Omega_k \cdot \nabla (H_i) H_j + \sigma^t_{ge} H_i H_j dV + \sum_{\gamma^+} |\Omega_k \cdot n_{\gamma}| \int_{\gamma} H_i H_j dA \right] \\ & - \sum_{\gamma^-} |\Omega_k \cdot n_{\gamma}| \sum_{p} \psi_{gk, e(\gamma),p } \int_{\gamma} H_i H_p^{(e(\gamma))}dA\end{split}\]

where \(\gamma\) denotes a face of element \(e\), with \(n_{\gamma}\) the outward normal of that face, and \(e(\gamma)\) the neighboring element on the other side of face \(\gamma\). \(\gamma^+\) and \(\gamma^-\) denote the faces where the dot product with the angle direction is positive (outflow) and negative (inflow) respectively.

The required geometric factors and finite elements are provided through the SpatialData object.

Inversion

The inverse can be computed through a so-called sweep. Since a topological ordering can be computed for the elements based on the angle direction, the transport equation can be solved element by element in that order. For each element, one then computes:

\[\begin{split}G_{e, ij} &= \int_{e} - \Omega_k \cdot \nabla (H_i) H_j + \sigma^t_{ge} H_i H_j dV + \sum_{\gamma^+} |\Omega_k \cdot n_{\gamma}| \int_{\gamma} H_i H_j dA \\ G^{-1}_{e, ij} &= \left( G_{e, ij} \right)^{-1} \\ D_{e \gamma, ij} &= |\Omega_k \cdot n_{\gamma}| \int_{\gamma} H_i H_j^{(e(\gamma))}dA \\\end{split}\]

These can be computed in parallel since they do not depend on the angular fluxes, or even stored a priori (although this is expensive in memory).

The ‘sweep’ then consists of iterating over the elements in topological order, and for each element collapsing (\(b_{ej}\) is the right-hand side for element \(e\), \(a_{ej}\) the unknown angular fluxes for element e):

\[a_{ei} = \sum_j G_{e,ij}^{-1} \left( b_{ej} + \sum_{\gamma^-} \sum_{p} D_{e \gamma,jp} a_{e(\gamma),p} \right)\]

This is the only truly sequential part of the transport operator inversion.

This inversion can be implemented in various ways, e.g., using vmap over angles, batched over angles, using Pallas, etc. These are implemented in subclasses.

Transposition

The transpose of the transport operator is obtained by simply reversing the angle (including the topological sort information). This is implemented in transpose(). See here for a detailed proof of this fact.

The transpose of the inverse is obtained by the inverse of the transpose, i.e., by performing a sweep in reverse topological order with reversed angles.

data: TransportData

Transport Data

strategy: SingleAngleForwardStrategy

How to apply the forward operator.

inversion_strategy: SingleAngleInversionStrategy

How to invert the operator

property n_angles: int

Number of angles in the transport operator. Here is just one.

class TransportOperatorMultiAngle(data: TransportData, strategy: MultiAngleForwardStrategy, inversion_strategy: MultiAngleInversionStrategy)[source]

Bases: AbstractTransportOperator

Transport operator for multiple angles.

Mathematical description

The transport operator for multiple angles maps angular fluxes to transported angular fluxes for a specific energy group \(g\) and multiple angles \(\Omega_k\):

\[\begin{split}&\mathbf{T}_{g} : \mathcal{D} \otimes \mathcal{F} \to \mathcal{D} \otimes \mathcal{F} \\ &\mathbf{T}_g = \bigoplus_{k} \mathbf{T}_{gk}\end{split}\]

where \(\mathcal{F} = \mathcal{E}\otimes \mathcal{B}\) is the space of elements and basis functions, and \(\mathcal{D}\) is the space of discrete angles. The \(\mathbf{T}_{gk}\) are the single-angle transport operators defined in TransportOperatorSingleAngle.

Therefore, we have for an angular flux \(\psi_g \in \mathcal{D} \otimes \mathcal{F}\) with components \(\psi_{gk} \in \mathcal{F}\) for each angle \(k\):

\[(\mathbf{T}_{g}[\psi_{g}])_k = T_{gk}[\psi_{gk}]\]

Inversion

The inverse can be computed by inverting each single-angle transport operator independently due to the block-diagonal structure:

\[\mathbf{T}_g^{-1} = \bigoplus_{k} \mathbf{T}_{gk}^{-1}\]

Transpose

The transpose of the multiple-angle transport operator is obtained by transposing each single-angle transport operator, as in TransportOperatorSingleAngle.

Structure

Since this is by definition a block-diagonal operator with the single-angle transport operators on the diagonal, this operator does not couple angles and can be computed independently for each angle (and thus in parallel)

data: TransportData

Transport Data

strategy: MultiAngleForwardStrategy

How to apply the forward operator.

inversion_strategy: MultiAngleInversionStrategy

How to invert the operator

property n_angles: int

Number of angles in the transport operator.

to_single_angle(idx)[source]

Derive the genuine single-angle transport operator for angle idx, slicing the angle and the sweep ordering in lockstep and using the natural single-angle inversion named by inversion_strategy.

Return type:

TransportOperatorSingleAngle

class SingleAngleForwardStrategy[source]

Bases: Strategy

abstractmethod forward(data: TransportData, x)[source]
classmethod native_multi_angle()[source]

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 forward.

Return type:

Type[MultiAngleForwardStrategy]

class MultiAngleForwardStrategy[source]

Bases: Strategy

How to apply the forward operator over the angle axis. Consumes a batched TransportData (angle axis leading). Mirrors MultiAngleInversionStrategy.

abstractmethod forward(data: TransportData, x)[source]
abstractmethod single_angle_strategy()[source]

The single-angle SingleAngleForwardStrategy this angle lift corresponds to.

Return type:

SingleAngleForwardStrategy

class ComposedMultiAngleForwardStrategy(inner: SingleAngleForwardStrategy)[source]

Bases: MultiAngleForwardStrategy

A multi-angle forward strategy that is a combinator over a single-angle strategy.

inner: SingleAngleForwardStrategy
single_angle_strategy()[source]

The single-angle SingleAngleForwardStrategy this angle lift corresponds to.

Return type:

SingleAngleForwardStrategy

class AbstractInversionStrategy[source]

Bases: Module

How to sweep-invert the operator, and the sweep-data precondition that inversion declares. Shared by the single- and multi-angle hierarchies, which differ only in whether data carries an angle axis – so this is what a transport operator types its inversion strategy against.

abstractmethod invert(data: TransportData, b)[source]
required_capabilities()[source]

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

Return type:

Tuple[type, ...]

build_default_sweep_data(geometric_domain: GeometricDomain, angle: Array)[source]

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.

Return type:

SweepData

accepts(sweep_data: SweepData)[source]

Accepts any sweep data that provides every capability in required_capabilities().

Return type:

bool

check_compatible(data: TransportData)[source]

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

class SingleAngleInversionStrategy[source]

Bases: AbstractInversionStrategy

How to sweep-invert the single-angle operator.

abstractmethod classmethod vmappable()[source]

Whether an outer vmap over this inversion vectorises.

Return type:

bool

classmethod native_multi_angle()[source]

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.

Return type:

Type[MultiAngleInversionStrategy]

class MultiAngleInversionStrategy[source]

Bases: AbstractInversionStrategy

How to sweep-invert over the angle axis. Consumes a batched TransportData (angle axis leading).

abstractmethod invert(data: TransportData, b)[source]

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

abstractmethod single_angle_strategy()[source]

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

Return type:

SingleAngleInversionStrategy

class ComposedMultiAngleInversionStrategy(inner: SingleAngleInversionStrategy)[source]

Bases: MultiAngleInversionStrategy

A multi-angle inversion strategy that is a combinator over a single-angle strategy.

inner: SingleAngleInversionStrategy
single_angle_strategy()[source]

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

Return type:

SingleAngleInversionStrategy

required_capabilities()[source]

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

Return type:

Tuple[type, ...]

build_default_sweep_data(geometric_domain, angle)[source]

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.

Return type:

SweepData

check_compatible(data: TransportData)[source]

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

class VmapForward[source]

Bases: SingleAngleForwardStrategy

Forward via vmap over all elements.

forward(data: TransportData, x)[source]
class BatchedForward(batch_size: int = None)[source]

Bases: SingleAngleForwardStrategy

Forward via batched vmap + fori_loop. batch_size=None reproduces the non-batched version.

batch_size: int = None
forward(data: TransportData, x)[source]
class VmapSweep[source]

Bases: SingleAngleInversionStrategy

Inverse sweep with all per-element matrices computed via vmap (high memory).

invert(data: TransportData, b)[source]
classmethod vmappable()[source]

Whether an outer vmap over this inversion vectorises.

class BatchedSweep(batch_size: int = None)[source]

Bases: SingleAngleInversionStrategy

Inverse sweep over batches of elements. batch_size=None reproduces the non-batched version.

batch_size: int = None
invert(data: TransportData, b)[source]
classmethod vmappable()[source]

Whether an outer vmap over this inversion vectorises.

class VmapAnglesInversion(inner: SingleAngleInversionStrategy)[source]

Bases: ComposedMultiAngleInversionStrategy

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

invert(data: TransportData, b)[source]

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

class ScanAnglesInversion(inner: SingleAngleInversionStrategy)[source]

Bases: ComposedMultiAngleInversionStrategy

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

invert(data: TransportData, b)[source]

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

class VmapAnglesForward(inner: SingleAngleForwardStrategy)[source]

Bases: ComposedMultiAngleForwardStrategy

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

forward(data: TransportData, x)[source]
class ScanAnglesForward(inner: SingleAngleForwardStrategy)[source]

Bases: ComposedMultiAngleForwardStrategy

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

forward(data: TransportData, x)[source]
class VmapPallasSweep[source]

Bases: SingleAngleInversionStrategy

Inverse sweep with all per-element matrices via vmap, using Pallas for the matrix ops.

invert(data: TransportData, b)[source]
classmethod vmappable()[source]

Whether an outer vmap over this inversion vectorises.

class BatchedPallasSweep(batch_size: int = None)[source]

Bases: SingleAngleInversionStrategy

Inverse sweep over batches of elements, using Pallas for the matrix ops.

batch_size: int = None
invert(data: TransportData, b)[source]
classmethod vmappable()[source]

Whether an outer vmap over this inversion vectorises.

class WarpSequentialSweep[source]

Bases: SingleAngleInversionStrategy

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

invert(data: TransportData, b)[source]
check_compatible(data)[source]

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

classmethod vmappable()[source]

Whether an outer vmap over this inversion vectorises.

classmethod native_multi_angle()[source]

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)[source]

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)[source]
classmethod required_capabilities()[source]

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

classmethod vmappable()[source]

Whether an outer vmap over this inversion vectorises.

build_default_sweep_data(geometric_domain, angle)[source]

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()[source]

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)[source]

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

class WarpSequentialAngles[source]

Bases: MultiAngleInversionStrategy

Native multi-angle inverse: fused sequential Warp sweep.

single_angle_strategy()[source]

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

invert(data: TransportData, b)[source]

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

check_compatible(data)[source]

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

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

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()[source]

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

invert(data: TransportData, b)[source]

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

required_capabilities()[source]

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)[source]

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)[source]

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

Subpackages

Submodules