jax_sn.operators.mass_operator.mass_operator_strategies module

Strategy objects for the mass operator.

The mass operator applies a block-diagonal, per-element mass matrix scaled by the volume Jacobian determinant. The implementation of that contraction is selected by a strategy object rather than by subclassing, mirroring the transport, moment and discrete operators:

• MassStrategy – how to apply the (batched) block-diagonal mass matrix.

The mass operator is symmetric, so each strategy implements a single forward (forward == transpose); there is no separate transpose method and no transpose_bool. A new implementation is a single small class, not a subclass of the operator. The kernels live here next to the strategies that dispatch to them, so the operator module depends only on the (array-free) strategy classes.

class VmapMass

Bases: MassStrategy

Uses a vmap of all elements to compute the mass operator. Alone, this creates a temporary. However, in a combined operator (e.g. SingleGroupOperator), this temporary can often be re-used, such that it is not necessarily any worse than LaxMass (which does not create temporaries).

forward(data: MassData, x)

class LaxMass

Bases: MassStrategy

Batched matrix multiplication implementation of the mass operator using jax.lax. This implementation does not create temporaries and is not particularly less or more performant than VmapMass.

forward(data: MassData, x)

total_mass_kernel_lax(fem_matrices: FEMMatrices, precomputed_geometric_data: PrecomputedGeometricData, values: Array)