clebsch gordan submodule - implementation of TorchScript interface

python/rascaline-torch/rascaline/torch/utils/__init__.py

@@ -1,5 +1,6 @@
 import os
 
+from . import clebsch_gordan
 from .power_spectrum import PowerSpectrum
 
 
@@ -10,4 +11,4 @@
 Path containing the CMake configuration files for the underlying C library
 """
 
-__all__ = ["PowerSpectrum"]
+__all__ = ["PowerSpectrum", "clebsch_gordan"]

python/rascaline-torch/rascaline/torch/utils/clebsch_gordan.py

@@ -0,0 +1,73 @@
+import importlib
+import os
+import sys
+from typing import Any
+
+import torch
+from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap
+
+import rascaline.utils.clebsch_gordan
+
+
+# For details what is happening here take a look an `rascaline.torch.calculators`.
+
+# Step 1: create te `_classes` module as an empty module
+spec = importlib.util.spec_from_loader(
+    "rascaline.torch.utils.clebsch_gordan._classes",
+    loader=None,
+)
+module = importlib.util.module_from_spec(spec)
+# This module only exposes a handful of things, defined here. Any changes here MUST also
+# be made to the `metatensor/operations/_classes.py` file, which is used in non
+# TorchScript mode.
+module.__dict__["Labels"] = Labels
+module.__dict__["TensorBlock"] = TensorBlock
+module.__dict__["TensorMap"] = TensorMap
+module.__dict__["LabelsEntry"] = LabelsEntry
+module.__dict__["torch_jit_is_scripting"] = torch.jit.is_scripting
+module.__dict__["torch_jit_annotate"] = torch.jit.annotate
+module.__dict__["torch_jit_export"] = torch.jit.export
+module.__dict__["TorchTensor"] = torch.Tensor
+module.__dict__["TorchModule"] = torch.nn.Module
+module.__dict__["TorchScriptClass"] = torch.ScriptClass
+module.__dict__["Array"] = torch.Tensor
+
+
+def is_labels(obj: Any):
+    return isinstance(obj, Labels)
+
+
+if os.environ.get("RASCALINE_IMPORT_FOR_SPHINX") is None:
+    is_labels = torch.jit.script(is_labels)
+
+module.__dict__["is_labels"] = is_labels
+
+
+def check_isinstance(obj, ty):
+    if isinstance(ty, torch.ScriptClass):
+        # This branch is taken when `ty` is a custom class (TensorMap, …). since `ty` is
+        # an instance of `torch.ScriptClass` and not a class itself, there is no way to
+        # check if obj is an "instance" of this class, so we always return True and hope
+        # for the best. Most errors should be caught by the TorchScript compiler anyway.
+        return True
+    else:
+        assert isinstance(ty, type)
+        return isinstance(obj, ty)
+
+
+module.__dict__["check_isinstance"] = check_isinstance
+
+# register the module in sys.modules, so future import find it directly
+sys.modules[spec.name] = module
+
+
+# Step 2: create a module named `rascaline.torch.utils.clebsch_gordan` using code from
+# `rascaline.utils.clebsch_gordan`
+spec = importlib.util.spec_from_file_location(
+    "rascaline.torch.utils.clebsch_gordan",
+    rascaline.utils.clebsch_gordan.__file__,
+)
+
+module = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = module
+spec.loader.exec_module(module)

python/rascaline-torch/tests/utils/correlate_density.py

@@ -0,0 +1,109 @@
+# -*- coding: utf-8 -*-
+import io
+import os
+from typing import Any, List
+
+import ase.io
+import metatensor.torch
+import pytest
+import torch
+from metatensor.torch import Labels, TensorBlock, TensorMap  # noqa
+
+import rascaline.torch
+from rascaline.torch.utils.clebsch_gordan.correlate_density import DensityCorrelations
+
+
+DATA_ROOT = os.path.join(os.path.dirname(__file__), "data")
+
+
+def is_tensor_map(obj: Any):
+    return isinstance(obj, TensorMap)
+
+
+is_tensor_map = torch.jit.script(is_tensor_map)
+
+SPHEX_HYPERS = {
+    "cutoff": 2.5,  # Angstrom
+    "max_radial": 3,  # Exclusive
+    "max_angular": 3,  # Inclusive
+    "atomic_gaussian_width": 0.2,
+    "radial_basis": {"Gto": {}},
+    "cutoff_function": {"ShiftedCosine": {"width": 0.5}},
+    "center_atom_weight": 1.0,
+}
+
+
+def h2o_isolated():
+    return ase.io.read(os.path.join(DATA_ROOT, "h2o_isolated.xyz"), ":")
+
+
+def spherical_expansion(frames: List[ase.Atoms]):
+    """Returns a rascaline SphericalExpansion"""
+    calculator = rascaline.torch.SphericalExpansion(**SPHEX_HYPERS)
+    return calculator.compute(rascaline.torch.systems_to_torch(frames))
+
+
+# copy of def test_correlate_density_angular_selection(
+@pytest.mark.parametrize(
+    "selected_keys",
+    [
+        None,
+        Labels(
+            names=["spherical_harmonics_l"], values=torch.tensor([1, 3]).reshape(-1, 1)
+        ),
+    ],
+)
+@pytest.mark.parametrize("skip_redundant", [True, False])
+def test_torch_script_correlate_density_angular_selection(
+    selected_keys: Labels,
+    skip_redundant: bool,
+):
+    """
+    Tests that the correct angular channels are output based on the specified
+    ``selected_keys``.
+    """
+    frames = h2o_isolated()
+    nu_1 = spherical_expansion(frames)
+    correlation_order = 2
+    corr_calculator = DensityCorrelations(
+        max_angular=SPHEX_HYPERS["max_angular"] * correlation_order,
+        correlation_order=correlation_order,
+        angular_cutoff=None,
+        selected_keys=selected_keys,
+        skip_redundant=skip_redundant,
+    )
+
+    ref_nu_2 = corr_calculator.compute(nu_1)
+    scripted_corr_calculator = torch.jit.script(corr_calculator)
+
+    # Test compute
+    scripted_nu_2 = scripted_corr_calculator.compute(nu_1)
+    assert metatensor.torch.equal_metadata(scripted_nu_2, ref_nu_2)
+    assert metatensor.torch.allclose(scripted_nu_2, ref_nu_2)
+
+    # Test compute_metadata
+    scripted_nu_2 = scripted_corr_calculator.compute_metadata(nu_1)
+    assert metatensor.torch.equal_metadata(scripted_nu_2, ref_nu_2)
+
+    # Test if properties are accesible
+    assert isinstance(corr_calculator.correlation_order, int)
+    assert isinstance(corr_calculator.selected_keys, list)
+    assert isinstance(corr_calculator.skip_redundant, list)
+    assert isinstance(corr_calculator.output_selection, list)
+    assert isinstance(corr_calculator.arrays_backend, str)
+    assert isinstance(corr_calculator.cg_backend, str)
+    assert is_tensor_map(corr_calculator.cg_coeffs)
+
+
+def test_save_load():
+    corr_calculator = DensityCorrelations(
+        max_angular=2,
+        correlation_order=2,
+        angular_cutoff=1,
+    )
+    scripted_correlate_density = torch.jit.script(corr_calculator)
+    buffer = io.BytesIO()
+    torch.jit.save(scripted_correlate_density, buffer)
+    buffer.seek(0)
+    torch.jit.load(buffer)
+    buffer.close()

python/rascaline-torch/tests/utils/data/h2o_isolated.xyz

@@ -0,0 +1,5 @@
+3
+pbc="F F F"
+O        2.56633400       2.50000000       2.50370100
+H        1.97361700       1.73067300       2.47063400
+H        1.97361700       3.26932700       2.47063400

python/rascaline/rascaline/utils/__init__.py

@@ -1,6 +1,6 @@
 import os
 
-from .clebsch_gordan import correlate_density, correlate_density_metadata  # noqa
+from .clebsch_gordan import DensityCorrelations  # noqa
 from .power_spectrum import PowerSpectrum  # noqa
 from .splines import (  # noqa
     AtomicDensityBase,

python/rascaline/rascaline/utils/clebsch_gordan/__init__.py

@@ -1,7 +1,6 @@
-from .correlate_density import correlate_density, correlate_density_metadata  # noqa
+from .correlate_density import DensityCorrelations  # noqa
 
 
 __all__ = [
-    "correlate_density",
-    "correlate_density_metadata",
+    "DensityCorrelations",
 ]