Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement convenience calculator EquivariantPowerSpectrum #329

Closed
wants to merge 2 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion python/rascaline-torch/tests/utils/cartesian_spherical.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion python/rascaline-torch/tests/utils/cg_product.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down
2 changes: 1 addition & 1 deletion python/rascaline-torch/tests/utils/density_correlations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down
31 changes: 31 additions & 0 deletions python/rascaline-torch/tests/utils/equivariant_power_spectrum.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
# -*- coding: utf-8 -*-
jwa7 marked this conversation as resolved.
Show resolved Hide resolved
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)
1 change: 1 addition & 0 deletions python/rascaline/rascaline/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from .clebsch_gordan import ( # noqa
ClebschGordanProduct,
DensityCorrelations,
EquivariantPowerSpectrum,
calculate_cg_coefficients,
cartesian_to_spherical,
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Original file line number Diff line number Diff line change
@@ -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
jwa7 marked this conversation as resolved.
Show resolved Hide resolved
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,201 @@
"""
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):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
class EquivariantPowerSpectrum(TorchModule):
class EquivariantSoapPowerSpectrum(TorchModule):

Should we refer to SOAP here? it makes the class name a bit longer, but I could live with it

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm super convinced about this one, but happy for you to decide

"""
Computes an equivariant power spectrum descriptor, or a "lambda-SOAP".

For a full description of the hyper-parameters, see the corresponding
:ref:`documentation <soap-power-spectrum>`.
"""

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,
):
jwa7 marked this conversation as resolved.
Show resolved Hide resolved
"""
: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. 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
30 changes: 27 additions & 3 deletions python/rascaline/rascaline/utils/clebsch_gordan/_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading