metatomic-torch 0.1.0__py3-none-win_amd64.whl

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/_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/_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}