From 5a5b2de6f176afe4fd60ae88876f57dfcf9d6213 Mon Sep 17 00:00:00 2001 From: Joseph Abbott Date: Fri, 20 Sep 2024 13:00:10 +0100 Subject: [PATCH 1/2] First draft of `EquivariantPowerSpectrum` --- .../tests/utils/cartesian_spherical.py | 2 +- .../rascaline-torch/tests/utils/cg_product.py | 2 +- .../tests/utils/density_correlations.py | 2 +- .../tests/utils/equivariant_power_spectrum.py | 31 +++ python/rascaline/rascaline/utils/__init__.py | 1 + .../utils/clebsch_gordan/__init__.py | 1 + .../clebsch_gordan/_density_correlations.py | 2 +- .../_equivariant_power_spectrum.py | 202 ++++++++++++++++++ .../rascaline/utils/clebsch_gordan/_utils.py | 30 ++- .../tests/utils/equivariant_power_spectrum.py | 140 ++++++++++++ .../rascaline/tests/utils/power_spectrum.py | 2 +- 11 files changed, 407 insertions(+), 8 deletions(-) create mode 100644 python/rascaline-torch/tests/utils/equivariant_power_spectrum.py create mode 100644 python/rascaline/rascaline/utils/clebsch_gordan/_equivariant_power_spectrum.py create mode 100644 python/rascaline/tests/utils/equivariant_power_spectrum.py diff --git a/python/rascaline-torch/tests/utils/cartesian_spherical.py b/python/rascaline-torch/tests/utils/cartesian_spherical.py index dea2cbe96..0e6fb05ab 100644 --- a/python/rascaline-torch/tests/utils/cartesian_spherical.py +++ b/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 diff --git a/python/rascaline-torch/tests/utils/cg_product.py b/python/rascaline-torch/tests/utils/cg_product.py index f81921d51..1d022aef9 100644 --- a/python/rascaline-torch/tests/utils/cg_product.py +++ b/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 = { diff --git a/python/rascaline-torch/tests/utils/density_correlations.py b/python/rascaline-torch/tests/utils/density_correlations.py index 171d1fe41..f38382a12 100644 --- a/python/rascaline-torch/tests/utils/density_correlations.py +++ b/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 = { diff --git a/python/rascaline-torch/tests/utils/equivariant_power_spectrum.py b/python/rascaline-torch/tests/utils/equivariant_power_spectrum.py new file mode 100644 index 000000000..ea62b918a --- /dev/null +++ b/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) diff --git a/python/rascaline/rascaline/utils/__init__.py b/python/rascaline/rascaline/utils/__init__.py index cbff94460..26b6b5943 100644 --- a/python/rascaline/rascaline/utils/__init__.py +++ b/python/rascaline/rascaline/utils/__init__.py @@ -3,6 +3,7 @@ from .clebsch_gordan import ( # noqa ClebschGordanProduct, DensityCorrelations, + EquivariantPowerSpectrum, calculate_cg_coefficients, cartesian_to_spherical, ) diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py b/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py index 6ba9fac07..54ed5c359 100644 --- a/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py +++ b/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 diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/_density_correlations.py b/python/rascaline/rascaline/utils/clebsch_gordan/_density_correlations.py index fbabeed6f..214f828dc 100644 --- a/python/rascaline/rascaline/utils/clebsch_gordan/_density_correlations.py +++ b/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 diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/_equivariant_power_spectrum.py b/python/rascaline/rascaline/utils/clebsch_gordan/_equivariant_power_spectrum.py new file mode 100644 index 000000000..5dc5263e3 --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/_equivariant_power_spectrum.py @@ -0,0 +1,202 @@ +""" +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 `. + """ + + 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. Note: typically, + setting to true results in a speed up of the computation. 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 diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/_utils.py b/python/rascaline/rascaline/utils/clebsch_gordan/_utils.py index b02292d3d..e365315c5 100644 --- a/python/rascaline/rascaline/utils/clebsch_gordan/_utils.py +++ b/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 diff --git a/python/rascaline/tests/utils/equivariant_power_spectrum.py b/python/rascaline/tests/utils/equivariant_power_spectrum.py new file mode 100644 index 000000000..b3c21bd56 --- /dev/null +++ b/python/rascaline/tests/utils/equivariant_power_spectrum.py @@ -0,0 +1,140 @@ +from typing import List + +import metatensor +import numpy as np +import pytest +from metatensor import Labels, TensorBlock, TensorMap + +import rascaline +from rascaline.utils import PowerSpectrum +from rascaline.utils.clebsch_gordan import EquivariantPowerSpectrum + + +# Try to import some modules +ase = pytest.importorskip("ase") +import ase.io # noqa: E402, F811 + + +SPHEX_HYPERS_SMALL = { + "cutoff": 2.5, # Angstrom + "max_radial": 1, # Exclusive + "max_angular": 2, # Inclusive + "atomic_gaussian_width": 0.2, + "radial_basis": {"Gto": {}}, + "cutoff_function": {"ShiftedCosine": {"width": 0.5}}, + "center_atom_weight": 1.0, +} + +# ============ Helper functions ============ + + +def h2o_isolated(): + 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], + ], + ) + ] + + +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], + ) + ] + + +def power_spectrum(frames: List[ase.Atoms]): + """Returns a rascaline PowerSpectrum constructed from a + SphericalExpansion""" + return PowerSpectrum(rascaline.SphericalExpansion(**SPHEX_HYPERS_SMALL)).compute( + frames + ) + + +# ============ Test EquivariantPowerSpectrum vs PowerSpectrum ============ + + +@pytest.mark.parametrize("frames", [h2o_isolated(), h2o_periodic()]) +def test_equivariant_power_spectrum_vs_powerspectrum(frames): + """ + Tests for exact equivalence between the invariant block of a generated + EquivariantPowerSpectrum the Python implementation of PowerSpectrum in rascaline + utils. + """ + # Build a PowerSpectrum + ps_1 = power_spectrum(frames) + + # Build an EquivariantPowerSpectrum + ps_2 = EquivariantPowerSpectrum(*SPHEX_HYPERS_SMALL).compute( + frames, + selected_keys=Labels(names=["o3_lambda"], values=np.array([0]).reshape(-1, 1)), + ) + + # Manipulate metadata to match + ps_2 = ps_2.keys_to_properties(["neighbor_1_type", "neighbor_2_type"]) + keys = ps_2.keys.remove(name="o3_lambda") # redundant as all 0 + keys = keys.remove("o3_sigma") # redundant as all 1 + + blocks = [] + for block in ps_2.blocks(): + n_samples, n_props = block.values.shape[0], block.values.shape[2] + new_props = block.properties + new_props = new_props.remove(name="l_1") + new_props = new_props.rename(old="l_2", new="l") + blocks.append( + TensorBlock( + values=block.values.reshape((n_samples, n_props)), + samples=block.samples, + components=[], + properties=new_props, + ) + ) + ps_2 = TensorMap(keys=keys, blocks=blocks) + + assert metatensor.allclose(ps_1, ps_2) + + +@pytest.mark.parametrize("frames", [h2o_isolated(), h2o_periodic()]) +def test_equivariant_power_spectrum_neighbors_to_properties(frames): + """ + 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(*SPHEX_HYPERS_SMALL) + + # Compute the first. Move keys after CG step + powspec_1 = powspec_calc.compute( + frames, + neighbors_to_properties=False, + ) + 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( + frames, + 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) diff --git a/python/rascaline/tests/utils/power_spectrum.py b/python/rascaline/tests/utils/power_spectrum.py index f8dda1a20..6f4c27266 100644 --- a/python/rascaline/tests/utils/power_spectrum.py +++ b/python/rascaline/tests/utils/power_spectrum.py @@ -13,7 +13,7 @@ ase = pytest.importorskip("ase") -HYPERS = hypers = { +HYPERS = { "cutoff": 5.0, "max_radial": 6, "max_angular": 4, From 8e01eb2646c79726209bbc82027a759bb23ef906 Mon Sep 17 00:00:00 2001 From: Joseph Abbott Date: Fri, 20 Sep 2024 13:11:54 +0100 Subject: [PATCH 2/2] Remove speed up claim as it depends! --- .../utils/clebsch_gordan/_equivariant_power_spectrum.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/_equivariant_power_spectrum.py b/python/rascaline/rascaline/utils/clebsch_gordan/_equivariant_power_spectrum.py index 5dc5263e3..356a71955 100644 --- a/python/rascaline/rascaline/utils/clebsch_gordan/_equivariant_power_spectrum.py +++ b/python/rascaline/rascaline/utils/clebsch_gordan/_equivariant_power_spectrum.py @@ -97,8 +97,7 @@ def compute( 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. Note: typically, - setting to true results in a speed up of the computation. Defaults to false. + prior to performing the Clebsch Gordan product step. Defaults to false. :return: :py:class:`TensorMap`, the output equivariant power spectrum. """