clebsch gordan submodule - implementation of TorchScript interface

docs/src/references/api/python/utils/clebsch-gordan.rst

@@ -0,0 +1,5 @@
+Clebsch-Gordan products
+=======================
+
+.. autoclass:: rascaline.utils.DensityCorrelations
+    :members:

docs/src/references/api/python/utils/index.rst

@@ -11,3 +11,4 @@ Utility functions and classes that extend the core usage of rascaline.
     radial-basis
     power-spectrum
     splines
+    clebsch-gordan

docs/src/references/api/torch/utils/clebsch-gordan.rst

@@ -0,0 +1,5 @@
+Clebsch-Gordan products
+=======================
+
+.. autoclass:: rascaline.torch.utils.DensityCorrelations
+    :members:

docs/src/references/api/torch/utils/index.rst

@@ -8,3 +8,4 @@ Utility functions and classes that extend the core usage of rascaline-torch
     :maxdepth: 1
 
     power-spectrum
+    clebsch-gordan

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

@@ -9,12 +9,12 @@
 
 _load_library()
 
-from . import utils  # noqa
-from .calculator_base import CalculatorModule, register_autograd  # noqa
+from . import utils  # noqa: E402, F401
+from .calculator_base import CalculatorModule, register_autograd  # noqa: E402, F401
 
 # don't forget to also update `rascaline/__init__.py` and
 # `rascaline/torch/calculators.py` when modifying this file
-from .calculators import (  # noqa
+from .calculators import (  # noqa: E402, F401
     AtomicComposition,
     LodeSphericalExpansion,
     NeighborList,
@@ -24,7 +24,7 @@
     SphericalExpansion,
     SphericalExpansionByPair,
 )
-from .system import systems_to_torch  # noqa
+from .system import systems_to_torch  # noqa: E402, F401
 
 
 __all__ = [

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

@@ -0,0 +1,88 @@
+import importlib
+import os
+import sys
+from typing import Any
+
+import torch
+from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap
+
+import rascaline.utils
+
+from .calculator_base import CalculatorModule
+from .system import System
+
+
+_HERE = os.path.dirname(__file__)
+
+
+# For details what is happening here take a look an `rascaline.torch.calculators`.
+
+# create the `_backend` module as an empty module
+spec = importlib.util.spec_from_loader(
+    "rascaline.torch.utils._backend",
+    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
+module.__dict__["CalculatorBase"] = CalculatorModule
+module.__dict__["IntoSystem"] = System
+
+
+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)
+
+
+# register the module in sys.modules, so future import find it directly
+sys.modules[spec.name] = module
+
+# create a module named `rascaline.torch.utils` using code from
+# `rascaline.utils`
+spec = importlib.util.spec_from_file_location(
+    "rascaline.torch.utils", rascaline.utils.__file__
+)
+
+module = importlib.util.module_from_spec(spec)
+
+
+cmake_prefix_path = os.path.realpath(os.path.join(_HERE, "..", "lib", "cmake"))
+"""
+Path containing the CMake configuration files for the underlying C library
+"""
+
+module.__dict__["cmake_prefix_path"] = cmake_prefix_path
+
+# override `rascaline.torch.utils` (the module associated with the current file)
+# with the newly created module
+sys.modules[spec.name] = module
+spec.loader.exec_module(module)

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

@@ -0,0 +1,113 @@
+# -*- 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)
+
+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,
+}
+
+SELECTED_KEYS = Labels(
+    names=["spherical_harmonics_l"], values=torch.tensor([1, 3]).reshape(-1, 1)
+)
+
+
+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(**SPHERICAL_EXPANSION_HYPERS)
+    return calculator.compute(rascaline.torch.systems_to_torch(frames))
+
+
+# copy of def test_correlate_density_angular_selection(
+@pytest.mark.parametrize("selected_keys", [None, SELECTED_KEYS])
+@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=SPHERICAL_EXPANSION_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)
+
+
+def test_jit_save_load():
+    corr_calculator = DensityCorrelations(
+        max_angular=2,
+        correlation_order=2,
+        angular_cutoff=1,
+    )
+    scripted_correlate_density = torch.jit.script(corr_calculator)
+    with io.BytesIO() as buffer:
+        torch.jit.save(scripted_correlate_density, buffer)
+        buffer.seek(0)
+        torch.jit.load(buffer)
+        buffer.close()
+
+
+def test_save_load():
+    """Tests for saving and loading with cg_backend="python-dense",
+    which makes the DensityCorrelations object non-scriptable due to
+    a non-contiguous CG cache."""
+    corr_calculator = DensityCorrelations(
+        max_angular=2,
+        correlation_order=2,
+        angular_cutoff=1,
+        cg_backend="python-dense",
+    )
+    with io.BytesIO() as buffer:
+        torch.save(corr_calculator, buffer)
+        buffer.seek(0)
+        torch.load(buffer)
+        buffer.close()

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,