From b13937aa981f24bc0de05e0b8fbfd8b16140786c Mon Sep 17 00:00:00 2001 From: Matthias Kellner Date: Wed, 29 Jan 2025 12:51:55 +0100 Subject: [PATCH 1/4] adds sample selection to equivariant power spectrum --- .../_equivariant_power_spectrum.py | 20 +++++++++-- .../equivariant_power_spectrum.py | 36 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) 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..e090f07a6 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py @@ -4,6 +4,7 @@ import numpy as np import pytest from metatensor import Labels, TensorBlock, TensorMap +from metatensor.operations import equal, join from numpy.testing import assert_equal from featomic import SphericalExpansion @@ -136,6 +137,41 @@ def test_equivariant_power_spectrum_neighbors_to_properties(): metatensor.equal_raise(powspec_1, powspec_2) +def tests_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 = join([powspec_1, powspec_2], axis="samples", remove_tensor_name=True) + powspec_4 = powspec_calc.compute(frame, neighbors_to_properties=True) + + assert equal(powspec_3, powspec_4) + assert not equal(powspec_2, powspec_4) + assert not equal(powspec_1, powspec_4) + + def test_fill_types_option() -> None: """ Test that ``neighbor_types`` options adds arbitrary atomic neighbor types. From ce9c0e9bd2bd263c900e940479610bc504621ecf Mon Sep 17 00:00:00 2001 From: Matthias Kellner Date: Wed, 29 Jan 2025 13:00:13 +0100 Subject: [PATCH 2/4] changes operations import --- .../tests/clebsch_gordan/equivariant_power_spectrum.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py index e090f07a6..fc03aa2c5 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py @@ -4,7 +4,6 @@ import numpy as np import pytest from metatensor import Labels, TensorBlock, TensorMap -from metatensor.operations import equal, join from numpy.testing import assert_equal from featomic import SphericalExpansion @@ -164,12 +163,12 @@ def tests_sample_selection() -> None: frame, neighbors_to_properties=True, selected_samples=label_2nd ) - powspec_3 = join([powspec_1, powspec_2], axis="samples", remove_tensor_name=True) + 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 equal(powspec_3, powspec_4) - assert not equal(powspec_2, powspec_4) - assert not equal(powspec_1, powspec_4) + 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: From fa162836d450c6971e4cb1d241bee1d26c3deb28 Mon Sep 17 00:00:00 2001 From: Matthias Kellner <62180584+bananenpampe@users.noreply.github.com> Date: Wed, 29 Jan 2025 13:27:58 +0100 Subject: [PATCH 3/4] Update python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py Co-authored-by: Joseph W. Abbott --- .../featomic/tests/clebsch_gordan/equivariant_power_spectrum.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py index fc03aa2c5..8d7eb829a 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py @@ -136,7 +136,7 @@ def test_equivariant_power_spectrum_neighbors_to_properties(): metatensor.equal_raise(powspec_1, powspec_2) -def tests_sample_selection() -> None: +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. From 4ad7fa0a89c69f6156cba396884ed3fd5a02c0d3 Mon Sep 17 00:00:00 2001 From: Matthias Kellner Date: Thu, 30 Jan 2025 12:28:15 +0100 Subject: [PATCH 4/4] lint --- .../tests/clebsch_gordan/equivariant_power_spectrum.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py index fc03aa2c5..9b439cb0d 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum.py @@ -163,7 +163,9 @@ def tests_sample_selection() -> None: frame, neighbors_to_properties=True, selected_samples=label_2nd ) - powspec_3 = metatensor.join([powspec_1, powspec_2], axis="samples", remove_tensor_name=True) + 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)