Transform cartesian tensors into spherical ones

docs/src/references/api/python/utils/clebsch-gordan.rst

@@ -3,3 +3,8 @@ Clebsch-Gordan products
 
 .. autoclass:: rascaline.utils.DensityCorrelations
     :members:
+
+
+.. autofunction:: rascaline.utils.cartesian_to_spherical
+
+.. autofunction:: rascaline.utils.calculate_cg_coefficients

python/rascaline-torch/rascaline/torch/utils.py

@@ -38,6 +38,7 @@
 module.__dict__["Array"] = torch.Tensor
 module.__dict__["CalculatorBase"] = CalculatorModule
 module.__dict__["IntoSystem"] = System
+module.__dict__["BACKEND_IS_METATENSOR_TORCH"] = True
 
 
 def is_labels(obj: Any):

python/rascaline-torch/tests/utils/cartesian_spherical.py

@@ -0,0 +1,98 @@
+import pytest
+import torch
+from metatensor.torch import Labels, TensorBlock, TensorMap
+
+from rascaline.torch.utils.clebsch_gordan import cartesian_to_spherical
+
+
+@pytest.fixture
+def cartesian():
+    # the first block is completely symmetric
+    values_1 = torch.rand(10, 4, 3, 3, 3, 2, dtype=torch.float64)
+    values_1[:, :, 0, 1, 0, :] = values_1[:, :, 0, 0, 1, :]
+    values_1[:, :, 1, 0, 0, :] = values_1[:, :, 0, 0, 1, :]
+
+    values_1[:, :, 0, 2, 0, :] = values_1[:, :, 0, 0, 2, :]
+    values_1[:, :, 2, 0, 0, :] = values_1[:, :, 0, 0, 2, :]
+
+    values_1[:, :, 1, 0, 1, :] = values_1[:, :, 0, 1, 1, :]
+    values_1[:, :, 1, 1, 0, :] = values_1[:, :, 0, 1, 1, :]
+
+    values_1[:, :, 2, 0, 2, :] = values_1[:, :, 0, 2, 2, :]
+    values_1[:, :, 2, 2, 0, :] = values_1[:, :, 0, 2, 2, :]
+
+    values_1[:, :, 2, 1, 2, :] = values_1[:, :, 2, 2, 1, :]
+    values_1[:, :, 1, 2, 2, :] = values_1[:, :, 2, 2, 1, :]
+
+    values_1[:, :, 1, 2, 1, :] = values_1[:, :, 1, 1, 2, :]
+    values_1[:, :, 2, 1, 1, :] = values_1[:, :, 1, 1, 2, :]
+
+    values_1[:, :, 0, 2, 1, :] = values_1[:, :, 0, 1, 2, :]
+    values_1[:, :, 2, 0, 1, :] = values_1[:, :, 0, 1, 2, :]
+    values_1[:, :, 1, 0, 2, :] = values_1[:, :, 0, 1, 2, :]
+    values_1[:, :, 1, 2, 0, :] = values_1[:, :, 0, 1, 2, :]
+    values_1[:, :, 2, 1, 0, :] = values_1[:, :, 0, 1, 2, :]
+
+    block_1 = TensorBlock(
+        values=values_1,
+        samples=Labels.range("s", 10),
+        components=[
+            Labels.range("other", 4),
+            Labels.range("xyz_1", 3),
+            Labels.range("xyz_2", 3),
+            Labels.range("xyz_3", 3),
+        ],
+        properties=Labels.range("p", 2),
+    )
+
+    # second block does not have any specific symmetry
+    block_2 = TensorBlock(
+        values=torch.rand(12, 6, 3, 3, 3, 7, dtype=torch.float64),
+        samples=Labels.range("s", 12),
+        components=[
+            Labels.range("other", 6),
+            Labels.range("xyz_1", 3),
+            Labels.range("xyz_2", 3),
+            Labels.range("xyz_3", 3),
+        ],
+        properties=Labels.range("p", 7),
+    )
+
+    return TensorMap(Labels.range("key", 2), [block_1, block_2])
+
+
+def test_torch_script():
+    torch.jit.script(cartesian_to_spherical)
+
+
+def test_cartesian_to_spherical(cartesian):
+    # rank 1
+    spherical = cartesian_to_spherical(cartesian, components=["xyz_1"])
+
+    assert spherical.component_names == ["other", "o3_mu", "xyz_2", "xyz_3"]
+    assert spherical.keys.names == ["o3_lambda", "o3_sigma", "key"]
+    assert len(spherical.keys) == 2
+
+    # rank 2
+    spherical = cartesian_to_spherical(cartesian, components=["xyz_1", "xyz_2"])
+
+    assert spherical.component_names == ["other", "o3_mu", "xyz_3"]
+    assert spherical.keys.names == ["o3_lambda", "o3_sigma", "key"]
+    assert len(spherical.keys) == 5
+
+    # rank 3
+    spherical = cartesian_to_spherical(
+        cartesian, components=["xyz_1", "xyz_2", "xyz_3"]
+    )
+
+    assert spherical.component_names == ["other", "o3_mu"]
+    assert spherical.keys.names == [
+        "o3_lambda",
+        "o3_sigma",
+        "l_3",
+        "k_1",
+        "l_2",
+        "l_1",
+        "key",
+    ]
+    assert len(spherical.keys) == 10

python/rascaline/rascaline/utils/__init__.py

@@ -1,6 +1,10 @@
 import os
 
-from .clebsch_gordan import DensityCorrelations  # noqa
+from .clebsch_gordan import (  # noqa
+    DensityCorrelations,
+    calculate_cg_coefficients,
+    cartesian_to_spherical,
+)
 from .power_spectrum import PowerSpectrum  # noqa
 from .splines import (  # noqa
     AtomicDensityBase,

python/rascaline/rascaline/utils/_backend.py

@@ -39,6 +39,8 @@ class TorchScriptClass:
 
 Array = Union[np.ndarray, TorchTensor]
 
+BACKEND_IS_METATENSOR_TORCH = False
+
 __all__ = [
     "Array",
     "CalculatorBase",
@@ -53,4 +55,5 @@ class TorchScriptClass:
     "torch_jit_is_scripting",
     "torch_jit_export",
     "is_labels",
+    "BACKEND_IS_METATENSOR_TORCH",
 ]

python/rascaline/rascaline/utils/_dispatch.py

@@ -437,12 +437,25 @@ def imag(array):
     """
     Takes the imag part of the array
 
-    This function has the same behavior as
-    ``np.imag(array)`` or ``torch.imag(array)``.
+    This function has the same behavior as ``np.imag(array)`` or ``torch.imag(array)``.
     """
     if isinstance(array, TorchTensor):
         return torch.imag(array)
     elif isinstance(array, np.ndarray):
         return np.imag(array)
     else:
         raise TypeError(UNKNOWN_ARRAY_TYPE)
+
+
+def roll(array, shifts: List[int], axis: List[int]):
+    """
+    Roll array elements along a given axis.
+
+    This function has the same behavior as ``np.roll(array)`` or ``torch.roll(array)``.
+    """
+    if isinstance(array, TorchTensor):
+        return torch.roll(array, shifts=shifts, dims=axis)
+    elif isinstance(array, np.ndarray):
+        return np.roll(array, shift=shifts, axis=axis)
+    else:
+        raise TypeError(UNKNOWN_ARRAY_TYPE)

python/rascaline/rascaline/utils/clebsch_gordan/__init__.py

@@ -1 +1,3 @@
+from ._cartesian_spherical import cartesian_to_spherical  # noqa: F401
+from ._coefficients import calculate_cg_coefficients  # noqa: F401
 from ._correlate_density import DensityCorrelations  # noqa: F401