jax_sn.operators.transport_operator.implementations.simplex_warp_kernel module

Warp sweep kernels for simplex elements, parametrised on (n_basis, n_faces).

Warp’s fixed-size vector/matrix types stop at length 4, so a per-element basis vector must be a vecNd with N = n_basis <= 4 and the element matrix a matNNd. That is the only limit: the kernel body never mentions the spatial dimension or the element degree, so every simplex whose basis count fits is expressible. See SUPPORTED for the resulting set.

n_basis and n_faces are independent axes – a quadratic interval is (3, 2) while a linear triangle is (3, 3) – so build_kernels() is keyed on both.

The spatial types are not parametrised. Angles and face normals are 3-vectors at every spatial dimension, and vol_jacobians_inv/derivative_matrix are zero-padded to 3 on the host (see _pad_to_3()), which makes angle_ref a plain 3x3 solve whose trailing components vanish – identical to the truncation the JAX kernel does with angle[:vol_jacobian_inv.shape[-1]]. The padding is a no-op in 3-D, so the case that dominates runtime pays nothing for it, and it also buys 1-D, for which Warp has no 1x1 matrix type at all.

MAX_LENGTH = 4

Warp’s largest fixed-size vector/matrix type; bounds both n_basis and n_faces.

SUPPORTED = ((2, 2), (3, 2), (4, 2), (3, 3), (4, 4))

the linear, quadratic and cubic interval, the linear triangle and the linear tetrahedron.

Type:

The (n_basis, n_faces) pairs the simplex families reach within MAX_LENGTH

NDIM = 3

Angles and face normals are 3-vectors regardless of the spatial dimension.

INFLOW_EPSILON = float64(-1e-12)

Floating point differences can lead to JAX and Warp disagreeing on whether a face is inflow or outflow. Then, an element can access a non-computed face neighbour as inflow in Warp (trying to get non-computed elements). This leads to NaNs in the solution. Although, in principle, one could set the whole result array to zero before the sweep to avoid NaNs, this would not work for level-based parallelism, where the levels could in principle have two of these neighbours on the same level, which could lead to a race condition. Although extremely rare, we don’t want to introduce the possibility. This epsilon guards agains this by treaing faces with angle.face_normal > -INFLOW_EPSILON as outflow faces. Since the contribution is multiplied by this dot-product, even if it were physical, the contribution would be on the order of the INFLOW_EPSILON, which should thus be chosen to be on the order of floating point errors.

class SimplexWarpKernels(n_basis: int, n_faces: int, single_angle: Callable, multi_angle: Callable, level_single_angle: Callable, level_multi_angle: Callable)

Bases: NamedTuple

The four sweep entry points compiled for one (n_basis, n_faces). Selected at trace time and never traced through, so a plain tuple rather than a pytree node.

n_basis: int

Alias for field number 0

n_faces: int

Alias for field number 1

single_angle: Callable

Alias for field number 2

multi_angle: Callable

Alias for field number 3

level_single_angle: Callable

Alias for field number 4

level_multi_angle: Callable

Alias for field number 5

build_kernels(n_basis: int, n_faces: int)

Compile the sweep kernels for one element shape. Memoised, so a repeated call returns the same Warp module rather than recompiling; distinct shapes coexist as separate modules.

Return type:

SimplexWarpKernels