featomic-torch 0.6.0__py3-none-macosx_11_0_x86_64.whl

featomic/torch/__init__.py

@@ -0,0 +1,39 @@
+import importlib.metadata
+
+
+__version__ = importlib.metadata.version("featomic-torch")
+
+
+from ._c_lib import _load_library
+
+
+_load_library()
+
+from . import utils  # noqa: E402, F401
+from .calculator_base import CalculatorModule, register_autograd  # noqa: E402, F401
+
+# don't forget to also update `featomic/__init__.py` and
+# `featomic/torch/calculators.py` when modifying this file
+from .calculators import (  # noqa: E402, F401
+    AtomicComposition,
+    LodeSphericalExpansion,
+    NeighborList,
+    SoapPowerSpectrum,
+    SoapRadialSpectrum,
+    SortedDistances,
+    SphericalExpansion,
+    SphericalExpansionByPair,
+)
+from .system import systems_to_torch  # noqa: E402, F401
+
+
+__all__ = [
+    "AtomicComposition",
+    "LodeSphericalExpansion",
+    "NeighborList",
+    "SoapPowerSpectrum",
+    "SoapRadialSpectrum",
+    "SortedDistances",
+    "SphericalExpansion",
+    "SphericalExpansionByPair",
+]

featomic/torch/_build_versions.py

@@ -0,0 +1,5 @@
+# Autogenerated file, do not edit
+
+
+# version of featomic used when compiling this package
+BUILD_FEATOMIC_VERSION = '0.6.0'

featomic/torch/_c_lib.py

@@ -0,0 +1,138 @@
+import glob
+import os
+import re
+import sys
+from collections import namedtuple
+
+import metatensor.torch
+import torch
+
+import featomic
+
+from ._build_versions import BUILD_FEATOMIC_VERSION
+
+
+Version = namedtuple("Version", ["major", "minor", "patch"])
+
+
+def parse_version(version):
+    match = re.match(r"(\d+)\.(\d+)\.(\d+).*", version)
+    if match:
+        return Version(*map(int, match.groups()))
+    else:
+        raise ValueError("Invalid version string format")
+
+
+def version_compatible(actual, required):
+    actual = parse_version(actual)
+    required = parse_version(required)
+
+    if actual.major != required.major:
+        return False
+    elif actual.minor != required.minor:
+        return False
+    else:
+        return True
+
+
+if not version_compatible(featomic.__version__, BUILD_FEATOMIC_VERSION):
+    raise ImportError(
+        f"Trying to load featomic-torch with featomic v{featomic.__version__}, "
+        f"but it was compiled against featomic v{BUILD_FEATOMIC_VERSION}, which "
+        "is not ABI compatible"
+    )
+
+_HERE = os.path.realpath(os.path.dirname(__file__))
+
+
+def _lib_path():
+    torch_version = parse_version(torch.__version__)
+    install_prefix = os.path.join(
+        _HERE, f"torch-{torch_version.major}.{torch_version.minor}"
+    )
+
+    if os.path.exists(install_prefix):
+        if sys.platform.startswith("darwin"):
+            path = os.path.join(install_prefix, "lib", "libfeatomic_torch.dylib")
+            windows = False
+        elif sys.platform.startswith("linux"):
+            path = os.path.join(install_prefix, "lib", "libfeatomic_torch.so")
+            windows = False
+        elif sys.platform.startswith("win"):
+            path = os.path.join(install_prefix, "bin", "featomic_torch.dll")
+            windows = True
+        else:
+            raise ImportError("Unknown platform. Please edit this file")
+
+        if os.path.isfile(path):
+            if windows:
+                _check_dll(path)
+            return path
+        else:
+            raise ImportError("Could not find featomic_torch shared library at " + path)
+
+    # gather which torch version(s) the current install was built
+    # with to create the error message
+    existing_versions = []
+    for prefix in glob.glob(os.path.join(_HERE, "torch-*")):
+        existing_versions.append(os.path.basename(prefix)[6:])
+
+    if len(existing_versions) == 1:
+        raise ImportError(
+            f"Trying to load featomic-torch with torch v{torch.__version__}, "
+            f"but it was compiled against torch v{existing_versions[0]}, which "
+            "is not ABI compatible"
+        )
+    else:
+        all_versions = ", ".join(map(lambda version: f"v{version}", existing_versions))
+        raise ImportError(
+            f"Trying to load featomic-torch with torch v{torch.__version__}, "
+            f"we found builds for torch {all_versions}; which are not ABI compatible.\n"
+            "You can try to re-install from source with "
+            "`pip install featomic-torch --no-binary=featomic-torch`"
+        )
+
+
+def _check_dll(path):
+    """
+    Check if the DLL pointer size matches Python (32-bit or 64-bit)
+    """
+    import platform
+    import struct
+
+    IMAGE_FILE_MACHINE_I386 = 332
+    IMAGE_FILE_MACHINE_AMD64 = 34404
+
+    machine = None
+    with open(path, "rb") as fd:
+        header = fd.read(2).decode(encoding="utf-8", errors="strict")
+        if header != "MZ":
+            raise ImportError(path + " is not a DLL")
+        else:
+            fd.seek(60)
+            header = fd.read(4)
+            header_offset = struct.unpack("<L", header)[0]
+            fd.seek(header_offset + 4)
+            header = fd.read(2)
+            machine = struct.unpack("<H", header)[0]
+
+    arch = platform.architecture()[0]
+    if arch == "32bit":
+        if machine != IMAGE_FILE_MACHINE_I386:
+            raise ImportError("Python is 32-bit, but this DLL is not")
+    elif arch == "64bit":
+        if machine != IMAGE_FILE_MACHINE_AMD64:
+            raise ImportError("Python is 64-bit, but this DLL is not")
+    else:
+        raise ImportError("Could not determine pointer size of Python")
+
+
+def _load_library():
+    # Load featomic & metatensor-torch shared library in the process first, to ensure
+    # the featomic_torch shared library can find them
+    metatensor.torch._c_lib._load_library()
+
+    featomic._c_lib._get_library()
+
+    # load the C++ operators and custom classes
+    torch.ops.load_library(_lib_path())

featomic/torch/calculator_base.py

@@ -0,0 +1,176 @@
+from typing import List, Optional, Union
+
+import torch
+from metatensor.torch import Labels, TensorMap
+from metatensor.torch.atomistic import NeighborListOptions
+
+from .system import System
+
+
+CalculatorHolder = torch.classes.featomic.CalculatorHolder
+
+
+def register_autograd(
+    systems: Union[List[System], System],
+    precomputed: TensorMap,
+    forward_gradients: Optional[List[str]] = None,
+) -> TensorMap:
+    """
+    Register autograd nodes between ``system.positions`` and ``system.cell`` for each of
+    the systems and the values in the ``precomputed``
+    :py:class:`metatensor.torch.TensorMap`.
+
+    This is an advanced function must users should not need to use.
+
+    The autograd nodes ``backward`` function will use the gradients already stored in
+    ``precomputed``, meaning that if any of the system's positions ``requires_grad``,
+    ``precomputed`` must contain ``"positions"`` gradients. Similarly, if any of the
+    system's cell ``requires_grad``, ``precomputed`` must contain ``"cell"`` gradients.
+
+    :param systems: list of system used to compute ``precomputed``
+    :param precomputed: precomputed :py:class:`metatensor.torch.TensorMap`
+    :param forward_gradients: which gradients to keep in the output, defaults to None
+    """
+    if forward_gradients is None:
+        forward_gradients = []
+
+    if not isinstance(systems, list):
+        systems = [systems]
+
+    return torch.ops.featomic.register_autograd(systems, precomputed, forward_gradients)
+
+
+class CalculatorModule(torch.nn.Module):
+    """
+    This is the base class for calculators in featomic-torch, providing the
+    :py:meth:`CalculatorModule.compute` function and integration with
+    :py:class:`torch.nn.Module`.
+
+    One can initialize a py:class:`CalculatorModule` in two ways: either directly with
+    the registered name and JSON parameter string (which are documented in the
+    :ref:`userdoc-calculators`); or through one of the child class documented below.
+
+    :param name: name used to register this calculator
+    :param parameters: JSON parameter string for the calculator
+    """
+
+    def __init__(self, name: str, parameters: str):
+        """"""
+        # empty docstring here for the docs to render corectly
+        super().__init__()
+        self._c_name = name
+        self._c = CalculatorHolder(name=name, parameters=parameters)
+
+    @property
+    def name(self) -> str:
+        """name of this calculator"""
+        return self._c.name
+
+    @property
+    def c_name(self) -> str:
+        """name used to register & create this calculator"""
+        return self._c_name
+
+    @property
+    def parameters(self) -> str:
+        """parameters (formatted as JSON) used to create this calculator"""
+        return self._c.parameters
+
+    @property
+    def cutoffs(self) -> List[float]:
+        """all the radial cutoffs used by this calculator's neighbors lists"""
+        return self._c.cutoffs
+
+    def requested_neighbor_lists(self) -> List[NeighborListOptions]:
+        options = []
+        for cutoff in self.cutoffs:
+            options.append(
+                NeighborListOptions(
+                    cutoff=cutoff,
+                    full_list=False,
+                    # we will re-filter the NL when converting to featomic internal
+                    # type, so we don't need the engine to pre-filter it for us
+                    strict=False,
+                    requestor="featomic",
+                )
+            )
+        return options
+
+    def compute(
+        self,
+        systems: Union[System, List[System]],
+        gradients: Optional[List[str]] = None,
+        use_native_system: bool = True,
+        selected_samples: Optional[Union[Labels, TensorMap]] = None,
+        selected_properties: Optional[Union[Labels, TensorMap]] = None,
+        selected_keys: Optional[Labels] = None,
+    ) -> TensorMap:
+        """Runs a calculation with this calculator on the given ``systems``.
+
+        .. seealso::
+
+            :py:func:`featomic.calculators.CalculatorBase.compute` for more information
+            on the different parameters of this function.
+
+        :param systems: single system or list of systems on which to run the
+            calculation. If any of the systems' ``positions`` or ``cell`` has
+            ``requires_grad`` set to ``True``, then the corresponding gradients are
+            computed and registered as a custom node in the computational graph, to
+            allow backward propagation of the gradients later.
+
+        :param gradients: List of forward gradients to keep in the output. If this is
+            ``None`` or an empty list ``[]``, no gradients are kept in the output. Some
+            gradients might still be computed at runtime to allow for backward
+            propagation.
+
+        :param use_native_system: This can only be ``True``, and is here for
+            compatibility with the same parameter on
+            :py:meth:`featomic.calculators.CalculatorBase.compute`.
+
+        :param selected_samples: Set of samples on which to run the calculation, with
+            the same meaning as in
+            :py:func:`featomic.calculators.CalculatorBase.compute`.
+
+        :param selected_properties: Set of properties to compute, with the same meaning
+            as in :py:func:`featomic.calculators.CalculatorBase.compute`.
+
+        :param selected_keys: Selection for the keys to include in the output, with the
+            same meaning as in :py:func:`featomic.calculators.CalculatorBase.compute`.
+        """
+        if gradients is None:
+            gradients = []
+
+        if not isinstance(systems, list):
+            systems = [systems]
+
+        # We have this parameter to have the same API as featomic.
+        if not use_native_system:
+            raise ValueError("only `use_native_system=True` is supported")
+
+        options = torch.classes.featomic.CalculatorOptions()
+        options.gradients = gradients
+        options.selected_samples = selected_samples
+        options.selected_properties = selected_properties
+        options.selected_keys = selected_keys
+
+        return self._c.compute(systems=systems, options=options)
+
+    def forward(
+        self,
+        systems: List[System],
+        gradients: Optional[List[str]] = None,
+        use_native_system: bool = True,
+        selected_samples: Optional[Union[Labels, TensorMap]] = None,
+        selected_properties: Optional[Union[Labels, TensorMap]] = None,
+        selected_keys: Optional[Labels] = None,
+    ) -> TensorMap:
+        """forward just calls :py:meth:`CalculatorModule.compute`"""
+
+        return self.compute(
+            systems=systems,
+            gradients=gradients,
+            use_native_system=use_native_system,
+            selected_samples=selected_samples,
+            selected_properties=selected_properties,
+            selected_keys=selected_keys,
+        )

featomic/torch/calculators.py

@@ -0,0 +1,52 @@
+import importlib
+import sys
+
+import featomic.calculators
+
+from .calculator_base import CalculatorModule
+
+
+#                       CAREFUL ADVENTURER, HERE BE DRAGONS!
+#
+#                                         \||/
+#                                         |  @___oo
+#                               /\  /\   / (__,,,,|
+#                              ) /^\) ^\/ _)
+#                              )   /^\/   _)
+#                              )   _ /  / _)
+#                          /\  )/\/ ||  | )_)
+#                         <  >      |(,,) )__)
+#                          ||      /    \)___)\
+#                          | \____(      )___) )___
+#                           \______(_______;;; __;;;
+#
+#
+# This module tries to re-use code from `featomic.calculators`, which contains a more
+# user-friendly interface to the different calculator. At the C-API level calculators
+# are just defined by a name & JSON parameter string. `featomic.calculators` defines
+# one class for each name and set the `__init__` parameters with the top-level keys of
+# the JSON parameters.
+#
+# To achieve this, we import the module in a special mode with `importlib`, defining a
+# global variable `CalculatorBase` which is pointing to `CalculatorModule`. Then,
+# `featomic.calculators` checks if `CalculatorBase` is defined and otherwise imports it
+# from `featomic.calculator_base`.
+#
+# This means the same code is used to define two versions of each class: one will be
+# used in `featomic` and have a base class of `featomic.CalculatorBase`, and one in
+# `featomic.torch` with base classes `featomic.torch.CalculatorModule` and
+# `torch.nn.Module`.
+
+
+spec = importlib.util.spec_from_file_location(
+    # create a module with this name
+    "featomic.torch.calculators",
+    # using the code from there
+    featomic.calculators.__file__,
+)
+module = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = module
+
+module.__dict__["CalculatorBase"] = CalculatorModule
+
+spec.loader.exec_module(module)

featomic/torch/clebsch_gordan.py

@@ -0,0 +1,89 @@
+import importlib
+import os
+import sys
+from typing import Any
+
+import metatensor.torch
+import torch
+
+import featomic.utils
+
+from .calculator_base import CalculatorModule
+from .system import System
+
+
+# For details what is happening here take a look an `featomic.torch.calculators`.
+
+# create the `_backend` module as an empty module
+spec = importlib.util.spec_from_loader(
+    "featomic.torch.clebsch_gordan._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 `featomic/clebsch_gordan/_backend.py` file, which is used in
+# non-TorchScript mode.
+module.__dict__["BACKEND_IS_METATENSOR_TORCH"] = True
+
+module.__dict__["Labels"] = metatensor.torch.Labels
+module.__dict__["TensorBlock"] = metatensor.torch.TensorBlock
+module.__dict__["TensorMap"] = metatensor.torch.TensorMap
+module.__dict__["LabelsEntry"] = metatensor.torch.LabelsEntry
+
+module.__dict__["CalculatorBase"] = CalculatorModule
+module.__dict__["IntoSystem"] = System
+
+module.__dict__["TorchTensor"] = torch.Tensor
+module.__dict__["TorchModule"] = torch.nn.Module
+module.__dict__["TorchScriptClass"] = torch.ScriptClass
+module.__dict__["Array"] = torch.Tensor
+module.__dict__["DType"] = torch.dtype
+module.__dict__["Device"] = torch.device
+
+module.__dict__["torch_jit_is_scripting"] = torch.jit.is_scripting
+module.__dict__["torch_jit_export"] = torch.jit.export
+
+if os.environ.get("METATENSOR_IMPORT_FOR_SPHINX", "0") == "0":
+    module.__dict__["operations"] = metatensor.torch.operations
+else:
+    # FIXME: we can remove this hack once metatensor-operations v0.2.4 is released
+    module.__dict__["operations"] = None
+
+
+def is_labels(obj: Any):
+    return isinstance(obj, metatensor.torch.Labels)
+
+
+if os.environ.get("FEATOMIC_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 `featomic.torch.clebsch_gordan` using code from
+# `featomic.clebsch_gordan`
+spec = importlib.util.spec_from_file_location(
+    "featomic.torch.clebsch_gordan", featomic.clebsch_gordan.__file__
+)
+
+module = importlib.util.module_from_spec(spec)
+
+# override `featomic.torch.clebsch_gordan` (the module associated with the current file)
+# with the newly created module
+sys.modules[spec.name] = module
+spec.loader.exec_module(module)

featomic/torch/system.py

@@ -0,0 +1,106 @@
+from typing import List, Optional, Sequence, overload
+
+import numpy as np
+import torch
+from metatensor.torch.atomistic import System
+
+import featomic
+
+
+@overload
+def systems_to_torch(
+    systems: featomic.systems.IntoSystem,
+    positions_requires_grad: Optional[bool] = None,
+    cell_requires_grad: Optional[bool] = None,
+) -> System:
+    pass
+
+
+@overload
+def systems_to_torch(
+    systems: Sequence[featomic.systems.IntoSystem],
+    positions_requires_grad: Optional[bool] = None,
+    cell_requires_grad: Optional[bool] = None,
+) -> List[System]:
+    pass
+
+
+def systems_to_torch(
+    systems,
+    positions_requires_grad=None,
+    cell_requires_grad=None,
+) -> List[System]:
+    """
+    Convert a arbitrary system to metatensor's atomistic
+    :py:class:`metatensor.torch.atomistic.System`, putting all the data in
+    :py:class:`torch.Tensor` and making the overall object compatible with TorchScript.
+
+    :param system: any system supported by featomic. If this is an iterable of system,
+        this function converts them all and returns a list of converted systems.
+
+    :param positions_requires_grad: The value of ``requires_grad`` on the output
+        ``positions``. If ``None`` and the positions of the input is already a
+        :py:class:`torch.Tensor`, ``requires_grad`` is kept the same. Otherwise it is
+        initialized to ``False``.
+
+    :param cell_requires_grad: The value of ``requires_grad`` on the output ``cell``. If
+        ``None`` and the positions of the input is already a :py:class:`torch.Tensor`,
+        ``requires_grad`` is kept the same. Otherwise it is initialized to ``False``.
+    """
+
+    try:
+        return _system_to_torch(systems, positions_requires_grad, cell_requires_grad)
+    except TypeError:
+        # try iterating over the systems
+        return [
+            _system_to_torch(system, positions_requires_grad, cell_requires_grad)
+            for system in systems
+        ]
+
+
+def _system_to_torch(system, positions_requires_grad, cell_requires_grad):
+    if not _is_torch_system(system):
+        system = featomic.systems.wrap_system(system)
+        system = System(
+            types=torch.tensor(system.types()),
+            positions=torch.tensor(system.positions()),
+            cell=torch.tensor(system.cell()),
+            pbc=(
+                torch.tensor([False, False, False])
+                if np.all(system.cell() == 0.0)
+                else torch.tensor([True, True, True])
+            ),
+        )
+
+    if positions_requires_grad is not None:
+        system.positions.requires_grad_(positions_requires_grad)
+
+    if cell_requires_grad is not None:
+        system.cell.requires_grad_(cell_requires_grad)
+
+    return system
+
+
+def _is_torch_system(system):
+    if not isinstance(system, torch.ScriptObject):
+        return False
+
+    torch_version_tuple = tuple(map(int, torch.__version__.split(".")[:2]))
+    if torch_version_tuple >= (2, 1):
+        return system._type().name() == "System"
+
+    # For older torch version, we check that we have the right properties
+    properties = system._properties()
+    if len(properties) != 3:
+        return False
+
+    if properties[0].name != "species":
+        return False
+
+    if properties[1].name != "positions":
+        return False
+
+    if properties[2].name != "cell":
+        return False
+
+    return True

featomic/torch/torch-2.1/include/featomic/torch/autograd.hpp

@@ -0,0 +1,58 @@
+// IWYU pragma: private; include "featomic/torch.hpp"
+
+#ifndef FEATOMIC_TORCH_AUTOGRAD_HPP
+#define FEATOMIC_TORCH_AUTOGRAD_HPP
+
+#include <ATen/core/ivalue.h>
+#include <torch/autograd.h>
+
+#include <metatensor/torch.hpp>
+
+#include "featomic/torch/exports.h"
+
+namespace featomic_torch {
+
+/// Custom torch::autograd::Function integrating featomic with torch autograd.
+///
+/// This is a bit more complex than your typical autograd because there is some
+/// impedance mismatch between featomic and torch. Most of it should be taken
+/// care of by the `compute` function below.
+class FEATOMIC_TORCH_EXPORT FeatomicAutograd: public torch::autograd::Function<FeatomicAutograd> {
+public:
+    /// Register a pseudo node in Torch's computational graph going from
+    /// `all_positions` and `all_cell` to the values in `block`; using the
+    /// pre-computed gradients in `block`.
+    ///
+    /// If `all_positions.requires_grad` is True, `block` must have a
+    /// `"positions"` gradient; and `systems_start` should contain the index of
+    /// the first atom of each system in `all_positions`.
+    ///
+    /// If `all_cells.requires_grad` is True, `block` must have a `"cell"`
+    /// gradient, and the block samples must contain a `"stucture"` dimension.
+    ///
+    /// This function returns a vector with one element corresponding to
+    /// `block.values`, which should be left unused. It is only there to make
+    /// sure torch registers a `grad_fn` for the tensors stored inside the
+    /// TensorBlock (the values in the TensorBlock are references to the ones
+    /// returned by this function, so when a `grad_fn` is added to one, it is
+    /// also added to the other).
+    static std::vector<torch::Tensor> forward(
+        torch::autograd::AutogradContext *ctx,
+        torch::Tensor all_positions,
+        torch::Tensor all_cells,
+        torch::IValue systems_start,
+        metatensor_torch::TorchTensorBlock block
+    );
+
+    /// Backward step: get the gradients of some quantity `A` w.r.t. the outputs
+    /// of `forward`; and compute the gradients of the same quantity `A` w.r.t.
+    /// the inputs of `forward` (i.e. cell and positions).
+    static std::vector<torch::Tensor> backward(
+        torch::autograd::AutogradContext *ctx,
+        std::vector<torch::Tensor> grad_outputs
+    );
+};
+
+}
+
+#endif