UMA / eSCN compatibility¶

e3nn_mlx exposes the factorized S2-grid surface that eSCN-style models — including FairChem’s UMA — build their spherical grid matrices from.

The compatibility surface¶

ToS2Grid and FromS2Grid expose two constant matrices each:

Attribute	Shape
`sha`	`(res_alpha, 2*lmax + 1)`
`shb`	`(2lmax + 1, res_beta, (lmax + 1)*2)`

sha is the real Fourier basis along longitude, ordered m = -l … +l, and is available directly as e3nn_mlx.o3.spherical_harmonics_alpha(). shb is the latitude factor, carrying the transform’s normalization.

The dense grid matrices are recovered with the contractions eSCN uses:

to_grid_matrix = mx.einsum("mbi,am->bai", to_grid.shb, to_grid.sha)
from_grid_matrix = mx.einsum("am,mbi->bai", from_grid.sha, from_grid.shb)

Both have shape (res_beta, res_alpha, (lmax + 1)**2).

These are constants, not parameters: to_grid.parameters() and from_grid.parameters() are both empty, and mx.compile works on either transform.

Grids where `res_alpha < 2*lmax + 1`¶

res_alpha may be smaller than 2*lmax + 1. UMA relies on this: it works in an m-truncated subspace, keeping only coefficients with |m| <= mmax, and sizes the longitude grid for mmax rather than for lmax.

The K16L6 configuration is the motivating case:

lmax, mmax = 6, 2
res_beta, res_alpha = 2 * (lmax + 1), 2 * mmax + 1   # 14, 5

to_grid = o3.ToS2Grid(lmax, (res_beta, res_alpha), normalization="integral")
from_grid = o3.FromS2Grid((res_beta, res_alpha), lmax, normalization="integral")

to_grid.sha.shape    # (5, 13)
to_grid.shb.shape    # (13, 14, 49)

What does not round-trip, by construction¶

Such a grid cannot invert the full (lmax + 1)**2 representation, and is not expected to. Five longitude samples resolve five independent Fourier modes, so a 14x5 grid cannot carry 49 independent coefficients. That is a property of sampling, not a limitation of this implementation.

The guarantee is on the truncated subspace instead: coefficients satisfying

|m| <= mmax

round-trip through the restricted matrices. For lmax=6, mmax=2 that is 29 of the 49 coefficients, and the measured round-trip error is order 1e-6.

This behaviour is intentional. tests/test_uma_s2_compat.py pins both halves: the truncated round-trip must hold, and the full-basis round-trip is not asserted.

How the factorization is built¶

The beta factor is obtained by projecting the canonical e3nn_mlx.o3.spherical_harmonics() onto the alpha basis, evaluated on a fully sampled reference longitude grid of 2*lmax + 1 points — independent of the requested res_alpha. On an undersampled grid the alpha modes alias onto one another, and projecting there would mix them arbitrarily.

Each coefficient belongs to exactly one m, so only that row of shb is populated. The sparsity is deliberate: it prevents a truncated grid from folding one m into another.

FromS2Grid builds its inverse by spherical quadrature rather than a matrix inverse, because no inverse exists for an undersampled grid. On a fully sampled grid the quadrature is exact and reproduces the previous behaviour.

What this library does not provide¶

UMA-specific machinery stays in the UMA port, not here: reduced (l,m) indexing (CoefficientMapping, to_m), SO(2) convolutions, Wigner caching, graph construction, model blocks, and checkpoint loading. e3nn_mlx supplies the generic mathematics — spherical harmonics, S2 transforms, sha/shb, irreps and tensor products.

Verification¶

The matrices are checked against upstream e3nn fixtures in tests/reference_data/uma_s2_reference.npz, generated by tests/reference_generation/generate_uma_s2.py. Agreement of dense_to and dense_from is the release gate, since those are what UMA consumes; the measured relative difference is order 1e-6, set by float32 evaluation of the harmonics.

Fixtures are generated in the project’s pinned reference environment — the same versions the [reference] extra declares — so they are reproducible from the repository alone:

uv venv .venv-reference --python 3.12
VIRTUAL_ENV=.venv-reference uv pip install \
    -r tests/reference_generation/requirements.txt
.venv-reference/bin/python tests/reference_generation/generate_uma_s2.py

tests/reference_data/uma_s2_manifest.json records the versions actually used, so a regeneration with anything else is visible.

Normalization is pinned absolutely¶

UMA uses integral throughout, but this transform serves component and norm as well, so the fixtures also cover those — including lmax_in != lmax, which drives the bandwidth correction applied when projecting.

Those cases compare the matrices themselves rather than a round trip, because a round trip cannot detect this class of error: a scaling mistake in synthesis that is mirrored in analysis cancels exactly. Dropping the bandwidth correction, or keying the per-degree scale off the wrong degree, both fail these tests.