From e5bc669119e9a049ca0c4cc6b9a6d704aaea8207 Mon Sep 17 00:00:00 2001
From: ppegolo <paolo.pegolo@epfl.ch>
Date: Tue, 21 Jan 2025 18:20:20 +0100
Subject: [PATCH 01/21] Start fixing CG utils

---
 .../clebsch_gordan/_cartesian_spherical.py    |  3 +
 .../featomic/clebsch_gordan/_coefficients.py  | 57 ++++++++++++-------
 2 files changed, 40 insertions(+), 20 deletions(-)

diff --git a/python/featomic/featomic/clebsch_gordan/_cartesian_spherical.py b/python/featomic/featomic/clebsch_gordan/_cartesian_spherical.py
index 32fc7eba2..f849be677 100644
--- a/python/featomic/featomic/clebsch_gordan/_cartesian_spherical.py
+++ b/python/featomic/featomic/clebsch_gordan/_cartesian_spherical.py
@@ -245,6 +245,9 @@ def cartesian_to_spherical(
     #           [..., 3, 5, ...] => [..., 3, ...] (o3_lambda=1, o3_sigma=+1)
     #                            => [..., 5, ...] (o3_lambda=2, o3_sigma=-1)
     #                            => [..., 7, ...] (o3_lambda=3, o3_sigma=+1)
+
+    tensor = TensorMap(tensor.keys, new_blocks)
+
     if cg_coefficients is None:
         if torch_jit_is_scripting():
             raise ValueError(
diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py
index e7d94f336..c5ffa0637 100644
--- a/python/featomic/featomic/clebsch_gordan/_coefficients.py
+++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py
@@ -198,17 +198,27 @@ def _build_dense_cg_coeff_dict(
                     dtype=complex_like.dtype,
                 )
 
-                real_cg = (r2c[l1] @ complex_cg.reshape(2 * l1 + 1, -1)).reshape(
-                    complex_cg.shape
-                )
+                # real_cg = (r2c[l1] @ complex_cg.reshape(2 * l1 + 1, -1)).reshape(
+                #     complex_cg.shape
+                # )
+                real_cg = (complex_cg.permute(2, 1, 0) @ r2c[l1]).swapaxes(
+                    1, 2
+                )  # this has shape 2*L+1, 2*l1+1, 2*l2+1)
 
-                real_cg = real_cg.swapaxes(0, 1)
-                real_cg = (r2c[l2] @ real_cg.reshape(2 * l2 + 1, -1)).reshape(
-                    real_cg.shape
-                )
-                real_cg = real_cg.swapaxes(0, 1)
+                # real_cg = real_cg.swapaxes(0, 1)
+                # real_cg = (r2c[l2] @ real_cg.reshape(2 * l2 + 1, -1)).reshape(
+                #     real_cg.shape
+                # )
+
+                real_cg = (
+                    real_cg @ r2c[l2]
+                )  # This has shape (2*o3_lambda+1, 2*l1+1, 2*l2+1)
+
+                # real_cg = real_cg.swapaxes(0, 1)
+
+                # real_cg = real_cg @ c2r[o3_lambda]
 
-                real_cg = real_cg @ c2r[o3_lambda]
+                real_cg = real_cg.permute(1, 2, 0) @ c2r[o3_lambda].T
 
                 if (l1 + l2 + o3_lambda) % 2 == 0:
                     cg_l1l2lam_dense = _dispatch.real(real_cg)
@@ -366,7 +376,7 @@ def _real2complex(o3_lambda: int, like: Array) -> Array:
             # Positive part
             result[o3_lambda + m, o3_lambda + m] = inv_sqrt_2 * ((-1) ** m)
 
-    return result
+    return result.T
 
 
 def _complex2real(o3_lambda: int, like) -> Array:
@@ -438,18 +448,19 @@ def cg_couple(
             for o3_lambda in o3_lambdas
         ]
     elif cg_backend == "python-dense":
+
         results = []
 
-        n_samples = array.shape[0]
-        n_properties = array.shape[3]
+        # n_samples = array.shape[0]
+        # n_properties = array.shape[3]
 
-        array = array.swapaxes(1, 3)
-        array = array.reshape(n_samples * n_properties, 2 * l2 + 1, 2 * l1 + 1)
+        # array = array.swapaxes(1, 3)
+        # array = array.reshape(n_samples * n_properties, 2 * l2 + 1, 2 * l1 + 1)
 
         for o3_lambda in o3_lambdas:
             result = _cg_couple_dense(array, o3_lambda, cg_coefficients)
-            result = result.reshape(n_samples, n_properties, -1)
-            result = result.swapaxes(1, 2)
+            # result = result.reshape(n_samples, n_properties, -1)
+            # result = result.swapaxes(1, 2)
             results.append(result)
 
         return results
@@ -519,7 +530,10 @@ def _cg_couple_dense(
     :param cg_coefficients: CG coefficients as returned by
         :py:func:`calculate_cg_coefficients` with ``cg_backed="python-dense"``
     """
-    assert len(array.shape) == 3
+    # assert len(array.shape) == 3
+
+    # l1 = (array.shape[1] - 1) // 2
+    # l2 = (array.shape[2] - 1) // 2
 
     l1 = (array.shape[1] - 1) // 2
     l2 = (array.shape[2] - 1) // 2
@@ -527,13 +541,16 @@ def _cg_couple_dense(
     cg_l1l2lam = cg_coefficients.block({"l1": l1, "l2": l2, "lambda": o3_lambda}).values
 
     # [samples, l1, l2] => [samples, (l1 l2)]
-    array = array.reshape(-1, (2 * l1 + 1) * (2 * l2 + 1))
+    # array = array.reshape(-1, (2 * l1 + 1) * (2 * l2 + 1))
 
     # [l1, l2, lambda] -> [(l1 l2), lambda]
-    cg_l1l2lam = cg_l1l2lam.reshape(-1, 2 * o3_lambda + 1)
+    # cg_l1l2lam = cg_l1l2lam.reshape(-1, 2 * o3_lambda + 1)
 
     # [samples, (l1 l2)] @ [(l1 l2), lambda] => [samples, lambda]
-    return array @ cg_l1l2lam
+    # return array @ cg_l1l2lam
+    import torch
+
+    return torch.einsum("smnp,mnM->sMp", array, cg_l1l2lam[0, ..., 0])
 
 
 # ======================================================================= #

From 1159c6f024e58bc31061132d5d1b5dc58b438702 Mon Sep 17 00:00:00 2001
From: Joseph Abbott <joseph.william.abbott@gmail.com>
Date: Wed, 22 Jan 2025 12:10:53 +0100
Subject: [PATCH 02/21] dispatch operations

---
 .../clebsch_gordan/_cartesian_spherical.py    |  1 +
 .../featomic/clebsch_gordan/_coefficients.py  | 63 ++++---------------
 .../featomic/clebsch_gordan/_dispatch.py      | 22 +++++++
 3 files changed, 36 insertions(+), 50 deletions(-)

diff --git a/python/featomic/featomic/clebsch_gordan/_cartesian_spherical.py b/python/featomic/featomic/clebsch_gordan/_cartesian_spherical.py
index f849be677..2c9b678d0 100644
--- a/python/featomic/featomic/clebsch_gordan/_cartesian_spherical.py
+++ b/python/featomic/featomic/clebsch_gordan/_cartesian_spherical.py
@@ -246,6 +246,7 @@ def cartesian_to_spherical(
     #                            => [..., 5, ...] (o3_lambda=2, o3_sigma=-1)
     #                            => [..., 7, ...] (o3_lambda=3, o3_sigma=+1)
 
+    # Use the rolled block values as the starting point for higher-rank tensors
     tensor = TensorMap(tensor.keys, new_blocks)
 
     if cg_coefficients is None:
diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py
index c5ffa0637..0d904921f 100644
--- a/python/featomic/featomic/clebsch_gordan/_coefficients.py
+++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py
@@ -197,28 +197,15 @@ def _build_dense_cg_coeff_dict(
                     device=device,
                     dtype=complex_like.dtype,
                 )
-
-                # real_cg = (r2c[l1] @ complex_cg.reshape(2 * l1 + 1, -1)).reshape(
-                #     complex_cg.shape
-                # )
-                real_cg = (complex_cg.permute(2, 1, 0) @ r2c[l1]).swapaxes(
+                # TODO: comment to explain what is happening here and the shapes
+                # involved
+                real_cg = (_dispatch.permute(complex_cg, [2, 1, 0]) @ r2c[l1]).swapaxes(
                     1, 2
-                )  # this has shape 2*L+1, 2*l1+1, 2*l2+1)
-
-                # real_cg = real_cg.swapaxes(0, 1)
-                # real_cg = (r2c[l2] @ real_cg.reshape(2 * l2 + 1, -1)).reshape(
-                #     real_cg.shape
-                # )
-
-                real_cg = (
-                    real_cg @ r2c[l2]
-                )  # This has shape (2*o3_lambda+1, 2*l1+1, 2*l2+1)
+                )  # this has shape (2*L+1, 2*l1+1, 2*l2+1)
 
-                # real_cg = real_cg.swapaxes(0, 1)
+                real_cg = real_cg @ r2c[l2]  # this has shape (2*L+1, 2*l1+1, 2*l2+1)
 
-                # real_cg = real_cg @ c2r[o3_lambda]
-
-                real_cg = real_cg.permute(1, 2, 0) @ c2r[o3_lambda].T
+                real_cg = _dispatch.permute(real_cg, [1, 2, 0]) @ c2r[o3_lambda].T
 
                 if (l1 + l2 + o3_lambda) % 2 == 0:
                     cg_l1l2lam_dense = _dispatch.real(real_cg)
@@ -448,22 +435,10 @@ def cg_couple(
             for o3_lambda in o3_lambdas
         ]
     elif cg_backend == "python-dense":
-
-        results = []
-
-        # n_samples = array.shape[0]
-        # n_properties = array.shape[3]
-
-        # array = array.swapaxes(1, 3)
-        # array = array.reshape(n_samples * n_properties, 2 * l2 + 1, 2 * l1 + 1)
-
-        for o3_lambda in o3_lambdas:
-            result = _cg_couple_dense(array, o3_lambda, cg_coefficients)
-            # result = result.reshape(n_samples, n_properties, -1)
-            # result = result.swapaxes(1, 2)
-            results.append(result)
-
-        return results
+        return [
+            _cg_couple_dense(array, o3_lambda, cg_coefficients)
+            for o3_lambda in o3_lambdas
+        ]
 
     else:
         raise ValueError(
@@ -530,27 +505,15 @@ def _cg_couple_dense(
     :param cg_coefficients: CG coefficients as returned by
         :py:func:`calculate_cg_coefficients` with ``cg_backed="python-dense"``
     """
-    # assert len(array.shape) == 3
-
-    # l1 = (array.shape[1] - 1) // 2
-    # l2 = (array.shape[2] - 1) // 2
+    assert len(array.shape) == 3
 
     l1 = (array.shape[1] - 1) // 2
     l2 = (array.shape[2] - 1) // 2
 
     cg_l1l2lam = cg_coefficients.block({"l1": l1, "l2": l2, "lambda": o3_lambda}).values
 
-    # [samples, l1, l2] => [samples, (l1 l2)]
-    # array = array.reshape(-1, (2 * l1 + 1) * (2 * l2 + 1))
-
-    # [l1, l2, lambda] -> [(l1 l2), lambda]
-    # cg_l1l2lam = cg_l1l2lam.reshape(-1, 2 * o3_lambda + 1)
-
-    # [samples, (l1 l2)] @ [(l1 l2), lambda] => [samples, lambda]
-    # return array @ cg_l1l2lam
-    import torch
-
-    return torch.einsum("smnp,mnM->sMp", array, cg_l1l2lam[0, ..., 0])
+    # TODO: use something more efficient than an einsum
+    return _dispatch.einsum("smnp,mnM->sMp", array, cg_l1l2lam[0, ..., 0])
 
 
 # ======================================================================= #
diff --git a/python/featomic/featomic/clebsch_gordan/_dispatch.py b/python/featomic/featomic/clebsch_gordan/_dispatch.py
index dc097df56..c015a9d4a 100644
--- a/python/featomic/featomic/clebsch_gordan/_dispatch.py
+++ b/python/featomic/featomic/clebsch_gordan/_dispatch.py
@@ -63,6 +63,28 @@ def concatenate(arrays: List[TorchTensor], axis: int):
         raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
+def einsum(equation: str, *operands):
+    if isinstance(operands[0], TorchTensor):
+        return torch.einsum(equation, *operands)
+    elif isinstance(operands[0], np.ndarray):
+        return np.einsum(equation, *operands)
+    else:
+        raise TypeError(UNKNOWN_ARRAY_TYPE)
+
+
+def permute(array, axes: List[int]):
+    """
+    Permute the axes of the given ``array``. For numpy and torch, ``transpose`` and
+    ``permute`` are equivalent operations.
+    """
+    if isinstance(array, TorchTensor):
+        return torch.permute(array, axes)
+    elif isinstance(array, np.ndarray):
+        return np.transpose(array, axes)
+    else:
+        raise TypeError(UNKNOWN_ARRAY_TYPE)
+
+
 def empty_like(array, shape: Optional[List[int]] = None, requires_grad: bool = False):
     """
     Create an uninitialized array, with the given ``shape``, and similar dtype, device

From bd75c9d8c83a9fa138007382791995441f43f4c3 Mon Sep 17 00:00:00 2001
From: Joseph Abbott <joseph.william.abbott@gmail.com>
Date: Wed, 22 Jan 2025 13:16:37 +0100
Subject: [PATCH 03/21] Use tensordot instead of einsum

---
 .../featomic/clebsch_gordan/_coefficients.py  |  6 +++--
 .../featomic/clebsch_gordan/_dispatch.py      | 23 +++++++++++--------
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py
index 0d904921f..23697ffa7 100644
--- a/python/featomic/featomic/clebsch_gordan/_coefficients.py
+++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py
@@ -512,8 +512,10 @@ def _cg_couple_dense(
 
     cg_l1l2lam = cg_coefficients.block({"l1": l1, "l2": l2, "lambda": o3_lambda}).values
 
-    # TODO: use something more efficient than an einsum
-    return _dispatch.einsum("smnp,mnM->sMp", array, cg_l1l2lam[0, ..., 0])
+    return _dispatch.permute(
+        _dispatch.tensordot(array, cg_l1l2lam[0, ..., 0], axes=([1, 2], [0, 1])),
+        [0, 2, 1]
+    )
 
 
 # ======================================================================= #
diff --git a/python/featomic/featomic/clebsch_gordan/_dispatch.py b/python/featomic/featomic/clebsch_gordan/_dispatch.py
index c015a9d4a..5aee626c0 100644
--- a/python/featomic/featomic/clebsch_gordan/_dispatch.py
+++ b/python/featomic/featomic/clebsch_gordan/_dispatch.py
@@ -63,15 +63,6 @@ def concatenate(arrays: List[TorchTensor], axis: int):
         raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
-def einsum(equation: str, *operands):
-    if isinstance(operands[0], TorchTensor):
-        return torch.einsum(equation, *operands)
-    elif isinstance(operands[0], np.ndarray):
-        return np.einsum(equation, *operands)
-    else:
-        raise TypeError(UNKNOWN_ARRAY_TYPE)
-
-
 def permute(array, axes: List[int]):
     """
     Permute the axes of the given ``array``. For numpy and torch, ``transpose`` and
@@ -83,6 +74,20 @@ def permute(array, axes: List[int]):
         return np.transpose(array, axes)
     else:
         raise TypeError(UNKNOWN_ARRAY_TYPE)
+    
+def tensordot(array_1, array_2, axes: Union[int, List[int]]):
+    """
+    Compute tensor dot product along specified axes for arrays.
+
+    This function has the same behavior as ``np.tensordot(array_1, array_2, axes)`` and
+    ``torch.tensordot(array_1, array_2, axes)``.
+    """
+    if isinstance(array_1, TorchTensor):
+        return torch.tensordot(array_1, array_2, axes)
+    elif isinstance(array_1, np.ndarray):
+        return np.tensordot(array_1, array_2, axes)
+    else:
+        raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
 def empty_like(array, shape: Optional[List[int]] = None, requires_grad: bool = False):

From 3a7e1834c25dc775f11877292c7b55d213d1a925 Mon Sep 17 00:00:00 2001
From: Joseph Abbott <joseph.william.abbott@gmail.com>
Date: Wed, 22 Jan 2025 15:30:42 +0100
Subject: [PATCH 04/21] Passing `cartesian_to_spherical` tests for L=2 and 3
 components

---
 .../featomic/clebsch_gordan/_coefficients.py  |   3 +-
 .../clebsch_gordan/cartesian_spherical.py     | 178 ++++++++++++++++++
 2 files changed, 180 insertions(+), 1 deletion(-)

diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py
index 23697ffa7..fc00b3e6e 100644
--- a/python/featomic/featomic/clebsch_gordan/_coefficients.py
+++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py
@@ -505,7 +505,8 @@ def _cg_couple_dense(
     :param cg_coefficients: CG coefficients as returned by
         :py:func:`calculate_cg_coefficients` with ``cg_backed="python-dense"``
     """
-    assert len(array.shape) == 3
+    # TODO: fix this!
+    # assert len(array.shape) == 3
 
     l1 = (array.shape[1] - 1) // 2
     l2 = (array.shape[2] - 1) // 2
diff --git a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
index cffdfc685..59a1bb3da 100644
--- a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
+++ b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 from metatensor import Labels, TensorBlock, TensorMap
+from metatensor import operations
 
 from featomic.clebsch_gordan import cartesian_to_spherical
 
@@ -198,3 +199,180 @@ def test_cartesian_to_spherical_errors(cartesian):
             components=["xyz_1", "xyz_2", "xyz_3"],
             keep_l_in_keys=False,
         )
+
+def _l2_components_from_matrix(A):
+    """
+    *The* reference equations for projecting a cartesian rank 2 tensor to the
+    irreducible spherical component with lambda = 2.
+    """
+    A = A.reshape(3, 3)
+
+    l2_A = np.empty((5,))
+    l2_A[0] = (A[0, 1] + A[1, 0]) / 2.0
+    l2_A[1] = (A[1, 2] + A[2, 1]) / 2.0
+    l2_A[2] = (2.0 * A[2, 2] - A[0, 0] - A[1, 1]) / ((2.0) * np.sqrt(3.0))
+    l2_A[3] = (A[0, 2] + A[2, 0]) / 2.0
+    l2_A[4] = (A[0, 0] - A[1, 1]) / 2.0
+
+    return l2_A * np.sqrt(2)
+
+def _l3_components_from_matrix(A):
+    """
+    *The* reference equations for projecting a cartesian rank 3 tensor to the
+    irreducible spherical component with lambda = 3.
+    """
+    A = A.reshape(3, 3, 3)
+
+    l3_A = np.empty((7,))
+
+    # -3
+    l3_A[0] = (A[0, 1, 0] + A[1, 0, 0] + A[0, 0, 1] - A[1, 1, 1]) / 2.0
+
+    # -2
+    l3_A[1] = (
+        A[0, 1, 2] + A[1, 0, 2] + A[1, 2, 0] + A[2, 1, 0] + A[0, 2, 1] + A[2, 0, 1]
+    ) / np.sqrt(6)
+
+    # -1
+    l3_A[2] = (
+        4.0 * A[1, 2, 2]
+        + 4.0 * A[2, 1, 2]
+        + 4.0 * A[2, 2, 1]
+        - 3.0 * A[1, 1, 1]
+        - A[0, 0, 1]
+        - A[0, 1, 0]
+        - A[1, 0, 0]
+    ) / np.sqrt(60)
+
+    # 0
+    l3_A[3] = (
+        2.0 * A[2, 2, 2]
+        - A[0, 2, 0]
+        - A[2, 0, 0]
+        - A[0, 0, 2]
+        - A[1, 2, 1]
+        - A[2, 1, 1]
+        - A[1, 1, 2]
+    ) / np.sqrt(10)
+
+    # 1
+    l3_A[4] = (
+        4.0 * A[0, 2, 2]
+        + 4.0 * A[2, 0, 2]
+        + 4.0 * A[2, 2, 0]
+        - 3.0 * A[0, 0, 0]
+        - A[1, 1, 0]
+        - A[0, 1, 1]
+        - A[1, 0, 1]
+    ) / np.sqrt(60)
+
+    # 2
+    l3_A[5] = (
+        A[0, 0, 2] + A[0, 2, 0] + A[2, 0, 0] - A[1, 1, 2] - A[1, 2, 1] - A[2, 1, 1]
+    ) / np.sqrt(6)
+
+    # 3
+    l3_A[6] = (A[0, 0, 0] - A[1, 1, 0] - A[0, 1, 1] - A[1, 0, 1]) / 2.0
+
+    return l3_A 
+
+@pytest.mark.parametrize("cg_backend", ["python-dense", "python-sparse"])
+def test_cartesian_to_spherical_rank_2_by_equation(cg_backend):
+    """
+    Tests cartesian_to_spherical for a random rank-2 tensor by comparing the result to
+    the result from *the* reference equation.
+    """
+    # Build the reference lambda = 2 component
+    random_rank_2_arr = np.random.rand(100, 3, 3, 1)
+    l2_reference = TensorMap(
+        keys=Labels(["o3_lambda", "o3_sigma"], np.array([[2, 1]])),
+        blocks=[
+            TensorBlock(
+                samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
+                components=[
+                    Labels(["o3_mu"], np.arange(-2, 3).reshape(-1, 1))
+                ],
+                properties=Labels(["_"], np.array([[0]])),
+                values=np.stack(
+                    [_l2_components_from_matrix(A[..., 0]) for A in random_rank_2_arr]
+                ).reshape(random_rank_2_arr.shape[0], 5, 1),
+            )
+        ],
+    )
+
+    # Build the cartesian tensor and do cartesian to spherical
+    rank_2_input_cart = TensorMap(
+        keys=Labels.single(),
+        blocks=[
+            TensorBlock(
+                values=random_rank_2_arr,
+                samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
+                components=[
+                    Labels([a], np.arange(3).reshape(-1, 1)) for a in ["xyz1", "xyz2"]
+                ],
+                properties=Labels(["_"], np.array([[0]])),
+            )
+        ]
+    )
+    rank_2_input_sph = cartesian_to_spherical(rank_2_input_cart, ["xyz1", "xyz2"])
+
+    # Extract the lambda = 2 component
+    l2_input = operations.drop_blocks(
+        operations.remove_dimension(rank_2_input_sph, "keys", "_"),
+        keys=Labels(["o3_lambda"], np.array([[0], [1]])),
+    )
+
+    assert operations.equal_metadata(l2_input, l2_reference)
+    assert operations.allclose(l2_input, l2_reference)
+
+
+@pytest.mark.parametrize("cg_backend", ["python-dense", "python-sparse"])
+def test_cartesian_to_spherical_rank_3_by_equation(cg_backend):
+    """
+    Tests cartesian_to_spherical for a random rank-2 tensor by comparing the result to
+    the result from *the* reference equation.
+    """
+    # Build the reference lambda = 2 component
+    random_rank_3_arr = np.random.rand(100, 3, 3, 3, 1)
+    l3_reference = TensorMap(
+        keys=Labels(["o3_lambda", "o3_sigma"], np.array([[3, 1]])),
+        blocks=[
+            TensorBlock(
+                samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
+                components=[
+                    Labels(["o3_mu"], np.arange(-3, 4).reshape(-1, 1))
+                ],
+                properties=Labels(["_"], np.array([[0]])),
+                values=np.stack(
+                    [_l3_components_from_matrix(A[..., 0]) for A in random_rank_3_arr]
+                ).reshape(random_rank_3_arr.shape[0], 7, 1),
+            )
+        ],
+    )
+
+    # Build the cartesian tensor and do cartesian to spherical
+    rank_3_input_cart = TensorMap(
+        keys=Labels.single(),
+        blocks=[
+            TensorBlock(
+                values=random_rank_3_arr,
+                samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
+                components=[
+                    Labels([a], np.arange(3).reshape(-1, 1)) for a in ["xyz1", "xyz2", "xyz3"]
+                ],
+                properties=Labels(["_"], np.array([[0]])),
+            )
+        ]
+    )
+    rank_3_input_sph = cartesian_to_spherical(rank_3_input_cart, ["xyz1", "xyz2", "xyz3"])
+
+    # Extract the lambda = 3 component
+    l3_input = operations.drop_blocks(
+        operations.remove_dimension(rank_3_input_sph, "keys", "_"),
+        keys=Labels(["o3_lambda"], np.array([[0], [1], [2]])),
+    )
+    for dim in ["l_3", "k_1", "l_2", "l_1"]:
+        l3_input = operations.remove_dimension(l3_input, "keys", dim)
+
+    assert operations.equal_metadata(l3_input, l3_reference)
+    assert operations.allclose(l3_input, l3_reference)

From 80e95f98badf877119d58dbd7dbd30d462615437 Mon Sep 17 00:00:00 2001
From: Joseph Abbott <joseph.william.abbott@gmail.com>
Date: Wed, 22 Jan 2025 15:49:19 +0100
Subject: [PATCH 05/21] Add sparse vs dense test

---
 .../clebsch_gordan/cartesian_spherical.py     |  8 +++-
 .../tests/clebsch_gordan/cg_product.py        | 42 +++++++++++++++++--
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
index 59a1bb3da..e32b7b4ff 100644
--- a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
+++ b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
@@ -314,7 +314,9 @@ def test_cartesian_to_spherical_rank_2_by_equation(cg_backend):
             )
         ]
     )
-    rank_2_input_sph = cartesian_to_spherical(rank_2_input_cart, ["xyz1", "xyz2"])
+    rank_2_input_sph = cartesian_to_spherical(
+        rank_2_input_cart, ["xyz1", "xyz2"], cg_backend=cg_backend
+    )
 
     # Extract the lambda = 2 component
     l2_input = operations.drop_blocks(
@@ -364,7 +366,9 @@ def test_cartesian_to_spherical_rank_3_by_equation(cg_backend):
             )
         ]
     )
-    rank_3_input_sph = cartesian_to_spherical(rank_3_input_cart, ["xyz1", "xyz2", "xyz3"])
+    rank_3_input_sph = cartesian_to_spherical(
+        rank_3_input_cart, ["xyz1", "xyz2", "xyz3"], cg_backend=cg_backend
+    )
 
     # Extract the lambda = 3 component
     l3_input = operations.drop_blocks(
diff --git a/python/featomic/tests/clebsch_gordan/cg_product.py b/python/featomic/tests/clebsch_gordan/cg_product.py
index 59743c7f9..3eb94f2cc 100644
--- a/python/featomic/tests/clebsch_gordan/cg_product.py
+++ b/python/featomic/tests/clebsch_gordan/cg_product.py
@@ -85,14 +85,16 @@ def spherical_expansion_by_pair(frames: List[ase.Atoms]):
     calculator = featomic.SphericalExpansionByPair(**SPHEX_HYPERS)
     return calculator.compute(frames)
 
-
-def test_keys_are_matched():
+@pytest.mark.parametrize("cg_backend", ["python-sparse", "python-dense"])
+def test_keys_are_matched(cg_backend):
     """
     Tests that key dimensions named the same in two tensors are matched.
     """
     # Set up
     frames = h2o_isolated()
-    calculator = ClebschGordanProduct(max_angular=MAX_ANGULAR * 2)
+    calculator = ClebschGordanProduct(
+        max_angular=MAX_ANGULAR * 2, cg_backend=cg_backend
+    )
 
     # Compute lambda-SOAP
     density = spherical_expansion(frames)
@@ -431,3 +433,37 @@ def test_device_dtype(dtype, device):
         o3_lambda_1_new_name="l_1",
         o3_lambda_2_new_name="l_2",
     )
+
+
+def test_dense_sparse_agree():
+    """
+    Tests that the max_angular is sufficient to correlate the two tensors when not large
+    enough to cover MAX_ANGULAR, but when an angular cutoff is applied.
+    """
+    frames = h2o_isolated()
+    density = spherical_expansion(frames)
+
+    results = []
+    for cg_backend in ["python-sparse", "python-dense"]:
+    
+
+        # max_angular to be twice as big here if not using an angular cutoff
+        calculator = ClebschGordanProduct(
+            max_angular=MAX_ANGULAR,
+            cg_backend=cg_backend,
+        )
+        results.append(
+            calculator.compute(
+                metatensor.rename_dimension(density, "properties", "n", "n_1"),
+                metatensor.rename_dimension(density, "properties", "n", "n_2"),
+                o3_lambda_1_new_name="l_1",
+                o3_lambda_2_new_name="l_2",
+                selected_keys=Labels(
+                    names=["o3_lambda"],
+                    values=np.arange(MAX_ANGULAR + 1).reshape(-1, 1),
+                ),
+            )
+        )
+
+    assert metatensor.equal_metadata(results[0], results[1])
+    assert metatensor.allclose(results[0], results[1])

From f6df3d774bc9be0b7c10ab2fd3c9c503c55a0754 Mon Sep 17 00:00:00 2001
From: ppegolo <paolo.pegolo@epfl.ch>
Date: Wed, 22 Jan 2025 15:50:06 +0100
Subject: [PATCH 06/21] Change input and output shapes of `_cg_couple_dense` to
 the previous version

---
 .../featomic/clebsch_gordan/_coefficients.py  | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py
index fc00b3e6e..899249a63 100644
--- a/python/featomic/featomic/clebsch_gordan/_coefficients.py
+++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py
@@ -435,6 +435,13 @@ def cg_couple(
             for o3_lambda in o3_lambdas
         ]
     elif cg_backend == "python-dense":
+
+        n_samples = array.shape[0]
+        n_properties = array.shape[3]
+
+        array = array.swapaxes(1, 3)
+        array = array.reshape(n_samples * n_properties, 2 * l2 + 1, 2 * l1 + 1)
+
         return [
             _cg_couple_dense(array, o3_lambda, cg_coefficients)
             for o3_lambda in o3_lambdas
@@ -505,18 +512,14 @@ def _cg_couple_dense(
     :param cg_coefficients: CG coefficients as returned by
         :py:func:`calculate_cg_coefficients` with ``cg_backed="python-dense"``
     """
-    # TODO: fix this!
-    # assert len(array.shape) == 3
+    assert len(array.shape) == 3
 
-    l1 = (array.shape[1] - 1) // 2
-    l2 = (array.shape[2] - 1) // 2
+    l1 = (array.shape[2] - 1) // 2
+    l2 = (array.shape[1] - 1) // 2
 
     cg_l1l2lam = cg_coefficients.block({"l1": l1, "l2": l2, "lambda": o3_lambda}).values
 
-    return _dispatch.permute(
-        _dispatch.tensordot(array, cg_l1l2lam[0, ..., 0], axes=([1, 2], [0, 1])),
-        [0, 2, 1]
-    )
+    return _dispatch.tensordot(array, cg_l1l2lam[0, ..., 0], axes=([2, 1], [0, 1]))
 
 
 # ======================================================================= #

From da8b93eaa924dbf6c2ac483d3b764d86aed9dca8 Mon Sep 17 00:00:00 2001
From: ppegolo <paolo.pegolo@epfl.ch>
Date: Wed, 22 Jan 2025 17:37:01 +0100
Subject: [PATCH 07/21] Fix agreement between dense and sparse CG coupling

---
 .../featomic/clebsch_gordan/_cg_product.py    |  3 +-
 .../featomic/clebsch_gordan/_coefficients.py  | 36 +++++++++++++------
 .../featomic/clebsch_gordan/_dispatch.py      |  3 +-
 .../clebsch_gordan/cartesian_spherical.py     | 23 ++++++------
 .../tests/clebsch_gordan/cg_product.py        |  3 +-
 .../tests/clebsch_gordan/coefficients.py      |  5 +--
 6 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/python/featomic/featomic/clebsch_gordan/_cg_product.py b/python/featomic/featomic/clebsch_gordan/_cg_product.py
index 9d1ac4e0b..9ad32a86a 100644
--- a/python/featomic/featomic/clebsch_gordan/_cg_product.py
+++ b/python/featomic/featomic/clebsch_gordan/_cg_product.py
@@ -117,8 +117,7 @@ def __init__(
 
         if max_angular < 0:
             raise ValueError(
-                f"Given `max_angular={max_angular}` negative. "
-                "Must be greater equal 0."
+                f"Given `max_angular={max_angular}` negative. Must be greater equal 0."
             )
         self._max_angular = max_angular
 
diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py
index 899249a63..40294176e 100644
--- a/python/featomic/featomic/clebsch_gordan/_coefficients.py
+++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py
@@ -435,17 +435,28 @@ def cg_couple(
             for o3_lambda in o3_lambdas
         ]
     elif cg_backend == "python-dense":
-
         n_samples = array.shape[0]
         n_properties = array.shape[3]
 
-        array = array.swapaxes(1, 3)
-        array = array.reshape(n_samples * n_properties, 2 * l2 + 1, 2 * l1 + 1)
+        # Move properties next to samples
+        array = _dispatch.permute(array, [0, 3, 1, 2])
 
-        return [
-            _cg_couple_dense(array, o3_lambda, cg_coefficients)
-            for o3_lambda in o3_lambdas
-        ]
+        # Reshape to [n_samples * n_properties, 2*l1 + 1, 2*l2 + 1]
+        array = array.reshape(n_samples * n_properties, 2 * l1 + 1, 2 * l2 + 1)
+
+        results = []
+        for o3_lambda in o3_lambdas:
+            result = _cg_couple_dense(array, o3_lambda, cg_coefficients)
+
+            # Separate samples and properties
+            result = result.reshape(n_samples, n_properties, -1)
+
+            # Back to TensorBlock-like axes
+            result = result.swapaxes(1, 2)
+
+            results.append(result)
+
+        return results
 
     else:
         raise ValueError(
@@ -507,19 +518,22 @@ def _cg_couple_dense(
     degree ``o3_lambda``) using CG coefficients. This is a "dense" implementation, using
     all CG coefficients at the same time.
 
-    :param array: input array, we expect a shape of ``[samples, 2*l1 + 1, 2*l2 + 1]``
+    :param array: input array, we expect a shape of ``[samples * properties, 2*l1 + 1,
+        2*l2 + 1]``
     :param o3_lambda: value of lambda for the output spherical harmonic
     :param cg_coefficients: CG coefficients as returned by
         :py:func:`calculate_cg_coefficients` with ``cg_backed="python-dense"``
+
+    :return: array of shape ``[samples * properties, 2 * o3_lambda + 1]``
     """
     assert len(array.shape) == 3
 
-    l1 = (array.shape[2] - 1) // 2
-    l2 = (array.shape[1] - 1) // 2
+    l1 = (array.shape[1] - 1) // 2
+    l2 = (array.shape[2] - 1) // 2
 
     cg_l1l2lam = cg_coefficients.block({"l1": l1, "l2": l2, "lambda": o3_lambda}).values
 
-    return _dispatch.tensordot(array, cg_l1l2lam[0, ..., 0], axes=([2, 1], [0, 1]))
+    return _dispatch.tensordot(array, cg_l1l2lam[0, ..., 0], axes=([2, 1], [1, 0]))
 
 
 # ======================================================================= #
diff --git a/python/featomic/featomic/clebsch_gordan/_dispatch.py b/python/featomic/featomic/clebsch_gordan/_dispatch.py
index 5aee626c0..49bd31b3c 100644
--- a/python/featomic/featomic/clebsch_gordan/_dispatch.py
+++ b/python/featomic/featomic/clebsch_gordan/_dispatch.py
@@ -74,7 +74,8 @@ def permute(array, axes: List[int]):
         return np.transpose(array, axes)
     else:
         raise TypeError(UNKNOWN_ARRAY_TYPE)
-    
+
+
 def tensordot(array_1, array_2, axes: Union[int, List[int]]):
     """
     Compute tensor dot product along specified axes for arrays.
diff --git a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
index e32b7b4ff..81300f3a1 100644
--- a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
+++ b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
@@ -1,7 +1,6 @@
 import numpy as np
 import pytest
-from metatensor import Labels, TensorBlock, TensorMap
-from metatensor import operations
+from metatensor import Labels, TensorBlock, TensorMap, operations
 
 from featomic.clebsch_gordan import cartesian_to_spherical
 
@@ -200,6 +199,7 @@ def test_cartesian_to_spherical_errors(cartesian):
             keep_l_in_keys=False,
         )
 
+
 def _l2_components_from_matrix(A):
     """
     *The* reference equations for projecting a cartesian rank 2 tensor to the
@@ -216,6 +216,7 @@ def _l2_components_from_matrix(A):
 
     return l2_A * np.sqrt(2)
 
+
 def _l3_components_from_matrix(A):
     """
     *The* reference equations for projecting a cartesian rank 3 tensor to the
@@ -274,7 +275,8 @@ def _l3_components_from_matrix(A):
     # 3
     l3_A[6] = (A[0, 0, 0] - A[1, 1, 0] - A[0, 1, 1] - A[1, 0, 1]) / 2.0
 
-    return l3_A 
+    return l3_A
+
 
 @pytest.mark.parametrize("cg_backend", ["python-dense", "python-sparse"])
 def test_cartesian_to_spherical_rank_2_by_equation(cg_backend):
@@ -289,9 +291,7 @@ def test_cartesian_to_spherical_rank_2_by_equation(cg_backend):
         blocks=[
             TensorBlock(
                 samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
-                components=[
-                    Labels(["o3_mu"], np.arange(-2, 3).reshape(-1, 1))
-                ],
+                components=[Labels(["o3_mu"], np.arange(-2, 3).reshape(-1, 1))],
                 properties=Labels(["_"], np.array([[0]])),
                 values=np.stack(
                     [_l2_components_from_matrix(A[..., 0]) for A in random_rank_2_arr]
@@ -312,7 +312,7 @@ def test_cartesian_to_spherical_rank_2_by_equation(cg_backend):
                 ],
                 properties=Labels(["_"], np.array([[0]])),
             )
-        ]
+        ],
     )
     rank_2_input_sph = cartesian_to_spherical(
         rank_2_input_cart, ["xyz1", "xyz2"], cg_backend=cg_backend
@@ -341,9 +341,7 @@ def test_cartesian_to_spherical_rank_3_by_equation(cg_backend):
         blocks=[
             TensorBlock(
                 samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
-                components=[
-                    Labels(["o3_mu"], np.arange(-3, 4).reshape(-1, 1))
-                ],
+                components=[Labels(["o3_mu"], np.arange(-3, 4).reshape(-1, 1))],
                 properties=Labels(["_"], np.array([[0]])),
                 values=np.stack(
                     [_l3_components_from_matrix(A[..., 0]) for A in random_rank_3_arr]
@@ -360,11 +358,12 @@ def test_cartesian_to_spherical_rank_3_by_equation(cg_backend):
                 values=random_rank_3_arr,
                 samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
                 components=[
-                    Labels([a], np.arange(3).reshape(-1, 1)) for a in ["xyz1", "xyz2", "xyz3"]
+                    Labels([a], np.arange(3).reshape(-1, 1))
+                    for a in ["xyz1", "xyz2", "xyz3"]
                 ],
                 properties=Labels(["_"], np.array([[0]])),
             )
-        ]
+        ],
     )
     rank_3_input_sph = cartesian_to_spherical(
         rank_3_input_cart, ["xyz1", "xyz2", "xyz3"], cg_backend=cg_backend
diff --git a/python/featomic/tests/clebsch_gordan/cg_product.py b/python/featomic/tests/clebsch_gordan/cg_product.py
index 3eb94f2cc..33071ed29 100644
--- a/python/featomic/tests/clebsch_gordan/cg_product.py
+++ b/python/featomic/tests/clebsch_gordan/cg_product.py
@@ -85,6 +85,7 @@ def spherical_expansion_by_pair(frames: List[ase.Atoms]):
     calculator = featomic.SphericalExpansionByPair(**SPHEX_HYPERS)
     return calculator.compute(frames)
 
+
 @pytest.mark.parametrize("cg_backend", ["python-sparse", "python-dense"])
 def test_keys_are_matched(cg_backend):
     """
@@ -445,8 +446,6 @@ def test_dense_sparse_agree():
 
     results = []
     for cg_backend in ["python-sparse", "python-dense"]:
-    
-
         # max_angular to be twice as big here if not using an angular cutoff
         calculator = ClebschGordanProduct(
             max_angular=MAX_ANGULAR,
diff --git a/python/featomic/tests/clebsch_gordan/coefficients.py b/python/featomic/tests/clebsch_gordan/coefficients.py
index 3c36a6d19..44fa37249 100644
--- a/python/featomic/tests/clebsch_gordan/coefficients.py
+++ b/python/featomic/tests/clebsch_gordan/coefficients.py
@@ -29,9 +29,10 @@ def complex_to_real_matrix(sph):
 
     matrix = _complex2real(ell, sph)
 
-    real = sph @ matrix
+    real = matrix @ sph.T
+    real = real.T
 
-    assert np.linalg.norm(np.imag(real)) < 1e-15
+    assert np.allclose(np.imag(real), 0)
     return np.real(real)
 
 

From c157c231aa31058d1695891211731435854b7c62 Mon Sep 17 00:00:00 2001
From: Joseph Abbott <joseph.william.abbott@gmail.com>
Date: Wed, 22 Jan 2025 17:54:24 +0100
Subject: [PATCH 08/21] Write an equivariance test for `cartesian_to_spherical`

---
 .../clebsch_gordan/cartesian_spherical.py     | 74 +++++++++++++++++++
 1 file changed, 74 insertions(+)

diff --git a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
index 81300f3a1..84d4d0bda 100644
--- a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
+++ b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
@@ -4,6 +4,8 @@
 
 from featomic.clebsch_gordan import cartesian_to_spherical
 
+from .rotations import cartesian_rotation, WignerDReal
+
 
 @pytest.fixture
 def cartesian():
@@ -379,3 +381,75 @@ def test_cartesian_to_spherical_rank_3_by_equation(cg_backend):
 
     assert operations.equal_metadata(l3_input, l3_reference)
     assert operations.allclose(l3_input, l3_reference)
+
+@pytest.mark.parametrize("cg_backend", ["python-dense", "python-sparse"])
+def test_cartesian_to_spherical_equivariance(cg_backend):
+
+    # Define a random cartesian rank 2 tensor
+    random_rank_2_arr = np.random.rand(100, 3, 3, 1)
+
+    # Define some rotation angles, the cartesian rotation matrix, and the Wigner
+    # matrices
+    angles = np.random.randn(3) * np.pi
+    R = cartesian_rotation(angles)
+    wigner = WignerDReal(max_angular=2, angles=angles)
+
+    # Rotate the cartesian tensor
+    random_rank_2_arr_rot = (
+            (
+            random_rank_2_arr.copy().reshape(100, 3, 3) @ R.T
+        ).transpose(0, 2, 1) @ R.T
+    ).transpose(0, 2, 1).reshape(100, 3, 3, 1)
+
+    # Build the Cartesian TMs and do cartesian to spherical
+    # Build the cartesian tensor and do cartesian to spherical
+    rank_2_input_cart = TensorMap(
+        keys=Labels.single(),
+        blocks=[
+            TensorBlock(
+                values=random_rank_2_arr,
+                samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
+                components=[
+                    Labels([a], np.arange(3).reshape(-1, 1)) for a in ["xyz1", "xyz2"]
+                ],
+                properties=Labels(["_"], np.array([[0]])),
+            )
+        ],
+    )
+    rank_2_input_sph = cartesian_to_spherical(
+        rank_2_input_cart, ["xyz1", "xyz2"], cg_backend=cg_backend
+    )
+    rank_2_input_cart_rot = TensorMap(
+        keys=Labels.single(),
+        blocks=[
+            TensorBlock(
+                values=random_rank_2_arr_rot,
+                samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
+                components=[
+                    Labels([a], np.arange(3).reshape(-1, 1)) for a in ["xyz1", "xyz2"]
+                ],
+                properties=Labels(["_"], np.array([[0]])),
+            )
+        ],
+    )
+    rank_2_input_sph_rot = cartesian_to_spherical(
+        rank_2_input_cart_rot, ["xyz1", "xyz2"], cg_backend=cg_backend
+    )
+
+    # Extract the lambda = 2 components
+    l2_input = operations.drop_blocks(
+        operations.remove_dimension(rank_2_input_sph, "keys", "_"),
+        keys=Labels(["o3_lambda"], np.array([[0], [1]])),
+    )
+    l2_input_rot = operations.drop_blocks(
+        operations.remove_dimension(rank_2_input_sph_rot, "keys", "_"),
+        keys=Labels(["o3_lambda"], np.array([[0], [1]])),
+    )
+
+    # Rotate the unrotated L=2 component
+    l2_input_original_rot = wigner.transform_tensormap_so3(l2_input)
+
+    print(l2_input_rot[0].values[0], l2_input_original_rot[0].values[0])
+
+    assert operations.equal_metadata(l2_input_rot, l2_input_original_rot)
+    assert operations.allclose(l2_input_rot, l2_input_original_rot)

From 7ce47542f86124cb8773058c11bba20428685d67 Mon Sep 17 00:00:00 2001
From: Joseph Abbott <joseph.william.abbott@gmail.com>
Date: Wed, 22 Jan 2025 18:36:32 +0100
Subject: [PATCH 09/21] Fix torchscipt error

---
 python/featomic/featomic/clebsch_gordan/_dispatch.py |  2 +-
 python/featomic/featomic/clebsch_gordan/_utils.py    |  2 +-
 .../tests/clebsch_gordan/cartesian_spherical.py      | 12 ++++++------
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/python/featomic/featomic/clebsch_gordan/_dispatch.py b/python/featomic/featomic/clebsch_gordan/_dispatch.py
index 49bd31b3c..43d9a8d01 100644
--- a/python/featomic/featomic/clebsch_gordan/_dispatch.py
+++ b/python/featomic/featomic/clebsch_gordan/_dispatch.py
@@ -76,7 +76,7 @@ def permute(array, axes: List[int]):
         raise TypeError(UNKNOWN_ARRAY_TYPE)
 
 
-def tensordot(array_1, array_2, axes: Union[int, List[int]]):
+def tensordot(array_1, array_2, axes: List[List[int]]):
     """
     Compute tensor dot product along specified axes for arrays.
 
diff --git a/python/featomic/featomic/clebsch_gordan/_utils.py b/python/featomic/featomic/clebsch_gordan/_utils.py
index 10c23f29f..b8d6ea8e0 100644
--- a/python/featomic/featomic/clebsch_gordan/_utils.py
+++ b/python/featomic/featomic/clebsch_gordan/_utils.py
@@ -365,7 +365,7 @@ def _compute_labels_full_cartesian_product(
     # Check for no shared labels dimensions
     for name in labels_1.names:
         assert name not in labels_2.names, (
-            "`labels_1` and `labels_2` must not have a" " dimension ({name}) in common"
+            "`labels_1` and `labels_2` must not have a dimension ({name}) in common"
         )
     # Create the new labels names by concatenating the names of the two input labels
     labels_names: List[str] = labels_1.names + labels_2.names
diff --git a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
index 84d4d0bda..a0bf7045a 100644
--- a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
+++ b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
@@ -4,7 +4,7 @@
 
 from featomic.clebsch_gordan import cartesian_to_spherical
 
-from .rotations import cartesian_rotation, WignerDReal
+from .rotations import WignerDReal, cartesian_rotation
 
 
 @pytest.fixture
@@ -382,9 +382,9 @@ def test_cartesian_to_spherical_rank_3_by_equation(cg_backend):
     assert operations.equal_metadata(l3_input, l3_reference)
     assert operations.allclose(l3_input, l3_reference)
 
+
 @pytest.mark.parametrize("cg_backend", ["python-dense", "python-sparse"])
 def test_cartesian_to_spherical_equivariance(cg_backend):
-
     # Define a random cartesian rank 2 tensor
     random_rank_2_arr = np.random.rand(100, 3, 3, 1)
 
@@ -396,10 +396,10 @@ def test_cartesian_to_spherical_equivariance(cg_backend):
 
     # Rotate the cartesian tensor
     random_rank_2_arr_rot = (
-            (
-            random_rank_2_arr.copy().reshape(100, 3, 3) @ R.T
-        ).transpose(0, 2, 1) @ R.T
-    ).transpose(0, 2, 1).reshape(100, 3, 3, 1)
+        ((random_rank_2_arr.copy().reshape(100, 3, 3) @ R.T).transpose(0, 2, 1) @ R.T)
+        .transpose(0, 2, 1)
+        .reshape(100, 3, 3, 1)
+    )
 
     # Build the Cartesian TMs and do cartesian to spherical
     # Build the cartesian tensor and do cartesian to spherical

From e6638055528d4a6da1a1d4816c3fd06886b3aa3d Mon Sep 17 00:00:00 2001
From: ppegolo <paolo.pegolo@epfl.ch>
Date: Thu, 23 Jan 2025 09:20:00 +0100
Subject: [PATCH 10/21] Lint

---
 python/featomic/tests/log.py            | 2 +-
 python/featomic/tests/systems/ase.py    | 2 +-
 python/scripts/generate-declarations.py | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/featomic/tests/log.py b/python/featomic/tests/log.py
index ae613af58..2b848fccc 100644
--- a/python/featomic/tests/log.py
+++ b/python/featomic/tests/log.py
@@ -35,7 +35,7 @@ def record_log_events(level, message):
     calculator.compute(SystemForTests())
 
     message = (
-        "featomic::calculators::dummy_calculator -- " "log-test-info: test info message"
+        "featomic::calculators::dummy_calculator -- log-test-info: test info message"
     )
     event = (FEATOMIC_LOG_LEVEL_INFO, message)
     assert event in recorded_events
diff --git a/python/featomic/tests/systems/ase.py b/python/featomic/tests/systems/ase.py
index fc74633d9..bda3288b5 100644
--- a/python/featomic/tests/systems/ase.py
+++ b/python/featomic/tests/systems/ase.py
@@ -128,7 +128,7 @@ def test_partial_pbc():
     atoms.pbc = [True, True, False]
 
     message = (
-        "different periodic boundary conditions on different axis " "are not supported"
+        "different periodic boundary conditions on different axis are not supported"
     )
     with pytest.raises(ValueError, match=message):
         AseSystem(atoms)
diff --git a/python/scripts/generate-declarations.py b/python/scripts/generate-declarations.py
index 22f19da7d..99f9db2ed 100755
--- a/python/scripts/generate-declarations.py
+++ b/python/scripts/generate-declarations.py
@@ -143,7 +143,7 @@ def funcdecl_to_ctypes(type, ndpointer=False):
     restype = type_to_ctypes(type.type, ndpointer)
     args = [type_to_ctypes(t.type, ndpointer) for t in type.args.params]
 
-    return f'CFUNCTYPE({restype}, {", ".join(args)})'
+    return f"CFUNCTYPE({restype}, {', '.join(args)})"
 
 
 def type_to_ctypes(type, ndpointer=False):

From 071440a18a1c3884ba75dc52211208798e1ddecd Mon Sep 17 00:00:00 2001
From: Joseph Abbott <joseph.william.abbott@gmail.com>
Date: Thu, 23 Jan 2025 10:24:09 +0100
Subject: [PATCH 11/21] Update changelog

---
 featomic/CHANGELOG.md | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/featomic/CHANGELOG.md b/featomic/CHANGELOG.md
index 769b6ffa6..e5c6f964f 100644
--- a/featomic/CHANGELOG.md
+++ b/featomic/CHANGELOG.md
@@ -17,6 +17,15 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows
 ### Removed
 -->
 
+### Changed
+
+- Bug fixes in `featomic.clebsch_gordan`, see issue
+  [#370](https://github.com/metatensor/featomic/issues/370):
+    - fixed `cartesian_to_spherical` transformation for tensors of rank greater than 2.
+      Previously, axes weren't properly rolled to `yzx`.
+    - added a missing transformation in the real to complex spherical harmonics
+      transformation, used in the calculation of CG coefficients.
+
 ## [Version 0.6.0](https://github.com/metatensor/featomic/releases/tag/featomic-v0.6.0) - 2024-12-20
 
 ### Added

From 5fcc2c40d9743107be1f5718b4a7e2d97c4f61b4 Mon Sep 17 00:00:00 2001
From: ppegolo <paolo.pegolo@epfl.ch>
Date: Wed, 29 Jan 2025 12:51:51 +0100
Subject: [PATCH 12/21] Fix negative sign convention in cg_couple

---
 .../featomic/featomic/clebsch_gordan/_coefficients.py  | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py
index 40294176e..873da1512 100644
--- a/python/featomic/featomic/clebsch_gordan/_coefficients.py
+++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py
@@ -503,7 +503,11 @@ def _cg_couple_sparse(
         m2 = int(m1m2mu[1])
         mu = int(m1m2mu[2])
         # Broadcast arrays, multiply together and with CG coeff
-        output[mu, :, :] += arrays[str((m1, m2))] * cg_l1l2lam.values[i, 0]
+        output[mu, :, :] += (
+            arrays[str((m1, m2))]
+            * (-1) ** (l1 + l2 + o3_lambda)
+            * cg_l1l2lam.values[i, 0]
+        )
 
     return output.swapaxes(0, 1)
 
@@ -531,7 +535,9 @@ def _cg_couple_dense(
     l1 = (array.shape[1] - 1) // 2
     l2 = (array.shape[2] - 1) // 2
 
-    cg_l1l2lam = cg_coefficients.block({"l1": l1, "l2": l2, "lambda": o3_lambda}).values
+    cg_l1l2lam = (-1) ** (l1 + l2 + o3_lambda) * cg_coefficients.block(
+        {"l1": l1, "l2": l2, "lambda": o3_lambda}
+    ).values
 
     return _dispatch.tensordot(array, cg_l1l2lam[0, ..., 0], axes=([2, 1], [1, 0]))
 

From 2ba147852ccc41b4d6dfb8a57015cb709ec9181f Mon Sep 17 00:00:00 2001
From: ppegolo <paolo.pegolo@epfl.ch>
Date: Wed, 29 Jan 2025 14:13:57 +0100
Subject: [PATCH 13/21] Add test for lambda=1, sigma=-1 to check that the
 adopted convention (page 254 Sakurai) is enforced

---
 .../clebsch_gordan/cartesian_spherical.py     | 42 +++++++++++++++----
 1 file changed, 33 insertions(+), 9 deletions(-)

diff --git a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
index a0bf7045a..f0934e868 100644
--- a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
+++ b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
@@ -202,6 +202,21 @@ def test_cartesian_to_spherical_errors(cartesian):
         )
 
 
+def _l1_components_from_matrix(A):
+    """
+    *The* reference equations for projecting a cartesian rank 2 tensor to the
+    irreducible spherical component with lambda = 1 and sigma = -1.
+    """
+    A = A.reshape(3, 3)
+
+    l1_A = np.empty((3,))
+    l1_A[0] = A[2, 0] - A[0, 2]  # y
+    l1_A[1] = A[0, 1] - A[1, 0]  # z
+    l1_A[2] = A[1, 2] - A[2, 1]  # x
+
+    return -l1_A / np.sqrt(2)
+
+
 def _l2_components_from_matrix(A):
     """
     *The* reference equations for projecting a cartesian rank 2 tensor to the
@@ -286,11 +301,20 @@ def test_cartesian_to_spherical_rank_2_by_equation(cg_backend):
     Tests cartesian_to_spherical for a random rank-2 tensor by comparing the result to
     the result from *the* reference equation.
     """
-    # Build the reference lambda = 2 component
+    # Build the reference (lambda, sigma) = (1, -1) and (lambda, sigma) = (2, 1)
+    # components
     random_rank_2_arr = np.random.rand(100, 3, 3, 1)
-    l2_reference = TensorMap(
-        keys=Labels(["o3_lambda", "o3_sigma"], np.array([[2, 1]])),
+    l1_and_l2_reference = TensorMap(
+        keys=Labels(["o3_lambda", "o3_sigma"], np.array([[1, -1], [2, 1]])),
         blocks=[
+            TensorBlock(
+                samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
+                components=[Labels(["o3_mu"], np.arange(-1, 2).reshape(-1, 1))],
+                properties=Labels(["_"], np.array([[0]])),
+                values=np.stack(
+                    [_l1_components_from_matrix(A[..., 0]) for A in random_rank_2_arr]
+                ).reshape(random_rank_2_arr.shape[0], 3, 1),
+            ),
             TensorBlock(
                 samples=Labels(["system"], np.arange(100).reshape(-1, 1)),
                 components=[Labels(["o3_mu"], np.arange(-2, 3).reshape(-1, 1))],
@@ -298,7 +322,7 @@ def test_cartesian_to_spherical_rank_2_by_equation(cg_backend):
                 values=np.stack(
                     [_l2_components_from_matrix(A[..., 0]) for A in random_rank_2_arr]
                 ).reshape(random_rank_2_arr.shape[0], 5, 1),
-            )
+            ),
         ],
     )
 
@@ -320,14 +344,14 @@ def test_cartesian_to_spherical_rank_2_by_equation(cg_backend):
         rank_2_input_cart, ["xyz1", "xyz2"], cg_backend=cg_backend
     )
 
-    # Extract the lambda = 2 component
-    l2_input = operations.drop_blocks(
+    # Extract the lambda = 1 and lambda = 2 components
+    l1_and_l2_input = operations.drop_blocks(
         operations.remove_dimension(rank_2_input_sph, "keys", "_"),
-        keys=Labels(["o3_lambda"], np.array([[0], [1]])),
+        keys=Labels(["o3_lambda"], np.array([[0]])),
     )
 
-    assert operations.equal_metadata(l2_input, l2_reference)
-    assert operations.allclose(l2_input, l2_reference)
+    assert operations.equal_metadata(l1_and_l2_input, l1_and_l2_reference)
+    assert operations.allclose(l1_and_l2_input, l1_and_l2_reference)
 
 
 @pytest.mark.parametrize("cg_backend", ["python-dense", "python-sparse"])

From 97202f490f10e1356c486ba3e788313d6a370649 Mon Sep 17 00:00:00 2001
From: Joseph Abbott <joseph.william.abbott@gmail.com>
Date: Wed, 29 Jan 2025 15:26:26 +0100
Subject: [PATCH 14/21] Address review comments

---
 featomic/CHANGELOG.md                         | 11 +---
 .../featomic/clebsch_gordan/_coefficients.py  | 55 +++++++++++++++----
 .../clebsch_gordan/cartesian_spherical.py     | 49 ++++++++---------
 3 files changed, 71 insertions(+), 44 deletions(-)

diff --git a/featomic/CHANGELOG.md b/featomic/CHANGELOG.md
index e5c6f964f..b7f96c9d2 100644
--- a/featomic/CHANGELOG.md
+++ b/featomic/CHANGELOG.md
@@ -17,14 +17,9 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows
 ### Removed
 -->
 
-### Changed
-
-- Bug fixes in `featomic.clebsch_gordan`, see issue
-  [#370](https://github.com/metatensor/featomic/issues/370):
-    - fixed `cartesian_to_spherical` transformation for tensors of rank greater than 2.
-      Previously, axes weren't properly rolled to `yzx`.
-    - added a missing transformation in the real to complex spherical harmonics
-      transformation, used in the calculation of CG coefficients.
+### Fixed
+- Fixed `featomic.clebsch_gordan.cartesian_to_spherical` transformation for tensors of rank greater than 2 (#371)
+- Fixed the calculation of Clebsch-Gordan coefficients, use all across the `featomic.clebsch_gordan` module (#371)
 
 ## [Version 0.6.0](https://github.com/metatensor/featomic/releases/tag/featomic-v0.6.0) - 2024-12-20
 
diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py
index 873da1512..62dd0e42c 100644
--- a/python/featomic/featomic/clebsch_gordan/_coefficients.py
+++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py
@@ -197,20 +197,53 @@ def _build_dense_cg_coeff_dict(
                     device=device,
                     dtype=complex_like.dtype,
                 )
-                # TODO: comment to explain what is happening here and the shapes
-                # involved
-                real_cg = (_dispatch.permute(complex_cg, [2, 1, 0]) @ r2c[l1]).swapaxes(
-                    1, 2
-                )  # this has shape (2*L+1, 2*l1+1, 2*l2+1)
-
-                real_cg = real_cg @ r2c[l2]  # this has shape (2*L+1, 2*l1+1, 2*l2+1)
-
-                real_cg = _dispatch.permute(real_cg, [1, 2, 0]) @ c2r[o3_lambda].T
+                # We want to create a CG coefficient for a real representation, using
+                # the expression:
+                #
+                #  C_{l1m1l2m2}^{lm}
+                #       = \sum_{m'm1'm2'} ...
+                #             ... U_{mm'}^{l} C_{l1m1'm2m2'}^{lm'} ...
+                #             ... U^{\dagger}_{m1'm1}^{l1} U^{\dagger}_{m2'm2}^{l2}
+                #
+                # where:
+                #       U is the "c2r" transformation matrix below
+                #       U^{\dagger} is the "r2c" transformation matrix below
+                #
+                # 0) We have a CG coefficient of shape:
+                #       complex_cg.shape = (2*l1+1, 2*l2+1, 2*L+1)
+                #    and transormation matrices of shape:
+                #       c2r[L].shape = (2*L+1, 2*L+1)
+                #       r2c[L].shape = (2*L+1, 2*L+1)
+                #
+                # 1) take the matrix product:
+                #       \sum_{m1'} C_{l1m1'm2m2'}^{lm'} U^{\dagger}_{m1'm1}^{l1}
+                #    this requires some permuting of axes of the objects involved to
+                #    ensure proper matching for matmul. i.e. permute axes of the complex
+                #    CG coefficient from: (m1, m2, m) -> (m, m2, m1)
+                first_step = _dispatch.permute(complex_cg, [2, 1, 0]) @ r2c[l1]
+                # The result `first_step` has shape (2*L+1, 2*l2+1, 2*l1+1)
+                #
+                # 2) take the matrix product:
+                #       \sum_{m2'} ...
+                #          ... (first_step)_{m',m1,m2'}^{l,l1} U^{\dagger}_{m2'm2}^{l2}
+                first_step_swap = first_step.swapaxes(1, 2)
+                # The result `first_step_swap` has shape (2*L+1, 2*l1+1, 2*l2+1)
+                second_step = first_step_swap @ r2c[l2]
+                # The result `second_step` has shape (2*L+1, 2*l1+1, 2*l2+1)
+                #
+                # 3) take the matrix product:
+                #       \sum_{m'} U_{mm'}^{l} (second_step)_{m',m1,m2}^{l,l1,l2}
+                second_step_swap = _dispatch.permute(second_step, [1, 0, 2])
+                # The result `second_step_swap` has shape (2*l1+1, 2*L+1, 2*l2+1)
+                third_step = c2r[o3_lambda] @ second_step_swap
+                # The result `third_step` has shape (2*l1+1, 2*L+1, 2*l2+1)
+                third_step_swap = _dispatch.permute(third_step, [0, 2, 1])
+                # The result `third_step_swap` has shape (2*l1+1, 2*l2+1, 2*L+1)
 
                 if (l1 + l2 + o3_lambda) % 2 == 0:
-                    cg_l1l2lam_dense = _dispatch.real(real_cg)
+                    cg_l1l2lam_dense = _dispatch.real(third_step_swap)
                 else:
-                    cg_l1l2lam_dense = _dispatch.imag(real_cg)
+                    cg_l1l2lam_dense = _dispatch.imag(third_step_swap)
 
                 coeff_dict[(l1, l2, o3_lambda)] = _dispatch.to(
                     cg_l1l2lam_dense,
diff --git a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
index f0934e868..16d805d74 100644
--- a/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
+++ b/python/featomic/tests/clebsch_gordan/cartesian_spherical.py
@@ -1,6 +1,7 @@
+import metatensor as mts
 import numpy as np
 import pytest
-from metatensor import Labels, TensorBlock, TensorMap, operations
+from metatensor import Labels, TensorBlock, TensorMap
 
 from featomic.clebsch_gordan import cartesian_to_spherical
 
@@ -243,15 +244,15 @@ def _l3_components_from_matrix(A):
 
     l3_A = np.empty((7,))
 
-    # -3
+    # o3_mu = -3
     l3_A[0] = (A[0, 1, 0] + A[1, 0, 0] + A[0, 0, 1] - A[1, 1, 1]) / 2.0
 
-    # -2
+    # o3_mu = -2
     l3_A[1] = (
         A[0, 1, 2] + A[1, 0, 2] + A[1, 2, 0] + A[2, 1, 0] + A[0, 2, 1] + A[2, 0, 1]
     ) / np.sqrt(6)
 
-    # -1
+    # o3_mu = -1
     l3_A[2] = (
         4.0 * A[1, 2, 2]
         + 4.0 * A[2, 1, 2]
@@ -262,7 +263,7 @@ def _l3_components_from_matrix(A):
         - A[1, 0, 0]
     ) / np.sqrt(60)
 
-    # 0
+    # o3_mu = 0
     l3_A[3] = (
         2.0 * A[2, 2, 2]
         - A[0, 2, 0]
@@ -273,7 +274,7 @@ def _l3_components_from_matrix(A):
         - A[1, 1, 2]
     ) / np.sqrt(10)
 
-    # 1
+    # o3_mu = 1
     l3_A[4] = (
         4.0 * A[0, 2, 2]
         + 4.0 * A[2, 0, 2]
@@ -284,12 +285,12 @@ def _l3_components_from_matrix(A):
         - A[1, 0, 1]
     ) / np.sqrt(60)
 
-    # 2
+    # o3_mu = 2
     l3_A[5] = (
         A[0, 0, 2] + A[0, 2, 0] + A[2, 0, 0] - A[1, 1, 2] - A[1, 2, 1] - A[2, 1, 1]
     ) / np.sqrt(6)
 
-    # 3
+    # o3_mu = 3
     l3_A[6] = (A[0, 0, 0] - A[1, 1, 0] - A[0, 1, 1] - A[1, 0, 1]) / 2.0
 
     return l3_A
@@ -345,13 +346,13 @@ def test_cartesian_to_spherical_rank_2_by_equation(cg_backend):
     )
 
     # Extract the lambda = 1 and lambda = 2 components
-    l1_and_l2_input = operations.drop_blocks(
-        operations.remove_dimension(rank_2_input_sph, "keys", "_"),
+    l1_and_l2_input = mts.drop_blocks(
+        mts.remove_dimension(rank_2_input_sph, "keys", "_"),
         keys=Labels(["o3_lambda"], np.array([[0]])),
     )
 
-    assert operations.equal_metadata(l1_and_l2_input, l1_and_l2_reference)
-    assert operations.allclose(l1_and_l2_input, l1_and_l2_reference)
+    assert mts.equal_metadata(l1_and_l2_input, l1_and_l2_reference)
+    assert mts.allclose(l1_and_l2_input, l1_and_l2_reference)
 
 
 @pytest.mark.parametrize("cg_backend", ["python-dense", "python-sparse"])
@@ -396,15 +397,15 @@ def test_cartesian_to_spherical_rank_3_by_equation(cg_backend):
     )
 
     # Extract the lambda = 3 component
-    l3_input = operations.drop_blocks(
-        operations.remove_dimension(rank_3_input_sph, "keys", "_"),
+    l3_input = mts.drop_blocks(
+        mts.remove_dimension(rank_3_input_sph, "keys", "_"),
         keys=Labels(["o3_lambda"], np.array([[0], [1], [2]])),
     )
     for dim in ["l_3", "k_1", "l_2", "l_1"]:
-        l3_input = operations.remove_dimension(l3_input, "keys", dim)
+        l3_input = mts.remove_dimension(l3_input, "keys", dim)
 
-    assert operations.equal_metadata(l3_input, l3_reference)
-    assert operations.allclose(l3_input, l3_reference)
+    assert mts.equal_metadata(l3_input, l3_reference)
+    assert mts.allclose(l3_input, l3_reference)
 
 
 @pytest.mark.parametrize("cg_backend", ["python-dense", "python-sparse"])
@@ -461,19 +462,17 @@ def test_cartesian_to_spherical_equivariance(cg_backend):
     )
 
     # Extract the lambda = 2 components
-    l2_input = operations.drop_blocks(
-        operations.remove_dimension(rank_2_input_sph, "keys", "_"),
+    l2_input = mts.drop_blocks(
+        mts.remove_dimension(rank_2_input_sph, "keys", "_"),
         keys=Labels(["o3_lambda"], np.array([[0], [1]])),
     )
-    l2_input_rot = operations.drop_blocks(
-        operations.remove_dimension(rank_2_input_sph_rot, "keys", "_"),
+    l2_input_rot = mts.drop_blocks(
+        mts.remove_dimension(rank_2_input_sph_rot, "keys", "_"),
         keys=Labels(["o3_lambda"], np.array([[0], [1]])),
     )
 
     # Rotate the unrotated L=2 component
     l2_input_original_rot = wigner.transform_tensormap_so3(l2_input)
 
-    print(l2_input_rot[0].values[0], l2_input_original_rot[0].values[0])
-
-    assert operations.equal_metadata(l2_input_rot, l2_input_original_rot)
-    assert operations.allclose(l2_input_rot, l2_input_original_rot)
+    assert mts.equal_metadata(l2_input_rot, l2_input_original_rot)
+    assert mts.allclose(l2_input_rot, l2_input_original_rot)

From d97f44901cbbe381ff92f2187a6a13270f6824b6 Mon Sep 17 00:00:00 2001
From: Joseph Abbott <joseph.william.abbott@gmail.com>
Date: Wed, 29 Jan 2025 15:41:36 +0100
Subject: [PATCH 15/21] Update docstring

---
 .../featomic/clebsch_gordan/_coefficients.py        | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py
index 62dd0e42c..3741526e7 100644
--- a/python/featomic/featomic/clebsch_gordan/_coefficients.py
+++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py
@@ -424,10 +424,10 @@ def cg_couple(
     Go from an uncoupled product basis that behave like a product of spherical harmonics
     to a coupled basis that behaves like a single spherical harmonic.
 
-    The ``array`` shape should be ``(n_samples, 2 * l1 + 1, 2 * l2 + 1, n_q)``.
-    ``n_samples`` is the number of samples, and ``n_q`` is the number of properties.
+    The ``array`` shape should be ``(n_s, 2 * l1 + 1, 2 * l2 + 1, n_q)``.
+    ``n_s`` is the number of samples, and ``n_q`` is the number of properties.
 
-    This function will output a list of arrays, whose shape will be ``[n_samples, (2 *
+    This function will output a list of arrays, whose shape will be ``[n_s, (2 *
     o3_lambda+1), n_q]``, with the requested ``o3_lambdas``.
 
     These arrays will contain the result of projecting from a product of spherical
@@ -441,7 +441,7 @@ def cg_couple(
     The operation is dispatched such that numpy arrays or torch tensors are
     automatically handled.
 
-    :param array: input array with shape ``[n_samples, 2 * l1 + 1, 2 * l2 + 1, n_q]``
+    :param array: input array with shape ``[n_s, 2 * l1 + 1, 2 * l2 + 1, n_q]``
     :param o3_lambdas: list of degrees of spherical harmonics to compute
     :param cg_coefficients: CG coefficients as returned by
         :py:func:`calculate_cg_coefficients` with the same ``cg_backed`` given to this
@@ -555,13 +555,12 @@ def _cg_couple_dense(
     degree ``o3_lambda``) using CG coefficients. This is a "dense" implementation, using
     all CG coefficients at the same time.
 
-    :param array: input array, we expect a shape of ``[samples * properties, 2*l1 + 1,
-        2*l2 + 1]``
+    :param array: input array, we expect a shape of ``[N, 2*l1 + 1, 2*l2 + 1]``
     :param o3_lambda: value of lambda for the output spherical harmonic
     :param cg_coefficients: CG coefficients as returned by
         :py:func:`calculate_cg_coefficients` with ``cg_backed="python-dense"``
 
-    :return: array of shape ``[samples * properties, 2 * o3_lambda + 1]``
+    :return: array of shape ``[N, 2 * o3_lambda + 1]``
     """
     assert len(array.shape) == 3
 

From 141cf2dd0405af0ccbfb96d7755960e5dc2616ad Mon Sep 17 00:00:00 2001
From: Joseph Abbott <joseph.william.abbott@gmail.com>
Date: Wed, 29 Jan 2025 16:21:46 +0100
Subject: [PATCH 16/21] =?UTF-8?q?Minor=20change=20=C3=A0=20la=20Michelange?=
 =?UTF-8?q?lo?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 python/featomic/tests/clebsch_gordan/rotations.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/featomic/tests/clebsch_gordan/rotations.py b/python/featomic/tests/clebsch_gordan/rotations.py
index db67bc79d..39ae4e90e 100644
--- a/python/featomic/tests/clebsch_gordan/rotations.py
+++ b/python/featomic/tests/clebsch_gordan/rotations.py
@@ -99,6 +99,7 @@ def __init__(self, max_angular: int, angles: Sequence[float] = None):
         # Randomly generate Euler angles between 0 and 2 pi if none are provided
         if angles is None:
             angles = np.random.uniform(size=(3)) * 2 * np.pi
+            angles[1] = angles[1]/2
         self.angles = angles
         self.rotation = cartesian_rotation(angles)
 

From c95291275df7b7821748201d49f667d8f3e40074 Mon Sep 17 00:00:00 2001
From: ppegolo <paolo.pegolo@epfl.ch>
Date: Thu, 30 Jan 2025 11:21:41 +0100
Subject: [PATCH 17/21] Linting

---
 python/featomic/tests/clebsch_gordan/rotations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/featomic/tests/clebsch_gordan/rotations.py b/python/featomic/tests/clebsch_gordan/rotations.py
index 39ae4e90e..ccadc353e 100644
--- a/python/featomic/tests/clebsch_gordan/rotations.py
+++ b/python/featomic/tests/clebsch_gordan/rotations.py
@@ -99,7 +99,7 @@ def __init__(self, max_angular: int, angles: Sequence[float] = None):
         # Randomly generate Euler angles between 0 and 2 pi if none are provided
         if angles is None:
             angles = np.random.uniform(size=(3)) * 2 * np.pi
-            angles[1] = angles[1]/2
+            angles[1] = angles[1] / 2
         self.angles = angles
         self.rotation = cartesian_rotation(angles)
 

From 109aad9323a84dd3dae140ad663faca559a18749 Mon Sep 17 00:00:00 2001
From: "Joseph W. Abbott" <joseph.william.abbott@gmail.com>
Date: Fri, 31 Jan 2025 12:02:00 +0100
Subject: [PATCH 18/21] Update featomic/CHANGELOG.md

Co-authored-by: Guillaume Fraux <guillaume.fraux@epfl.ch>
---
 featomic/CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/featomic/CHANGELOG.md b/featomic/CHANGELOG.md
index b7f96c9d2..1d5e968ab 100644
--- a/featomic/CHANGELOG.md
+++ b/featomic/CHANGELOG.md
@@ -19,7 +19,7 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows
 
 ### Fixed
 - Fixed `featomic.clebsch_gordan.cartesian_to_spherical` transformation for tensors of rank greater than 2 (#371)
-- Fixed the calculation of Clebsch-Gordan coefficients, use all across the `featomic.clebsch_gordan` module (#371)
+- Fixed the calculation of Clebsch-Gordan coefficients, used all across the `featomic.clebsch_gordan` module (#371)
 
 ## [Version 0.6.0](https://github.com/metatensor/featomic/releases/tag/featomic-v0.6.0) - 2024-12-20
 

From 44945a030734e7d2bf4a95a16712a49258750bbe Mon Sep 17 00:00:00 2001
From: "Joseph W. Abbott" <joseph.william.abbott@gmail.com>
Date: Fri, 31 Jan 2025 12:02:10 +0100
Subject: [PATCH 19/21] Update
 python/featomic/featomic/clebsch_gordan/_coefficients.py

Co-authored-by: Guillaume Fraux <guillaume.fraux@epfl.ch>
---
 python/featomic/featomic/clebsch_gordan/_coefficients.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py
index 3741526e7..bc134d63f 100644
--- a/python/featomic/featomic/clebsch_gordan/_coefficients.py
+++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py
@@ -474,7 +474,6 @@ def cg_couple(
         # Move properties next to samples
         array = _dispatch.permute(array, [0, 3, 1, 2])
 
-        # Reshape to [n_samples * n_properties, 2*l1 + 1, 2*l2 + 1]
         array = array.reshape(n_samples * n_properties, 2 * l1 + 1, 2 * l2 + 1)
 
         results = []

From 2ad555dae25dea1763ac2c8e63e45df33e204ecd Mon Sep 17 00:00:00 2001
From: ppegolo <paolo.pegolo@epfl.ch>
Date: Fri, 31 Jan 2025 17:08:46 +0100
Subject: [PATCH 20/21] Explicit mention of Euler angles

---
 python/featomic/tests/clebsch_gordan/rotations.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/python/featomic/tests/clebsch_gordan/rotations.py b/python/featomic/tests/clebsch_gordan/rotations.py
index ccadc353e..d8efba30f 100644
--- a/python/featomic/tests/clebsch_gordan/rotations.py
+++ b/python/featomic/tests/clebsch_gordan/rotations.py
@@ -96,10 +96,12 @@ def __init__(self, max_angular: int, angles: Sequence[float] = None):
         :param angles: Sequence[float], the alpha, beta, gamma Euler angles, in radians.
         """
         self.max_angular = max_angular
-        # Randomly generate Euler angles between 0 and 2 pi if none are provided
+        # Randomly generate Euler angles between if none are provided
         if angles is None:
-            angles = np.random.uniform(size=(3)) * 2 * np.pi
-            angles[1] = angles[1] / 2
+            psi = 2 * np.pi * np.random.uniform(size=1)  # psi is between 0 and 2pi
+            theta = np.pi * np.random.uniform(size=1)  # theta is between 0 and pi
+            phi = 2 * np.pi * np.random.uniform(size=1)  # phi is between 0 and 2pi
+            angles = np.array([psi, theta, phi])
         self.angles = angles
         self.rotation = cartesian_rotation(angles)
 

From 593918b35615c6aee98e9d39e07935398b1fd0ba Mon Sep 17 00:00:00 2001
From: ppegolo <paolo.pegolo@epfl.ch>
Date: Fri, 31 Jan 2025 17:52:12 +0100
Subject: [PATCH 21/21] Fix shape of angles

---
 python/featomic/tests/clebsch_gordan/rotations.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/featomic/tests/clebsch_gordan/rotations.py b/python/featomic/tests/clebsch_gordan/rotations.py
index d8efba30f..ba0cc83f2 100644
--- a/python/featomic/tests/clebsch_gordan/rotations.py
+++ b/python/featomic/tests/clebsch_gordan/rotations.py
@@ -98,9 +98,9 @@ def __init__(self, max_angular: int, angles: Sequence[float] = None):
         self.max_angular = max_angular
         # Randomly generate Euler angles between if none are provided
         if angles is None:
-            psi = 2 * np.pi * np.random.uniform(size=1)  # psi is between 0 and 2pi
-            theta = np.pi * np.random.uniform(size=1)  # theta is between 0 and pi
-            phi = 2 * np.pi * np.random.uniform(size=1)  # phi is between 0 and 2pi
+            psi = 2 * np.pi * np.random.uniform()  # psi is between 0 and 2pi
+            theta = np.pi * np.random.uniform()  # theta is between 0 and pi
+            phi = 2 * np.pi * np.random.uniform()  # phi is between 0 and 2pi
             angles = np.array([psi, theta, phi])
         self.angles = angles
         self.rotation = cartesian_rotation(angles)