From 67754f3a67a45226c7b09e19d96120de974dd891 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jan 2025 16:29:12 +0100 Subject: [PATCH 01/33] Explicit reshaping (required if n_s==0) --- python/featomic/featomic/clebsch_gordan/_coefficients.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py index e7d94f336..924b13a11 100644 --- a/python/featomic/featomic/clebsch_gordan/_coefficients.py +++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py @@ -701,7 +701,7 @@ def _cg_tensor_product_dense( # output shape is [(samples q p), lambda] output = _cg_couple_dense(tensor_product, o3_lambda, cg_coefficients) # => [samples, (q p), lambda] - output = output.reshape(n_s, (n_p * n_q), -1) + output = output.reshape(n_s, (n_p * n_q), 2 * o3_lambda + 1) # => [samples, lambda, (q p)] output = _dispatch.swapaxes(output, 1, 2) result.append(output) From dca005e299a119a8ed72641bac520cc9e12e8d05 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jan 2025 16:29:29 +0100 Subject: [PATCH 02/33] Add equivariant power spectrum by pair --- .../_equivariant_power_spectrum_by_pair.py | 361 ++++++++++++++++++ 1 file changed, 361 insertions(+) create mode 100644 python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py new file mode 100644 index 000000000..781fd9e6f --- /dev/null +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -0,0 +1,361 @@ +""" +This module provides a convenience calculator for computing a two-center equivariant +power spectrum, or equivariant power spectrum by pair +""" + +import json +from typing import List, Optional, Union + +from . import _dispatch +from ._backend import ( + CalculatorBase, + Device, + DType, + IntoSystem, + Labels, + TensorMap, + TorchModule, + operations, +) +from ._cg_product import ClebschGordanProduct +from ._density_correlations import _filter_redundant_keys + + +class EquivariantPowerSpectrumByPair(TorchModule): + r""" + Computes a general equivariant power spectrum by pair descriptor of two calculators, + the second being a calculator by pair. + + Example + ------- + + TODO: update docstring + + As an example we calculate the equivariant power spectrum by pair for a spherical + expansion and a spherical expansion by pair for a NaCl crystal. + + >>> import featomic + >>> import ase + + Construct the NaCl crystal + + >>> atoms = ase.Atoms( + ... symbols="NaCl", + ... positions=[[0, 0, 0], [0.5, 0.5, 0.5]], + ... pbc=True, + ... cell=[1, 1, 1], + ... ) + + Define the hyper parameters for the short-range spherical expansion + + >>> spex_hypers = { + ... "cutoff": { + ... "radius": 3.0, + ... "smoothing": {"type": "ShiftedCosine", "width": 0.5}, + ... }, + ... "density": { + ... "type": "Gaussian", + ... "width": 0.3, + ... }, + ... "basis": { + ... "type": "TensorProduct", + ... "max_angular": 2, + ... "radial": {"type": "Gto", "max_radial": 5}, + ... }, + ... } + + Define the hyper parameters for the spherical expansion by pair + + >>> spex_by_pair_hypers = { + ... "cutoff": { + ... "radius": 5.0, + ... "smoothing": {"type": "ShiftedCosine", "width": 0.5}, + ... }, + ... "density": { + ... "type": "Gaussian", + ... "width": 0.3, + ... }, + ... "basis": { + ... "type": "TensorProduct", + ... "max_angular": 1, + ... "radial": {"type": "Gto", "max_radial": 5}, + ... }, + ... } + + Construct the calculators + + >>> spex_calculator = featomic.SphericalExpansion(**spex_hypers) + >>> spex_by_pair_calculator = featomic.LodeSphericalExpansion(**spex_by_pair_hypers) + + Construct the power spectrum by pair calculators and compute the spherical expansion + + >>> calculator = featomic.clebsch_gordan.EquivariantPowerSpectrumByPair( + ... spex_calculator, spex_by_pair_calculator + ... ) + >>> power_spectrum_by_pair = calculator.compute(atoms, neighbors_to_properties=True) + + The resulting equivariants are stored as :py:class:`metatensor.TensorMap` as for any + other calculator. The keys contain the symmetry information: + + >>> power_spectrum_by_pair.keys + Labels( + o3_lambda o3_sigma first_atom_type second_atom_type + 0 1 11 11 + 1 1 11 11 + 0 1 11 17 + 1 1 11 17 + 1 -1 11 11 + 2 1 11 11 + 1 -1 11 17 + 2 1 11 17 + 0 1 17 11 + 1 1 17 11 + 0 1 17 17 + 1 1 17 17 + 1 -1 17 11 + 2 1 17 11 + 1 -1 17 17 + 2 1 17 17 + ) + + The block properties contain the angular order of the combined blocks ("l_1", + "l_2"), along with the neighbor type of the full spherical expansion and the radial + channel indices. + + >>> power_spectrum_by_pair[0].properties.names + ['l_1', 'l_2', 'neighbor_1_type', 'n_1', 'n_2'] + + .. seealso:: + An equivariant power spectrum calculator for single-center descriptors can be + found at :py:class:`featomic.clebsch_gordan.EquivariantPowerSpectrum`. + """ + + def __init__( + self, + calculator_1: CalculatorBase, + calculator_2: CalculatorBase, + neighbor_types: Optional[List[int]] = None, + *, + dtype: Optional[DType] = None, + device: Optional[Device] = None, + ): + """ + + Constructs the equivariant power spectrum by pair calculator. + + :param calculator_1: first calculator that computes a density descriptor, either + a :py:class:`featomic.SphericalExpansion` or + :py:class:`featomic.LodeSphericalExpansion`. + :param calculator_2: second calculator that computes a density by pair + descriptor, it must be :py:class:`featomic.SphericalExpansionByPair` for + now. + :param neighbor_types: List of ``"neighbor_type"`` to use in the properties of + the output. This option might be useful when running the calculation on + subset of a whole dataset and trying to join along the ``sample`` dimension + after the calculation. If ``None``, blocks are filled with + ``"neighbor_type"`` found in the systems. This parameter is only used if + ``neighbors_to_properties=True`` is passed to the :py:meth:`compute` method. + :param dtype: the scalar type to use to store coefficients + :param device: the computational device to use for calculations. + """ + + super().__init__() + self.calculator_1 = calculator_1 + self.calculator_2 = calculator_2 + self.neighbor_types = neighbor_types + self.dtype = dtype + self.device = device + + supported_calculators = ["lode_spherical_expansion", "spherical_expansion"] + + if self.calculator_1.c_name not in supported_calculators: + raise ValueError( + f"Only [{', '.join(supported_calculators)}] are supported for " + f"`calculator_1`, got '{self.calculator_1.c_name}'" + ) + + parameters_1 = json.loads(calculator_1.parameters) + + # For the moment, only spherical expansion by pair is supported for calculator_2 + supported_calculators = ["spherical_expansion_by_pair"] + if self.calculator_2.c_name not in supported_calculators: + raise ValueError( + f"Only [{', '.join(supported_calculators)}] are supported for " + f"`calculator_2`, got '{self.calculator_2.c_name}'" + ) + + parameters_2 = json.loads(calculator_2.parameters) + if parameters_1["basis"]["type"] != "TensorProduct": + raise ValueError("only 'TensorProduct' basis is supported for calculator_1") + + if parameters_2["basis"]["type"] != "TensorProduct": + raise ValueError("only 'TensorProduct' basis is supported for calculator_2") + + self._cg_product = ClebschGordanProduct( + max_angular=parameters_1["basis"]["max_angular"] + + parameters_2["basis"]["max_angular"], + cg_backend=None, + keys_filter=_filter_redundant_keys, + arrays_backend=None, + dtype=dtype, + device=device, + ) + + @property + def name(self): + """Name of this calculator.""" + return "EquivariantPowerSpectrumByPair" + + def compute( + self, + systems: Union[IntoSystem, List[IntoSystem]], + selected_keys: Optional[Labels] = None, + neighbors_to_properties: bool = False, + ) -> TensorMap: + """ + Computes an equivariant power spectrum by pair, that can be thought of as a + "Lambda-SOAP by pair" when doing a correlation of the SOAP density with a + spherical expansion by pair. + + First computes a :py:class:`SphericalExpansion` density descriptor of body order + 2. TODO: is is not body order 1? + + Before performing the Clebsch-Gordan tensor product, the spherical expansion + density can be densified by moving the key dimension "neighbor_type" to the + block properties. This is controlled by the ``neighbors_to_properties`` + parameter. Depending on the specific systems descriptors are being computed for, + the sparsity or density of the density can affect the computational cost of the + Clebsch-Gordan tensor product. + + If ``neighbors_to_properties=True`` and ``neighbor_types`` have been passed to + the constructor, property dimensions are created for all of these global atom + types when moving the key dimension to properties. This ensures that the output + properties dimension is of consistent size across all systems passed in + ``systems``. + + Finally a single Clebsch-Gordan tensor product is taken to produce a body order + 3 equivariant power spectrum by pair. + + :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_by_pair( + 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_by_pair( + systems=systems, + selected_keys=selected_keys, + neighbors_to_properties=neighbors_to_properties, + compute_metadata=True, + ) + + def _equivariant_power_spectrum_by_pair( + 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_1 = self.calculator_1.compute(systems) + # Rename "center_type" dimension to match "first_atom_type" + density_1 = operations.rename_dimension( + density_1, "keys", "center_type", "first_atom_type" + ) + # Rename "neighbor_type" dimension so they are correlated + density_1 = operations.rename_dimension( + density_1, "keys", "neighbor_type", "neighbor_1_type" + ) + # Rename samples to match the pair + density_1 = operations.rename_dimension( + density_1, "samples", "atom", "first_atom" + ) + # Rename properties so they are correlated + density_1 = operations.rename_dimension(density_1, "properties", "n", "n_1") + + # Compute pair density + density_2 = self.calculator_2.compute(systems) + + # Rename properties so they are correlated + density_2 = operations.rename_dimension(density_2, "properties", "n", "n_2") + + if neighbors_to_properties: + if self.neighbor_types is None: # just move neighbor type + keys_to_move_1 = "neighbor_1_type" + else: # use the user-specified types + values = _dispatch.list_to_array( + array=density_1.keys.values, + data=[[t] for t in self.neighbor_types], + ) + keys_to_move_1 = Labels(names="neighbor_1_type", values=values) + + density_1 = density_1.keys_to_properties(keys_to_move_1) + + # 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 From ceb9c67b1069a021a54412d55dce732dbcd440a2 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jan 2025 16:30:12 +0100 Subject: [PATCH 03/33] Add equivariant power spectrum by pair in init --- python/featomic/featomic/clebsch_gordan/__init__.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/featomic/featomic/clebsch_gordan/__init__.py b/python/featomic/featomic/clebsch_gordan/__init__.py index d2f805482..3e8c5b404 100644 --- a/python/featomic/featomic/clebsch_gordan/__init__.py +++ b/python/featomic/featomic/clebsch_gordan/__init__.py @@ -3,4 +3,7 @@ from ._coefficients import calculate_cg_coefficients # noqa: F401 from ._density_correlations import DensityCorrelations # noqa: F401 from ._equivariant_power_spectrum import EquivariantPowerSpectrum # noqa: F401 +from ._equivariant_power_spectrum_by_pair import ( + EquivariantPowerSpectrumByPair, +) # noqa: F401 from ._power_spectrum import PowerSpectrum # noqa: F401 From e95fdba7a2e44c0f591afa8bb31f1671728ec4ee Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jan 2025 16:37:44 +0100 Subject: [PATCH 04/33] Lint --- python/featomic/featomic/clebsch_gordan/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/featomic/featomic/clebsch_gordan/__init__.py b/python/featomic/featomic/clebsch_gordan/__init__.py index 3e8c5b404..058329b83 100644 --- a/python/featomic/featomic/clebsch_gordan/__init__.py +++ b/python/featomic/featomic/clebsch_gordan/__init__.py @@ -3,7 +3,7 @@ from ._coefficients import calculate_cg_coefficients # noqa: F401 from ._density_correlations import DensityCorrelations # noqa: F401 from ._equivariant_power_spectrum import EquivariantPowerSpectrum # noqa: F401 -from ._equivariant_power_spectrum_by_pair import ( - EquivariantPowerSpectrumByPair, +from ._equivariant_power_spectrum_by_pair import ( # noqa: F401 + EquivariantPowerSpectrumByPair, # noqa: F401 ) # noqa: F401 from ._power_spectrum import PowerSpectrum # noqa: F401 From ab91b175612dd2abc97cc0206bf464c54f56fb6f Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jan 2025 19:01:21 +0100 Subject: [PATCH 05/33] Permute and rename dimensions to comply with the conventions of the cg_product calculator --- .../_equivariant_power_spectrum_by_pair.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index 781fd9e6f..a4dcaad1e 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -358,4 +358,23 @@ def _equivariant_power_spectrum_by_pair( # Move the CG combination info keys to properties pow_spec = pow_spec.keys_to_properties(["l_1", "l_2"]) + if neighbors_to_properties: + return operations.permute_dimensions( + operations.rename_dimension( + operations.rename_dimension( + operations.rename_dimension( + pow_spec, "properties", "neighbor_1_type", "n_1_" + ), + "properties", + "n_1", + "neighbor_1_type", + ), + "properties", + "n_1_", + "n_1", + ), + "properties", + [0, 1, 3, 2, 4], + ) + return pow_spec From 750769a54deab2128a8e3188a3f976932e158b2b Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 30 Jan 2025 19:01:43 +0100 Subject: [PATCH 06/33] Add tests for equivariant_power_spectrum_by_pair --- .../equivariant_power_spectrum_by_pair.py | 166 ++++++++++++++++++ 1 file changed, 166 insertions(+) create mode 100644 python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py new file mode 100644 index 000000000..8ede44aaf --- /dev/null +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py @@ -0,0 +1,166 @@ +# from typing import List + +import metatensor +import numpy as np +import pytest +from metatensor import Labels # , TensorBlock, TensorMap +from numpy.testing import assert_equal + +from featomic import SphericalExpansion, SphericalExpansionByPair +from featomic.clebsch_gordan import ( + EquivariantPowerSpectrum, + EquivariantPowerSpectrumByPair, +) + + +# Try to import some modules +ase = pytest.importorskip("ase") +import ase.io # noqa: E402, F811 + + +SPHEX_HYPERS_SMALL = { + "cutoff": { + "radius": 2.5, + "smoothing": {"type": "ShiftedCosine", "width": 0.5}, + }, + "density": { + "type": "Gaussian", + "width": 0.2, + }, + "basis": { + "type": "TensorProduct", + "max_angular": 2, + "radial": {"type": "Gto", "max_radial": 1}, + }, +} + +# ============ Helper functions ============ + + +def h2o_periodic(): + return [ + ase.Atoms( + symbols=["O", "H", "H"], + positions=[ + [2.56633400, 2.50000000, 2.50370100], + [1.97361700, 1.73067300, 2.47063400], + [1.97361700, 3.26932700, 2.47063400], + ], + cell=[5, 5, 5], + pbc=[True, True, True], + ) + ] + + +# ============ Test EquivariantPowerSpectrum vs PowerSpectrum ============ + + +def test_equivariant_power_spectrum_vs_equivariant_power_spectrum_by_pair(): + """ + TODO + Tests for exact equivalence between EquivariantPowerSpectrumByPair and + EquivariantPowerSpectrum after metadata manipulation and reduction over samples. + """ + + # Build an EquivariantPowerSpectrum + ps = EquivariantPowerSpectrum(SphericalExpansion(**SPHEX_HYPERS_SMALL)).compute( + h2o_periodic(), + selected_keys=Labels(names=["o3_lambda"], values=np.array([0]).reshape(-1, 1)), + neighbors_to_properties=False, + ) + + # Build an EquivariantPowerSpectrumByPair + ps_by_pair = EquivariantPowerSpectrumByPair( + SphericalExpansion(**SPHEX_HYPERS_SMALL), + SphericalExpansionByPair(**SPHEX_HYPERS_SMALL), + ).compute( + h2o_periodic(), + selected_keys=Labels(names=["o3_lambda"], values=np.array([0]).reshape(-1, 1)), + neighbors_to_properties=False, + ) + + # Manipulate metadata to match + reduced_ps_by_pair = metatensor.rename_dimension( + ps_by_pair, "keys", "first_atom_type", "center_type" + ) + reduced_ps_by_pair = metatensor.rename_dimension( + reduced_ps_by_pair, + "keys", + "second_atom_type", + "neighbor_2_type", + ) + reduced_ps_by_pair = metatensor.rename_dimension( + reduced_ps_by_pair, + "samples", + "first_atom", + "atom", + ) + # Sum over `second_atom` and `cell_shift` samples + reduced_ps_by_pair = metatensor.sum_over_samples( + reduced_ps_by_pair, + ["second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"], + ) + + for k, b1 in ps.items(): + b2 = reduced_ps_by_pair.block(k) + if b1.values.shape != b2.values.shape: + assert b2.values.shape[0] == 0 + assert np.allclose(np.linalg.norm(b1.values), np.linalg.norm(b2.values)) + else: + assert np.allclose(b1.values, b2.values) + + +# def test_equivariant_power_spectrum_by_pair_neighbors_to_properties(): +# """ +# TODO +# Tests that computing an EquivariantPowerSpectrum is equivalent when passing +# `neighbors_to_properties` as both True and False (after metadata manipulation). +# """ +# # Build an EquivariantPowerSpectrum +# powspec_calc = EquivariantPowerSpectrum(SphericalExpansion(**SPHEX_HYPERS_SMALL)) + +# # Compute the first. Move keys after CG step +# powspec_1 = powspec_calc.compute( +# h2o_periodic(), +# neighbors_to_properties=False, +# ) +# powspec_1 = powspec_1.keys_to_properties(["neighbor_1_type", "neighbor_2_type"]) + +# # Compute the second. Move keys before the CG step +# powspec_2 = powspec_calc.compute( +# h2o_periodic(), +# neighbors_to_properties=True, +# ) + +# # Permute properties dimensions to match ``powspec_1`` and sort +# powspec_2 = metatensor.sort( +# metatensor.permute_dimensions(powspec_2, "properties", [2, 4, 0, 1, 3, 5]) +# ) + +# # Check equivalent +# metatensor.equal_metadata_raise(powspec_1, powspec_2) +# metatensor.equal_raise(powspec_1, powspec_2) + + +def test_fill_types_option() -> None: + """ + Test that ``neighbor_types`` options adds arbitrary atomic neighbor types. + """ + + frames = [ + ase.Atoms("H", positions=np.zeros([1, 3])), + ase.Atoms("O", positions=np.zeros([1, 3])), + ] + + neighbor_types = [1, 8, 10] + calculator = EquivariantPowerSpectrumByPair( + calculator_1=SphericalExpansion(**SPHEX_HYPERS_SMALL), + calculator_2=SphericalExpansionByPair(**SPHEX_HYPERS_SMALL), + neighbor_types=neighbor_types, + ) + + descriptor = calculator.compute(frames, neighbors_to_properties=True) + + print(descriptor[0].properties["neighbor_1_type"]) + + assert_equal(np.unique(descriptor[0].properties["neighbor_1_type"]), neighbor_types) From 6b407a4748b73eda0d57f0f3cfbe7caf6056dc2f Mon Sep 17 00:00:00 2001 From: ppegolo Date: Fri, 31 Jan 2025 11:37:47 +0100 Subject: [PATCH 07/33] Fix docs --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index a4dcaad1e..27c49d038 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -85,7 +85,9 @@ class EquivariantPowerSpectrumByPair(TorchModule): Construct the calculators >>> spex_calculator = featomic.SphericalExpansion(**spex_hypers) - >>> spex_by_pair_calculator = featomic.LodeSphericalExpansion(**spex_by_pair_hypers) + >>> spex_by_pair_calculator = featomic.SphericalExpansionByPair( + ... **spex_by_pair_hypers + ... ) Construct the power spectrum by pair calculators and compute the spherical expansion From 79d2d6648753548cf4b5666d5b34af21080b257a Mon Sep 17 00:00:00 2001 From: ppegolo Date: Fri, 31 Jan 2025 13:55:09 +0100 Subject: [PATCH 08/33] Add samples_selection argument --- .../_equivariant_power_spectrum_by_pair.py | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index 27c49d038..f4026349c 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -29,8 +29,6 @@ class EquivariantPowerSpectrumByPair(TorchModule): Example ------- - TODO: update docstring - As an example we calculate the equivariant power spectrum by pair for a spherical expansion and a spherical expansion by pair for a NaCl crystal. @@ -212,6 +210,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: """ @@ -241,6 +240,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. @@ -250,6 +254,7 @@ def compute( return self._equivariant_power_spectrum_by_pair( systems=systems, selected_keys=selected_keys, + selected_samples=selected_samples, neighbors_to_properties=neighbors_to_properties, compute_metadata=False, ) @@ -258,6 +263,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: """ @@ -271,6 +277,7 @@ def forward( return self.compute( systems=systems, selected_keys=selected_keys, + selected_samples=selected_samples, neighbors_to_properties=neighbors_to_properties, ) @@ -278,6 +285,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: """ @@ -298,6 +306,7 @@ def _equivariant_power_spectrum_by_pair( self, systems: Union[IntoSystem, List[IntoSystem]], selected_keys: Optional[Labels], + selected_samples: Optional[Labels], neighbors_to_properties: bool, compute_metadata: bool, ) -> TensorMap: @@ -305,7 +314,9 @@ def _equivariant_power_spectrum_by_pair( 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 + ) # Rename "center_type" dimension to match "first_atom_type" density_1 = operations.rename_dimension( density_1, "keys", "center_type", "first_atom_type" @@ -322,7 +333,9 @@ def _equivariant_power_spectrum_by_pair( density_1 = operations.rename_dimension(density_1, "properties", "n", "n_1") # Compute pair density - density_2 = self.calculator_2.compute(systems) + density_2 = self.calculator_2.compute( + systems, selected_samples=selected_samples + ) # Rename properties so they are correlated density_2 = operations.rename_dimension(density_2, "properties", "n", "n_2") From 4dd23fb325fabbcd9010596b9bd8af0d59949c6e Mon Sep 17 00:00:00 2001 From: ppegolo Date: Fri, 31 Jan 2025 14:09:43 +0100 Subject: [PATCH 09/33] Add tests for sample_selection --- .../_equivariant_power_spectrum_by_pair.py | 5 ++ .../equivariant_power_spectrum_by_pair.py | 71 ++++++++++--------- 2 files changed, 44 insertions(+), 32 deletions(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index f4026349c..1ca8d49fe 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -333,6 +333,11 @@ def _equivariant_power_spectrum_by_pair( density_1 = operations.rename_dimension(density_1, "properties", "n", "n_1") # Compute pair density + if selected_samples is not None: + if "atom" in selected_samples.names: + new_names = selected_samples.names + new_names[selected_samples.names.index("atom")] = "first_atom" + selected_samples = Labels(new_names, selected_samples.values) density_2 = self.calculator_2.compute( systems, selected_samples=selected_samples ) diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py index 8ede44aaf..41b987bda 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py @@ -52,12 +52,11 @@ def h2o_periodic(): ] -# ============ Test EquivariantPowerSpectrum vs PowerSpectrum ============ +# =========== Test EquivariantPowerSpectrumByPair vs EquivariantPowerSpectrum ========== def test_equivariant_power_spectrum_vs_equivariant_power_spectrum_by_pair(): """ - TODO Tests for exact equivalence between EquivariantPowerSpectrumByPair and EquivariantPowerSpectrum after metadata manipulation and reduction over samples. """ @@ -110,36 +109,44 @@ def test_equivariant_power_spectrum_vs_equivariant_power_spectrum_by_pair(): assert np.allclose(b1.values, b2.values) -# def test_equivariant_power_spectrum_by_pair_neighbors_to_properties(): -# """ -# TODO -# Tests that computing an EquivariantPowerSpectrum is equivalent when passing -# `neighbors_to_properties` as both True and False (after metadata manipulation). -# """ -# # Build an EquivariantPowerSpectrum -# powspec_calc = EquivariantPowerSpectrum(SphericalExpansion(**SPHEX_HYPERS_SMALL)) - -# # Compute the first. Move keys after CG step -# powspec_1 = powspec_calc.compute( -# h2o_periodic(), -# neighbors_to_properties=False, -# ) -# powspec_1 = powspec_1.keys_to_properties(["neighbor_1_type", "neighbor_2_type"]) - -# # Compute the second. Move keys before the CG step -# powspec_2 = powspec_calc.compute( -# h2o_periodic(), -# neighbors_to_properties=True, -# ) - -# # Permute properties dimensions to match ``powspec_1`` and sort -# powspec_2 = metatensor.sort( -# metatensor.permute_dimensions(powspec_2, "properties", [2, 4, 0, 1, 3, 5]) -# ) - -# # Check equivalent -# metatensor.equal_metadata_raise(powspec_1, powspec_2) -# metatensor.equal_raise(powspec_1, powspec_2) +def test_sample_selection() -> None: + """Tests that the sample selection works as expected. + By first computing the equivariant power spectrum by pair 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_by_pair_calc = EquivariantPowerSpectrumByPair( + SphericalExpansion(**SPHEX_HYPERS_SMALL), + SphericalExpansionByPair(**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_by_pair_calc.compute( + frame, neighbors_to_properties=True, selected_samples=label_1st + ) + + powspec_2 = powspec_by_pair_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_by_pair_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: From 0af50eaf968371ed8462f70cfdc17c2f33f23bb4 Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Fri, 31 Jan 2025 14:46:26 +0100 Subject: [PATCH 10/33] Update python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py index 41b987bda..4705964d7 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py @@ -3,7 +3,7 @@ import metatensor import numpy as np import pytest -from metatensor import Labels # , TensorBlock, TensorMap +from metatensor import Labels from numpy.testing import assert_equal from featomic import SphericalExpansion, SphericalExpansionByPair From 28096a26a5fb6153158d775e6393faea3f2dd0bb Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Fri, 31 Jan 2025 14:46:40 +0100 Subject: [PATCH 11/33] Update python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py index 4705964d7..c1c80c07f 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py @@ -1,4 +1,3 @@ -# from typing import List import metatensor import numpy as np From 05333eaa22828a5e50304e995ef654772ab3bb93 Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Fri, 31 Jan 2025 14:47:12 +0100 Subject: [PATCH 12/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index 1ca8d49fe..c9b7e3a9c 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -219,7 +219,6 @@ def compute( spherical expansion by pair. First computes a :py:class:`SphericalExpansion` density descriptor of body order - 2. TODO: is is not body order 1? Before performing the Clebsch-Gordan tensor product, the spherical expansion density can be densified by moving the key dimension "neighbor_type" to the From bbd9b920cfeea22c01c75e17aec041e02f7e5c3c Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Fri, 31 Jan 2025 14:47:22 +0100 Subject: [PATCH 13/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index c9b7e3a9c..a03c3d3f4 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -248,7 +248,7 @@ def compute( 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: :py:class:`TensorMap`, the output equivariant power spectrum by pair. """ return self._equivariant_power_spectrum_by_pair( systems=systems, From c07fa8e9ea39dfd8eae5ecdca39316f6fbd3e35e Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Fri, 31 Jan 2025 14:47:29 +0100 Subject: [PATCH 14/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index a03c3d3f4..ff70cac6a 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -166,7 +166,8 @@ def __init__( self.dtype = dtype self.device = device - supported_calculators = ["lode_spherical_expansion", "spherical_expansion"] + supported_calculators_1 = ["lode_spherical_expansion", "spherical_expansion"] + supported_calculators_2 = ["spherical_expansion_by_pair"] if self.calculator_1.c_name not in supported_calculators: raise ValueError( From ec0f9026ae8e753e93980a25cae5adb6db67eb07 Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Fri, 31 Jan 2025 14:47:36 +0100 Subject: [PATCH 15/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index ff70cac6a..c867e5048 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -178,7 +178,6 @@ def __init__( parameters_1 = json.loads(calculator_1.parameters) # For the moment, only spherical expansion by pair is supported for calculator_2 - supported_calculators = ["spherical_expansion_by_pair"] if self.calculator_2.c_name not in supported_calculators: raise ValueError( f"Only [{', '.join(supported_calculators)}] are supported for " From 2790e8dd3980d766f05915166d997fac657b5bdb Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Fri, 31 Jan 2025 14:48:00 +0100 Subject: [PATCH 16/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index c867e5048..4dd2d7780 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -178,7 +178,7 @@ def __init__( parameters_1 = json.loads(calculator_1.parameters) # For the moment, only spherical expansion by pair is supported for calculator_2 - if self.calculator_2.c_name not in supported_calculators: + if self.calculator_2.c_name not in supported_calculators_2: raise ValueError( f"Only [{', '.join(supported_calculators)}] are supported for " f"`calculator_2`, got '{self.calculator_2.c_name}'" From f67a95d8c71dc62606ed3b0af5527da7dd18c2a8 Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Fri, 31 Jan 2025 14:48:14 +0100 Subject: [PATCH 17/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index 4dd2d7780..789b16364 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -169,7 +169,7 @@ def __init__( supported_calculators_1 = ["lode_spherical_expansion", "spherical_expansion"] supported_calculators_2 = ["spherical_expansion_by_pair"] - if self.calculator_1.c_name not in supported_calculators: + if self.calculator_1.c_name not in supported_calculators_1: raise ValueError( f"Only [{', '.join(supported_calculators)}] are supported for " f"`calculator_1`, got '{self.calculator_1.c_name}'" From 7f0b4103f5421c66daeda5f14484c140f52e50a0 Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Fri, 31 Jan 2025 14:48:19 +0100 Subject: [PATCH 18/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index 789b16364..f99e23c32 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -180,7 +180,7 @@ def __init__( # For the moment, only spherical expansion by pair is supported for calculator_2 if self.calculator_2.c_name not in supported_calculators_2: raise ValueError( - f"Only [{', '.join(supported_calculators)}] are supported for " + f"Only [{', '.join(supported_calculators_2)}] are supported for " f"`calculator_2`, got '{self.calculator_2.c_name}'" ) From 150d911a61d102cf893e1f2aff69f7685fff9181 Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Fri, 31 Jan 2025 14:49:03 +0100 Subject: [PATCH 19/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index f99e23c32..52886e594 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -171,7 +171,7 @@ def __init__( if self.calculator_1.c_name not in supported_calculators_1: raise ValueError( - f"Only [{', '.join(supported_calculators)}] are supported for " + f"Only [{', '.join(supported_calculators_1)}] are supported for " f"`calculator_1`, got '{self.calculator_1.c_name}'" ) From a17abc22d356d8906279612422e9d21f668a9b33 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Fri, 31 Jan 2025 14:53:20 +0100 Subject: [PATCH 20/33] Linting --- .../tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py index c1c80c07f..7885392f6 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py @@ -1,8 +1,7 @@ - import metatensor import numpy as np import pytest -from metatensor import Labels +from metatensor import Labels from numpy.testing import assert_equal from featomic import SphericalExpansion, SphericalExpansionByPair From 3fa60d1431d7cd70a0876eff8dcdcdb1a7dfa622 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Fri, 31 Jan 2025 14:53:26 +0100 Subject: [PATCH 21/33] Add torch tests --- .../equivariant_power_spectrum_by_pair.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 python/featomic_torch/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py diff --git a/python/featomic_torch/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic_torch/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py new file mode 100644 index 000000000..93fa17382 --- /dev/null +++ b/python/featomic_torch/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py @@ -0,0 +1,38 @@ +import io + +import torch + +from featomic.torch import SphericalExpansion, SphericalExpansionByPair +from featomic.torch.clebsch_gordan import EquivariantPowerSpectrumByPair + + +SPHEX_HYPERS_SMALL = { + "cutoff": { + "radius": 5.5, + "smoothing": {"type": "ShiftedCosine", "width": 0.5}, + }, + "density": { + "type": "Gaussian", + "width": 0.2, + }, + "basis": { + "type": "TensorProduct", + "max_angular": 6, + "radial": {"type": "Gto", "max_radial": 4}, + }, +} + + +def test_jit_save_load(): + calculator = torch.jit.script( + EquivariantPowerSpectrumByPair( + SphericalExpansion(**SPHEX_HYPERS_SMALL), + SphericalExpansionByPair(**SPHEX_HYPERS_SMALL), + dtype=torch.float64, + ) + ) + + with io.BytesIO() as buffer: + torch.jit.save(calculator, buffer) + buffer.seek(0) + torch.jit.load(buffer) From 0cbc4bca641e8f1d19330b573a4d0254a1a4c512 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Fri, 31 Jan 2025 14:53:36 +0100 Subject: [PATCH 22/33] Update docs --- docs/src/references/api/python/clebsch-gordan.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/src/references/api/python/clebsch-gordan.rst b/docs/src/references/api/python/clebsch-gordan.rst index 1fc3b1edc..6c20e0c86 100644 --- a/docs/src/references/api/python/clebsch-gordan.rst +++ b/docs/src/references/api/python/clebsch-gordan.rst @@ -4,6 +4,9 @@ Clebsch-Gordan products .. autoclass:: featomic.clebsch_gordan.EquivariantPowerSpectrum :members: +.. autoclass:: featomic.clebsch_gordan.EquivariantPowerSpectrumByPair + :members: + .. autoclass:: featomic.clebsch_gordan.PowerSpectrum :members: From 6832f34db1f7dd95d3374d0554ab263d1dadb2a8 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Fri, 31 Jan 2025 14:57:50 +0100 Subject: [PATCH 23/33] Update changelog --- featomic/CHANGELOG.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/featomic/CHANGELOG.md b/featomic/CHANGELOG.md index 769b6ffa6..0996b6d05 100644 --- a/featomic/CHANGELOG.md +++ b/featomic/CHANGELOG.md @@ -17,6 +17,11 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows ### Removed --> +### Added + +- `clebsch_gordan.EquivariantPowerSpectrumByPair` calculator for two-center equivariant + descriptors. Similar API to `clebsch_gordan.EquivariantPowerSpectrum` (#378) + ## [Version 0.6.0](https://github.com/metatensor/featomic/releases/tag/featomic-v0.6.0) - 2024-12-20 ### Added From 3c254390d87541b79999060f52f5b35478883c90 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Mon, 3 Feb 2025 13:51:44 +0100 Subject: [PATCH 24/33] Fix bug in _compute_labels_full_cartesian_product --- .../_equivariant_power_spectrum_by_pair.py | 19 ------------------- .../featomic/clebsch_gordan/_utils.py | 13 +++++++------ 2 files changed, 7 insertions(+), 25 deletions(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index 52886e594..6656046f7 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -377,23 +377,4 @@ def _equivariant_power_spectrum_by_pair( # Move the CG combination info keys to properties pow_spec = pow_spec.keys_to_properties(["l_1", "l_2"]) - if neighbors_to_properties: - return operations.permute_dimensions( - operations.rename_dimension( - operations.rename_dimension( - operations.rename_dimension( - pow_spec, "properties", "neighbor_1_type", "n_1_" - ), - "properties", - "n_1", - "neighbor_1_type", - ), - "properties", - "n_1_", - "n_1", - ), - "properties", - [0, 1, 3, 2, 4], - ) - return pow_spec diff --git a/python/featomic/featomic/clebsch_gordan/_utils.py b/python/featomic/featomic/clebsch_gordan/_utils.py index b8d6ea8e0..dc10094e6 100644 --- a/python/featomic/featomic/clebsch_gordan/_utils.py +++ b/python/featomic/featomic/clebsch_gordan/_utils.py @@ -364,9 +364,9 @@ def _compute_labels_full_cartesian_product( """ # Check for no shared labels dimensions for name in labels_1.names: - assert name not in labels_2.names, ( - "`labels_1` and `labels_2` must not have a dimension ({name}) in common" - ) + assert ( + name not in labels_2.names + ), "`labels_1` and `labels_2` must not have a dimension ({name}) in common" # Create the new labels names by concatenating the names of the two input labels labels_names: List[str] = labels_1.names + labels_2.names @@ -374,16 +374,17 @@ def _compute_labels_full_cartesian_product( # [[i, j] for i in range(len(labels_1.values)) for j in # range(len(labels_2.values))] # [0, 1, 2], [0, 1] -> [[0, 1], [0, 2], [1, 0], [1, 1], [2, 0], [2, 1]] + labels_1_labels_2_product_idx = _dispatch.cartesian_prod( - _dispatch.int_range_like(0, len(labels_2.values), like=labels_2.values), _dispatch.int_range_like(0, len(labels_1.values), like=labels_1.values), + _dispatch.int_range_like(0, len(labels_2.values), like=labels_2.values), ) return Labels( names=labels_names, values=_dispatch.int_array_like( [ - _dispatch.to_int_list(labels_2.values[indices[0]]) - + _dispatch.to_int_list(labels_1.values[indices[1]]) + _dispatch.to_int_list(labels_1.values[indices[0]]) + + _dispatch.to_int_list(labels_2.values[indices[1]]) for indices in labels_1_labels_2_product_idx ], like=labels_1.values, From 7f4db79ff53b702564713864229309c59ab602aa Mon Sep 17 00:00:00 2001 From: ppegolo Date: Mon, 3 Feb 2025 16:03:10 +0100 Subject: [PATCH 25/33] Allow for more flexibility in _match_samples_of_blocks --- .../featomic/clebsch_gordan/_utils.py | 32 ++++++++++++++++--- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/python/featomic/featomic/clebsch_gordan/_utils.py b/python/featomic/featomic/clebsch_gordan/_utils.py index dc10094e6..6a07fc52c 100644 --- a/python/featomic/featomic/clebsch_gordan/_utils.py +++ b/python/featomic/featomic/clebsch_gordan/_utils.py @@ -333,15 +333,39 @@ def _match_samples_of_blocks( block_2.samples.values[:, dims_2][:, None] == block_1.samples.values, axis=2, ) - )[1].tolist() + ) - # Build new block and return + # Build new block_1 block_1 = TensorBlock( - values=block_1.values[matches], - samples=block_2.samples, + values=block_1.values[matches[1].tolist()], + samples=Labels( + block_2.samples.names, block_2.samples.values[matches[0].tolist()] + ), components=block_1.components, properties=block_1.properties, ) + + # Check if we need to slice block_2 samples due to samples in block_2 not being + # matched in block_1 + if ( + len(matches[0]) == len(block_2.samples) + and not _dispatch.all( + matches[0] + == _dispatch.int_range_like( + 0, len(block_2.samples), like=block_2.samples.values + ) + ) + ) or len(matches[0]) != len(block_2.samples): + # Build new block_2 + block_2 = TensorBlock( + values=block_2.values[matches[0].tolist()], + samples=Labels( + block_2.samples.names, block_2.samples.values[matches[0].tolist()] + ), + components=block_2.components, + properties=block_2.properties, + ) + if swapped: return block_2, block_1 From 18f10dd58ce432d1cff9357815c313c8083a3d11 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Mon, 3 Feb 2025 16:03:43 +0100 Subject: [PATCH 26/33] Add test for `neighbors_to_properties` in `EquivariantPowerSpectrumByPair` --- .../equivariant_power_spectrum_by_pair.py | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py index 7885392f6..051a0c2b1 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py @@ -147,6 +147,41 @@ def test_sample_selection() -> None: assert not metatensor.equal(powspec_1, powspec_4) +def test_equivariant_power_spectrum_neighbors_to_properties(): + """ + Tests that computing an EquivariantPowerSpectrumByPair is equivalent when passing + `neighbors_to_properties` as both True and False (after metadata manipulation). + """ + # Build an EquivariantPowerSpectrum + powspec_calc = EquivariantPowerSpectrumByPair( + SphericalExpansion(**SPHEX_HYPERS_SMALL), + SphericalExpansionByPair(**SPHEX_HYPERS_SMALL), + ) + + # Compute the first. Move keys after CG step + powspec_1 = powspec_calc.compute( + h2o_periodic(), + neighbors_to_properties=False, + ) + powspec_1 = powspec_1.keys_to_properties(["neighbor_1_type"]) + + # Compute the second. Move keys before the CG step + powspec_2 = powspec_calc.compute( + h2o_periodic(), + neighbors_to_properties=True, + ) + + # Permute properties dimensions to match ``powspec_1`` and sort + powspec_2 = metatensor.sort( + metatensor.permute_dimensions(powspec_2, "properties", [2, 0, 1, 3, 4]) + ) + + # Check equivalent + powspec_1 = metatensor.sort(powspec_1) + metatensor.equal_metadata_raise(powspec_1, powspec_2) + metatensor.allclose_raise(powspec_1, powspec_2) + + def test_fill_types_option() -> None: """ Test that ``neighbor_types`` options adds arbitrary atomic neighbor types. From 4e279cbc18463a564c635a336ffccc2e3bcfb9c1 Mon Sep 17 00:00:00 2001 From: Paolo Pegolo Date: Mon, 3 Feb 2025 16:05:01 +0100 Subject: [PATCH 27/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py Co-authored-by: Joseph W. Abbott --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index 6656046f7..0f7d82ff0 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -219,6 +219,8 @@ def compute( spherical expansion by pair. First computes a :py:class:`SphericalExpansion` density descriptor of body order + 2. Then a :py:class:`SphericalExpansionByPair` two-center density descriptor of + body order 2 is computed. Before performing the Clebsch-Gordan tensor product, the spherical expansion density can be densified by moving the key dimension "neighbor_type" to the From c456fffdd866ca04bcedc220da95b106a458517e Mon Sep 17 00:00:00 2001 From: ppegolo Date: Mon, 3 Feb 2025 18:13:53 +0100 Subject: [PATCH 28/33] Lint --- python/featomic/featomic/clebsch_gordan/_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/featomic/featomic/clebsch_gordan/_utils.py b/python/featomic/featomic/clebsch_gordan/_utils.py index 6a07fc52c..1117e7141 100644 --- a/python/featomic/featomic/clebsch_gordan/_utils.py +++ b/python/featomic/featomic/clebsch_gordan/_utils.py @@ -388,9 +388,9 @@ def _compute_labels_full_cartesian_product( """ # Check for no shared labels dimensions for name in labels_1.names: - assert ( - name not in labels_2.names - ), "`labels_1` and `labels_2` must not have a dimension ({name}) in common" + assert name not in labels_2.names, ( + "`labels_1` and `labels_2` must not have a dimension ({name}) in common" + ) # Create the new labels names by concatenating the names of the two input labels labels_names: List[str] = labels_1.names + labels_2.names From e55d4555a3fd8bab43f2ab2a39614eab179fd069 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Thu, 6 Feb 2025 17:52:51 +0100 Subject: [PATCH 29/33] Change hypers to avoid empty blocks in tests for equivariant power spectrum by pair --- .../equivariant_power_spectrum_by_pair.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py index 051a0c2b1..dcbef8f44 100644 --- a/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py @@ -32,6 +32,21 @@ }, } +SPHEX_HYPERS_LARGE = { + "cutoff": { + "radius": 5.0, + "smoothing": {"type": "ShiftedCosine", "width": 0.5}, + }, + "density": { + "type": "Gaussian", + "width": 0.2, + }, + "basis": { + "type": "TensorProduct", + "max_angular": 2, + "radial": {"type": "Gto", "max_radial": 1}, + }, +} # ============ Helper functions ============ @@ -44,7 +59,7 @@ def h2o_periodic(): [1.97361700, 1.73067300, 2.47063400], [1.97361700, 3.26932700, 2.47063400], ], - cell=[5, 5, 5], + cell=[3, 3, 3], pbc=[True, True, True], ) ] @@ -154,8 +169,8 @@ def test_equivariant_power_spectrum_neighbors_to_properties(): """ # Build an EquivariantPowerSpectrum powspec_calc = EquivariantPowerSpectrumByPair( - SphericalExpansion(**SPHEX_HYPERS_SMALL), - SphericalExpansionByPair(**SPHEX_HYPERS_SMALL), + SphericalExpansion(**SPHEX_HYPERS_LARGE), + SphericalExpansionByPair(**SPHEX_HYPERS_LARGE), ) # Compute the first. Move keys after CG step From a4c8ccb7f76394c880c2e1cfaa7880074076da62 Mon Sep 17 00:00:00 2001 From: ppegolo Date: Mon, 10 Feb 2025 11:16:59 +0100 Subject: [PATCH 30/33] Fix TorchScript error --- .../featomic/featomic/clebsch_gordan/_utils.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/python/featomic/featomic/clebsch_gordan/_utils.py b/python/featomic/featomic/clebsch_gordan/_utils.py index 1117e7141..2a2321b8f 100644 --- a/python/featomic/featomic/clebsch_gordan/_utils.py +++ b/python/featomic/featomic/clebsch_gordan/_utils.py @@ -7,7 +7,7 @@ from typing import List, Tuple from . import _coefficients, _dispatch -from ._backend import Labels, TensorBlock, TensorMap +from ._backend import Array, Labels, TensorBlock, TensorMap # ======================================== # @@ -328,7 +328,7 @@ def _match_samples_of_blocks( # Broadcast the values of block_1 along the samples dimensions to match those of # block_2 dims_2 = [block_2.samples.names.index(name) for name in block_1.samples.names] - matches: List[int] = _dispatch.where( + matches: List[Array] = _dispatch.where( _dispatch.all( block_2.samples.values[:, dims_2][:, None] == block_1.samples.values, axis=2, @@ -337,10 +337,8 @@ def _match_samples_of_blocks( # Build new block_1 block_1 = TensorBlock( - values=block_1.values[matches[1].tolist()], - samples=Labels( - block_2.samples.names, block_2.samples.values[matches[0].tolist()] - ), + values=block_1.values[matches[1]], + samples=Labels(block_2.samples.names, block_2.samples.values[matches[0]]), components=block_1.components, properties=block_1.properties, ) @@ -358,10 +356,8 @@ def _match_samples_of_blocks( ) or len(matches[0]) != len(block_2.samples): # Build new block_2 block_2 = TensorBlock( - values=block_2.values[matches[0].tolist()], - samples=Labels( - block_2.samples.names, block_2.samples.values[matches[0].tolist()] - ), + values=block_2.values[matches[0]], + samples=Labels(block_2.samples.names, block_2.samples.values[matches[0]]), components=block_2.components, properties=block_2.properties, ) From 4121634f21e0648904202626f9f5de300e453289 Mon Sep 17 00:00:00 2001 From: "Joseph W. Abbott" Date: Mon, 10 Feb 2025 12:09:15 +0100 Subject: [PATCH 31/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index 0f7d82ff0..a2e60175a 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -299,6 +299,7 @@ def compute_metadata( return self._equivariant_power_spectrum_by_pair( systems=systems, selected_keys=selected_keys, + selected_samples=selected_samples, neighbors_to_properties=neighbors_to_properties, compute_metadata=True, ) From 0545c010ae5897bd98bc383df76de90392ebcd88 Mon Sep 17 00:00:00 2001 From: "Joseph W. Abbott" Date: Mon, 10 Feb 2025 12:09:23 +0100 Subject: [PATCH 32/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index a2e60175a..0d76eddf2 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -313,7 +313,7 @@ def _equivariant_power_spectrum_by_pair( compute_metadata: bool, ) -> TensorMap: """ - Computes the equivariant power spectrum, either fully or just metadata + Computes the equivariant power spectrum by pair, either fully or just metadata """ # Compute density density_1 = self.calculator_1.compute( From 26df6aaf30d243c4aef1c967ad8da1e3dd38ae98 Mon Sep 17 00:00:00 2001 From: "Joseph W. Abbott" Date: Mon, 10 Feb 2025 12:09:30 +0100 Subject: [PATCH 33/33] Update python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py --- .../clebsch_gordan/_equivariant_power_spectrum_by_pair.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py index 0d76eddf2..d54f26fe3 100644 --- a/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -226,7 +226,7 @@ def compute( density can be densified by moving the key dimension "neighbor_type" to the block properties. This is controlled by the ``neighbors_to_properties`` parameter. Depending on the specific systems descriptors are being computed for, - the sparsity or density of the density can affect the computational cost of the + the sparsity of the spherical expansion can affect the computational cost of the Clebsch-Gordan tensor product. If ``neighbors_to_properties=True`` and ``neighbor_types`` have been passed to