Adds sample selection to equivariant power spectrum (#376)

Author: bananenpampe

python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum.py

@@ -217,6 +217,7 @@ def compute(
         self,
         systems: Union[IntoSystem, List[IntoSystem]],
         selected_keys: Optional[Labels] = None,
+        selected_samples: Optional[Labels] = None,
         neighbors_to_properties: bool = False,
     ) -> TensorMap:
         """
@@ -245,6 +246,11 @@ def compute(
         :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 selected_samples: :py:class:`Labels`, Set of samples on which to run the
+            calculation. Use ``None`` to run the calculation on all samples in
+            the systems (this is the default). Gets passed to ``calculator_1``  and
+            ``calculator_2``, therefore requiring that both calculators support sample
+            selection.
         :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.
@@ -254,6 +260,7 @@ def compute(
         return self._equivariant_power_spectrum(
             systems=systems,
             selected_keys=selected_keys,
+            selected_samples=selected_samples,
             neighbors_to_properties=neighbors_to_properties,
             compute_metadata=False,
         )
@@ -262,6 +269,7 @@ def forward(
         self,
         systems: Union[IntoSystem, List[IntoSystem]],
         selected_keys: Optional[Labels] = None,
+        selected_samples: Optional[Labels] = None,
         neighbors_to_properties: bool = False,
     ) -> TensorMap:
         """
@@ -275,13 +283,15 @@ def forward(
         return self.compute(
             systems=systems,
             selected_keys=selected_keys,
+            selected_samples=selected_samples,
             neighbors_to_properties=neighbors_to_properties,
         )
 
     def compute_metadata(
         self,
         systems: Union[IntoSystem, List[IntoSystem]],
         selected_keys: Optional[Labels] = None,
+        selected_samples: Optional[Labels] = None,
         neighbors_to_properties: bool = False,
     ) -> TensorMap:
         """
@@ -294,6 +304,7 @@ def compute_metadata(
         return self._equivariant_power_spectrum(
             systems=systems,
             selected_keys=selected_keys,
+            selected_samples=selected_samples,
             neighbors_to_properties=neighbors_to_properties,
             compute_metadata=True,
         )
@@ -302,19 +313,24 @@ def _equivariant_power_spectrum(
         self,
         systems: Union[IntoSystem, List[IntoSystem]],
         selected_keys: Optional[Labels],
+        selected_samples: Optional[Labels],
         neighbors_to_properties: bool,
         compute_metadata: bool,
     ) -> TensorMap:
         """
         Computes the equivariant power spectrum, either fully or just metadata
         """
         # Compute density
-        density_1 = self.calculator_1.compute(systems)
+        density_1 = self.calculator_1.compute(
+            systems, selected_samples=selected_samples
+        )
 
         if self.calculator_2 is None:
             density_2 = density_1
         else:
-            density_2 = self.calculator_2.compute(systems)
+            density_2 = self.calculator_2.compute(
+                systems, selected_samples=selected_samples
+            )
 
         # Rename "neighbor_type" dimension so they are correlated
         density_1 = operations.rename_dimension(

python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py

@@ -136,6 +136,43 @@ def test_equivariant_power_spectrum_neighbors_to_properties():
     metatensor.equal_raise(powspec_1, powspec_2)
 
 
+def test_sample_selection() -> None:
+    """Tests that the sample selection works as expected.
+    By first computing the powerspectruim for all atoms in H2O
+    Then first for atom 1 and then atom 2  and 3.
+    Their join should be identical to computing it for all atoms.
+    """
+
+    frame = h2o_periodic()
+
+    powspec_calc = EquivariantPowerSpectrum(SphericalExpansion(**SPHEX_HYPERS_SMALL))
+
+    label_1st = metatensor.Labels(
+        ["system", "atom"], np.array([[0, 0]], dtype=np.int32)
+    )
+
+    label_2nd = metatensor.Labels(
+        ["system", "atom"], np.array([[0, 1], [0, 2]], dtype=np.int32)
+    )
+
+    powspec_1 = powspec_calc.compute(
+        frame, neighbors_to_properties=True, selected_samples=label_1st
+    )
+
+    powspec_2 = powspec_calc.compute(
+        frame, neighbors_to_properties=True, selected_samples=label_2nd
+    )
+
+    powspec_3 = metatensor.join(
+        [powspec_1, powspec_2], axis="samples", remove_tensor_name=True
+    )
+    powspec_4 = powspec_calc.compute(frame, neighbors_to_properties=True)
+
+    assert metatensor.equal(powspec_3, powspec_4)
+    assert not metatensor.equal(powspec_2, powspec_4)
+    assert not metatensor.equal(powspec_1, powspec_4)
+
+
 def test_fill_types_option() -> None:
     """
     Test that ``neighbor_types`` options adds arbitrary atomic neighbor types.