diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum.py index 6354f8356..79d6f52f2 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum.py +++ b/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,6 +283,7 @@ def forward( return self.compute( systems=systems, selected_keys=selected_keys, + selected_samples=selected_samples, neighbors_to_properties=neighbors_to_properties, ) @@ -282,6 +291,7 @@ 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,6 +313,7 @@ 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: @@ -309,12 +321,16 @@ def _equivariant_power_spectrum( 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( diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py index 1f47efe64..6f2ca3a77 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py +++ b/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.