Implementation of EquivariantPowerSpectrum and EquivariantPowerSpectrumByPair convenience calculators

python/rascaline/rascaline/utils/__init__.py

@@ -2,6 +2,7 @@
 
 from .clebsch_gordan import (  # noqa
     DensityCorrelations,
+    CorrelateTensorWithDensity,
     calculate_cg_coefficients,
     cartesian_to_spherical,
 )

python/rascaline/rascaline/utils/_backend.py

@@ -4,6 +4,7 @@
 from metatensor import Labels, LabelsEntry, TensorBlock, TensorMap, operations
 
 from ..calculator_base import CalculatorBase
+from ..calculators import SphericalExpansion, SphericalExpansionByPair
 from ..systems import IntoSystem
 
 

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

@@ -1,3 +1,5 @@
 from ._cartesian_spherical import cartesian_to_spherical  # noqa: F401
 from ._coefficients import calculate_cg_coefficients  # noqa: F401
 from ._correlate_density import DensityCorrelations  # noqa: F401
+from ._correlate_tensors import CorrelateTensorWithDensity  # noqa: F401
+from ._lambda_soap import LambdaSoap, LambdaSoapXSphExByPair  # noqa: F401

python/rascaline/rascaline/utils/clebsch_gordan/_correlate_density.py

@@ -135,6 +135,7 @@ def __init__(
         keep_l_in_keys: Optional[Union[bool, List[bool]]] = False,
         arrays_backend: Optional[str] = None,
         cg_backend: Optional[str] = None,
+        cg_coefficients: Optional[TensorMap] = None,
     ) -> None:
         super().__init__()
         if arrays_backend is None:
@@ -174,11 +175,14 @@ def __init__(
                 "Must be greater equal 0."
             )
         self._max_angular = max_angular
-        self._cg_coefficients = _coefficients.calculate_cg_coefficients(
-            lambda_max=self._max_angular,
-            cg_backend=self._cg_backend,
-            use_torch=(arrays_backend == "torch"),
-        )
+        if cg_coefficients is None:
+            self._cg_coefficients = _coefficients.calculate_cg_coefficients(
+                lambda_max=self._max_angular,
+                cg_backend=self._cg_backend,
+                use_torch=(arrays_backend == "torch"),
+            )
+        else:  # using a pre-computed set of CG coefficients
+            self._cg_coefficients = cg_coefficients
 
         # Check inputs
         if body_order <= 2:
@@ -274,7 +278,6 @@ def _correlate_density(
 
         # Add "order_nu" to the keys metadata
         density = operations.insert_dimension(density, "keys", 0, "order_nu", 1)
-        # density = _utils.standardize_keys(density)  # standardize metadata
 
         # Check metadata
         assert density.keys.names[:3] == [