Fix cartesian_to_spherical for rank > 2 tensors and real_to_complex transformation

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

python/featomic/featomic/clebsch_gordan/_cartesian_spherical.py

@@ -245,6 +245,10 @@ 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)
+
+    # Use the rolled block values as the starting point for higher-rank tensors
+    tensor = TensorMap(tensor.keys, new_blocks)
+
     if cg_coefficients is None:
         if torch_jit_is_scripting():
             raise ValueError(

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
 

python/featomic/featomic/clebsch_gordan/_coefficients.py

@@ -197,18 +197,15 @@ 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 = (r2c[l1] @ complex_cg.reshape(2 * l1 + 1, -1)).reshape(
-                    complex_cg.shape
-                )
-
-                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 @ r2c[l2]  # this has shape (2*L+1, 2*l1+1, 2*l2+1)
 
-                real_cg = real_cg @ c2r[o3_lambda]
+                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)
@@ -366,7 +363,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 +435,25 @@ 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)
+        # 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 = []
         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
@@ -514,10 +518,13 @@ 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
 
@@ -526,14 +533,7 @@ 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))
-
-    # [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
+    return _dispatch.tensordot(array, cg_l1l2lam[0, ..., 0], axes=([2, 1], [1, 0]))
 
 
 # ======================================================================= #

python/featomic/featomic/clebsch_gordan/_dispatch.py

@@ -63,6 +63,34 @@ def concatenate(arrays: List[TorchTensor], axis: int):
         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 tensordot(array_1, array_2, axes: List[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):
     """
     Create an uninitialized array, with the given ``shape``, and similar dtype, device

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