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: diff --git a/featomic/CHANGELOG.md b/featomic/CHANGELOG.md index 1d5e968ab..049ca5131 100644 --- a/featomic/CHANGELOG.md +++ b/featomic/CHANGELOG.md @@ -21,6 +21,11 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows - Fixed `featomic.clebsch_gordan.cartesian_to_spherical` transformation for tensors of rank greater than 2 (#371) - Fixed the calculation of Clebsch-Gordan coefficients, used all across the `featomic.clebsch_gordan` module (#371) +### 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 diff --git a/python/featomic/featomic/clebsch_gordan/__init__.py b/python/featomic/featomic/clebsch_gordan/__init__.py index d2f805482..058329b83 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 ( # noqa: F401 + EquivariantPowerSpectrumByPair, # noqa: F401 +) # noqa: F401 from ._power_spectrum import PowerSpectrum # noqa: F401 diff --git a/python/featomic/featomic/clebsch_gordan/_coefficients.py b/python/featomic/featomic/clebsch_gordan/_coefficients.py index bc134d63f..b3dd47187 100644 --- a/python/featomic/featomic/clebsch_gordan/_coefficients.py +++ b/python/featomic/featomic/clebsch_gordan/_coefficients.py @@ -738,7 +738,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) 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..d54f26fe3 --- /dev/null +++ b/python/featomic/featomic/clebsch_gordan/_equivariant_power_spectrum_by_pair.py @@ -0,0 +1,383 @@ +""" +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 + ------- + + 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.SphericalExpansionByPair( + ... **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_1 = ["lode_spherical_expansion", "spherical_expansion"] + supported_calculators_2 = ["spherical_expansion_by_pair"] + + if self.calculator_1.c_name not in supported_calculators_1: + raise ValueError( + f"Only [{', '.join(supported_calculators_1)}] 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 + if self.calculator_2.c_name not in supported_calculators_2: + raise ValueError( + f"Only [{', '.join(supported_calculators_2)}] 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, + selected_samples: 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. 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 + block properties. This is controlled by the ``neighbors_to_properties`` + parameter. Depending on the specific systems descriptors are being computed for, + 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 + 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 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. + + :return: :py:class:`TensorMap`, the output equivariant power spectrum by pair. + """ + 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, + ) + + def forward( + self, + systems: Union[IntoSystem, List[IntoSystem]], + selected_keys: Optional[Labels] = None, + selected_samples: 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, + selected_samples=selected_samples, + neighbors_to_properties=neighbors_to_properties, + ) + + def compute_metadata( + self, + systems: Union[IntoSystem, List[IntoSystem]], + selected_keys: Optional[Labels] = None, + selected_samples: Optional[Labels] = None, + neighbors_to_properties: bool = False, + ) -> TensorMap: + """ + 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, + selected_samples=selected_samples, + 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], + selected_samples: Optional[Labels], + neighbors_to_properties: bool, + compute_metadata: bool, + ) -> TensorMap: + """ + Computes the equivariant power spectrum by pair, either fully or just metadata + """ + # Compute density + 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" + ) + # 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 + 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 + ) + + # 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 diff --git a/python/featomic/featomic/clebsch_gordan/_utils.py b/python/featomic/featomic/clebsch_gordan/_utils.py index b8d6ea8e0..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,20 +328,40 @@ 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, ) - )[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]], + samples=Labels(block_2.samples.names, block_2.samples.values[matches[0]]), 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]], + samples=Labels(block_2.samples.names, block_2.samples.values[matches[0]]), + components=block_2.components, + properties=block_2.properties, + ) + if swapped: return block_2, block_1 @@ -374,16 +394,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, 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..dcbef8f44 --- /dev/null +++ b/python/featomic/tests/clebsch_gordan/equivariant_power_spectrum_by_pair.py @@ -0,0 +1,221 @@ +import metatensor +import numpy as np +import pytest +from metatensor import Labels +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}, + }, +} + +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 ============ + + +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=[3, 3, 3], + pbc=[True, True, True], + ) + ] + + +# =========== Test EquivariantPowerSpectrumByPair vs EquivariantPowerSpectrum ========== + + +def test_equivariant_power_spectrum_vs_equivariant_power_spectrum_by_pair(): + """ + 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_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_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_LARGE), + SphericalExpansionByPair(**SPHEX_HYPERS_LARGE), + ) + + # 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. + """ + + 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) 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)