jax_sn.operators.mass_operator

class MassOperator(data: MassData, additional_shape: Tuple[int], strategy: MassStrategy)[source]

Bases: SymmetricDataStrategyLinearOperator

Mass operator. Computes for a given vector \(\{ v_j \}\) the mass matrix map:

\[a_i = \sum_{j} \int H_i H_j dV v_j\]

Is arbitrarily batched over additional leading dimensions.

Mathematical description

The non-batched version computes the mass matrix for each element and basis function:

\[\overline{\mathbf{H}} : F \to F\]

where in this case:

\[F : \{ \text{all elements} \} \otimes \{ \text{all basis functions per element} \}\]

and since all elements are independent, the mass matrix is block-diagonal with each block representing the mass matrix for one element.

Its definition is (\(e,p\) is element index, i,j basis function index):

\[\overline{\mathbf{H}}[\{ v_{ei} \}] = a_{ej} = \sum_{i} J_{vq} \int_{B} H_i H_j dV v_{ei}\]

where \(J_{vq}\) is the volume Jacobian determinant for element \(e\), and \(B\) is the reference element. Note this integral is already computed in jax_sn.solution_domain.FEMMatrices.mass_matrix, and the Jacobian determinants are precomputed and stored jax_sn.solution_domain.PrecomputedGeometricData.vol_jacobians_det.

Transposition

This operator is symmetric since the mass matrix is symmetric. Transposition thus returns itself. The strategy therefore implements a single apply() (forward == transpose); there is no transpose_bool.

Structure

Batched versions operate over additional leading dimensions, e.g.:

\[\begin{split}&\mathbf{H} : \mathcal{K} \otimes \mathcal{F} \to \mathcal{K} \otimes \mathcal{F}\\ &\mathbf{H} = \mathbf{I}_{\mathcal{K}} \otimes \overline{\mathbf{H}}\end{split}\]

where \(\mathcal{K}\) is an arbitrary additional space (e.g. angles).

In essence, this just is a batched matrix product (and is implemented like that as well).

data: MassData

MassData containing FEM matrices and precomputed geometric data. [differentiable]

additional_shape: Tuple[int]

Additional leading shape to batch over. Used for consistent in/out structure definitions.

strategy: MassStrategy

How to apply the mass contraction. See mass_operator_strategies.

class MassData(fem_matrices: FEMMatrices, precomputed_geometric_data: PrecomputedGeometricData)[source]

Bases: BaseOperatorData

fem_matrices: FEMMatrices
precomputed_geometric_data: PrecomputedGeometricData
class MassStrategy[source]

Bases: SymmetricStrategy

How to apply the mass operator. Symmetric, so transpose == forward.

abstractmethod forward(data: MassData, x)[source]
class VmapMass[source]

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

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

Bases: MassStrategy

Fused Warp kernel, supporting n_basis <= MAX_NBASIS (=4). Forward == transpose; differentiable.

forward(operator, x)[source]

Submodules