Implement convenience calculator EquivariantPowerSpectrum

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

@@ -2,7 +2,7 @@
 import torch
 from metatensor.torch import Labels, TensorBlock, TensorMap
 
-from rascaline.torch.utils.clebsch_gordan import cartesian_to_spherical
+from rascaline.torch.utils import cartesian_to_spherical
 
 
 @pytest.fixture

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

@@ -8,7 +8,7 @@
 from metatensor.torch.atomistic import System
 
 import rascaline.torch
-from rascaline.torch.utils.clebsch_gordan import ClebschGordanProduct
+from rascaline.torch.utils import ClebschGordanProduct
 
 
 SPHERICAL_EXPANSION_HYPERS = {

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

@@ -8,7 +8,7 @@
 from metatensor.torch.atomistic import System
 
 import rascaline.torch
-from rascaline.torch.utils.clebsch_gordan import DensityCorrelations
+from rascaline.torch.utils import DensityCorrelations
 
 
 SPHERICAL_EXPANSION_HYPERS = {

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

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+import io
+
+import torch
+
+from rascaline.torch.utils import EquivariantPowerSpectrum
+
+
+SPHERICAL_EXPANSION_HYPERS = {
+    "cutoff": 2.5,
+    "max_radial": 3,
+    "max_angular": 3,
+    "atomic_gaussian_width": 0.2,
+    "radial_basis": {"Gto": {}},
+    "cutoff_function": {"ShiftedCosine": {"width": 0.5}},
+    "center_atom_weight": 1.0,
+}
+
+
+def test_jit_save_load():
+    calculator = torch.jit.script(
+        EquivariantPowerSpectrum(
+            **SPHERICAL_EXPANSION_HYPERS,
+            dtype=torch.float64,
+        )
+    )
+
+    with io.BytesIO() as buffer:
+        torch.jit.save(calculator, buffer)
+        buffer.seek(0)
+        torch.jit.load(buffer)

python/rascaline/rascaline/utils/__init__.py

@@ -3,6 +3,7 @@
 from .clebsch_gordan import (  # noqa
     ClebschGordanProduct,
     DensityCorrelations,
+    EquivariantPowerSpectrum,
     calculate_cg_coefficients,
     cartesian_to_spherical,
 )

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

@@ -2,3 +2,4 @@
 from ._cg_product import ClebschGordanProduct  # noqa: F401
 from ._coefficients import calculate_cg_coefficients  # noqa: F401
 from ._density_correlations import DensityCorrelations  # noqa: F401
+from ._equivariant_power_spectrum import EquivariantPowerSpectrum  # noqa: F401

python/rascaline/rascaline/utils/clebsch_gordan/_density_correlations.py

@@ -1,5 +1,5 @@
 """
-This module provides convenience calculators for preforming density correlations, i.e.
+This module provides a convenience calculator for preforming density correlations, i.e.
 the (iterative) CG tensor products of density (body order 2) tensors.
 
 All of these calculators wrap the :py:class:`ClebschGordanProduct` class, handling the

python/rascaline/rascaline/utils/clebsch_gordan/_equivariant_power_spectrum.py

@@ -0,0 +1,201 @@
+"""
+This module provides a convenience calculator for computing a single-center equivariant
+power spectrum.
+"""
+
+from typing import List, Optional, Union
+
+from ...calculators import SphericalExpansion
+from ...systems import IntoSystem
+from .._backend import Device, DType, Labels, TensorMap, TorchModule, operations
+from ._cg_product import ClebschGordanProduct
+from ._density_correlations import _filter_redundant_keys
+
+
+class EquivariantPowerSpectrum(TorchModule):
+    """
+    Computes an equivariant power spectrum descriptor, or a "lambda-SOAP".
+
+    For a full description of the hyper-parameters, see the corresponding
+    :ref:`documentation <soap-power-spectrum>`.
+    """
+
+    def __init__(
+        self,
+        cutoff,
+        max_radial,
+        max_angular,
+        atomic_gaussian_width,
+        center_atom_weight,
+        radial_basis,
+        cutoff_function,
+        radial_scaling=None,
+        *,
+        dtype: Optional[DType] = None,
+        device: Optional[Device] = None,
+    ):
+        """
+        :param spherical_expansion_hypers: :py:class:`dict` containing the
+            hyper-parameters used to initialize a :py:class:`SphericalExpansion` for
+            computing the initial density.
+        :param atom_types: :py:class:`list` of :py:class:`str`, the global atom types
+            to compute neighbor correlations for. Ensures consistent global properties
+            dimensions.
+        :param dtype: the scalar type to use to store coefficients
+        :param device: the computational device to use for calculations. This must be
+            ``"cpu"`` if ``array_backend="numpy"``.
+        """
+
+        super().__init__()
+
+        parameters = {
+            "cutoff": cutoff,
+            "max_radial": max_radial,
+            "max_angular": max_angular,
+            "atomic_gaussian_width": atomic_gaussian_width,
+            "center_atom_weight": center_atom_weight,
+            "radial_basis": radial_basis,
+            "cutoff_function": cutoff_function,
+        }
+
+        if radial_scaling is not None:
+            parameters["radial_scaling"] = radial_scaling
+
+        self._spherical_expansion = SphericalExpansion(**parameters)
+        self._cg_product = ClebschGordanProduct(
+            max_angular=max_angular * 2,
+            cg_backend=None,
+            keys_filter=_filter_redundant_keys,
+            arrays_backend=None,
+            dtype=dtype,
+            device=device,
+        )
+
+    def compute(
+        self,
+        systems: Union[IntoSystem, List[IntoSystem]],
+        *,
+        selected_keys: Optional[Labels] = None,
+        neighbors_to_properties: bool = False,
+    ) -> TensorMap:
+        """
+        Computes an equivariant power spectrum, or "lambda-SOAP".
+
+        First computes a :py:class:`SphericalExpansion` density descriptor of body order
+        2.
+
+        The key dimension 'neighbor_type' is then moved to properties so that they are
+        correlated. The global atom types passed in the constructor are taken into
+        account so that a consistent output properties dimension is achieved regardless
+        of the atomic composition of the systems passed in ``systems``.
+
+        Finally a single Clebsch-Gordan tensor product is taken to produce a body order
+        3 equivariant power spectrum, or "lambda-SOAP".
+
+        :param selected_keys: :py:class:`Labels`, the output keys to computed. If
+            ``None``, all keys are computed. Subsets of key dimensions can be passed to
+            compute output blocks that match in these dimensions.
+        :param neighbors_to_properties: :py:class:`bool`, if true, densifies the
+            spherical expansion by moving key dimension "neighbor_type" to properties
+            prior to performing the Clebsch Gordan product step. Defaults to false.
+
+        :return: :py:class:`TensorMap`, the output equivariant power spectrum.
+        """
+        return self._equivariant_power_spectrum(
+            systems=systems,
+            selected_keys=selected_keys,
+            neighbors_to_properties=neighbors_to_properties,
+            compute_metadata=False,
+        )
+
+    def forward(
+        self,
+        systems: Union[IntoSystem, List[IntoSystem]],
+        *,
+        selected_keys: Optional[Labels] = None,
+        neighbors_to_properties: bool = False,
+    ) -> TensorMap:
+        """
+        Calls the :py:meth:`compute` method.
+
+        This is intended for :py:class:`torch.nn.Module` compatibility, and should be
+        ignored in pure Python mode.
+
+        See :py:meth:`compute` for a full description of the parameters.
+        """
+        return self.compute(
+            systems=systems,
+            selected_keys=selected_keys,
+            neighbors_to_properties=neighbors_to_properties,
+        )
+
+    def compute_metadata(
+        self,
+        systems: Union[IntoSystem, List[IntoSystem]],
+        *,
+        selected_keys: Optional[Labels] = None,
+        neighbors_to_properties: bool = False,
+    ) -> TensorMap:
+        """
+        Returns the metadata-only :py:class:`TensorMap` that would be output by the
+        function :py:meth:`compute` for the same calculator under the same settings,
+        without performing the actual Clebsch-Gordan tensor products in the second step.
+
+        See :py:meth:`compute` for a full description of the parameters.
+        """
+        return self._equivariant_power_spectrum(
+            systems=systems,
+            selected_keys=selected_keys,
+            neighbors_to_properties=neighbors_to_properties,
+            compute_metadata=True,
+        )
+
+    def _equivariant_power_spectrum(
+        self,
+        systems: Union[IntoSystem, List[IntoSystem]],
+        selected_keys: Optional[Labels],
+        neighbors_to_properties: bool,
+        compute_metadata: bool,
+    ) -> TensorMap:
+        """
+        Computes the equivariant power spectrum, either fully or just metadata
+        """
+        # Compute density
+        density = self._spherical_expansion.compute(systems)
+
+        # Rename "neighbor_type" dimension so they are correlated
+        density_1 = operations.rename_dimension(
+            density, "keys", "neighbor_type", "neighbor_1_type"
+        )
+        density_2 = operations.rename_dimension(
+            density, "keys", "neighbor_type", "neighbor_2_type"
+        )
+        density_1 = operations.rename_dimension(density_1, "properties", "n", "n_1")
+        density_2 = operations.rename_dimension(density_2, "properties", "n", "n_2")
+
+        if neighbors_to_properties:
+            density_1 = density_1.keys_to_properties("neighbor_1_type")
+            density_2 = density_2.keys_to_properties("neighbor_2_type")
+
+        # Compute the power spectrum
+        if compute_metadata:
+            pow_spec = self._cg_product.compute_metadata(
+                tensor_1=density_1,
+                tensor_2=density_2,
+                o3_lambda_1_new_name="l_1",
+                o3_lambda_2_new_name="l_2",
+                selected_keys=selected_keys,
+            )
+        else:
+            pow_spec = self._cg_product.compute(
+                tensor_1=density_1,
+                tensor_2=density_2,
+                o3_lambda_1_new_name="l_1",
+                o3_lambda_2_new_name="l_2",
+                selected_keys=selected_keys,
+            )
+
+        # Move the CG combination info keys to properties
+        pow_spec = pow_spec.keys_to_properties(["l_1", "l_2"])
+
+        return pow_spec

python/rascaline/rascaline/utils/clebsch_gordan/_utils.py

@@ -285,10 +285,34 @@ def _match_samples_of_blocks(
 
     TODO: implement for samples dimensions that are not a subset of the other.
     """
-    # If the number of dimensions are the same, check they are equivalent and return
+    # The number of dimensions are the same
     if len(block_1.samples.names) == len(block_2.samples.names):
-        if not block_1.samples == block_2.samples:
-            raise ValueError("Samples dimensions of the two blocks are not equivalent.")
+        if block_1.samples == block_2.samples:  # nothing needs to be done
+            return block_1, block_2
+
+        # Find the union of samples and broadcast both blocks along samples axis
+        new_blocks = []
+        union, mapping_1, mapping_2 = block_1.samples.union_and_mapping(block_2.samples)
+        for block, mapping in [(block_1, mapping_1), (block_2, mapping_2)]:
+            new_block_vals = _dispatch.zeros_like(
+                block.values, (len(union), *block.values.shape[1:])
+            )
+            new_block_vals[mapping] = block.values
+            new_block = TensorBlock(
+                values=new_block_vals,
+                samples=union,
+                components=block.components,
+                properties=block.properties,
+            )
+            new_blocks.append(new_block)
+
+        block_1, block_2 = new_blocks
+
+        return block_1, block_2
+
+    # Otherwise, the samples dimensions of one block is assumed (and checked) to be a
+    # subset of the other.
+    assert len(block_1.samples.names) != len(block_2.samples.names)
 
     # First find the block with fewer dimensions. Reorder to have this block on the
     # 'left' for simplicity, but record the original order for the final output