e3j.core.SymmetricContraction¶

class e3j.core.SymmetricContraction(source_irreps: str, correlation: int, keep_irrep_out: str, num_species: int, num_channels: int, layout: str | Layout, l_max: int | None = None, config: Config | None = None)¶

Contracts the power expansion of E3-features with index-wise weights,

B(A, i) = W[i] @ (A + (A ⊗ A) + … + A**(⊗ ν))

where:

ν is the correlation order (typically lower than 4),
A : (N, C, D) is an array of D-dimensional E3-features,
i : (N,) is the vector of species indices 0 <= i < S,
W stores one (S, C, m_out, m_in) weight per block of linear_blocks, passed explicitly to __call__().

Species indices outside [0, S) are undefined behaviour: they are not checked, and the result may differ between implementations.

References

See MACE, eq.10.

MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields. Batatia, Kovacs, Simm, Ortner & Csanyi, 2022.

Methods

__call__(features, indices, weight)

Power expansion of features, contracted with index-wise weights.

Attributes

`expansion`	Direct sum of equivariant powers.
`linear`	Index-wise linear mixing the sorted powers.
`linear_blocks`	Per-(l, p) blocks of the index-wise linear, in weight order.
`mosaic_gpu_params`	Static metadata of the fused Mosaic GPU kernel.
`mosaic_tpu_params`	Static metadata of the fused Mosaic TPU kernel.
`permutation`	Groups equivalent irreps before the output linear transform.
`source`
`target`

__call__(features: Array, indices: Array, weight: list[Array]) → Array¶

Power expansion of features, contracted with index-wise weights.

Parameters:

features (jnp.ndarray) – (..., C, dim) or (..., dim, C) features depending on layout.
indices (jnp.ndarray) – (...)-array of species indices in [0, S); out-of-range indices are undefined behaviour.
weight (list[jnp.ndarray]) – one (S, C, m_out, m_in) array per block of linear_blocks.

Returns:

target features, in layout.

Return type:

jnp.ndarray

property expansion: PowerExpansion¶: Direct sum of equivariant powers.

property linear: LinearIndexwise¶: Index-wise linear mixing the sorted powers.

property linear_blocks: list[LinearBlock]¶: Per-(l, p) blocks of the index-wise linear, in weight order.

property mosaic_gpu_params: PallasMosaicGPUSymmetricContractionParams¶: Static metadata of the fused Mosaic GPU kernel.

property mosaic_tpu_params: PallasMosaicTPUSymmetricContractionParams¶: Static metadata of the fused Mosaic TPU kernel.

property permutation: Permutation¶: Groups equivalent irreps before the output linear transform.