From 2f356895b702877ecd8efc4e06747d619ac2ae16 Mon Sep 17 00:00:00 2001
From: "Rose K. Cersonsky" <47536110+rosecers@users.noreply.github.com>
Date: Sat, 6 Jun 2026 11:56:03 -0700
Subject: [PATCH 1/8] Updating some things that break metatrain

---
 .../ellipsoidal_density_projection.py         | 22 +++++---
 anisoap/representations/radial_basis.py       | 53 +++++++++++++------
 2 files changed, 52 insertions(+), 23 deletions(-)

diff --git a/anisoap/representations/ellipsoidal_density_projection.py b/anisoap/representations/ellipsoidal_density_projection.py
index cd67b87..ee0810e 100644
--- a/anisoap/representations/ellipsoidal_density_projection.py
+++ b/anisoap/representations/ellipsoidal_density_projection.py
@@ -740,12 +740,8 @@ def requested_neighbor_lists(self) -> List[Any]:
         try:
             from metatomic.torch import NeighborListOptions
         except Exception:
-            try:
-                from metatomic.torch.system import NeighborListOptions
-            except Exception as exc:  # pragma: no cover
-                raise ImportError(
-                    "metatomic.torch is required for requested_neighbor_lists"
-                ) from exc
+            from metatomic.torch.system import NeighborListOptions
+            
         if self._neighbor_list_options is None:
             try:
                 self._neighbor_list_options = NeighborListOptions(
@@ -1003,7 +999,14 @@ def power_spectrum_features_from_tensormap(
                 device=device,
                 dtype=torch.long,
             )
-            dense.index_copy_(0, rows, vals)
+
+            block_dense = dense[:, col : col + dim]
+            block_dense.index_copy_(0, rows, vals)
+            dense = torch.cat(
+                [dense[:, :col], block_dense, dense[:, col + dim :]],
+                dim=1,
+            )
+
             col += dim
 
         samples = (
@@ -1054,7 +1057,9 @@ def power_spectrum_features(
             nu2, aggregate_by_system=aggregate_by_system
         )
 
-    def power_spectrum_feature_tensor_map(self, **kwargs: Any) -> TensorMap:
+    def power_spectrum_feature_tensor_map(
+        self, *, normalize: bool = True, **kwargs: Any
+    ) -> TensorMap:
         """Return a single-block per-atom feature TensorMap for AniSOAP-BPNN.
 
         The block layout is ``samples=['system', 'atom']`` and
@@ -1080,6 +1085,7 @@ def power_spectrum_feature_tensor_map(self, **kwargs: Any) -> TensorMap:
             atom_indices=graph.atom_indices,
             rotations=graph.rotations,
             ellipsoid_lengths=graph.ellipsoid_lengths,
+            normalize=normalize,
         )
         features, _ = self.power_spectrum_features_from_tensormap(
             nu2, target_samples=target_samples
diff --git a/anisoap/representations/radial_basis.py b/anisoap/representations/radial_basis.py
index 4f9517d..db4d702 100644
--- a/anisoap/representations/radial_basis.py
+++ b/anisoap/representations/radial_basis.py
@@ -6,6 +6,7 @@
 
 import numpy as np
 import torch
+import math
 from metatensor import TensorMap
 from scipy.special import (
     gamma,
@@ -45,15 +46,12 @@ def inverse_matrix_sqrt(matrix: torch.Tensor, rcond=1e-8, tol=1e-3) -> torch.Ten
     # matrices. The old NumPy path used this only as a diagnostic; torch high-order
     # GTO slices can be badly conditioned even when the retained eigenspace is valid.
     if tol is not None:
-        try:
-            matrix2 = torch.linalg.pinv(result @ result)
-            err = torch.linalg.norm(matrix - matrix2)
-            if torch.isfinite(err) and err > tol:
-                raise ValueError(
-                    f"Incurred Numerical Imprecision {torch.linalg.norm(matrix-matrix2)= :.8f}"
-                )
-        except RuntimeError:
-            warnings.warn("Could not run inverse_matrix_sqrt reconstruction check")
+        matrix2 = torch.linalg.pinv(result @ result)
+        err = torch.linalg.norm(matrix - matrix2)
+        if torch.isfinite(err) and err > tol:
+            raise ValueError(
+                f"Incurred Numerical Imprecision {torch.linalg.norm(matrix-matrix2)= :.8f}"
+            )
 
     return result.to(dtype=original_dtype, device=original_device)
 
@@ -191,7 +189,16 @@ def gto_prefactor(n, sigma):
         The normalization constant
 
     """
-    return np.sqrt(1 / gto_square_norm(n, sigma))
+    n = torch.as_tensor(n, dtype=torch.float64)
+    sigma = torch.as_tensor(sigma, device=n.device, dtype=torch.float64)
+
+    log_square_norm = (
+        math.log(0.5)
+        + (2.0 * n + 3.0) * torch.log(sigma)
+        + torch.lgamma(n + 1.5)
+    )
+
+    return torch.exp(-0.5 * log_square_norm)
 
 
 def gto_overlap(n, m, sigma_n, sigma_m):
@@ -225,12 +232,28 @@ def gto_overlap(n, m, sigma_n, sigma_m):
         overlap of the two normalized GTOs
 
     """
-    N_n = gto_prefactor(n, sigma_n)
-    N_m = gto_prefactor(m, sigma_m)
-    n_eff = (n + m) / 2
-    sigma_eff = np.sqrt(2 * sigma_n**2 * sigma_m**2 / (sigma_n**2 + sigma_m**2))
-    return N_n * N_m * gto_square_norm(n_eff, sigma_eff)
+    n = torch.as_tensor(n, dtype=torch.float64)
+    m = torch.as_tensor(m, device=n.device, dtype=torch.float64)
+    sigma_n = torch.as_tensor(sigma_n, device=n.device, dtype=torch.float64)
+    sigma_m = torch.as_tensor(sigma_m, device=n.device, dtype=torch.float64)
+
+    prefactor_n = gto_prefactor(n, sigma_n)
+    prefactor_m = gto_prefactor(m, sigma_m)
+
+    n_eff = (n + m) / 2.0
+    sigma_eff = torch.sqrt(
+        2.0 * sigma_n**2 * sigma_m**2 / (sigma_n**2 + sigma_m**2)
+    )
+
+    log_overlap = (
+        torch.log(prefactor_n)
+        + torch.log(prefactor_m)
+        + math.log(0.5)
+        + (2.0 * n_eff + 3.0) * torch.log(sigma_eff)
+        + torch.lgamma(n_eff + 1.5)
+    )
 
+    return torch.exp(log_overlap)
 
 def monomial_square_norm(n, r_cut):
     """

From 5648d054b8fab5f1628dc26ad6ec25f34acb6d94 Mon Sep 17 00:00:00 2001
From: "Rose K. Cersonsky" <47536110+rosecers@users.noreply.github.com>
Date: Sat, 6 Jun 2026 12:03:17 -0700
Subject: [PATCH 2/8] Removing recursion

---
 .../ellipsoidal_density_projection.py         | 192 +++++++++++-------
 anisoap/representations/radial_basis.py       |  11 +-
 .../test_torch_correctness.py                 |   6 +-
 tests/unit-tests/test_batched_moments.py      |  78 +++++++
 4 files changed, 209 insertions(+), 78 deletions(-)
 create mode 100644 tests/unit-tests/test_batched_moments.py

diff --git a/anisoap/representations/ellipsoidal_density_projection.py b/anisoap/representations/ellipsoidal_density_projection.py
index ee0810e..2ec5395 100644
--- a/anisoap/representations/ellipsoidal_density_projection.py
+++ b/anisoap/representations/ellipsoidal_density_projection.py
@@ -41,84 +41,138 @@
 )
 
 
-def compute_moments(A: torch.Tensor, a: torch.Tensor, maxdeg: int) -> torch.Tensor:
-    r"""Differentiable trivariate Gaussian moments.
+def _moment_index_maps(maxdeg: int, device=None):
+    """Build compact monomial maps for all exponents with total degree <= maxdeg."""
+    exponents_list: List[Tuple[int, int, int]] = []
+    for degree in range(maxdeg + 1):
+        for n0 in range(degree + 1):
+            for n1 in range(degree + 1 - n0):
+                n2 = degree - n0 - n1
+                exponents_list.append((n0, n1, n2))
+
+    index = {exp: i for i, exp in enumerate(exponents_list)}
+    exponents = torch.tensor(exponents_list, device=device, dtype=torch.long)
+    degrees = exponents.sum(dim=1)
+
+    parent = torch.full((len(exponents_list),), -1, device=device, dtype=torch.long)
+    direction = torch.full((len(exponents_list),), -1, device=device, dtype=torch.long)
+    decrement = torch.full(
+        (len(exponents_list), 3), -1, device=device, dtype=torch.long
+    )
+
+    for i, (n0, n1, n2) in enumerate(exponents_list):
+        if n0 + n1 + n2 == 0:
+            continue
+        if n0 > 0:
+            k = 0
+            p = (n0 - 1, n1, n2)
+        elif n1 > 0:
+            k = 1
+            p = (n0, n1 - 1, n2)
+        else:
+            k = 2
+            p = (n0, n1, n2 - 1)
+
+        parent[i] = index[p]
+        direction[i] = k
+        p_list = list(p)
+        for j in range(3):
+            if p_list[j] > 0:
+                q = p_list.copy()
+                q[j] -= 1
+                decrement[i, j] = index[tuple(q)]
+
+    return exponents, degrees, parent, direction, decrement
+
+
+def compute_moments_batched(
+    A: torch.Tensor,
+    a: torch.Tensor,
+    maxdeg: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    r"""Batched unnormalized trivariate Gaussian raw moments.
 
     Computes moments of ``exp(-1/2 (x-a)^T A (x-a))`` up to total polynomial
-    degree ``maxdeg``. This replaces the original Rust ``compute_moments`` call
-    on the torch path.
+    degree ``maxdeg``. The result is compact: only exponents with total degree
+    <= maxdeg are stored.
+
+    Returns
+    -------
+    moments
+        Shape ``(batch, n_monomials)``.
+    exponents
+        Shape ``(n_monomials, 3)``; each row is ``(n0, n1, n2)``.
     """
     A = torch.as_tensor(A)
+    if A.ndim == 2:
+        A = A.reshape(1, 3, 3)
     a = torch.as_tensor(a, device=A.device, dtype=A.dtype)
-    if A.shape != (3, 3):
-        raise ValueError(f"A must have shape (3, 3), got {tuple(A.shape)}")
-    if a.shape != (3,):
-        raise ValueError(f"a must have shape (3,), got {tuple(a.shape)}")
+    if a.ndim == 1:
+        a = a.reshape(1, 3)
+    if A.shape[-2:] != (3, 3):
+        raise ValueError(f"A must have shape (..., 3, 3), got {tuple(A.shape)}")
+    if a.shape[-1] != 3:
+        raise ValueError(f"a must have shape (..., 3), got {tuple(a.shape)}")
+    if A.shape[0] != a.shape[0]:
+        raise ValueError("A and a must have the same batch dimension")
     if maxdeg < 0:
         raise ValueError("maxdeg must be non-negative")
 
-    device, dtype = A.device, A.dtype
+    device = A.device
+    dtype = A.dtype
+    batch = A.shape[0]
+    exponents, degrees, parent, direction, decrement = _moment_index_maps(
+        maxdeg, device=device
+    )
+
     cov = torch.linalg.inv(A)
-    norm = torch.as_tensor(
-        (2.0 * math.pi) ** 1.5, device=device, dtype=dtype
-    ) / torch.sqrt(torch.linalg.det(A))
-
-    # Store normalized raw moments first; multiply by the Gaussian integral at end.
-    M: Dict[Tuple[int, int, int], torch.Tensor] = {
-        (0, 0, 0): torch.ones((), device=device, dtype=dtype)
-    }
-    if maxdeg >= 1:
-        M[(1, 0, 0)] = a[0]
-        M[(0, 1, 0)] = a[1]
-        M[(0, 0, 1)] = a[2]
-    if maxdeg >= 2:
-        M[(2, 0, 0)] = cov[0, 0] + a[0] * a[0]
-        M[(0, 2, 0)] = cov[1, 1] + a[1] * a[1]
-        M[(0, 0, 2)] = cov[2, 2] + a[2] * a[2]
-        M[(1, 1, 0)] = cov[0, 1] + a[0] * a[1]
-        M[(0, 1, 1)] = cov[1, 2] + a[1] * a[2]
-        M[(1, 0, 1)] = cov[0, 2] + a[0] * a[2]
-
-    def get(i: int, j: int, k: int) -> torch.Tensor:
-        if i < 0 or j < 0 or k < 0 or i + j + k > maxdeg:
-            return torch.zeros((), device=device, dtype=dtype)
-        return M.get((i, j, k), torch.zeros((), device=device, dtype=dtype))
-
-    # Isserlis/Stein recurrence: E[X_p f(X)] = mu_p E[f] + sum_q Sigma_pq E[df/dx_q]
-    for degree in range(2, maxdeg):
-        updates: Dict[Tuple[int, int, int], torch.Tensor] = {}
-        for n0 in range(degree + 1):
-            for n1 in range(degree + 1 - n0):
-                n2 = degree - n0 - n1
-                base = get(n0, n1, n2)
-                updates[(n0 + 1, n1, n2)] = (
-                    a[0] * base
-                    + cov[0, 0] * n0 * get(n0 - 1, n1, n2)
-                    + cov[0, 1] * n1 * get(n0, n1 - 1, n2)
-                    + cov[0, 2] * n2 * get(n0, n1, n2 - 1)
+    sign, logabsdet = torch.linalg.slogdet(A)
+    if not bool((sign > 0).all()):
+        raise ValueError("Gaussian precision matrices must be positive definite")
+    norm = torch.exp(1.5 * math.log(2.0 * math.pi) - 0.5 * logabsdet)
+
+    moments = torch.zeros((batch, exponents.shape[0]), device=device, dtype=dtype)
+    moments[:, 0] = 1.0
+
+    for degree in range(1, maxdeg + 1):
+        ids = torch.nonzero(degrees == degree, as_tuple=False).reshape(-1)
+        p = parent[ids]
+        k = direction[ids]
+        values = a[:, k] * moments[:, p]
+        parent_exponents = exponents[p]
+
+        for j in range(3):
+            dec = decrement[ids, j]
+            valid = dec >= 0
+            if bool(valid.any()):
+                coeff = parent_exponents[valid, j].to(dtype=dtype)
+                values[:, valid] = values[:, valid] + (
+                    coeff.reshape(1, -1) * cov[:, k[valid], j] * moments[:, dec[valid]]
                 )
-                if n0 == 0:
-                    updates[(n0, n1 + 1, n2)] = (
-                        a[1] * base
-                        + cov[1, 0] * n0 * get(n0 - 1, n1, n2)
-                        + cov[1, 1] * n1 * get(n0, n1 - 1, n2)
-                        + cov[1, 2] * n2 * get(n0, n1, n2 - 1)
-                    )
-                    if n1 == 0:
-                        updates[(n0, n1, n2 + 1)] = (
-                            a[2] * base
-                            + cov[2, 0] * n0 * get(n0 - 1, n1, n2)
-                            + cov[2, 1] * n1 * get(n0, n1 - 1, n2)
-                            + cov[2, 2] * n2 * get(n0, n1, n2 - 1)
-                        )
-        M.update(updates)
-
-    out = torch.zeros((maxdeg + 1, maxdeg + 1, maxdeg + 1), device=device, dtype=dtype)
-    if M:
-        idx = torch.tensor(list(M.keys()), device=device, dtype=torch.long).T
-        vals = torch.stack(list(M.values())) * norm
-        out = out.index_put(tuple(idx), vals, accumulate=False)
-    return out
+
+        moments[:, ids] = values
+
+    return norm.reshape(-1, 1) * moments, exponents
+
+
+def _compact_moments_to_cube(
+    moments: torch.Tensor,
+    exponents: torch.Tensor,
+    maxdeg: int,
+) -> torch.Tensor:
+    cube = torch.zeros(
+        (moments.shape[0], maxdeg + 1, maxdeg + 1, maxdeg + 1),
+        device=moments.device,
+        dtype=moments.dtype,
+    )
+    cube[:, exponents[:, 0], exponents[:, 1], exponents[:, 2]] = moments
+    return cube
+
+
+def compute_moments(A: torch.Tensor, a: torch.Tensor, maxdeg: int) -> torch.Tensor:
+    r"""Compatibility wrapper returning the historical dense moment cube."""
+    moments, exponents = compute_moments_batched(A, a, maxdeg)
+    return _compact_moments_to_cube(moments, exponents, maxdeg)[0]
 
 
 @dataclass
@@ -741,7 +795,7 @@ def requested_neighbor_lists(self) -> List[Any]:
             from metatomic.torch import NeighborListOptions
         except Exception:
             from metatomic.torch.system import NeighborListOptions
-            
+
         if self._neighbor_list_options is None:
             try:
                 self._neighbor_list_options = NeighborListOptions(
diff --git a/anisoap/representations/radial_basis.py b/anisoap/representations/radial_basis.py
index db4d702..5770389 100644
--- a/anisoap/representations/radial_basis.py
+++ b/anisoap/representations/radial_basis.py
@@ -1,3 +1,4 @@
+import math
 import warnings
 from typing import (
     Any,
@@ -6,7 +7,6 @@
 
 import numpy as np
 import torch
-import math
 from metatensor import TensorMap
 from scipy.special import (
     gamma,
@@ -193,9 +193,7 @@ def gto_prefactor(n, sigma):
     sigma = torch.as_tensor(sigma, device=n.device, dtype=torch.float64)
 
     log_square_norm = (
-        math.log(0.5)
-        + (2.0 * n + 3.0) * torch.log(sigma)
-        + torch.lgamma(n + 1.5)
+        math.log(0.5) + (2.0 * n + 3.0) * torch.log(sigma) + torch.lgamma(n + 1.5)
     )
 
     return torch.exp(-0.5 * log_square_norm)
@@ -241,9 +239,7 @@ def gto_overlap(n, m, sigma_n, sigma_m):
     prefactor_m = gto_prefactor(m, sigma_m)
 
     n_eff = (n + m) / 2.0
-    sigma_eff = torch.sqrt(
-        2.0 * sigma_n**2 * sigma_m**2 / (sigma_n**2 + sigma_m**2)
-    )
+    sigma_eff = torch.sqrt(2.0 * sigma_n**2 * sigma_m**2 / (sigma_n**2 + sigma_m**2))
 
     log_overlap = (
         torch.log(prefactor_n)
@@ -255,6 +251,7 @@ def gto_overlap(n, m, sigma_n, sigma_m):
 
     return torch.exp(log_overlap)
 
+
 def monomial_square_norm(n, r_cut):
     """
     Compute the square norm of monomials (inner product of itself over R^3).
diff --git a/tests/integration-tests/test_torch_correctness.py b/tests/integration-tests/test_torch_correctness.py
index abc7777..9c6f3dc 100644
--- a/tests/integration-tests/test_torch_correctness.py
+++ b/tests/integration-tests/test_torch_correctness.py
@@ -1,4 +1,5 @@
 import metatensor.torch as mts
+from pathlib import Path
 import torch
 import pytest
 import numpy as np
@@ -19,7 +20,8 @@ def test_benzene_correctness_5frames(self):
         """
         lmax = 9
         nmax = 6
-        frames = read("./notebooks/ellipsoids.xyz", ":5")
+        repo_root = Path(__file__).resolve().parents[1]
+        frames = read(repo_root / "../notebooks" / "ellipsoids.xyz", ":5")
 
         a1, a2, a3 = 4.0, 4.0, 0.5
         for frame in frames:
@@ -44,6 +46,6 @@ def test_benzene_correctness_5frames(self):
         x_anisoap_torch = calculator.power_spectrum(frames)
 
         x_anisoap_numpy = np.load(
-            "./tests/integration-tests/benzene_numpy_impl_5frames.npy"
+            repo_root / "integration-tests/benzene_numpy_impl_5frames.npy"
         )
         assert_allclose(x_anisoap_torch, x_anisoap_numpy, rtol=0, atol=1e-2)
diff --git a/tests/unit-tests/test_batched_moments.py b/tests/unit-tests/test_batched_moments.py
new file mode 100644
index 0000000..7a01fb5
--- /dev/null
+++ b/tests/unit-tests/test_batched_moments.py
@@ -0,0 +1,78 @@
+import torch
+
+from anisoap.representations.ellipsoidal_density_projection import (
+    compute_moments,
+    compute_moments_batched,
+)
+
+
+def test_compute_moments_batched_matches_single():
+    dtype = torch.float64
+
+    A = torch.tensor(
+        [
+            [[2.0, 0.1, 0.0], [0.1, 1.7, 0.2], [0.0, 0.2, 1.4]],
+            [[1.5, -0.05, 0.1], [-0.05, 2.2, 0.0], [0.1, 0.0, 1.8]],
+        ],
+        dtype=dtype,
+    )
+    centers = torch.tensor(
+        [
+            [0.2, -0.1, 0.4],
+            [-0.3, 0.5, 0.1],
+        ],
+        dtype=dtype,
+    )
+    maxdeg = 5
+
+    batched, exponents = compute_moments_batched(A, centers, maxdeg)
+
+    for i in range(A.shape[0]):
+        single_cube = compute_moments(A[i], centers[i], maxdeg)
+        expected = torch.stack(
+            [
+                single_cube[int(px), int(py), int(pz)]
+                for px, py, pz in exponents.detach().cpu().tolist()
+            ]
+        ).to(dtype=dtype)
+
+        torch.testing.assert_close(
+            batched[i],
+            expected,
+            rtol=1e-10,
+            atol=1e-10,
+        )
+
+
+def test_compute_moments_batched_has_gradients():
+    dtype = torch.float64
+
+    raw = torch.tensor(
+        [
+            [[1.5, 0.1, 0.0], [0.1, 1.4, 0.2], [0.0, 0.2, 1.8]],
+            [[1.7, -0.1, 0.1], [-0.1, 1.9, 0.0], [0.1, 0.0, 1.6]],
+        ],
+        dtype=dtype,
+        requires_grad=True,
+    )
+    A = raw @ raw.transpose(-1, -2) + 0.5 * torch.eye(3, dtype=dtype)
+
+    centers = torch.tensor(
+        [
+            [0.2, -0.1, 0.4],
+            [-0.3, 0.5, 0.1],
+        ],
+        dtype=dtype,
+        requires_grad=True,
+    )
+
+    moments, _ = compute_moments_batched(A, centers, maxdeg=5)
+    loss = moments.square().sum()
+    loss.backward()
+
+    assert raw.grad is not None
+    assert centers.grad is not None
+    assert torch.isfinite(raw.grad).all()
+    assert torch.isfinite(centers.grad).all()
+    assert raw.grad.abs().sum() > 0
+    assert centers.grad.abs().sum() > 0

From 4b373dd1789cd286541a2bcf6aa26b7ee6833653 Mon Sep 17 00:00:00 2001
From: "Rose K. Cersonsky" <47536110+rosecers@users.noreply.github.com>
Date: Sat, 6 Jun 2026 13:08:42 -0700
Subject: [PATCH 3/8] Adding a benchmark

---
 benchmarks/benchmark_moment_generation.py | 277 ++++++++++++++++++++++
 1 file changed, 277 insertions(+)
 create mode 100644 benchmarks/benchmark_moment_generation.py

diff --git a/benchmarks/benchmark_moment_generation.py b/benchmarks/benchmark_moment_generation.py
new file mode 100644
index 0000000..ad61c12
--- /dev/null
+++ b/benchmarks/benchmark_moment_generation.py
@@ -0,0 +1,277 @@
+#!/usr/bin/env python
+"""Benchmark AniSOAP Gaussian moment generation schemes.
+
+Run from the AniSOAP repo root, or from any environment where ``anisoap`` is
+importable:
+
+    python benchmark_moment_generation.py --device cpu --dtype float64
+    python benchmark_moment_generation.py --device cuda --dtype float32
+
+This compares:
+  1. legacy/single ``compute_moments`` called once per Gaussian;
+  2. batched ``compute_moments_batched``;
+  3. optionally ``torch.compile(compute_moments_batched)``.
+
+The benchmark also checks numerical agreement between legacy and batched output
+for a small subset of Gaussians.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import math
+import statistics
+import time
+from typing import Callable, Optional, Tuple
+
+import torch
+
+from anisoap.representations.ellipsoidal_density_projection import compute_moments
+
+try:
+    from anisoap.representations.ellipsoidal_density_projection import (
+        compute_moments_batched,
+    )
+except ImportError as exc:  # pragma: no cover
+    raise ImportError(
+        "Could not import compute_moments_batched. Apply the batched-moments "
+        "patch first."
+    ) from exc
+
+
+def parse_dtype(name: str) -> torch.dtype:
+    if name == "float32":
+        return torch.float32
+    if name == "float64":
+        return torch.float64
+    raise ValueError(f"Unsupported dtype: {name}")
+
+
+def make_spd_precision_matrices(
+    n: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int,
+) -> torch.Tensor:
+    gen = torch.Generator(device="cpu")
+    gen.manual_seed(seed)
+
+    raw = torch.randn((n, 3, 3), generator=gen, dtype=dtype).to(device)
+    eye = torch.eye(3, device=device, dtype=dtype).expand(n, 3, 3)
+
+    # SPD with moderate condition numbers; this avoids benchmarking numerical
+    # pathologies instead of moment generation.
+    A = raw @ raw.transpose(-1, -2) + 0.75 * eye
+    return 0.5 * (A + A.transpose(-1, -2))
+
+
+def make_centers(
+    n: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int,
+) -> torch.Tensor:
+    gen = torch.Generator(device="cpu")
+    gen.manual_seed(seed + 1)
+    return torch.randn((n, 3), generator=gen, dtype=dtype).to(device) * 0.5
+
+
+def sync(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+
+def time_callable(
+    fn: Callable[[], object],
+    *,
+    device: torch.device,
+    warmup: int,
+    repeats: int,
+) -> Tuple[float, float, float]:
+    for _ in range(warmup):
+        fn()
+    sync(device)
+
+    times = []
+    for _ in range(repeats):
+        gc.collect()
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
+        sync(device)
+        t0 = time.perf_counter()
+        fn()
+        sync(device)
+        times.append(time.perf_counter() - t0)
+
+    return min(times), statistics.median(times), max(times)
+
+
+def legacy_loop(A: torch.Tensor, centers: torch.Tensor, maxdeg: int) -> torch.Tensor:
+    # compute_moments returns a dense cube per Gaussian. This intentionally uses
+    # the old public interface to measure the current per-Gaussian overhead.
+    chunks = []
+    for i in range(A.shape[0]):
+        chunks.append(compute_moments(A[i], centers[i], maxdeg).reshape(1, -1))
+    return torch.cat(chunks, dim=0)
+
+
+def legacy_loop_valid_only(
+    A: torch.Tensor,
+    centers: torch.Tensor,
+    maxdeg: int,
+    exponents: torch.Tensor,
+) -> torch.Tensor:
+    cubes = legacy_loop(A, centers, maxdeg)
+    side = maxdeg + 1
+    linear = exponents[:, 0] * side * side + exponents[:, 1] * side + exponents[:, 2]
+    return cubes[:, linear]
+
+
+def check_correctness(
+    A: torch.Tensor,
+    centers: torch.Tensor,
+    maxdeg: int,
+    n_check: int,
+    rtol: float,
+    atol: float,
+) -> None:
+    n_check = min(n_check, A.shape[0])
+    A_small = A[:n_check]
+    c_small = centers[:n_check]
+
+    batched, exponents = compute_moments_batched(A_small, c_small, maxdeg)
+    legacy = legacy_loop_valid_only(A_small, c_small, maxdeg, exponents)
+
+    torch.testing.assert_close(batched, legacy, rtol=rtol, atol=atol)
+
+
+def bytes_for(t: torch.Tensor) -> int:
+    return t.numel() * t.element_size()
+
+
+def maybe_compile(fn: Callable) -> Optional[Callable]:
+    compile_fn = getattr(torch, "compile", None)
+    if compile_fn is None:
+        return None
+    try:
+        return compile_fn(fn, fullgraph=False)
+    except Exception:
+        return None
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n", type=int, default=512, help="number of Gaussians")
+    parser.add_argument("--maxdeg", type=int, default=8)
+    parser.add_argument("--device", type=str, default="cpu")
+    parser.add_argument("--dtype", type=str, default="float64", choices=["float32", "float64"])
+    parser.add_argument("--warmup", type=int, default=3)
+    parser.add_argument("--repeats", type=int, default=10)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--check", type=int, default=8, help="number of entries used for correctness check")
+    parser.add_argument("--no-compile", action="store_true")
+    parser.add_argument("--skip-legacy", action="store_true", help="skip slow legacy loop timing")
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+    dtype = parse_dtype(args.dtype)
+
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA requested but torch.cuda.is_available() is False")
+
+    A = make_spd_precision_matrices(args.n, device=device, dtype=dtype, seed=args.seed)
+    centers = make_centers(args.n, device=device, dtype=dtype, seed=args.seed)
+
+    print("AniSOAP moment generation benchmark")
+    print(f"  n        = {args.n}")
+    print(f"  maxdeg   = {args.maxdeg}")
+    print(f"  device   = {device}")
+    print(f"  dtype    = {dtype}")
+    print(f"  repeats  = {args.repeats}")
+    print()
+
+    # Correctness check uses valid monomials only.
+    check_correctness(
+        A,
+        centers,
+        args.maxdeg,
+        n_check=args.check,
+        rtol=5e-5 if dtype is torch.float32 else 1e-10,
+        atol=5e-6 if dtype is torch.float32 else 1e-10,
+    )
+    print(f"correctness: batched matches legacy on first {min(args.check, args.n)} Gaussians")
+    print()
+
+    results = []
+
+    batched_out, exponents = compute_moments_batched(A, centers, args.maxdeg)
+    sync(device)
+    n_valid = exponents.shape[0]
+    dense_cube = (args.maxdeg + 1) ** 3
+    print(f"valid monomials = {n_valid}")
+    print(f"dense cube size = {dense_cube}")
+    print(f"batched output memory = {bytes_for(batched_out) / 1024**2:.3f} MiB")
+    print()
+
+    def batched_fn() -> torch.Tensor:
+        out, _ = compute_moments_batched(A, centers, args.maxdeg)
+        return out
+
+    t_min, t_med, t_max = time_callable(
+        batched_fn, device=device, warmup=args.warmup, repeats=args.repeats
+    )
+    results.append(("batched", t_min, t_med, t_max))
+
+    if not args.no_compile:
+        compiled_batched = maybe_compile(compute_moments_batched)
+        if compiled_batched is not None:
+            def compiled_fn() -> torch.Tensor:
+                out, _ = compiled_batched(A, centers, args.maxdeg)
+                return out
+
+            t_min, t_med, t_max = time_callable(
+                compiled_fn, device=device, warmup=args.warmup, repeats=args.repeats
+            )
+            results.append(("batched torch.compile", t_min, t_med, t_max))
+        else:
+            print("torch.compile unavailable or failed to initialize; skipping compiled benchmark")
+            print()
+
+    if not args.skip_legacy:
+        def legacy_fn() -> torch.Tensor:
+            return legacy_loop_valid_only(A, centers, args.maxdeg, exponents)
+
+        t_min, t_med, t_max = time_callable(
+            legacy_fn, device=device, warmup=max(1, args.warmup // 2), repeats=args.repeats
+        )
+        results.append(("legacy loop", t_min, t_med, t_max))
+
+    print("timings")
+    print("  scheme                    min [ms]   median [ms]   max [ms]   speedup vs legacy")
+    legacy_med = None
+    for name, _, med, _ in results:
+        if name == "legacy loop":
+            legacy_med = med
+            break
+
+    for name, t_min, t_med, t_max in results:
+        if legacy_med is None or name == "legacy loop":
+            speedup = "--"
+        else:
+            speedup = f"{legacy_med / t_med:8.2f}x"
+        print(
+            f"  {name:<24} {1e3*t_min:9.3f} {1e3*t_med:12.3f} {1e3*t_max:9.3f}   {speedup}"
+        )
+
+    print()
+    print("notes")
+    print("  - legacy loop measures public compute_moments called once per Gaussian")
+    print("  - batched output stores only valid total-degree monomials")
+    print("  - use --skip-legacy for large n/maxdeg where the old loop is too slow")
+
+
+if __name__ == "__main__":
+    main()

From 0454dcda86140cc4f0afaf0abd5793b1c6da4ba2 Mon Sep 17 00:00:00 2001
From: "Rose K. Cersonsky" <47536110+rosecers@users.noreply.github.com>
Date: Sat, 6 Jun 2026 13:18:58 -0700
Subject: [PATCH 4/8] linters

---
 benchmarks/benchmark_moment_generation.py | 34 ++++++++++++++++++-----
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/benchmarks/benchmark_moment_generation.py b/benchmarks/benchmark_moment_generation.py
index ad61c12..de90569 100644
--- a/benchmarks/benchmark_moment_generation.py
+++ b/benchmarks/benchmark_moment_generation.py
@@ -167,13 +167,22 @@ def main() -> None:
     parser.add_argument("--n", type=int, default=512, help="number of Gaussians")
     parser.add_argument("--maxdeg", type=int, default=8)
     parser.add_argument("--device", type=str, default="cpu")
-    parser.add_argument("--dtype", type=str, default="float64", choices=["float32", "float64"])
+    parser.add_argument(
+        "--dtype", type=str, default="float64", choices=["float32", "float64"]
+    )
     parser.add_argument("--warmup", type=int, default=3)
     parser.add_argument("--repeats", type=int, default=10)
     parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--check", type=int, default=8, help="number of entries used for correctness check")
+    parser.add_argument(
+        "--check",
+        type=int,
+        default=8,
+        help="number of entries used for correctness check",
+    )
     parser.add_argument("--no-compile", action="store_true")
-    parser.add_argument("--skip-legacy", action="store_true", help="skip slow legacy loop timing")
+    parser.add_argument(
+        "--skip-legacy", action="store_true", help="skip slow legacy loop timing"
+    )
     args = parser.parse_args()
 
     device = torch.device(args.device)
@@ -202,7 +211,9 @@ def main() -> None:
         rtol=5e-5 if dtype is torch.float32 else 1e-10,
         atol=5e-6 if dtype is torch.float32 else 1e-10,
     )
-    print(f"correctness: batched matches legacy on first {min(args.check, args.n)} Gaussians")
+    print(
+        f"correctness: batched matches legacy on first {min(args.check, args.n)} Gaussians"
+    )
     print()
 
     results = []
@@ -228,6 +239,7 @@ def batched_fn() -> torch.Tensor:
     if not args.no_compile:
         compiled_batched = maybe_compile(compute_moments_batched)
         if compiled_batched is not None:
+
             def compiled_fn() -> torch.Tensor:
                 out, _ = compiled_batched(A, centers, args.maxdeg)
                 return out
@@ -237,20 +249,28 @@ def compiled_fn() -> torch.Tensor:
             )
             results.append(("batched torch.compile", t_min, t_med, t_max))
         else:
-            print("torch.compile unavailable or failed to initialize; skipping compiled benchmark")
+            print(
+                "torch.compile unavailable or failed to initialize; skipping compiled benchmark"
+            )
             print()
 
     if not args.skip_legacy:
+
         def legacy_fn() -> torch.Tensor:
             return legacy_loop_valid_only(A, centers, args.maxdeg, exponents)
 
         t_min, t_med, t_max = time_callable(
-            legacy_fn, device=device, warmup=max(1, args.warmup // 2), repeats=args.repeats
+            legacy_fn,
+            device=device,
+            warmup=max(1, args.warmup // 2),
+            repeats=args.repeats,
         )
         results.append(("legacy loop", t_min, t_med, t_max))
 
     print("timings")
-    print("  scheme                    min [ms]   median [ms]   max [ms]   speedup vs legacy")
+    print(
+        "  scheme                    min [ms]   median [ms]   max [ms]   speedup vs legacy"
+    )
     legacy_med = None
     for name, _, med, _ in results:
         if name == "legacy loop":

From b11e898d5f5931068bbe80995ee3c3e53d37d9a2 Mon Sep 17 00:00:00 2001
From: "Rose K. Cersonsky" <47536110+rosecers@users.noreply.github.com>
Date: Mon, 8 Jun 2026 11:34:57 -0500
Subject: [PATCH 5/8] adapting to metatran

---
 .../ellipsoidal_density_projection.py         | 172 ++++++++++++++++--
 1 file changed, 161 insertions(+), 11 deletions(-)

diff --git a/anisoap/representations/ellipsoidal_density_projection.py b/anisoap/representations/ellipsoidal_density_projection.py
index 2ec5395..0ebf383 100644
--- a/anisoap/representations/ellipsoidal_density_projection.py
+++ b/anisoap/representations/ellipsoidal_density_projection.py
@@ -581,10 +581,15 @@ def pairwise_ellip_expansion(
                 )
                 keys.append(key)
 
+    if len(keys) == 0:
+        key_values = torch.empty((0, 3), device=device, dtype=torch.int32)
+    else:
+        key_values = torch.as_tensor(keys, device=device, dtype=torch.int32)
+
     return TensorMap(
         keys=Labels(
             ["types_center", "types_neighbor", "angular_channel"],
-            torch.as_tensor(keys, device=device, dtype=torch.int32),
+            key_values,
         ),
         blocks=blocks,
     )
@@ -1111,6 +1116,30 @@ def power_spectrum_features(
             nu2, aggregate_by_system=aggregate_by_system
         )
 
+    def _feature_size(self, *, device=None, dtype=None) -> int:
+        """Infer the dense power-spectrum feature dimension from one dummy edge."""
+        if self.shape is not None:
+            return int(self.shape)
+
+        if dtype is None:
+            dtype = self.dtype
+        if device is None:
+            device = torch.device("cpu")
+
+        center_type = self.species[0] if self.species is not None else 0
+        dummy_features = self.power_spectrum_feature_tensor_map(
+            R_ij=torch.zeros((1, 3), device=device, dtype=dtype),
+            centers=torch.tensor([0], device=device, dtype=torch.long),
+            neighbors=torch.tensor([0], device=device, dtype=torch.long),
+            species=torch.tensor([center_type], device=device, dtype=torch.long),
+            structures=torch.tensor([0], device=device, dtype=torch.long),
+            atom_indices=torch.tensor([0], device=device, dtype=torch.long),
+            rotations=torch.eye(3, device=device, dtype=dtype).reshape(1, 3, 3),
+            ellipsoid_lengths=torch.ones((1, 3), device=device, dtype=dtype),
+        )
+        self.shape = int(dummy_features.block(0).values.shape[1])
+        return self.shape
+
     def power_spectrum_feature_tensor_map(
         self, *, normalize: bool = True, **kwargs: Any
     ) -> TensorMap:
@@ -1128,6 +1157,68 @@ def power_spectrum_feature_tensor_map(
                 dim=1,
             ).to(device=graph.R_ij.device, dtype=torch.int32),
         )
+
+        if graph.R_ij.shape[0] == 0:
+            if self.shape is None:
+                self.shape = self._feature_size()
+
+            all_species = (
+                self.species
+                if self.species is not None
+                else sorted(
+                    int(x) for x in torch.unique(graph.species).detach().cpu().tolist()
+                )
+            )
+
+            blocks = []
+            keys = []
+
+            for center_type in all_species:
+                mask = graph.species == int(center_type)
+                if not bool(mask.any()):
+                    continue
+
+                sample_values = torch.stack(
+                    [
+                        graph.structures[mask].to(
+                            device=graph.R_ij.device, dtype=torch.int32
+                        ),
+                        graph.atom_indices[mask].to(
+                            device=graph.R_ij.device, dtype=torch.int32
+                        ),
+                    ],
+                    dim=1,
+                )
+
+                blocks.append(
+                    TensorBlock(
+                        values=torch.zeros(
+                            (sample_values.shape[0], self.shape),
+                            device=graph.R_ij.device,
+                            dtype=graph.R_ij.dtype,
+                        ),
+                        samples=Labels(["system", "atom"], sample_values),
+                        components=[],
+                        properties=Labels(
+                            ["property"],
+                            torch.arange(
+                                self.shape,
+                                device=graph.R_ij.device,
+                                dtype=torch.int32,
+                            ).reshape(-1, 1),
+                        ),
+                    )
+                )
+                keys.append((int(center_type),))
+
+            return TensorMap(
+                keys=Labels(
+                    ["center_type"],
+                    torch.as_tensor(keys, device=graph.R_ij.device, dtype=torch.int32),
+                ),
+                blocks=blocks,
+            )
+
         # Reuse graph tensors to avoid reconstructing systems/frames.
         nu2 = self.power_spectrum(
             mean_over_samples=False,
@@ -1146,26 +1237,85 @@ def power_spectrum_feature_tensor_map(
         )
 
         self.shape = int(features.shape[1])
-        return TensorMap(
-            keys=Labels(
-                ["_"], torch.tensor([[0]], device=features.device, dtype=torch.int32)
-            ),
-            blocks=[
+
+        blocks = []
+        keys = []
+
+        all_species = (
+            self.species
+            if self.species is not None
+            else sorted(
+                int(x) for x in torch.unique(graph.species).detach().cpu().tolist()
+            )
+        )
+
+        for center_type in all_species:
+            mask = graph.species == int(center_type)
+            if not bool(mask.any()):
+                continue
+
+            sample_values = torch.stack(
+                [
+                    graph.structures[mask].to(
+                        device=features.device, dtype=torch.int32
+                    ),
+                    graph.atom_indices[mask].to(
+                        device=features.device, dtype=torch.int32
+                    ),
+                ],
+                dim=1,
+            )
+
+            target_rows = {
+                tuple(int(v) for v in row.detach().cpu().tolist()): idx
+                for idx, row in enumerate(target_samples.values)
+            }
+            row_indices = torch.tensor(
+                [
+                    target_rows[tuple(int(v) for v in row.detach().cpu().tolist())]
+                    for row in sample_values
+                ],
+                device=features.device,
+                dtype=torch.long,
+            )
+
+            blocks.append(
                 TensorBlock(
-                    values=features,
-                    samples=target_samples,
+                    values=features.index_select(0, row_indices),
+                    samples=Labels(["system", "atom"], sample_values),
                     components=[],
                     properties=Labels(
                         ["property"],
                         torch.arange(
-                            features.shape[1], device=features.device, dtype=torch.int32
+                            features.shape[1],
+                            device=features.device,
+                            dtype=torch.int32,
                         ).reshape(-1, 1),
                     ),
                 )
-            ],
+            )
+            keys.append((int(center_type),))
+
+        return TensorMap(
+            keys=Labels(
+                ["center_type"],
+                torch.as_tensor(keys, device=features.device, dtype=torch.int32),
+            ),
+            blocks=blocks,
         )
 
-    def forward(self, **kwargs: Any) -> TensorMap:
+    def forward(
+        self,
+        R_ij: torch.Tensor,
+        centers: torch.Tensor,
+        neighbors: torch.Tensor,
+        species: torch.Tensor,
+        structures: torch.Tensor,
+        atom_indices: torch.Tensor,
+        rotations: torch.Tensor,
+        ellipsoid_lengths: torch.Tensor,
+        normalize: bool = True,
+    ) -> TensorMap:
         """Default module output for AniSOAP-BPNN: per-atom scalar feature map."""
         return self.power_spectrum_feature_tensor_map(**kwargs)
 

From 734d81a27477b8bed302ab285c2b13185832c893 Mon Sep 17 00:00:00 2001
From: "Rose K. Cersonsky" <47536110+rosecers@users.noreply.github.com>
Date: Mon, 8 Jun 2026 11:37:02 -0500
Subject: [PATCH 6/8] adapting to metatrain

---
 .../test_torch_gradient_and_bpnn_adapter.py   | 34 ++++++++++++++-----
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/tests/unit-tests/test_torch_gradient_and_bpnn_adapter.py b/tests/unit-tests/test_torch_gradient_and_bpnn_adapter.py
index dd1ff57..f96616d 100644
--- a/tests/unit-tests/test_torch_gradient_and_bpnn_adapter.py
+++ b/tests/unit-tests/test_torch_gradient_and_bpnn_adapter.py
@@ -173,20 +173,38 @@ def test_anisoap_bpnn_feature_tensormap_layout_and_gradients():
     assert hasattr(features, "keys")
     assert hasattr(features, "blocks")
     assert hasattr(features, "block")
-    assert len(features.blocks()) == 1
+    assert list(features.keys.names) == ["center_type"]
+    assert features.keys.values.ndim == 2
+    assert features.keys.values.shape[1] == 1
 
-    assert list(features.keys.names) == ["_"]
-    assert features.keys.values.shape == (1, 1)
+    all_values = []
+    all_samples = []
 
-    block = features.block(0)
+    for _, block in features.items():
+        assert list(block.samples.names) == ["system", "atom"]
+        assert block.components == []
+        assert list(block.properties.names) == ["property"]
 
-    assert list(block.samples.names) == ["system", "atom"]
-    assert block.samples.values.shape == (2, 2)
+        all_values.append(block.values)
+        all_samples.append(block.samples.values)
+
+    samples = torch.cat(all_samples, dim=0)
+    values = torch.cat(all_values, dim=0)
+
+    assert samples.shape == (2, 2)
     assert torch.equal(
-        block.samples.values,
-        torch.tensor([[0, 0], [0, 1]], dtype=torch.int32),
+        samples,
+        torch.tensor([[0, 0], [0, 1]], dtype=torch.int32, device=samples.device),
     )
 
+    assert values.ndim == 2
+    assert values.shape[0] == 2
+    assert values.shape[1] > 0
+    assert calc.shape == values.shape[1]
+    assert torch.is_tensor(values)
+    assert values.requires_grad
+    assert torch.isfinite(values).all()
+
     assert block.components == []
     assert list(block.properties.names) == ["property"]
 

From 20c72e6d98b8acff60afbf1c8f8331211e4736a4 Mon Sep 17 00:00:00 2001
From: "Rose K. Cersonsky" <47536110+rosecers@users.noreply.github.com>
Date: Tue, 9 Jun 2026 11:49:20 -0500
Subject: [PATCH 7/8] Updating for metatrain

---
 .../ellipsoidal_density_projection.py         | 202 +++++++++++-------
 anisoap/representations/radial_basis.py       |  28 +--
 anisoap/utils/spherical_to_cartesian.py       |   5 +-
 3 files changed, 146 insertions(+), 89 deletions(-)

diff --git a/anisoap/representations/ellipsoidal_density_projection.py b/anisoap/representations/ellipsoidal_density_projection.py
index 0ebf383..eaf93d5 100644
--- a/anisoap/representations/ellipsoidal_density_projection.py
+++ b/anisoap/representations/ellipsoidal_density_projection.py
@@ -16,17 +16,17 @@
 
 import numpy as np
 import torch
-from metatensor.torch import (
-    Labels,
-    TensorBlock,
-    TensorMap,
-)
-
 from anisoap.representations.radial_basis import (
     GTORadialBasis,
     MonomialBasis,
+    _RadialBasis,
 )
 from anisoap.utils.spherical_to_cartesian import spherical_to_cartesian
+from metatensor.torch import (
+    Labels,
+    TensorBlock,
+    TensorMap,
+)
 
 from ..utils.metatensor_utils import (
     TorchClebschGordanReal,
@@ -467,10 +467,10 @@ def pairwise_ellip_expansion(
     atom_indices: torch.Tensor,
     rotation_matrices: torch.Tensor,
     ellipsoid_lengths: torch.Tensor,
-    sph_to_cart: Sequence[np.ndarray],
-    radial_basis: Any,
-    *,
-    types: Optional[Sequence[int]] = None,
+    sph_to_cart,
+    radial_basis,
+    types: List[int],
+    num_ns: List[int],
     normalize: bool = True,
 ) -> TensorMap:
     r"""Torch-native pairwise expansion ``<a n l m | rho_ij>``.
@@ -493,8 +493,12 @@ def pairwise_ellip_expansion(
     else:
         types = [int(x) for x in types]
 
-    num_ns = radial_basis.get_num_radial_functions()
-    maxdeg = int(np.max(np.arange(lmax + 1) + 2 * np.array(num_ns)))
+    # num_ns = radial_basis.get_num_radial_functions()
+    maxdeg = 0
+    for l in range(lmax + 1):
+        candidate = l + 2 * (int(num_ns[l]) - 1)
+        if candidate > maxdeg:
+            maxdeg = candidate
     scaled_sph_to_cart = []
     for l in range(lmax + 1):
         prefactor = math.sqrt((4.0 * math.pi) / (2 * l + 1))
@@ -595,6 +599,7 @@ def pairwise_ellip_expansion(
     )
 
 
+@torch.jit.ignore
 def contract_pairwise_feat(
     pair_ellip_feat: TensorMap, types: Sequence[int]
 ) -> TensorMap:
@@ -794,6 +799,7 @@ def __init__(
         self._cg = TorchClebschGordanReal(self.max_angular)
         self._neighbor_list_options = None
 
+    @torch.jit.ignore
     def requested_neighbor_lists(self) -> List[Any]:
         """Return the metatomic neighbor-list request for this descriptor."""
         try:
@@ -817,9 +823,8 @@ def requested_neighbor_lists(self) -> List[Any]:
 
     def _graph_from_inputs(
         self,
-        *,
-        systems: Optional[Sequence[Any]] = None,
-        frames: Optional[Sequence[Any]] = None,
+        systems=None,
+        frames=None,
         R_ij: Optional[torch.Tensor] = None,
         centers: Optional[torch.Tensor] = None,
         neighbors: Optional[torch.Tensor] = None,
@@ -866,11 +871,12 @@ def _graph_from_inputs(
             ),
         )
 
+    @torch.jit.ignore
     def pairwise_expansion(
         self,
-        frames: Optional[Sequence[Any]] = None,
+        frames=None,
         *,
-        systems: Optional[Sequence[Any]] = None,
+        systems=None,
         R_ij: Optional[torch.Tensor] = None,
         centers: Optional[torch.Tensor] = None,
         neighbors: Optional[torch.Tensor] = None,
@@ -918,13 +924,14 @@ def pairwise_expansion(
             self.radial_basis,
             types=types,
             normalize=normalize,
+            num_ns=self.radial_basis.get_num_radial_functions(),
         )
 
     def transform(
         self,
-        frames: Optional[Sequence[Any]] = None,
+        frames=None,
         *,
-        systems: Optional[Sequence[Any]] = None,
+        systems=None,
         R_ij: Optional[torch.Tensor] = None,
         centers: Optional[torch.Tensor] = None,
         neighbors: Optional[torch.Tensor] = None,
@@ -972,17 +979,19 @@ def transform(
             self.radial_basis,
             types=types,
             normalize=normalize,
+            num_ns=self.radial_basis.get_num_radial_functions(),
         )
         coeffs = contract_pairwise_feat(pairwise, types)
         if return_pairwise:
             return coeffs, pairwise
         return coeffs
 
+    @torch.jit.ignore
     def power_spectrum(
         self,
-        frames: Optional[Sequence[Any]] = None,
+        frames=None,
         *,
-        systems: Optional[Sequence[Any]] = None,
+        systems=None,
         mean_over_samples: bool = True,
         show_progress: bool = False,
         normalize: bool = True,
@@ -1107,6 +1116,7 @@ def power_spectrum_features_from_tensormap(
             )
         return dense, samples
 
+    @torch.jit.ignore
     def power_spectrum_features(
         self, aggregate_by_system: bool = False, **kwargs: Any
     ) -> Tuple[torch.Tensor, Labels]:
@@ -1128,20 +1138,31 @@ def _feature_size(self, *, device=None, dtype=None) -> int:
 
         center_type = self.species[0] if self.species is not None else 0
         dummy_features = self.power_spectrum_feature_tensor_map(
-            R_ij=torch.zeros((1, 3), device=device, dtype=dtype),
-            centers=torch.tensor([0], device=device, dtype=torch.long),
-            neighbors=torch.tensor([0], device=device, dtype=torch.long),
-            species=torch.tensor([center_type], device=device, dtype=torch.long),
-            structures=torch.tensor([0], device=device, dtype=torch.long),
-            atom_indices=torch.tensor([0], device=device, dtype=torch.long),
-            rotations=torch.eye(3, device=device, dtype=dtype).reshape(1, 3, 3),
-            ellipsoid_lengths=torch.ones((1, 3), device=device, dtype=dtype),
+            torch.zeros((1, 3), device=device, dtype=dtype),
+            torch.tensor([0], device=device, dtype=torch.long),
+            torch.tensor([0], device=device, dtype=torch.long),
+            torch.tensor([center_type], device=device, dtype=torch.long),
+            torch.tensor([0], device=device, dtype=torch.long),
+            torch.tensor([0], device=device, dtype=torch.long),
+            torch.eye(3, device=device, dtype=dtype).reshape(1, 3, 3),
+            torch.ones((1, 3), device=device, dtype=dtype),
+            True,
         )
         self.shape = int(dummy_features.block(0).values.shape[1])
         return self.shape
 
+    @torch.jit.export
     def power_spectrum_feature_tensor_map(
-        self, *, normalize: bool = True, **kwargs: Any
+        self,
+        R_ij: torch.Tensor,
+        centers: torch.Tensor,
+        neighbors: torch.Tensor,
+        species: torch.Tensor,
+        structures: torch.Tensor,
+        atom_indices: torch.Tensor,
+        rotations: torch.Tensor,
+        ellipsoid_lengths: torch.Tensor,
+        normalize: bool = True,
     ) -> TensorMap:
         """Return a single-block per-atom feature TensorMap for AniSOAP-BPNN.
 
@@ -1149,16 +1170,34 @@ def power_spectrum_feature_tensor_map(
         ``properties=['property']``, matching the SOAP-BPNN scalar descriptor
         interface.
         """
-        graph = self._graph_from_inputs(**kwargs)
+        R_ij = torch.as_tensor(R_ij)
+        device = R_ij.device
+        dtype = R_ij.dtype
+
+        centers = torch.as_tensor(centers, device=device, dtype=torch.long)
+        neighbors = torch.as_tensor(neighbors, device=device, dtype=torch.long)
+        species = torch.as_tensor(species, device=device, dtype=torch.long)
+        structures = torch.as_tensor(structures, device=device, dtype=torch.long)
+        atom_indices = torch.as_tensor(atom_indices, device=device, dtype=torch.long)
+        rotations = torch.as_tensor(rotations, device=device, dtype=dtype)
+        ellipsoid_lengths = torch.as_tensor(
+            ellipsoid_lengths,
+            device=device,
+            dtype=dtype,
+        )
+
         target_samples = Labels(
             ["system", "atom"],
             torch.stack(
-                [graph.structures.to(torch.int32), graph.atom_indices.to(torch.int32)],
+                [
+                    structures.to(dtype=torch.int32),
+                    atom_indices.to(dtype=torch.int32),
+                ],
                 dim=1,
-            ).to(device=graph.R_ij.device, dtype=torch.int32),
+            ),
         )
 
-        if graph.R_ij.shape[0] == 0:
+        if R_ij.shape[0] == 0:
             if self.shape is None:
                 self.shape = self._feature_size()
 
@@ -1166,26 +1205,22 @@ def power_spectrum_feature_tensor_map(
                 self.species
                 if self.species is not None
                 else sorted(
-                    int(x) for x in torch.unique(graph.species).detach().cpu().tolist()
+                    int(x) for x in torch.unique(species).detach().cpu().tolist()
                 )
             )
 
-            blocks = []
-            keys = []
+            blocks = torch.jit.annotate(List[TensorBlock], [])
+            keys = torch.jit.annotate(List[Tuple[int]], [])
 
             for center_type in all_species:
-                mask = graph.species == int(center_type)
+                mask = species == int(center_type)
                 if not bool(mask.any()):
                     continue
 
                 sample_values = torch.stack(
                     [
-                        graph.structures[mask].to(
-                            device=graph.R_ij.device, dtype=torch.int32
-                        ),
-                        graph.atom_indices[mask].to(
-                            device=graph.R_ij.device, dtype=torch.int32
-                        ),
+                        structures[mask].to(device=R_ij.device, dtype=torch.int32),
+                        atom_indices[mask].to(device=R_ij.device, dtype=torch.int32),
                     ],
                     dim=1,
                 )
@@ -1194,8 +1229,8 @@ def power_spectrum_feature_tensor_map(
                     TensorBlock(
                         values=torch.zeros(
                             (sample_values.shape[0], self.shape),
-                            device=graph.R_ij.device,
-                            dtype=graph.R_ij.dtype,
+                            device=R_ij.device,
+                            dtype=R_ij.dtype,
                         ),
                         samples=Labels(["system", "atom"], sample_values),
                         components=[],
@@ -1203,7 +1238,7 @@ def power_spectrum_feature_tensor_map(
                             ["property"],
                             torch.arange(
                                 self.shape,
-                                device=graph.R_ij.device,
+                                device=R_ij.device,
                                 dtype=torch.int32,
                             ).reshape(-1, 1),
                         ),
@@ -1214,23 +1249,42 @@ def power_spectrum_feature_tensor_map(
             return TensorMap(
                 keys=Labels(
                     ["center_type"],
-                    torch.as_tensor(keys, device=graph.R_ij.device, dtype=torch.int32),
+                    torch.as_tensor(keys, device=R_ij.device, dtype=torch.int32),
                 ),
                 blocks=blocks,
             )
 
-        # Reuse graph tensors to avoid reconstructing systems/frames.
-        nu2 = self.power_spectrum(
-            mean_over_samples=False,
-            R_ij=graph.R_ij,
-            centers=graph.centers,
-            neighbors=graph.neighbors,
-            species=graph.species,
-            structures=graph.structures,
-            atom_indices=graph.atom_indices,
-            rotations=graph.rotations,
-            ellipsoid_lengths=graph.ellipsoid_lengths,
+        if self.species is None:
+            raise RuntimeError(
+                "TorchScript path requires EllipsoidalDensityProjection.species to be set."
+            )
+        types = self.species
+
+        pairwise = pairwise_ellip_expansion(
+            self.max_angular,
+            R_ij,
+            centers,
+            neighbors,
+            species,
+            structures,
+            atom_indices,
+            rotations,
+            ellipsoid_lengths,
+            self.sph_to_cart,
+            self.radial_basis,
+            types=types,
             normalize=normalize,
+            num_ns=self.radial_basis.get_num_radial_functions(),
+        )
+
+        coeffs = contract_pairwise_feat(pairwise, types)
+        nu1 = standardize_keys(coeffs)
+        nu2 = cg_combine(
+            nu1,
+            nu1,
+            clebsch_gordan=self._cg,
+            lcut=0,
+            other_keys_match=["types_center"],
         )
         features, _ = self.power_spectrum_features_from_tensormap(
             nu2, target_samples=target_samples
@@ -1238,30 +1292,24 @@ def power_spectrum_feature_tensor_map(
 
         self.shape = int(features.shape[1])
 
-        blocks = []
-        keys = []
+        blocks = torch.jit.annotate(List[TensorBlock], [])
+        keys = torch.jit.annotate(List[Tuple[int]], [])
 
         all_species = (
             self.species
             if self.species is not None
-            else sorted(
-                int(x) for x in torch.unique(graph.species).detach().cpu().tolist()
-            )
+            else sorted(int(x) for x in torch.unique(species).detach().cpu().tolist())
         )
 
         for center_type in all_species:
-            mask = graph.species == int(center_type)
+            mask = species == int(center_type)
             if not bool(mask.any()):
                 continue
 
             sample_values = torch.stack(
                 [
-                    graph.structures[mask].to(
-                        device=features.device, dtype=torch.int32
-                    ),
-                    graph.atom_indices[mask].to(
-                        device=features.device, dtype=torch.int32
-                    ),
+                    structures[mask].to(device=features.device, dtype=torch.int32),
+                    atom_indices[mask].to(device=features.device, dtype=torch.int32),
                 ],
                 dim=1,
             )
@@ -1317,7 +1365,17 @@ def forward(
         normalize: bool = True,
     ) -> TensorMap:
         """Default module output for AniSOAP-BPNN: per-atom scalar feature map."""
-        return self.power_spectrum_feature_tensor_map(**kwargs)
+        return self.power_spectrum_feature_tensor_map(
+            R_ij=R_ij,
+            centers=centers,
+            neighbors=neighbors,
+            species=species,
+            structures=structures,
+            atom_indices=atom_indices,
+            rotations=rotations,
+            ellipsoid_lengths=ellipsoid_lengths,
+            normalize=normalize,
+        )
 
 
 __all__ = [
diff --git a/anisoap/representations/radial_basis.py b/anisoap/representations/radial_basis.py
index 5770389..ecfaeff 100644
--- a/anisoap/representations/radial_basis.py
+++ b/anisoap/representations/radial_basis.py
@@ -121,7 +121,7 @@ def orthonormalization_matrix(
     device=None,
     dtype=None,
 ) -> torch.Tensor:
-    """Return the original Lowdin radial orthonormalization matrix as torch."""
+    r"""Return the original Lowdin radial orthonormalization matrix as torch."""
     n_l = radial_basis.num_radial_functions[angular_channel]
     l_2n = angular_channel + 2 * torch.arange(n_l, device=device, dtype=torch.long)
     overlap = torch.as_tensor(radial_basis.overlap_matrix, device=device, dtype=dtype)
@@ -170,9 +170,9 @@ def gto_square_norm(n, sigma):
 
 
 def gto_prefactor(n, sigma):
-    """Computes the normalization prefactor of an unnormalized GTO.
+    r"""Computes the normalization prefactor of an unnormalized GTO.
 
-    This prefactor is simply :math:`\\frac{1}{\\sqrt{\\text{square_norm_area}}}`.
+    This prefactor is simply :math:`\frac{1}{\sqrt{\text{square_norm_area}}}`.
     Scaling a GTO by this prefactor will ensure that the GTO has square norm
     equal to 1.
 
@@ -253,7 +253,7 @@ def gto_overlap(n, m, sigma_n, sigma_m):
 
 
 def monomial_square_norm(n, r_cut):
-    """
+    r"""
     Compute the square norm of monomials (inner product of itself over R^3).
 
     Parameters
@@ -270,7 +270,7 @@ def monomial_square_norm(n, r_cut):
 
 
 def monomial_prefactor(n, r_cut):
-    """
+    r"""
     Computes the normalization prefactor of an unnormalized monomial basis.
     This prefactor is simply :math:`1/\sqrt{square\_norm\_area}`.
     Scaling a basis by this prefactor will ensure that the basis has square norm equal to 1.
@@ -324,7 +324,7 @@ def monomial_overlap(n, m, r_cut):
 
 
 class _RadialBasis:
-    """
+    r"""
     Class for precomputing and storing all results related to the radial basis.
     This helps to keep a cleaner main code by avoiding if-else clauses
     related to the radial basis.
@@ -387,7 +387,7 @@ def __init__(
 
     # Get number of radial functions
     def get_num_radial_functions(self):
-        """
+        r"""
         Output the number of radial basis functions considered per value of l.
         If max_angular and max_radial are both specified, then the list will contain repeated values of max_radial
         Otherwise, the outputted list will specify the number of radial basis functions per l, which may be automatically
@@ -403,7 +403,7 @@ def get_num_radial_functions(self):
         return self.num_radial_functions
 
     def plot_basis(self, n_r=100):
-        """
+        r"""
         Plot the normalized basis functions used in calculating the expansion
         coefficients
 
@@ -468,7 +468,7 @@ def compute_gaussian_parameters(self, r_ij, lengths, rotation_matrix):
         return gaussian_parameters(self, r_ij, lengths, rotation_matrix)
 
     def calc_overlap_matrix(self):
-        """
+        r"""
         Computes the overlap matrix for Monomnials over a fixed interval.
 
         The overlap matrix is a Gram matrix whose entries are the overlap:
@@ -500,7 +500,7 @@ def calc_overlap_matrix(self):
         return S
 
     def orthonormalize_basis(self, features: TensorMap):
-        """
+        r"""
         Apply an in-place orthonormalization on the features, using Lodwin Symmetric Orthonormalization.
         Each block in the features TensorMap uses a basis set of l + 2n, so we must take the appropriate slices of
         the overlap matrix to compute the orthonormalization matrix.
@@ -560,7 +560,7 @@ def orthonormalize_basis(self, features: TensorMap):
         return features
 
     def get_basis(self, rs):
-        """
+        r"""
         Evaluate orthonormalized monomial basis functions on mesh rs.
 
         If lmax and nmax defined, then the number of functions outputted is lmax*(nmax+1)
@@ -621,7 +621,7 @@ def get_basis(self, rs):
 
 
 class GTORadialBasis(_RadialBasis):
-    """
+    r"""
     A subclass of _RadialBasis that contains attributes and methods required for the GTO basis.
 
     The GTO basis of order n is defined to be :math:`R(r) = r^n e^{(-r^2/(2\sigma^2))}`.
@@ -665,7 +665,7 @@ def compute_gaussian_parameters(self, r_ij, lengths, rotation_matrix):
         return gaussian_parameters(self, r_ij, lengths, rotation_matrix)
 
     def calc_overlap_matrix(self):
-        """Computes the overlap matrix for GTOs.
+        r"""Computes the overlap matrix for GTOs.
 
         The overlap matrix is a Gram matrix whose entries are the overlap:
 
@@ -704,7 +704,7 @@ def calc_overlap_matrix(self):
         return S
 
     def orthonormalize_basis(self, features: TensorMap):
-        """Applies in-place orthonormalization on the features.
+        r"""Applies in-place orthonormalization on the features.
 
         Apply an in-place orthonormalization on the features, using Lodwin Symmetric
         Orthonormalization. Each block in the features TensorMap uses a GTO set
diff --git a/anisoap/utils/spherical_to_cartesian.py b/anisoap/utils/spherical_to_cartesian.py
index 1a1f5a2..1d29693 100644
--- a/anisoap/utils/spherical_to_cartesian.py
+++ b/anisoap/utils/spherical_to_cartesian.py
@@ -1,12 +1,11 @@
 import numpy as np
+from anisoap.utils import monomial_iterator
 from scipy.special import (
     comb,
     factorial,
     factorial2,
 )
 
-from anisoap.utils import monomial_iterator
-
 # Here we are implementing recurrence of the form R_{l}^m = prefact_minus1* z * T_{l-1} + prefact_minus2* r2 * T_{l-2}
 # where R_l^m is a solid harmonic, when expressed on a monomial basis - R_l^m = \sum_{n0+n1+n2=l} T_{l}[n0,n1,n2] x^n0 y^n1 z^n2
 # We will further define these coefficients with an additional n-dependent prefactor r^P2n}
@@ -83,7 +82,7 @@ def binom(n, k):
 
 
 def spherical_to_cartesian(lmax, num_ns):
-    """
+    r"""
     Finds the coefficients for the cartesian polynomial form of solid harmonics
     :math:`R_{lm} = \sqrt{(4\pi)/(2l+1)}*r^l*Y_{lm}`.  Note that our AniSOAP
     expansion does not contain the :math:`\sqrt{(4\pi)/(2l+1)}` prefactor, so in calculating

From 4194781aca41212d2a161e1b01645975b6f5a67f Mon Sep 17 00:00:00 2001
From: "Rose K. Cersonsky" <47536110+rosecers@users.noreply.github.com>
Date: Fri, 19 Jun 2026 16:21:32 -0500
Subject: [PATCH 8/8] Updating linters

---
 .../representations/ellipsoidal_density_projection.py | 11 ++++++-----
 anisoap/utils/spherical_to_cartesian.py               |  3 ++-
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/anisoap/representations/ellipsoidal_density_projection.py b/anisoap/representations/ellipsoidal_density_projection.py
index eaf93d5..8e9f224 100644
--- a/anisoap/representations/ellipsoidal_density_projection.py
+++ b/anisoap/representations/ellipsoidal_density_projection.py
@@ -16,17 +16,18 @@
 
 import numpy as np
 import torch
+from metatensor.torch import (
+    Labels,
+    TensorBlock,
+    TensorMap,
+)
+
 from anisoap.representations.radial_basis import (
     GTORadialBasis,
     MonomialBasis,
     _RadialBasis,
 )
 from anisoap.utils.spherical_to_cartesian import spherical_to_cartesian
-from metatensor.torch import (
-    Labels,
-    TensorBlock,
-    TensorMap,
-)
 
 from ..utils.metatensor_utils import (
     TorchClebschGordanReal,
diff --git a/anisoap/utils/spherical_to_cartesian.py b/anisoap/utils/spherical_to_cartesian.py
index 1d29693..b76ed49 100644
--- a/anisoap/utils/spherical_to_cartesian.py
+++ b/anisoap/utils/spherical_to_cartesian.py
@@ -1,11 +1,12 @@
 import numpy as np
-from anisoap.utils import monomial_iterator
 from scipy.special import (
     comb,
     factorial,
     factorial2,
 )
 
+from anisoap.utils import monomial_iterator
+
 # Here we are implementing recurrence of the form R_{l}^m = prefact_minus1* z * T_{l-1} + prefact_minus2* r2 * T_{l-2}
 # where R_l^m is a solid harmonic, when expressed on a monomial basis - R_l^m = \sum_{n0+n1+n2=l} T_{l}[n0,n1,n2] x^n0 y^n1 z^n2
 # We will further define these coefficients with an additional n-dependent prefactor r^P2n}