metatomic-torch 0.1.0__tar.gz

metatomic_torch-0.1.0/AUTHORS

@@ -0,0 +1,2 @@
+Guillaume Fraux
+Filippo Bigi

metatomic_torch-0.1.0/CMakeLists.txt

@@ -0,0 +1,88 @@
+# This file allow the python module in metatomic-torch to either use an
+# externally-provided version of the shared metatomic_torch library; or to build
+# the code from source and bundle the shared library inside the wheel.
+#
+# The first case is used when distributing the code in conda (since we have a
+# separate libmetatomic_torch package), the second one is used everywhere else
+# (for local development builds and for the PyPI distribution).
+
+cmake_minimum_required(VERSION 3.16)
+project(metatomic-torch-python NONE)
+
+option(METATOMIC_TORCH_PYTHON_USE_EXTERNAL_LIB "Force the usage of an external version of metatomic-torch" OFF)
+set(METATOMIC_TORCH_SOURCE_DIR "" CACHE PATH "Path to the sources of metatomic-torch")
+
+file(REMOVE ${CMAKE_INSTALL_PREFIX}/_external.py)
+
+set(REQUIRED_METATOMIC_TORCH_VERSION "0.1.0")
+if(${METATOMIC_TORCH_PYTHON_USE_EXTERNAL_LIB})
+    # when building a source checkout, update version to include git information
+    # this will not apply when building a sdist
+    if (EXISTS ${CMAKE_SOURCE_DIR}/../../metatomic-torch/cmake/dev-versions.cmake)
+        include(${CMAKE_SOURCE_DIR}/../../metatomic-torch/cmake/dev-versions.cmake)
+        create_development_version(
+            "${REQUIRED_METATOMIC_TORCH_VERSION}"
+            REQUIRED_METATOMIC_TORCH_VERSION
+            "metatomic-torch-v"
+        )
+    endif()
+
+    # strip any -dev/-rc suffix on the version since find_package does not support it
+    string(
+        REGEX REPLACE "([0-9]*)\\.([0-9]*)\\.([0-9]*).*" "\\1.\\2.\\3"
+        REQUIRED_METATOMIC_TORCH_VERSION
+        ${REQUIRED_METATOMIC_TORCH_VERSION}
+    )
+    find_package(metatomic_torch ${REQUIRED_METATOMIC_TORCH_VERSION} REQUIRED)
+
+    get_target_property(METATOMIC_TORCH_LOCATION metatomic_torch LOCATION)
+    message(STATUS "Using external metatomic-torch v${metatomic_torch_VERSION} at ${METATOMIC_TORCH_LOCATION}")
+
+    # Get the prefix to use as cmake_prefix_path when trying to load this
+    # version of the library again
+    get_filename_component(METATOMIC_TORCH_PREFIX "${METATOMIC_TORCH_LOCATION}" DIRECTORY)
+    get_filename_component(METATOMIC_TORCH_PREFIX "${METATOMIC_TORCH_PREFIX}" DIRECTORY)
+
+    file(WRITE ${CMAKE_INSTALL_PREFIX}/_external.py
+        "EXTERNAL_METATOMIC_TORCH_PATH = \"${METATOMIC_TORCH_LOCATION}\"\n\n"
+    )
+    file(APPEND ${CMAKE_INSTALL_PREFIX}/_external.py
+        "EXTERNAL_METATOMIC_TORCH_PREFIX = \"${METATOMIC_TORCH_PREFIX}\"\n"
+    )
+
+    install(CODE "message(STATUS \"nothing to install\")")
+else()
+    if ("${METATOMIC_TORCH_SOURCE_DIR}" STREQUAL "")
+        message(FATAL_ERROR
+            "Missing METATOMIC_TORCH_SOURCE_DIR, please specify where to \
+            find the source code for metatomic-torch"
+        )
+    endif()
+
+    message(STATUS "Using internal metatomic-torch from ${METATOMIC_TORCH_SOURCE_DIR}")
+
+    add_subdirectory("${METATOMIC_TORCH_SOURCE_DIR}" metatomic-torch)
+
+
+    if (LINUX OR APPLE)
+        if (LINUX)
+            set(rpath_origin "$ORIGIN")
+        elseif(APPLE)
+            set(rpath_origin "@loader_path")
+        endif()
+        find_package(Torch)
+
+        set(metatomic_install_rpath "${CMAKE_INSTALL_RPATH}")
+        # when loading the libraries from a Python installation:
+        # - $ORIGIN/../../../../torch/lib is where libtorch.so will be
+        # - $ORIGIN/../../../../metatensor/lib is where libmetatensor.so will be
+        # - $ORIGIN/../../../../metatensor/torch/torch-x.y/lib is where libmetatensor_torch.so will be
+        set(metatomic_install_rpath "${metatomic_install_rpath};${rpath_origin}/../../../../torch/lib")
+        set(metatomic_install_rpath "${metatomic_install_rpath};${rpath_origin}/../../../../metatensor/lib")
+        set(metatomic_install_rpath "${metatomic_install_rpath};${rpath_origin}/../../../../metatensor/torch/torch-${Torch_VERSION_MAJOR}.${Torch_VERSION_MINOR}/lib")
+
+        set_target_properties(
+            metatomic_torch PROPERTIES INSTALL_RPATH "${metatomic_install_rpath}"
+        )
+    endif()
+endif()

metatomic_torch-0.1.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2025, metatomic 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.

metatomic_torch-0.1.0/MANIFEST.in

@@ -0,0 +1,10 @@
+include pyproject.toml
+include CMakeLists.txt
+include AUTHORS
+include LICENSE
+
+include git_version_info
+
+include metatomic-torch-*.tar.gz
+
+recursive-include build-backend *.py

metatomic_torch-0.1.0/PKG-INFO

@@ -0,0 +1,39 @@
+Metadata-Version: 2.4
+Name: metatomic-torch
+Version: 0.1.0
+Summary: TODO
+Author: Guillaume Fraux, Filippo Bigi
+License-Expression: BSD-3-Clause
+Project-URL: homepage, https://docs.metatensor.org/metatomic/
+Project-URL: documentation, https://docs.metatensor.org/metatomic/
+Project-URL: repository, https://github.com/metatensor/metatomic
+Project-URL: changelog, https://docs.metatensor.org/metatomic/latest/torch/CHANGELOG.html
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+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>=2.1
+Requires-Dist: vesin
+Requires-Dist: metatensor-torch<0.8.0,>=0.7.0
+Requires-Dist: metatensor-operations<0.4.0,>=0.3.0
+Dynamic: author
+Dynamic: license-file
+Dynamic: requires-dist
+
+metatensor-torch
+================
+
+This package contains the TorchScript bindings to the core API of metatensor.

metatomic_torch-0.1.0/README.rst

@@ -0,0 +1,4 @@
+metatensor-torch
+================
+
+This package contains the TorchScript bindings to the core API of metatensor.

metatomic_torch-0.1.0/build-backend/backend.py

@@ -0,0 +1,33 @@
+# This is a custom Python build backend wrapping setuptool's to only depend on
+# torch/metatensor-torch when building the wheel and not the sdist
+import os
+
+from setuptools import build_meta
+
+
+ROOT = os.path.realpath(os.path.dirname(__file__))
+
+FORCED_TORCH_VERSION = os.environ.get("METATOMIC_TORCH_BUILD_WITH_TORCH_VERSION")
+if FORCED_TORCH_VERSION is not None:
+    TORCH_DEP = f"torch =={FORCED_TORCH_VERSION}"
+else:
+    TORCH_DEP = "torch >=2.1"
+
+# ==================================================================================== #
+#                   Build backend functions definition                                 #
+# ==================================================================================== #
+
+# Use the default version of these
+prepare_metadata_for_build_wheel = build_meta.prepare_metadata_for_build_wheel
+get_requires_for_build_sdist = build_meta.get_requires_for_build_sdist
+build_wheel = build_meta.build_wheel
+build_sdist = build_meta.build_sdist
+
+
+# Special dependencies to build the wheels
+def get_requires_for_build_wheel(config_settings=None):
+    defaults = build_meta.get_requires_for_build_wheel(config_settings)
+    return defaults + [
+        TORCH_DEP,
+        "metatensor-torch >=0.7.0,<0.8.0",
+    ]

metatomic_torch-0.1.0/git_version_info

@@ -0,0 +1,2 @@
+0
+git.47daea3

metatomic_torch-0.1.0/metatomic/torch/__init__.py

@@ -0,0 +1,49 @@
+import os
+
+import torch
+
+from ._c_lib import _load_library
+from .version import __version__  # noqa: F401
+
+
+if os.environ.get("METATOMIC_IMPORT_FOR_SPHINX", "0") != "0":
+    from .documentation import (
+        ModelCapabilities,
+        ModelEvaluationOptions,
+        ModelMetadata,
+        ModelOutput,
+        NeighborListOptions,
+        System,
+        check_atomistic_model,
+        load_model_extensions,
+        read_model_metadata,
+        register_autograd_neighbors,
+        unit_conversion_factor,
+    )
+
+else:
+    _load_library()
+
+    System = torch.classes.metatomic.System
+    NeighborListOptions = torch.classes.metatomic.NeighborListOptions
+
+    ModelOutput = torch.classes.metatomic.ModelOutput
+    ModelEvaluationOptions = torch.classes.metatomic.ModelEvaluationOptions
+    ModelCapabilities = torch.classes.metatomic.ModelCapabilities
+    ModelMetadata = torch.classes.metatomic.ModelMetadata
+
+    read_model_metadata = torch.ops.metatomic.read_model_metadata
+    load_model_extensions = torch.ops.metatomic.load_model_extensions
+    check_atomistic_model = torch.ops.metatomic.check_atomistic_model
+
+    register_autograd_neighbors = torch.ops.metatomic.register_autograd_neighbors
+    unit_conversion_factor = torch.ops.metatomic.unit_conversion_factor
+
+from .io import load_system, save  # noqa: F401
+from .model import (  # noqa: F401
+    AtomisticModel,
+    ModelInterface,
+    is_atomistic_model,
+    load_atomistic_model,  # noqa: F401
+)
+from .systems_to_torch import systems_to_torch  # noqa: F401

metatomic_torch-0.1.0/metatomic/torch/_c_lib.py

@@ -0,0 +1,134 @@
+import glob
+import importlib
+import os
+import sys
+
+import metatensor.torch
+import torch
+
+from .utils import parse_version, version_compatible
+from .version import __version__
+
+
+_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):
+        # check if we are using an externally-provided version of the shared library
+        external_path = os.path.join(install_prefix, "_external.py")
+        if os.path.exists(external_path):
+            spec = importlib.util.spec_from_file_location("_external", external_path)
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+            return module.EXTERNAL_METATOMIC_TORCH_PATH
+
+        if sys.platform.startswith("darwin"):
+            path = os.path.join(install_prefix, "lib", "libmetatomic_torch.dylib")
+            windows = False
+        elif sys.platform.startswith("linux"):
+            path = os.path.join(install_prefix, "lib", "libmetatomic_torch.so")
+            windows = False
+        elif sys.platform.startswith("win"):
+            path = os.path.join(install_prefix, "bin", "metatomic_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 metatomic_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 metatomic-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 metatomic-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 metatomic-torch --no-binary=metatomic-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 metatensor_torch shared library in the process first, to ensure
+    # the metatomic_torch shared library can find it
+    metatensor.torch._c_lib._load_library()
+
+    # load the C++ operators and custom classes
+    try:
+        torch.ops.load_library(_lib_path())
+    except Exception as e:
+        if "undefined symbol" in str(e):
+            file_name = os.path.basename(_lib_path())
+            raise ImportError(
+                f"{file_name} is not compatible with the current PyTorch "
+                "installation.\nThis can happen if PyTorch comes from one source "
+                "(pip, conda, custom), but metatomic-torch comes from a different "
+                "one.\nIn this case, you can try to compile metatomic-torch yourself "
+                "with `pip install --no-binary=metatomic-torch metatomic-torch`"
+            ) from e
+        else:
+            raise e
+
+    lib_version = torch.ops.metatomic.version()
+    if not version_compatible(lib_version, __version__):
+        raise ImportError(
+            f"Trying to load the Python package metatomic-torch v{__version__} "
+            f"with the incompatible metatomic-torch C++ library v{lib_version}"
+        )

metatomic_torch-0.1.0/metatomic/torch/_extensions.py

@@ -0,0 +1,202 @@
+import glob
+import hashlib
+import os
+import shutil
+import site
+import sys
+import warnings
+
+import metatensor.torch
+import torch
+
+from . import _c_lib
+
+
+METATOMIC_TORCH_LIB_PATH = _c_lib._lib_path()
+METATENSOR_TORCH_LIB_PATH = metatensor.torch._c_lib._lib_path()
+
+
+def _rascaline_lib_path():
+    # This is kept for backward compatibility, but rascaline is now named featomic.
+    # This code should be removed by the middle of 2025.
+    import rascaline
+
+    return [rascaline._c_lib._lib_path()]
+
+
+def _featomic_deps_path():
+    import featomic
+
+    deps_path = [featomic._c_lib._lib_path()]
+
+    libgomp_path = _find_openmp_dep("featomic_torch.libs")
+    if libgomp_path is not None:
+        deps_path.insert(0, libgomp_path)
+
+    return deps_path
+
+
+def _sphericart_deps_path():
+    import sphericart.torch
+
+    deps_path = []
+
+    libgomp_path = _find_openmp_dep("sphericart_torch.libs")
+    if libgomp_path is not None:
+        deps_path.append(libgomp_path)
+
+    if sys.platform.startswith("linux"):
+        # sphericart uses a separate library to get the CUDA stream corresponding to a
+        # tensor, see https://github.com/lab-cosmo/sphericart/pull/164
+        sphericart_torch_path = sphericart.torch._lib_path()
+        lib_dir = os.path.dirname(sphericart_torch_path)
+
+        cuda_stream_lib = os.path.join(lib_dir, "libsphericart_torch_cuda_stream.so")
+        if os.path.exists(cuda_stream_lib):
+            deps_path.append(cuda_stream_lib)
+
+    return deps_path
+
+
+def _find_openmp_dep(search_dir):
+    """
+    When building code that uses OpenMP on linux, we typically dynamically link to
+    libgomp. `cibuildwheel` then copies ``libgomp.so`` to
+    ``<wheel_name>.libs/libgomp-<hash>.so``, so we need to find and add this shared
+    library to the extensions dependencies.
+    """
+    if sys.platform.startswith("linux"):
+        libs_list = []
+
+        for prefix in site.getsitepackages():
+            libs_dir = os.path.join(prefix, search_dir)
+            if os.path.exists(libs_dir):
+                libs_list = glob.glob(os.path.join(libs_dir, "libgomp-*.so*"))
+                if len(libs_list) != 0:
+                    # found it!
+                    break
+
+        if len(libs_list) == 0:
+            warnings.warn(
+                f"No libgomp shared library found in '{search_dir}'. "
+                "This may cause issues when loading and running the model.",
+                stacklevel=2,
+            )
+        elif len(libs_list) > 1:
+            raise RuntimeError(
+                f"Multiple libgomp shared libraries found in '{search_dir}': "
+                f"{libs_list}. Try to re-install in a fresh environment."
+            )
+        else:  # len(libs_list) == 1
+            return libs_list[0]
+
+
+# Manual definition of which TorchScript extensions have their own dependencies. The
+# dependencies should be returned in the order they need to be loaded.
+EXTENSIONS_WITH_DEPENDENCIES = {
+    "rascaline_torch": _rascaline_lib_path,
+    "featomic_torch": _featomic_deps_path,
+    "sphericart_torch": _sphericart_deps_path,
+}
+
+
+def _collect_extensions(extensions_path):
+    """
+    Record the list of loaded TorchScript extensions (and their dependencies), to check
+    that they are also loaded when executing the model.
+    """
+    if extensions_path is not None:
+        if os.path.exists(extensions_path):
+            shutil.rmtree(extensions_path)
+        os.makedirs(extensions_path)
+        # TODO: the extensions are currently collected in a separate directory,
+        # should we store the files directly inside the model file? This would makes
+        # the model platform-specific but much more convenient (since the end user
+        # does not have to move a model around)
+
+    extensions = []
+    extensions_deps = []
+    for library in torch.ops.loaded_libraries:
+        assert os.path.exists(library)
+
+        # these should be provided by the simulation engine
+        if os.path.samefile(library, METATENSOR_TORCH_LIB_PATH):
+            continue
+        if os.path.samefile(library, METATOMIC_TORCH_LIB_PATH):
+            continue
+
+        path = _copy_extension(library, extensions_path)
+        info = _extension_info(library)
+        info["path"] = path
+        extensions.append(info)
+
+        for extra in EXTENSIONS_WITH_DEPENDENCIES.get(info["name"], lambda: [])():
+            path = _copy_extension(extra, extensions_path)
+            info = _extension_info(extra)
+            info["path"] = path
+            extensions_deps.append(info)
+
+    return extensions, extensions_deps
+
+
+def _copy_extension(full_path, extensions_dir):
+    full_path = os.path.realpath(full_path)
+
+    prefixes = site.getsitepackages()
+    if site.ENABLE_USER_SITE:
+        prefixes.append(site.getusersitepackages())
+
+    # this also takes care of extensions installed directly in the same prefix as
+    # Python, for example when installing a standalone libmetatensor_torch with conda.
+    prefixes.append(sys.prefix)
+
+    path = full_path
+    for prefix in prefixes:
+        prefix = os.path.realpath(prefix)
+        assert os.path.isabs(prefix)
+
+        # Remove any local prefix
+        if path.startswith(prefix):
+            path = os.path.relpath(path, prefix)
+            break
+
+    if extensions_dir is not None:
+        collect_path = os.path.realpath(os.path.join(extensions_dir, path))
+        if collect_path == path:
+            raise RuntimeError(
+                f"extensions directory '{extensions_dir}' would overwrite files, "
+                "you should set it to a local path instead"
+            )
+
+        if os.path.exists(collect_path):
+            raise RuntimeError(
+                f"more than one extension would be collected at {collect_path}"
+            )
+
+        os.makedirs(os.path.dirname(collect_path), exist_ok=True)
+        shutil.copyfile(full_path, collect_path)
+
+    return path
+
+
+def _extension_info(path):
+    # get the name of the library, excluding any shared object prefix/suffix
+    name = os.path.basename(path)
+    if name.startswith("lib"):
+        name = name[3:]
+
+    if name.endswith(".so"):
+        name = name[:-3]
+
+    if name.endswith(".dll"):
+        name = name[:-4]
+
+    if name.endswith(".dylib"):
+        name = name[:-6]
+
+    # Collect the hash of the extension shared library. We don't currently use
+    # this, but it would allow for binary-level reproducibility later.
+    with open(path, "rb") as fd:
+        sha256 = hashlib.sha256(fd.read()).hexdigest()
+
+    return {"name": name, "sha256": sha256}