featomic-torch 0.6.0__tar.gz

featomic_torch-0.6.0/AUTHORS

@@ -0,0 +1,6 @@
+Guillaume Fraux
+Philip Loche
+Sergei Kliavinek
+Kevin Kazuki Huguenin-Dumittan
+Davide Tisi
+Alexander Goscinski

featomic_torch-0.6.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2023, featomic developers
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

featomic_torch-0.6.0/MANIFEST.in

@@ -0,0 +1,7 @@
+include featomic-torch-cxx-*.tar.gz
+include git_version_info
+include build-backend/backend.py
+
+include pyproject.toml
+include AUTHORS
+include LICENSE

featomic_torch-0.6.0/PKG-INFO

@@ -0,0 +1,36 @@
+Metadata-Version: 2.1
+Name: featomic-torch
+Version: 0.6.0
+Summary: TorchScript bindings to featomic
+Author: Guillaume Fraux, Philip Loche, Sergei Kliavinek, Kevin Kazuki Huguenin-Dumittan, Davide Tisi, Alexander Goscinski
+License: BSD-3-Clause
+Project-URL: homepage, https://metatensor.github.io/featomic/latest/
+Project-URL: documentation, https://metatensor.github.io/featomic/latest/
+Project-URL: repository, https://github.com/metatensor/featomic
+Keywords: computational science,machine learning,molecular modeling,atomistic representations,torch
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Operating System :: POSIX
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Topic :: Scientific/Engineering :: Physics
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Description-Content-Type: text/x-rst
+License-File: LICENSE
+License-File: AUTHORS
+Requires-Dist: torch>=1.12
+Requires-Dist: metatensor-torch<0.7.0,>=0.6.0
+Requires-Dist: featomic<0.7.0,>=0.6.0
+
+featomic-torch
+===============
+
+This package contains the TorchScript bindings to featomic.

featomic_torch-0.6.0/README.rst

@@ -0,0 +1,4 @@
+featomic-torch
+===============
+
+This package contains the TorchScript bindings to featomic.

featomic_torch-0.6.0/build-backend/backend.py

@@ -0,0 +1,47 @@
+# this is a custom Python build backend wrapping setuptool's to add a build-time
+# dependencies to featomic, using the local version if it exists, and otherwise
+# falling back to the one on PyPI.
+import os
+
+from setuptools import build_meta
+
+
+ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), ".."))
+FEATOMIC_SRC = os.path.realpath(os.path.join(ROOT, "..", "featomic"))
+FORCED_FEATOMIC_VERSION = os.environ.get("FEATOMIC_TORCH_BUILD_WITH_FEATOMIC_VERSION")
+
+FEATOMIC_NO_LOCAL_DEPS = os.environ.get("FEATOMIC_NO_LOCAL_DEPS", "0") == "1"
+
+if FORCED_FEATOMIC_VERSION is not None:
+    # force a specific version for metatensor-core, this is used when checking the build
+    # from a sdist on a non-released version
+    FEATOMIC_DEP = f"featomic =={FORCED_FEATOMIC_VERSION}"
+
+elif not FEATOMIC_NO_LOCAL_DEPS and os.path.exists(FEATOMIC_SRC):
+    # we are building from a git checkout
+    FEATOMIC_DEP = f"featomic @ file://{FEATOMIC_SRC}"
+else:
+    # we are building from a sdist
+    FEATOMIC_DEP = "featomic >=0.6.0,<0.7.0"
+
+FORCED_TORCH_VERSION = os.environ.get("FEATOMIC_TORCH_BUILD_WITH_TORCH_VERSION")
+if FORCED_TORCH_VERSION is not None:
+    TORCH_DEP = f"torch =={FORCED_TORCH_VERSION}"
+else:
+    TORCH_DEP = "torch >=1.12"
+
+
+get_requires_for_build_sdist = build_meta.get_requires_for_build_sdist
+prepare_metadata_for_build_wheel = build_meta.prepare_metadata_for_build_wheel
+build_wheel = build_meta.build_wheel
+build_sdist = build_meta.build_sdist
+
+
+def get_requires_for_build_wheel(config_settings=None):
+    defaults = build_meta.get_requires_for_build_wheel(config_settings)
+    return defaults + [
+        "cmake",
+        TORCH_DEP,
+        "metatensor-torch >=0.6.0,<0.7.0",
+        FEATOMIC_DEP,
+    ]

featomic_torch-0.6.0/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-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-0.6.0/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-0.6.0/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-0.6.0/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-0.6.0/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-0.6.0/featomic/torch/utils.py

@@ -0,0 +1,19 @@
+import os
+
+import torch
+
+from ._c_lib import parse_version
+
+
+_HERE = os.path.dirname(__file__)
+
+
+_TORCH_VERSION = parse_version(torch.__version__)
+install_prefix = os.path.join(
+    _HERE, f"torch-{_TORCH_VERSION.major}.{_TORCH_VERSION.minor}"
+)
+
+cmake_prefix_path = os.path.join(install_prefix, "lib", "cmake")
+"""
+Path containing the CMake configuration files for the underlying C++ library
+"""

featomic_torch-0.6.0/featomic_torch.egg-info/PKG-INFO

@@ -0,0 +1,36 @@
+Metadata-Version: 2.1
+Name: featomic-torch
+Version: 0.6.0
+Summary: TorchScript bindings to featomic
+Author: Guillaume Fraux, Philip Loche, Sergei Kliavinek, Kevin Kazuki Huguenin-Dumittan, Davide Tisi, Alexander Goscinski
+License: BSD-3-Clause
+Project-URL: homepage, https://metatensor.github.io/featomic/latest/
+Project-URL: documentation, https://metatensor.github.io/featomic/latest/
+Project-URL: repository, https://github.com/metatensor/featomic
+Keywords: computational science,machine learning,molecular modeling,atomistic representations,torch
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Operating System :: POSIX
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Topic :: Scientific/Engineering :: Physics
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Description-Content-Type: text/x-rst
+License-File: LICENSE
+License-File: AUTHORS
+Requires-Dist: torch>=1.12
+Requires-Dist: metatensor-torch<0.7.0,>=0.6.0
+Requires-Dist: featomic<0.7.0,>=0.6.0
+
+featomic-torch
+===============
+
+This package contains the TorchScript bindings to featomic.

featomic_torch-0.6.0/featomic_torch.egg-info/SOURCES.txt

@@ -0,0 +1,22 @@
+AUTHORS
+LICENSE
+MANIFEST.in
+README.rst
+featomic-torch-cxx-0.6.0.tar.gz
+git_version_info
+pyproject.toml
+setup.py
+build-backend/backend.py
+featomic/torch/__init__.py
+featomic/torch/_c_lib.py
+featomic/torch/calculator_base.py
+featomic/torch/calculators.py
+featomic/torch/clebsch_gordan.py
+featomic/torch/system.py
+featomic/torch/utils.py
+featomic_torch.egg-info/PKG-INFO
+featomic_torch.egg-info/SOURCES.txt
+featomic_torch.egg-info/dependency_links.txt
+featomic_torch.egg-info/not-zip-safe
+featomic_torch.egg-info/requires.txt
+featomic_torch.egg-info/top_level.txt

featomic_torch-0.6.0/featomic_torch.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

featomic_torch-0.6.0/featomic_torch.egg-info/not-zip-safe

@@ -0,0 +1 @@
+

featomic_torch-0.6.0/featomic_torch.egg-info/requires.txt

@@ -0,0 +1,3 @@
+torch>=1.12
+metatensor-torch<0.7.0,>=0.6.0
+featomic<0.7.0,>=0.6.0

featomic_torch-0.6.0/featomic_torch.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+featomic
+featomic_torch

featomic_torch-0.6.0/git_version_info

@@ -0,0 +1,2 @@
+0
+git.575b7a8

featomic_torch-0.6.0/pyproject.toml

@@ -0,0 +1,62 @@
+[project]
+name = "featomic-torch"
+dynamic = ["version", "authors", "dependencies"]
+requires-python = ">=3.9"
+
+readme = "README.rst"
+license = {text = "BSD-3-Clause"}
+description = "TorchScript bindings to featomic"
+
+keywords = ["computational science", "machine learning", "molecular modeling", "atomistic representations", "torch"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: BSD License",
+    "Operating System :: POSIX",
+    "Operating System :: MacOS :: MacOS X",
+    "Operating System :: Microsoft :: Windows",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+    "Topic :: Scientific/Engineering :: Chemistry",
+    "Topic :: Scientific/Engineering :: Physics",
+    "Topic :: Software Development :: Libraries",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+
+[project.urls]
+homepage = "https://metatensor.github.io/featomic/latest/"
+documentation = "https://metatensor.github.io/featomic/latest/"
+repository = "https://github.com/metatensor/featomic"
+# changelog = "TODO"
+
+### ======================================================================== ###
+[build-system]
+requires = [
+    "setuptools",
+    "wheel",
+    "packaging",
+]
+
+# use a custom build backend to add a dependency on the right version of featomic
+build-backend = "backend"
+backend-path = ["build-backend"]
+
+[tool.setuptools]
+zip-safe = false
+
+[tool.setuptools.packages.find]
+include = ["featomic*"]
+namespaces = true
+
+### ======================================================================== ###
+
+[tool.pytest.ini_options]
+python_files = ["*.py"]
+testpaths = ["tests"]
+
+### ======================================================================== ###
+
+[tool.uv.pip]
+reinstall-package = ["featomic", "featomic-torch"]

featomic_torch-0.6.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

featomic_torch-0.6.0/setup.py

@@ -0,0 +1,379 @@
+import glob
+import os
+import subprocess
+import sys
+
+import packaging
+from setuptools import Extension, setup
+from setuptools.command.bdist_egg import bdist_egg
+from setuptools.command.build_ext import build_ext
+from setuptools.command.sdist import sdist
+from wheel.bdist_wheel import bdist_wheel
+
+
+ROOT = os.path.realpath(os.path.dirname(__file__))
+
+FEATOMIC_PYTHON_SRC = os.path.realpath(os.path.join(ROOT, "..", "featomic"))
+FEATOMIC_TORCH_SRC = os.path.realpath(os.path.join(ROOT, "..", "..", "featomic-torch"))
+
+
+class universal_wheel(bdist_wheel):
+    # When building the wheel, the `wheel` package assumes that if we have a
+    # binary extension then we are linking to `libpython.so`; and thus the wheel
+    # is only usable with a single python version. This is not the case for
+    # here, and the wheel will be compatible with any Python >=3. This is
+    # tracked in https://github.com/pypa/wheel/issues/185, but until then we
+    # manually override the wheel tag.
+    def get_tag(self):
+        tag = bdist_wheel.get_tag(self)
+        # tag[2:] contains the os/arch tags, we want to keep them
+        return ("py3", "none") + tag[2:]
+
+
+class cmake_ext(build_ext):
+    """Build the native library using cmake"""
+
+    def run(self):
+        import metatensor
+        import metatensor.torch
+        import torch
+
+        import featomic
+
+        source_dir = FEATOMIC_TORCH_SRC
+        build_dir = os.path.join(ROOT, "build", "cmake-build")
+        install_dir = os.path.join(os.path.realpath(self.build_lib), "featomic/torch")
+
+        os.makedirs(build_dir, exist_ok=True)
+
+        # Tell CMake where to find featomic & torch
+        cmake_prefix_path = [
+            featomic.utils.cmake_prefix_path,
+            metatensor.utils.cmake_prefix_path,
+            metatensor.torch.utils.cmake_prefix_path,
+            torch.utils.cmake_prefix_path,
+        ]
+
+        # Install the shared library in a prefix matching the torch version used to
+        # compile the code. This allows having multiple version of this shared library
+        # inside the wheel; and dynamically pick the right one.
+        torch_major, torch_minor, *_ = torch.__version__.split(".")
+        cmake_install_prefix = os.path.join(
+            install_dir, f"torch-{torch_major}.{torch_minor}"
+        )
+
+        cmake_options = [
+            "-DCMAKE_BUILD_TYPE=Release",
+            f"-DCMAKE_INSTALL_PREFIX={cmake_install_prefix}",
+            f"-DCMAKE_PREFIX_PATH={';'.join(cmake_prefix_path)}",
+        ]
+
+        # ==================================================================== #
+        # HACK: Torch cmake build system has a hard time finding CuDNN, so we
+        # help it by pointing it to the right files
+
+        # First try using the `nvidia.cudnn` package (dependency of torch on PyPI)
+        try:
+            import nvidia.cudnn
+
+            cudnn_root = os.path.dirname(nvidia.cudnn.__file__)
+        except ImportError:
+            # Otherwise try to find CuDNN inside PyTorch itself
+            cudnn_root = os.path.join(torch.utils.cmake_prefix_path, "..", "..")
+
+            cudnn_version = os.path.join(cudnn_root, "include", "cudnn_version.h")
+            if not os.path.exists(cudnn_version):
+                # create a minimal cudnn_version.h (with a made-up version),
+                # because it is not bundled together with the CuDNN shared
+                # library in PyTorch conda distribution, see
+                # https://github.com/pytorch/pytorch/issues/47743
+                with open(cudnn_version, "w") as fd:
+                    fd.write("#define CUDNN_MAJOR 8\n")
+                    fd.write("#define CUDNN_MINOR 5\n")
+                    fd.write("#define CUDNN_PATCHLEVEL 0\n")
+
+        cmake_options.append(f"-DCUDNN_INCLUDE_DIR={cudnn_root}/include")
+        cmake_options.append(f"-DCUDNN_LIBRARY={cudnn_root}/lib")
+        # do not warn if the two variables above aren't used
+        cmake_options.append("--no-warn-unused-cli")
+
+        # end of HACK
+        # ==================================================================== #
+
+        subprocess.run(
+            ["cmake", source_dir, *cmake_options],
+            cwd=build_dir,
+            check=True,
+        )
+        subprocess.run(
+            [
+                "cmake",
+                "--build",
+                build_dir,
+                "--parallel",
+                "--config",
+                "Release",
+                "--target",
+                "install",
+            ],
+            check=True,
+        )
+
+        with open(os.path.join(install_dir, "_build_versions.py"), "w") as fd:
+            fd.write("# Autogenerated file, do not edit\n\n\n")
+            # Store the version of featomic used to build featomic_torch, to give a
+            # nice error message to the user when trying to load the package
+            # with an older featomic version installed
+            fd.write(
+                "# version of featomic used when compiling this package\n"
+                f"BUILD_FEATOMIC_VERSION = '{featomic.__version__}'\n"
+            )
+
+
+class bdist_egg_disabled(bdist_egg):
+    """Disabled version of bdist_egg
+
+    Prevents setup.py install performing setuptools' default easy_install,
+    which it should never ever do.
+    """
+
+    def run(self):
+        sys.exit(
+            "Aborting implicit building of eggs. "
+            + "Use `pip install .` or `python setup.py bdist_wheel && pip "
+            + "install dist/metatensor-*.whl` to install from source."
+        )
+
+
+class sdist_generate_data(sdist):
+    """
+    Create a sdist with an additional generated files:
+        - `git_version_info`
+        - `featomic-torch-cxx-*.tar.gz`
+    """
+
+    def run(self):
+        n_commits, git_hash = git_version_info()
+        with open("git_version_info", "w") as fd:
+            fd.write(f"{n_commits}\n{git_hash}\n")
+
+        generate_cxx_tar()
+
+        # run original sdist
+        super().run()
+
+        os.unlink("git_version_info")
+        for path in glob.glob("featomic-torch-cxx-*.tar.gz"):
+            os.unlink(path)
+
+
+def generate_cxx_tar():
+    script = os.path.join(ROOT, "..", "..", "scripts", "package-featomic-torch.sh")
+    assert os.path.exists(script)
+
+    try:
+        output = subprocess.run(
+            ["bash", "--version"],
+            stderr=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            encoding="utf8",
+        )
+    except Exception as e:
+        raise RuntimeError("could not run `bash`, is it installed?") from e
+
+    output = subprocess.run(
+        ["bash", script, os.getcwd()],
+        stderr=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        encoding="utf8",
+    )
+    if output.returncode != 0:
+        stderr = output.stderr
+        stdout = output.stdout
+        raise RuntimeError(
+            "failed to collect C++ sources for Python sdist\n"
+            f"stdout:\n {stdout}\n\nstderr:\n {stderr}"
+        )
+
+
+def git_version_info():
+    """
+    If git is available and we are building from a checkout, get the number of commits
+    since the last tag & full hash of the code. Otherwise, this always returns (0, "").
+    """
+    TAG_PREFIX = "featomic-torch-v"
+
+    if os.path.exists("git_version_info"):
+        # we are building from a sdist, without git available, but the git
+        # version was recorded in the `git_version_info` file
+        with open("git_version_info") as fd:
+            n_commits = int(fd.readline().strip())
+            git_hash = fd.readline().strip()
+    else:
+        script = os.path.join(ROOT, "..", "..", "scripts", "git-version-info.py")
+        assert os.path.exists(script)
+
+        output = subprocess.run(
+            [sys.executable, script, TAG_PREFIX],
+            stderr=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            encoding="utf8",
+        )
+
+        if output.returncode != 0:
+            raise Exception(
+                "failed to get git version info.\n"
+                f"stdout: {output.stdout}\n"
+                f"stderr: {output.stderr}\n"
+            )
+        elif output.stderr:
+            print(output.stderr, file=sys.stderr)
+            n_commits = 0
+            git_hash = ""
+        else:
+            lines = output.stdout.splitlines()
+            n_commits = int(lines[0].strip())
+            git_hash = lines[1].strip()
+
+    return n_commits, git_hash
+
+
+def create_version_number(version):
+    version = packaging.version.parse(version)
+
+    n_commits, git_hash = git_version_info()
+    if n_commits != 0:
+        # `n_commits` will be non zero only if we have commits since the last tag. This
+        # mean we are in a pre-release of the next version. So we increase either the
+        # minor version number or the release candidate number (if we are closing up on
+        # a release)
+        if version.pre is not None:
+            assert version.pre[0] == "rc"
+            pre = ("rc", version.pre[1] + 1)
+            release = version.release
+        else:
+            major, minor, patch = version.release
+            release = (major, minor + 1, 0)
+            pre = None
+
+        # this is using a private API which is intended to become public soon:
+        # https://github.com/pypa/packaging/pull/698. In the mean time we'll
+        # use this
+        version._version = version._version._replace(release=release)
+        version._version = version._version._replace(pre=pre)
+        version._version = version._version._replace(dev=("dev", n_commits))
+        version._version = version._version._replace(local=(git_hash,))
+
+    return str(version)
+
+
+if __name__ == "__main__":
+    if sys.platform == "win32":
+        # On Windows, starting with PyTorch 2.3, the file shm.dll in torch has a
+        # dependency on mkl DLLs. When building the code using pip build isolation, pip
+        # installs the mkl package in a place where the os is not trying to load
+        #
+        # This is a very similar fix to https://github.com/pytorch/pytorch/pull/126095,
+        # except only applying when importing torch from a build-isolation virtual
+        # environment created by pip (`python -m build` does not seems to suffer from
+        # this).
+        import wheel
+
+        pip_virtualenv = os.path.realpath(
+            os.path.join(
+                os.path.dirname(wheel.__file__),
+                "..",
+                "..",
+                "..",
+                "..",
+            )
+        )
+        mkl_dll_dir = os.path.join(
+            pip_virtualenv,
+            "normal",
+            "Library",
+            "bin",
+        )
+
+        if os.path.exists(mkl_dll_dir):
+            os.add_dll_directory(mkl_dll_dir)
+
+        # End of Windows/MKL/PIP hack
+
+    if not os.path.exists(FEATOMIC_TORCH_SRC):
+        # we are building from a sdist, which should include featomic-torch
+        # sources as a tarball
+        tarballs = glob.glob(os.path.join(ROOT, "featomic-torch-cxx-*.tar.gz"))
+
+        if not len(tarballs) == 1:
+            raise RuntimeError(
+                "expected a single 'featomic-torch-cxx-*.tar.gz' file containing "
+                "featomic-torch C++ sources"
+            )
+
+        FEATOMIC_TORCH_SRC = os.path.realpath(tarballs[0])
+        subprocess.run(
+            ["cmake", "-E", "tar", "xf", FEATOMIC_TORCH_SRC],
+            cwd=ROOT,
+            check=True,
+        )
+
+        FEATOMIC_TORCH_SRC = ".".join(FEATOMIC_TORCH_SRC.split(".")[:-2])
+
+    with open(os.path.join(FEATOMIC_TORCH_SRC, "VERSION")) as fd:
+        version = create_version_number(fd.read().strip())
+
+    with open(os.path.join(ROOT, "AUTHORS")) as fd:
+        authors = fd.read().splitlines()
+
+    if authors[0].startswith(".."):
+        # handle "raw" symlink files (on Windows or from full repo tarball)
+        with open(os.path.join(ROOT, authors[0])) as fd:
+            authors = fd.read().splitlines()
+
+    try:
+        import torch
+
+        # if we have torch, we are building a wheel, which will only be compatible with
+        # a single torch version
+        torch_v_major, torch_v_minor, *_ = torch.__version__.split(".")
+        torch_version = f"== {torch_v_major}.{torch_v_minor}.*"
+    except ImportError:
+        # otherwise we are building a sdist
+        torch_version = ">= 1.12"
+
+    install_requires = [
+        f"torch {torch_version}",
+        "metatensor-torch >=0.6.0,<0.7.0",
+    ]
+
+    # when packaging a sdist for release, we should never use local dependencies
+    FEATOMIC_NO_LOCAL_DEPS = os.environ.get("FEATOMIC_NO_LOCAL_DEPS", "0") == "1"
+    if not FEATOMIC_NO_LOCAL_DEPS and os.path.exists(FEATOMIC_PYTHON_SRC):
+        # we are building from a git checkout
+        install_requires.append(f"featomic @ file://{FEATOMIC_PYTHON_SRC}")
+    else:
+        # we are building from a sdist/installing from a wheel
+        install_requires.append("featomic >=0.6.0,<0.7.0")
+
+    setup(
+        version=version,
+        author=", ".join(authors),
+        install_requires=install_requires,
+        ext_modules=[
+            Extension(name="featomic_torch", sources=[]),
+        ],
+        cmdclass={
+            "build_ext": cmake_ext,
+            "bdist_egg": bdist_egg if "bdist_egg" in sys.argv else bdist_egg_disabled,
+            "bdist_wheel": universal_wheel,
+            "sdist": sdist_generate_data,
+        },
+        package_data={
+            "featomic-torch": [
+                "featomic/torch*/bin/*",
+                "featomic/torch*/lib/*",
+                "featomic/torch*/include/*",
+            ]
+        },
+    )