Equivariant API Reference

Native E(3)-equivariant algebra and neural-network primitives used by the atomistic and quantum-chemistry models (SchNet, PaiNN, NequIP, QHNet). Provides irreducible representations, tensor products, gated nonlinearities, spherical harmonics, radial bases, and graph utilities.

Core algebra (irreps, tensor products, gates)

Native E(3)-equivariant building-block library (irreps, Clebsch-Gordan, etc.).

A dependency-free, JAX/Flax-NNX-native reimplementation of the equivariant primitives used by e3nn/e3nn-jax (Geiger & Smidt 2022, arXiv:2207.09453), shared across opifex's quantum-SciML families (interatomic potentials, equivariant Hamiltonian prediction).

CartesianLinear

CartesianLinear(*, in_channels: int, out_channels: int, rngs: Rngs)

Bases: Module

Equivariant channel-mixing of stacked rank-2 Cartesian tensors.

Given an input of shape (..., in_channels, 3, 3) the layer produces (..., out_channels, 3, 3) by mixing channels separately within each irreducible Cartesian part (scalar / vector / symmetric-traceless). Because each part carries a single O(3) irrep, an arbitrary learnable mixing over the channel axis -- applied independently per part -- commutes with the rotation action X -> R X R^T and is therefore equivariant (TensorNet uses the same per-irrep linear mixing, arXiv:2306.06482 Eq. 6).

The three parts get independent (in_channels, out_channels) weight matrices, so the layer is the Cartesian analogue of :class:~opifex.neural.equivariant.EquivariantLinear.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input Cartesian-tensor channels.	required
out_channels	int	Number of output channels.	required
rngs	Rngs	Random number generators (keyword-only); rngs.params() seeds the three weight matrices.	required

CartesianTensor

CartesianTensor(array: Float[Array, '... 3 3'])

A rank-2 Cartesian tensor feature transforming as X -> R X R^T.

The underlying array has shape (..., 3, 3) (arbitrary leading / channel dimensions). The object is a pytree whose single child is the array, so it flows through jit / grad / vmap unchanged. This is the TensorNet (arXiv:2306.06482) representation: an l <= 2 equivariant feature stored as one Cartesian tensor rather than a spherical-harmonic :class:~opifex.neural.equivariant.IrrepsArray.

shape property

shape: tuple[int, ...]

Shape of the underlying array.

rotate

rotate(rotation: Float[Array, '3 3']) -> CartesianTensor

Apply the equivariant action X -> R X R^T for a rotation R.

tensor_product

tensor_product(other: CartesianTensor) -> CartesianTensor

Equivariant tensor product via the matrix product X @ Y.

Equivariant because (R X R^T)(R Y R^T) = R (X Y) R^T (TensorNet, Eq. 5). No Clebsch-Gordan tensor is required.

decompose

decompose() -> tuple[Float[Array, ''], Float[Array, '... 3'], Float[Array, '... 3 3']]

Decompose into irreducible Cartesian parts (scalar, vector, symmetric).

Following TensorNet Eq. 1-4:

• scalar -- the trace tr(X) (an l = 0 invariant);
• vector -- the axial vector dual to the antisymmetric part (X - X^T)/2 (an l = 1 feature, rotating by R);
• symmetric -- the symmetric-traceless part (X + X^T)/2 - tr(X)/3 * I (an l = 2 feature, rotating by R . R^T).

Returns:

Type	Description
Float[Array, '']	A (scalar, vector, symmetric) triple with shapes
Float[Array, '... 3']	(...,), (..., 3) and (..., 3, 3).

from_parts classmethod

from_parts(scalar: Float[Array, ''], vector: Float[Array, '... 3'], symmetric: Float[Array, '... 3 3']) -> CartesianTensor

Assemble a Cartesian tensor from its irreducible parts.

Inverse of :meth:decompose: X = scalar/3 * I + A(vector) + symmetric where A(vector) is the antisymmetric matrix with axial vector vector and symmetric is the symmetric-traceless l = 2 part.

Parameters:

Name	Type	Description	Default
scalar	Float[Array, '']	The l = 0 trace, shape (...,).	required
vector	Float[Array, '... 3']	The l = 1 axial vector, shape (..., 3).	required
symmetric	Float[Array, '... 3 3']	The l = 2 symmetric-traceless tensor, shape (..., 3, 3).	required

Returns:

Type	Description
CartesianTensor	The reconstructed :class:CartesianTensor.

tree_flatten

tree_flatten() -> tuple[tuple[Float[Array, '... 3 3']], None]

Pytree flatten: the array is the single child, no aux-data.

tree_unflatten classmethod

tree_unflatten(_aux_data: None, children: tuple[Float[Array, '... 3 3']]) -> CartesianTensor

Pytree unflatten, bypassing shape re-validation for traced leaves.

Gate

Gate(irreps_in: Irreps | str, *, even_act: Callable[[Array], Array] = gelu, odd_act: Callable[[Array], Array] = soft_odd, gate_act: Callable[[Array], Array] = sigmoid)

Bases: Module

Configured wrapper around :func:gate for use in module stacks.

Stores the activation choices (and validates the input irreps on call) so the gate can be dropped into an nnx module sequence like a layer. It has no learnable parameters.

Parameters:

Name	Type	Description	Default
irreps_in	Irreps | str	Expected input layout (validated on call).	required
even_act	Callable[[Array], Array]	Activation for even scalars. Default :func:jax.nn.gelu.	gelu
odd_act	Callable[[Array], Array]	Activation for odd scalars. Default :func:soft_odd.	soft_odd
gate_act	Callable[[Array], Array]	Activation for the gate scalars. Default :func:jax.nn.sigmoid.	sigmoid

NormGate

NormGate(irreps: Irreps | str, *, rngs: Rngs)

Bases: Module

Norm-gated equivariant nonlinearity.

Unlike :func:gate -- which consumes the rightmost scalars as dedicated gates -- this gate drives every multiplicity from a learnable MLP of the input scalars concatenated with the per-irrep norms of the non-scalar channels (the gating signal is therefore rotation-invariant). The MLP output replaces the scalar channels and scales each non-scalar multiplicity, so the output layout equals the input layout. This is the nonlinearity used throughout the QHNet self / pair interaction layers (Yu et al. 2023).

Scaling an equivariant feature by an invariant gate is equivariant, so the module is rotation-equivariant; reusing :func:opifex.neural.equivariant.norm gives a NaN-safe gradient at zero vectors.

Parameters:

Name	Type	Description	Default
irreps	Irreps | str	Expected input layout (validated on call). Must carry at least one l = 0 scalar channel to drive the gates.	required
rngs	Rngs	Random number generators (keyword-only) seeding the MLP.	required

Raises:

Type	Description
ValueError	If irreps has no scalar channel.

Irrep

Irrep(degree: int | str | Irrep, parity: int | None = None)

A single irreducible representation (l, p) of O(3).

Construct from a string (Irrep("1o")) or from explicit degree and parity (Irrep(1, -1)).

dim property

dim: int

Dimension 2 l + 1 of the representation space.

Irreps

Irreps(value: Irreps | Irrep | str | tuple[tuple[int, Irrep], ...])

An ordered direct sum of mul x Irrep blocks (a steerable feature layout).

Construct from a string (Irreps("8x0e + 4x1o")), an :class:Irrep, another :class:Irreps, or a tuple of (mul, Irrep) blocks.

dim property

dim: int

Total dimension sum(mul * (2l + 1)) of the feature vector.

num_irreps property

num_irreps: int

Total number of irreps (multiplicities summed).

slices

slices() -> list[slice]

Return the contiguous slice into the feature axis for each block.

IrrepsArray

IrrepsArray(irreps: Irreps | Irrep | str, array: Float[Array, '... dim'])

A JAX array tagged with its :class:Irreps layout (an equivariant feature).

The array's last axis is the feature axis and must equal irreps.dim. The object is a pytree: the array is the single child (traced) and the :class:Irreps is static aux-data, so equivariant features pass through jit/grad/vmap unchanged.

shape property

shape: tuple[int, ...]

Shape of the underlying array.

ndim property

ndim: int

Number of dimensions of the underlying array.

chunks property

chunks: list[Float[Array, '... mul dim_l']]

Split the feature axis into per-block arrays of shape (..., mul, 2l+1).

tree_flatten

tree_flatten() -> tuple[tuple[Float[Array, '... dim']], Irreps]

Pytree flatten: array is the child, irreps is static aux-data.

tree_unflatten classmethod

tree_unflatten(aux_data: Irreps, children: tuple[Float[Array, '... dim']]) -> IrrepsArray

Pytree unflatten, bypassing dimension re-validation for traced leaves.

EquivariantLinear

EquivariantLinear(irreps_in: Irreps | str, irreps_out: Irreps | str, *, rngs: Rngs)

Bases: Module

An equivariant linear map between two :class:Irreps layouts.

For every output block, the layer sums learnable (mul_in, mul_out) mixings over all input blocks carrying the same irrep (l, p). Output blocks whose irrep is absent from the input are left at zero (no weights).

Parameters:

Name	Type	Description	Default
irreps_in	Irreps | str	Input layout.	required
irreps_out	Irreps | str	Output layout.	required
rngs	Rngs	Random number generators (keyword-only); rngs.params() seeds the weight blocks.	required

BesselBasis

BesselBasis(num_basis: int, cutoff: float, *, rngs: Rngs | None = None)

Bases: Module

Bessel radial basis b_n(r) = sqrt(2/r_c) sin(n pi r / r_c) / r.

Ported from MACE ../mace/mace/modules/radial.py:18 (eq. 7 of arXiv:2206.07697). The frequencies n pi / r_c for n = 1..num_basis are stored as a (non-trainable) buffer; the basis is a smooth, complete set of invariant radial features.

Parameters:

Name	Type	Description	Default
num_basis	int	Number of Bessel functions n = 1..num_basis.	required
cutoff	float	Cutoff radius r_c (must be positive).	required
rngs	Rngs | None	Unused (the basis has no learnable parameters); accepted for a uniform constructor signature across equivariant modules.	None

GaussianBasis

GaussianBasis(num_basis: int, cutoff: float, *, rngs: Rngs | None = None)

Bases: Module

Gaussian radial basis g_n(r) = exp(-(r - mu_n)^2 / (2 sigma^2)).

Ported from MACE ../mace/mace/modules/radial.py:88. Centres mu_n are evenly spaced on [0, r_c] and the width is the centre spacing (sigma = r_c / (num_basis - 1)), giving the coefficient coeff = -0.5 / sigma^2.

Parameters:

Name	Type	Description	Default
num_basis	int	Number of Gaussians (centres on [0, r_c]); must be >= 2 so the spacing is well defined.	required
cutoff	float	Cutoff radius r_c (positive); sets the centre range.	required
rngs	Rngs | None	Unused (no learnable parameters); accepted for a uniform constructor signature.	None

PiecewiseLinearBasis

PiecewiseLinearBasis(num_basis: int, cutoff: float, *, rngs: Rngs | None = None)

Bases: Module

Piecewise-linear (hat) radial basis on [0, cutoff].

Faithful to the isotropic basis of torch_harmonics's PiecewiseLinearFilterBasis (torch-harmonics/torch_harmonics/filter_basis.py), the canonical filter basis for discrete-continuous (DISCO) convolutions (Ocampo, Price & McEwen 2023, arXiv:2209.13603). The num_basis hat functions have collocation spacing dr = 2 * cutoff / (num_basis + 1) and half-width dr: phi_k(r) = max(0, 1 - |r - r_k| / dr) restricted to r <= cutoff, with centres r_k = k * dr (odd num_basis) or (k + 0.5) * dr (even). Each function is continuous and compactly supported, so no separate cutoff envelope is needed.

Parameters:

Name	Type	Description	Default
num_basis	int	Number of hat functions (>= 1).	required
cutoff	float	Support radius r_c (positive).	required
rngs	Rngs | None	Unused (the basis is parameter-free); accepted for interface uniformity.	None

SymmetricContraction

SymmetricContraction(irreps_in: Irreps | str, irreps_out: Irreps | str, *, correlation: int, num_species: int, num_channels: int, rngs: Rngs)

Bases: Module

Per-element symmetric contraction up to a given correlation (body order).

Maps a per-channel single-particle basis A of layout irreps_in to irreps_out features, channel-wise, summing the contributions of every tensor-power order 1 .. correlation. Weights are indexed by atomic element, so each species gets its own contraction (the MACE convention).

Parameters:

Name	Type	Description	Default
irreps_in	Irreps | str	Per-channel input irreps (multiplicity one each; the channel axis carries the multiplicity).	required
irreps_out	Irreps | str	Per-channel output irreps to keep.	required
correlation	int	Maximum tensor-power order nu (body order nu + 1).	required
num_species	int	Number of distinct atomic elements (the per-element weight axis).	required
num_channels	int	Number of feature channels.	required
rngs	Rngs	Random number generators (keyword-only) seeding the weights.	required

Raises:

Type	Description
ValueError	If correlation < 1.

ChannelwiseTensorProduct

ChannelwiseTensorProduct(irreps_in1: Irreps | str, irreps_in2: Irreps | str, irreps_out: Irreps | str, *, internal_weights: bool = True, rngs: Rngs)

Bases: Module

Channel-wise (e3nn "uuu") Clebsch-Gordan tensor product.

The two inputs and the output share one uniform multiplicity mul; each selection-rule-allowed path (i_1, i_2) -> i_3 whose output irrep is present in irreps_out couples the inputs channel-by-channel and carries a single (mul,) weight -- O(mul) parameters per path rather than the :class:FullyConnectedTensorProduct's O(mul^3). This is the connection QHNet's self / pair interaction layers use (../AIRS/OpenDFT/QHBench/QH9/ models/QHNet.py); the unweighted forward equals e3nn.elementwise_tensor_product.

The per-channel weights are the module's internal parameters by default, but a per-sample weights array (width :attr:weight_numel) can be supplied to :meth:__call__ -- the path QHNet uses to inject per-edge radial / inner-product modulation into the pair coupling.

Parameters:

Name	Type	Description	Default
irreps_in1	Irreps | str	First input layout (single uniform multiplicity).	required
irreps_in2	Irreps | str	Second input layout (same uniform multiplicity).	required
irreps_out	Irreps | str	Output layout; only paths producing an irrep present here are instantiated. Each output irrep should appear once.	required
internal_weights	bool	If True (default) allocate the per-channel path weights as module parameters; if False they must be supplied to :meth:__call__ every time.	True
rngs	Rngs	Random number generators (keyword-only); rngs.params() seeds the internal weights.	required

Raises:

Type	Description
ValueError	If the inputs and output do not share one multiplicity.

FullyConnectedTensorProduct

FullyConnectedTensorProduct(irreps_in1: Irreps | str, irreps_in2: Irreps | str, irreps_out: Irreps | str, *, rngs: Rngs)

Bases: Module

Fully-connected Clebsch-Gordan tensor product with learnable path weights.

Every selection-rule-allowed path (i_1, i_2) -> i_3 whose output irrep is present in irreps_out carries an independent (mul_1, mul_2, mul_3) weight; the contributions are summed per output block.

Parameters:

Name	Type	Description	Default
irreps_in1	Irreps | str	Layout of the first input.	required
irreps_in2	Irreps | str	Layout of the second input.	required
irreps_out	Irreps | str	Desired output layout; only paths producing an irrep that appears here are instantiated.	required
rngs	Rngs	Random number generators (keyword-only); rngs.params() seeds the path weights.	required

TensorProduct

Bases: Protocol

Protocol for an equivariant bilinear map of two :class:IrrepsArray s.

Concrete implementations (the dense Clebsch-Gordan :class:FullyConnectedTensorProduct, or a future Cartesian / cuEquivariance-backed variant) expose the same call signature so they are substitutable.

apply_scalar_weights

apply_scalar_weights(features: IrrepsArray, weights: Float[Array, '... num_irreps']) -> IrrepsArray

Scale each multiplicity of features by an invariant scalar weight.

Multiplying a steerable feature's multiplicity by a rotation-invariant scalar is equivariant, so this realises the radial / scalar-gated modulation used by the NequIP convolution and the equivariant Hamiltonian pair head (one weight per irrep, weights.shape[-1] == features.irreps.num_irreps).

Parameters:

Name	Type	Description	Default
features	IrrepsArray	The steerable feature to scale.	required
weights	Float[Array, '... num_irreps']	Per-multiplicity scalars of width features.irreps.num_irreps (multiplicity-major, matching the block order), broadcasting over the leading axes of features.	required

Returns:

Name	Type	Description
An	IrrepsArray	class:IrrepsArray with the same irreps as features, each
	IrrepsArray	multiplicity scaled by its weight.

inner_product

inner_product(x: IrrepsArray, y: IrrepsArray, *, normalize: bool = True) -> IrrepsArray

Return the per-multiplicity dot product of x and y as 0e scalars.

x and y must share their layout. Each mul x (l, p) block yields mul invariant scalars <x_u, y_u> (optionally divided by dim = 2l+1, the e3nn "component" path normalisation used by QHNet's InnerProduct).

Parameters:

Name	Type	Description	Default
x	IrrepsArray	First steerable feature.	required
y	IrrepsArray	Second steerable feature with y.irreps == x.irreps.	required
normalize	bool	If True (default) divide each block's dot product by its irrep dimension (component normalisation).	True

Returns:

Name	Type	Description
An	IrrepsArray	class:IrrepsArray of all-0e scalars, one per input multiplicity.

Raises:

Type	Description
ValueError	If x and y do not share the same irreps layout.

norm

norm(x: IrrepsArray, *, squared: bool = False) -> IrrepsArray

Return the per-irrep Euclidean norm of x as 0e scalars.

Each mul x (l, p) block contributes mul scalar norms, one per multiplicity, so the output layout replaces every block's irrep by 0e while keeping its multiplicity (e.g. 2x0e + 3x1o -> 2x0e + 3x0e).

Parameters:

Name	Type	Description	Default
x	IrrepsArray	Input steerable feature.	required
squared	bool	If True return the squared norm (no square root), which avoids the gradient singularity entirely.	False

Returns:

Name	Type	Description
An	IrrepsArray	class:IrrepsArray of all-0e scalars with the same leading shape
	IrrepsArray	as x and one scalar per input multiplicity.

rms_normalize

rms_normalize(x: IrrepsArray, *, eps: float = 1e-06) -> IrrepsArray

Scale each sample of x to unit root-mean-square (equivariant RMSNorm).

Divides the whole feature vector by the per-sample scalar sqrt(mean(x**2) + eps). That scalar is rotation-invariant (each irrep block's squared norm is preserved by rotation, so their mean is too), so dividing by it preserves the transformation law -- this is an equivariant analogue of RMSNorm. It bounds the feature magnitude before operations that square the features (the self / pair interaction tensor products), without which an unnormalised trunk's growing activations blow up the squared output.

Parameters:

Name	Type	Description	Default
x	IrrepsArray	Input steerable feature; the RMS is taken over its last (feature) axis.	required
eps	float	Numerical floor added to the mean square before the square root.	1e-06

Returns:

Name	Type	Description
An	IrrepsArray	class:IrrepsArray with the same layout as x, each sample scaled
	IrrepsArray	to unit RMS.

from_irreps_array

from_irreps_array(features: IrrepsArray) -> CartesianTensor

Convert an 1x0e + 1x1o + 1x2e IrrepsArray back to a Cartesian tensor.

Inverse of :func:to_irreps_array.

Parameters:

Name	Type	Description	Default
features	IrrepsArray	An :class:~opifex.neural.equivariant.IrrepsArray whose layout is exactly CARTESIAN_IRREPS.	required

Returns:

Type	Description
CartesianTensor	The reconstructed :class:CartesianTensor.

Raises:

Type	Description
ValueError	If features.irreps is not CARTESIAN_IRREPS.

to_irreps_array

to_irreps_array(tensor: CartesianTensor) -> IrrepsArray

Convert a :class:CartesianTensor to an 1x0e + 1x1o + 1x2e IrrepsArray.

The scalar (trace), vector (antisymmetric dual) and symmetric-traceless parts map to the l = 0 / l = 1 / l = 2 blocks respectively. The l = 2 block uses the change-of-basis derived from the canonical spherical-harmonic evaluator (see :func:_symmetric_to_l2_basis), so the result is exactly consistent with :func:opifex.geometry.algebra.wigner.wigner_d: a rotation of tensor rotates the blocks by D^l(R).

Parameters:

Name	Type	Description	Default
tensor	CartesianTensor	A rank-2 Cartesian tensor feature with leading shape (...,).	required

Returns:

Name	Type	Description
An	IrrepsArray	class:~opifex.neural.equivariant.IrrepsArray with layout
	IrrepsArray	CARTESIAN_IRREPS and array shape (..., 9).

radius_graph

radius_graph(positions: Float[Array, 'n 3'], cutoff: float, *, max_edges: int, self_loops: bool = False) -> tuple[Array, Array]

Build the radius-graph edge index for a point cloud.

Two nodes i and j are connected (with i the sender and j the receiver) when |positions[i] - positions[j]| < cutoff.

Parameters:

Name	Type	Description	Default
positions	Float[Array, 'n 3']	Node coordinates of shape (n, 3).	required
cutoff	float	Connection radius r_c (positive).	required
max_edges	int	Static upper bound on the number of returned edges; the output arrays have this length, padded with -1 for unused slots.	required
self_loops	bool	If True, include (i, i) edges. Default False.	False

Returns:

Type	Description
Array	A pair (senders, receivers) of integer arrays of shape
Array	(max_edges,); padding entries hold -1.

scatter_max

scatter_max(data: Float[Array, 'edges ...'], index: Array, *, num_segments: int) -> Float[Array, 'num_segments ...']

Segment-wise maximum of data addressed by index.

Parameters:

Name	Type	Description	Default
data	Float[Array, 'edges ...']	Per-edge values of shape (edges, ...).	required
index	Array	Destination node index per edge, shape (edges,).	required
num_segments	int	Number of output segments (nodes).	required

Returns:

Type	Description
Float[Array, 'num_segments ...']	Per-segment maximum of shape (num_segments, ...).

scatter_mean

scatter_mean(data: Float[Array, 'edges ...'], index: Array, *, num_segments: int) -> Float[Array, 'num_segments ...']

Average data within segments addressed by index.

Empty segments map to zero.

Parameters:

Name	Type	Description	Default
data	Float[Array, 'edges ...']	Per-edge values of shape (edges, ...).	required
index	Array	Destination node index per edge, shape (edges,).	required
num_segments	int	Number of output segments (nodes).	required

Returns:

Type	Description
Float[Array, 'num_segments ...']	Per-segment mean of shape (num_segments, ...).

scatter_sum

scatter_sum(data: Float[Array, 'edges ...'], index: Array, *, num_segments: int) -> Float[Array, 'num_segments ...']

Sum data into segments addressed by index.

output[index[e]] += data[e] for every edge e (negative indices, e.g. the -1 padding from :func:radius_graph, are dropped by segment_sum).

Parameters:

Name	Type	Description	Default
data	Float[Array, 'edges ...']	Per-edge values of shape (edges, ...).	required
index	Array	Destination node index per edge, shape (edges,).	required
num_segments	int	Number of output segments (nodes).	required

Returns:

Type	Description
Float[Array, 'num_segments ...']	Aggregated array of shape (num_segments, ...).

cosine_cutoff

cosine_cutoff(radius: Float[Array, ...], cutoff: float) -> Float[Array, ...]

Behler cosine cutoff envelope 0.5 (cos(pi r / r_c) + 1).

From Behler (J. Chem. Phys. 134, 074106, 2011), the original atom-centred symmetry-function cutoff. Decays smoothly from 1 at r = 0 to 0 at r = r_c (with vanishing first derivative there) and is masked to zero beyond r_c.

Parameters:

Name	Type	Description	Default
radius	Float[Array, ...]	Distances of shape (...).	required
cutoff	float	Cutoff radius r_c (positive).	required

Returns:

Type	Description
Float[Array, ...]	Envelope values of shape (...) in [0, 1].

polynomial_cutoff

polynomial_cutoff(radius: Float[Array, ...], cutoff: float, *, p: int = 6) -> Float[Array, ...]

Smooth polynomial cutoff envelope decaying from 1 to 0 on [0, r_c].

Ported from MACE ../mace/mace/modules/radial.py:113 (eq. 8 of arXiv:2206.07697); equivalent to e3nn_jax.poly_envelope. The envelope and its first p derivatives vanish at r_c::

f(r) = 1 - (p+1)(p+2)/2 (r/r_c)^p
         + p(p+2) (r/r_c)^{p+1}
         - p(p+1)/2 (r/r_c)^{p+2}

masked to zero for r >= r_c.

Parameters:

Name	Type	Description	Default
radius	Float[Array, ...]	Distances of shape (...).	required
cutoff	float	Cutoff radius r_c (positive).	required
p	int	Polynomial order controlling smoothness (default 6).	6

Returns:

Type	Description
Float[Array, ...]	Envelope values of shape (...) in [0, 1].

Spherical harmonics

Real spherical harmonics Y_l(r) for E(3)-equivariant networks.

A native, dependency-free port of the recursive spherical-harmonics algorithm of e3nn-jax (Geiger & Smidt 2022, arXiv:2207.09453; reference ../e3nn-jax/e3nn_jax/_src/spherical_harmonics/recursive.py). Higher degrees are built from lower ones by contracting two spherical-harmonic blocks with the real Clebsch-Gordan tensor::

Y_l = cste(l) / norm(l) * einsum('...i,...j,ijk->...k', Y_l1, Y_l2, C)

where l1 = l - 2**floor(log2(l-1)) and l2 = l - l1. The per-degree normalization constant norm(l) is computed numerically from the value of the contraction at the "north pole" index (rather than symbolically via sympy as in the reference), which is exact for the integer Clebsch-Gordan tables.

The output uses the same real-spherical-harmonic basis as :func:opifex.geometry.algebra.wigner.wigner_d, so equivariance Y_l(R r) = D^l(R) Y_l(r) holds. Y_0 is constant and Y_1 is proportional to the (normalized) direction. The implementation is jit/grad/vmap clean and handles batched inputs.

References

• ../e3nn-jax/e3nn_jax/_src/spherical_harmonics/recursive.py -- the recurrence, the l1/l2 split, and the integral / component normalization constants.
• ../e3nn-jax/e3nn_jax/_src/spherical_harmonics/__init__.py -- the normalize (project onto the sphere) and normalization options.

spherical_harmonics

spherical_harmonics(out: Irreps | int, vectors: Float[Array, '... 3'], *, normalize: bool = True, normalization: str = 'integral') -> IrrepsArray

Real spherical harmonics Y_l(r) for degrees l = 0..lmax.

The harmonics share the real basis of :func:opifex.geometry.algebra.wigner.wigner_d, so they are equivariant: Y_l(R r) = D^l(R) Y_l(r).

Parameters:

Name	Type	Description	Default
out	Irreps | int	Either an integer lmax (producing irreps 1x0e + 1x1o + ... + 1x{lmax}{parity} where each degree has the parity (-1)^l of a spherical harmonic) or an explicit :class:~opifex.neural.equivariant.Irreps layout.	required
vectors	Float[Array, '... 3']	Cartesian coordinates of shape (..., 3).	required
normalize	bool	If True (default), the input is projected onto the unit sphere before evaluating the polynomials.	True
normalization	str	One of "integral" (default), "component" or "norm" -- the convention for the per-degree scale.	'integral'

Returns:

Name	Type	Description
An	IrrepsArray	class:~opifex.neural.equivariant.IrrepsArray with the requested
	IrrepsArray	irreps and an array of shape (..., irreps.dim).

Raises:

Type	Description
ValueError	If normalization is unknown, vectors is not 3-dimensional in its last axis, or an explicit out layout is not a valid spherical-harmonic layout (multiplicity one, parity (-1)^l, ascending degrees).

Tensor products

Equivariant tensor product of two steerable features via Clebsch-Gordan.

The tensor product couples two :class:IrrepsArray s into a third by contracting each pair of input blocks with the real Clebsch-Gordan tensor. For input irreps ir_1 and ir_2, the selection rule |l_1 - l_2| <= l_3 <= l_1 + l_2 (with parity p_1 p_2) lists the allowed output irreps; the contraction einsum("...ui,...vj,ijk->...uvk", x_1, x_2, C) (with C the Clebsch-Gordan tensor) is equivariant by construction.

Ported from e3nn-jax (../e3nn-jax/e3nn_jax/_src/tensor_products.py:122 and the uvw connection of ../e3nn-jax/e3nn_jax/_src/legacy/core_tensor_product.py). The fully-connected variant gives every path (i_1, i_2) -> i_3 a learnable (mul_1, mul_2, mul_3) weight: einsum("uvw,ijk,...ui,...vj->...wk", w, C, x_1, x_2). Normalisation follows the e3nn default config (irrep_normalization="component", path_normalization="element"): the per-path scale alpha = dim(ir_3) / sum_paths(mul_1 mul_2) (cf. _normalize_instruction_path_weights) is folded into the weight initialisation standard deviation, so the forward pass uses the raw Clebsch-Gordan tensor.

TensorProduct

Bases: Protocol

Protocol for an equivariant bilinear map of two :class:IrrepsArray s.

Concrete implementations (the dense Clebsch-Gordan :class:FullyConnectedTensorProduct, or a future Cartesian / cuEquivariance-backed variant) expose the same call signature so they are substitutable.

FullyConnectedTensorProduct

FullyConnectedTensorProduct(irreps_in1: Irreps | str, irreps_in2: Irreps | str, irreps_out: Irreps | str, *, rngs: Rngs)

Bases: Module

Fully-connected Clebsch-Gordan tensor product with learnable path weights.

Every selection-rule-allowed path (i_1, i_2) -> i_3 whose output irrep is present in irreps_out carries an independent (mul_1, mul_2, mul_3) weight; the contributions are summed per output block.

Parameters:

Name	Type	Description	Default
irreps_in1	Irreps | str	Layout of the first input.	required
irreps_in2	Irreps | str	Layout of the second input.	required
irreps_out	Irreps | str	Desired output layout; only paths producing an irrep that appears here are instantiated.	required
rngs	Rngs	Random number generators (keyword-only); rngs.params() seeds the path weights.	required

ChannelwiseTensorProduct

ChannelwiseTensorProduct(irreps_in1: Irreps | str, irreps_in2: Irreps | str, irreps_out: Irreps | str, *, internal_weights: bool = True, rngs: Rngs)

Bases: Module

Channel-wise (e3nn "uuu") Clebsch-Gordan tensor product.

The two inputs and the output share one uniform multiplicity mul; each selection-rule-allowed path (i_1, i_2) -> i_3 whose output irrep is present in irreps_out couples the inputs channel-by-channel and carries a single (mul,) weight -- O(mul) parameters per path rather than the :class:FullyConnectedTensorProduct's O(mul^3). This is the connection QHNet's self / pair interaction layers use (../AIRS/OpenDFT/QHBench/QH9/ models/QHNet.py); the unweighted forward equals e3nn.elementwise_tensor_product.

The per-channel weights are the module's internal parameters by default, but a per-sample weights array (width :attr:weight_numel) can be supplied to :meth:__call__ -- the path QHNet uses to inject per-edge radial / inner-product modulation into the pair coupling.

Parameters:

Name	Type	Description	Default
irreps_in1	Irreps | str	First input layout (single uniform multiplicity).	required
irreps_in2	Irreps | str	Second input layout (same uniform multiplicity).	required
irreps_out	Irreps | str	Output layout; only paths producing an irrep present here are instantiated. Each output irrep should appear once.	required
internal_weights	bool	If True (default) allocate the per-channel path weights as module parameters; if False they must be supplied to :meth:__call__ every time.	True
rngs	Rngs	Random number generators (keyword-only); rngs.params() seeds the internal weights.	required

Raises:

Type	Description
ValueError	If the inputs and output do not share one multiplicity.

Radial bases and cutoffs

Radial bases and cutoff envelopes for E(3)-equivariant networks.

Interatomic-distance edge features are invariant scalars (they depend only on |r_i - r_j|), so these objects return plain arrays rather than :class:~opifex.neural.equivariant.IrrepsArray; downstream they are tagged as 0e channels. The formulae are ported from the MACE radial module (Batatia et al. 2022, "MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields", arXiv:2206.07697):

• :class:BesselBasis -- ../mace/mace/modules/radial.py:18 (eq. 7).
• :class:GaussianBasis -- ../mace/mace/modules/radial.py:88.
• :func:polynomial_cutoff -- ../mace/mace/modules/radial.py:113 (eq. 8), equivalent to e3nn_jax.poly_envelope.

The cosine envelope :func:cosine_cutoff follows Behler (J. Chem. Phys. 134, 074106, 2011), the original ACSF cutoff function.

BesselBasis

BesselBasis(num_basis: int, cutoff: float, *, rngs: Rngs | None = None)

Bases: Module

Bessel radial basis b_n(r) = sqrt(2/r_c) sin(n pi r / r_c) / r.

Ported from MACE ../mace/mace/modules/radial.py:18 (eq. 7 of arXiv:2206.07697). The frequencies n pi / r_c for n = 1..num_basis are stored as a (non-trainable) buffer; the basis is a smooth, complete set of invariant radial features.

Parameters:

Name	Type	Description	Default
num_basis	int	Number of Bessel functions n = 1..num_basis.	required
cutoff	float	Cutoff radius r_c (must be positive).	required
rngs	Rngs | None	Unused (the basis has no learnable parameters); accepted for a uniform constructor signature across equivariant modules.	None

GaussianBasis

GaussianBasis(num_basis: int, cutoff: float, *, rngs: Rngs | None = None)

Bases: Module

Gaussian radial basis g_n(r) = exp(-(r - mu_n)^2 / (2 sigma^2)).

Ported from MACE ../mace/mace/modules/radial.py:88. Centres mu_n are evenly spaced on [0, r_c] and the width is the centre spacing (sigma = r_c / (num_basis - 1)), giving the coefficient coeff = -0.5 / sigma^2.

Parameters:

Name	Type	Description	Default
num_basis	int	Number of Gaussians (centres on [0, r_c]); must be >= 2 so the spacing is well defined.	required
cutoff	float	Cutoff radius r_c (positive); sets the centre range.	required
rngs	Rngs | None	Unused (no learnable parameters); accepted for a uniform constructor signature.	None

PiecewiseLinearBasis

PiecewiseLinearBasis(num_basis: int, cutoff: float, *, rngs: Rngs | None = None)

Bases: Module

Piecewise-linear (hat) radial basis on [0, cutoff].

Faithful to the isotropic basis of torch_harmonics's PiecewiseLinearFilterBasis (torch-harmonics/torch_harmonics/filter_basis.py), the canonical filter basis for discrete-continuous (DISCO) convolutions (Ocampo, Price & McEwen 2023, arXiv:2209.13603). The num_basis hat functions have collocation spacing dr = 2 * cutoff / (num_basis + 1) and half-width dr: phi_k(r) = max(0, 1 - |r - r_k| / dr) restricted to r <= cutoff, with centres r_k = k * dr (odd num_basis) or (k + 0.5) * dr (even). Each function is continuous and compactly supported, so no separate cutoff envelope is needed.

Parameters:

Name	Type	Description	Default
num_basis	int	Number of hat functions (>= 1).	required
cutoff	float	Support radius r_c (positive).	required
rngs	Rngs | None	Unused (the basis is parameter-free); accepted for interface uniformity.	None

polynomial_cutoff

polynomial_cutoff(radius: Float[Array, ...], cutoff: float, *, p: int = 6) -> Float[Array, ...]

Smooth polynomial cutoff envelope decaying from 1 to 0 on [0, r_c].

Ported from MACE ../mace/mace/modules/radial.py:113 (eq. 8 of arXiv:2206.07697); equivalent to e3nn_jax.poly_envelope. The envelope and its first p derivatives vanish at r_c::

f(r) = 1 - (p+1)(p+2)/2 (r/r_c)^p
         + p(p+2) (r/r_c)^{p+1}
         - p(p+1)/2 (r/r_c)^{p+2}

masked to zero for r >= r_c.

Parameters:

Name	Type	Description	Default
radius	Float[Array, ...]	Distances of shape (...).	required
cutoff	float	Cutoff radius r_c (positive).	required
p	int	Polynomial order controlling smoothness (default 6).	6

Returns:

Type	Description
Float[Array, ...]	Envelope values of shape (...) in [0, 1].

cosine_cutoff

cosine_cutoff(radius: Float[Array, ...], cutoff: float) -> Float[Array, ...]

Behler cosine cutoff envelope 0.5 (cos(pi r / r_c) + 1).

From Behler (J. Chem. Phys. 134, 074106, 2011), the original atom-centred symmetry-function cutoff. Decays smoothly from 1 at r = 0 to 0 at r = r_c (with vanishing first derivative there) and is masked to zero beyond r_c.

Parameters:

Name	Type	Description	Default
radius	Float[Array, ...]	Distances of shape (...).	required
cutoff	float	Cutoff radius r_c (positive).	required

Returns:

Type	Description
Float[Array, ...]	Envelope values of shape (...) in [0, 1].

Graph utilities

Neighbour graph construction and segment-scatter for message passing.

These utilities turn a point cloud into the (senders, receivers) edge index used by equivariant message-passing networks, and aggregate per-edge messages back onto nodes.

• :func:radius_graph follows the dense pairwise-distance + fixed-size mask approach of ../e3nn-jax/e3nn_jax/_src/radius_graph.py (it uses jnp.where(mask, size=...) to return a statically shaped edge list). The neighbour-list concept is that of ../jax-md/jax_md/partition.py; unlike jax-md's cell-list partitioning, this implementation is a simple dense O(N^2) pairwise computation -- correct and jit-friendly for the small to medium molecules typical of interatomic-potential workloads, but not intended for very large N.

• :func:scatter_sum, :func:scatter_mean and :func:scatter_max wrap jax.ops.segment_* (cf. ../e3nn-jax/e3nn_jax/_src/scatter.py).

Static-shape contract: radius_graph returns edge arrays of fixed length max_edges (padded with -1), so the output shape does not depend on the data -- a requirement for jit. Callers must size max_edges to an upper bound on the true edge count; excess edges beyond max_edges are dropped.

radius_graph

radius_graph(positions: Float[Array, 'n 3'], cutoff: float, *, max_edges: int, self_loops: bool = False) -> tuple[Array, Array]

Build the radius-graph edge index for a point cloud.

Two nodes i and j are connected (with i the sender and j the receiver) when |positions[i] - positions[j]| < cutoff.

Parameters:

Name	Type	Description	Default
positions	Float[Array, 'n 3']	Node coordinates of shape (n, 3).	required
cutoff	float	Connection radius r_c (positive).	required
max_edges	int	Static upper bound on the number of returned edges; the output arrays have this length, padded with -1 for unused slots.	required
self_loops	bool	If True, include (i, i) edges. Default False.	False

Returns:

Type	Description
Array	A pair (senders, receivers) of integer arrays of shape
Array	(max_edges,); padding entries hold -1.

scatter_sum

scatter_sum(data: Float[Array, 'edges ...'], index: Array, *, num_segments: int) -> Float[Array, 'num_segments ...']

Sum data into segments addressed by index.

output[index[e]] += data[e] for every edge e (negative indices, e.g. the -1 padding from :func:radius_graph, are dropped by segment_sum).

Parameters:

Name	Type	Description	Default
data	Float[Array, 'edges ...']	Per-edge values of shape (edges, ...).	required
index	Array	Destination node index per edge, shape (edges,).	required
num_segments	int	Number of output segments (nodes).	required

Returns:

Type	Description
Float[Array, 'num_segments ...']	Aggregated array of shape (num_segments, ...).

scatter_mean

scatter_mean(data: Float[Array, 'edges ...'], index: Array, *, num_segments: int) -> Float[Array, 'num_segments ...']

Average data within segments addressed by index.

Empty segments map to zero.

Parameters:

Name	Type	Description	Default
data	Float[Array, 'edges ...']	Per-edge values of shape (edges, ...).	required
index	Array	Destination node index per edge, shape (edges,).	required
num_segments	int	Number of output segments (nodes).	required

Returns:

Type	Description
Float[Array, 'num_segments ...']	Per-segment mean of shape (num_segments, ...).

scatter_max

scatter_max(data: Float[Array, 'edges ...'], index: Array, *, num_segments: int) -> Float[Array, 'num_segments ...']

Segment-wise maximum of data addressed by index.

Parameters:

Name	Type	Description	Default
data	Float[Array, 'edges ...']	Per-edge values of shape (edges, ...).	required
index	Array	Destination node index per edge, shape (edges,).	required
num_segments	int	Number of output segments (nodes).	required

Returns:

Type	Description
Float[Array, 'num_segments ...']	Per-segment maximum of shape (num_segments, ...).