From 50298b43c2b697920442a7d078adb162a92ea7eb Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Tue, 17 Mar 2026 15:50:11 +0100 Subject: [PATCH] Split the ASE integration into it's own separate package --- .github/workflows/ase-tests.yml | 60 + AUTHORS | 2 + docs/src/engines/ase.rst | 33 +- docs/src/torch/reference/ase.rst | 29 +- metatomic-torch/CHANGELOG.md | 7 + pyproject.toml | 7 +- python/examples/2-running-ase-md.py | 2 +- python/examples/3-atomistic-model-with-nl.py | 2 +- python/examples/4-profiling.py | 2 +- python/metatomic_ase/AUTHORS | 4 + python/metatomic_ase/CHANGELOG.md | 20 + python/metatomic_ase/LICENSE | 1 + python/metatomic_ase/MANIFEST.in | 2 + python/metatomic_ase/README.md | 33 + python/metatomic_ase/pyproject.toml | 52 + python/metatomic_ase/setup.py | 154 ++ .../src/metatomic_ase/__init__.py | 8 + .../src/metatomic_ase/_calculator.py | 963 ++++++++++ .../src/metatomic_ase/_neighbors.py | 139 ++ .../src/metatomic_ase/_symmetry.py | 460 +++++ python/metatomic_ase/tests/__init__.py | 0 python/metatomic_ase/tests/_tests_utils.py | 28 + python/metatomic_ase/tests/calculator.py | 847 +++++++++ python/metatomic_ase/tests/neighbors.py | 111 ++ .../tests/symmetrized.py} | 51 +- .../metatomic/torch/__init__.py | 13 +- .../metatomic/torch/ase_calculator.py | 1553 +---------------- python/metatomic_torch/setup.py | 16 +- .../metatomic_torch/tests/ase_calculator.py | 972 +---------- python/metatomic_torch/tests/neighbors.py | 98 +- python/metatomic_torch/tests/utils.py | 23 +- python/metatomic_torchsim/README.md | 5 +- tox.ini | 130 +- 33 files changed, 3188 insertions(+), 2639 deletions(-) create mode 100644 .github/workflows/ase-tests.yml create mode 100644 python/metatomic_ase/AUTHORS create mode 100644 python/metatomic_ase/CHANGELOG.md create mode 120000 python/metatomic_ase/LICENSE create mode 100644 python/metatomic_ase/MANIFEST.in create mode 100644 python/metatomic_ase/README.md create mode 100644 python/metatomic_ase/pyproject.toml create mode 100644 python/metatomic_ase/setup.py create mode 100644 python/metatomic_ase/src/metatomic_ase/__init__.py create mode 100644 python/metatomic_ase/src/metatomic_ase/_calculator.py create mode 100644 python/metatomic_ase/src/metatomic_ase/_neighbors.py create mode 100644 python/metatomic_ase/src/metatomic_ase/_symmetry.py create mode 100644 python/metatomic_ase/tests/__init__.py create mode 100644 python/metatomic_ase/tests/_tests_utils.py create mode 100644 python/metatomic_ase/tests/calculator.py create mode 100644 python/metatomic_ase/tests/neighbors.py rename python/{metatomic_torch/tests/symmetrized_ase_calculator.py => metatomic_ase/tests/symmetrized.py} (93%) diff --git a/.github/workflows/ase-tests.yml b/.github/workflows/ase-tests.yml new file mode 100644 index 000000000..42247920b --- /dev/null +++ b/.github/workflows/ase-tests.yml @@ -0,0 +1,60 @@ +name: ASE integration tests + +on: + push: + branches: [main] + pull_request: + # Check all PR + +concurrency: + group: ase-tests-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +jobs: + tests: + runs-on: ubuntu-24.04 + name: ASE + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: setup Python + uses: actions/setup-python@v6 + with: + python-version: "3.13" + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.9 + with: + version: "v0.10.0" + + - name: Setup sccache environnement variables + run: | + echo "SCCACHE_GHA_ENABLED=true" >> $GITHUB_ENV + echo "RUSTC_WRAPPER=sccache" >> $GITHUB_ENV + echo "CMAKE_C_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV + echo "CMAKE_CXX_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV + + - name: install tests dependencies + run: | + python -m pip install --upgrade pip + python -m pip install tox coverage + + - name: run tests + run: tox -e ase-tests + env: + PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu + + - name: combine Python coverage files + shell: bash + run: | + coverage combine .tox/*/.coverage + coverage xml + + - name: upload to codecov.io + uses: codecov/codecov-action@v5 + with: + fail_ci_if_error: true + files: coverage.xml + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/AUTHORS b/AUTHORS index 9c13a1289..5b02179b7 100644 --- a/AUTHORS +++ b/AUTHORS @@ -1,2 +1,4 @@ Guillaume Fraux +Philip Loche Filippo Bigi +Qianjun Xu diff --git a/docs/src/engines/ase.rst b/docs/src/engines/ase.rst index 2bcdd8400..c43b1e56d 100644 --- a/docs/src/engines/ase.rst +++ b/docs/src/engines/ase.rst @@ -14,7 +14,7 @@ ASE Supported model outputs ^^^^^^^^^^^^^^^^^^^^^^^ -.. py:currentmodule:: metatomic.torch.ase_calculator +.. py:currentmodule:: metatomic_ase - the :ref:`energy `, non-conservative :ref:`forces ` and :ref:`stress ` @@ -26,16 +26,37 @@ Supported model outputs - for non-equivariant architectures like `PET `_, rotationally-averaged energies, forces, and stresses can be computed using - :py:class:`metatomic.torch.ase_calculator.SymmetrizedCalculator`. + :py:class:`metatomic_ase.SymmetrizedCalculator`. How to install the code ^^^^^^^^^^^^^^^^^^^^^^^ -The code is available in the ``metatomic-torch`` package, in the -:py:class:`metatomic.torch.ase_calculator.MetatomicCalculator` class. +The code is available in the ``metatomic-ase`` package, which can be installed +using ``pip install metatomic-ase``. How to use the code ^^^^^^^^^^^^^^^^^^^ -See the :ref:`corresponding tutorial `, and API -documentation of the :py:class:`MetatomicCalculator` class. +We offer two ASE calculators: :py:class:`metatomic_ase.MetatomicCalculator` is +the default one, and support all the features described above, while +:py:class:`metatomic_ase.SymmetrizedCalculator` is a wrapper around the former +that allows to compute rotationally-averaged energies, forces, and stresses for +non-equivariant architectures. Both calculators are designed to be used as +drop-in replacements for any ASE calculator, and can be used in any ASE +workflow. You can also check the :ref:`corresponding tutorial +`. + +.. _ase-integration-api: + +API documentation +----------------- + +.. _calculator: https://ase-lib.org/ase/calculators/calculators.html + +.. autoclass:: metatomic_ase.MetatomicCalculator + :show-inheritance: + :members: + +.. autoclass:: metatomic_ase.SymmetrizedCalculator + :show-inheritance: + :members: diff --git a/docs/src/torch/reference/ase.rst b/docs/src/torch/reference/ase.rst index ddb1f49a4..4dd244f91 100644 --- a/docs/src/torch/reference/ase.rst +++ b/docs/src/torch/reference/ase.rst @@ -1,23 +1,12 @@ Atomic Simulation Environment (ASE) integration =============================================== -.. py:currentmodule:: metatomic.torch - -The code in ``metatomic.torch.ase_calculator`` defines a class that -allows using a :py:class:`AtomisticModel` which predicts the energy and forces of a -system as an ASE `calculator`_; enabling the use of machine learning interatomic -potentials to drive calculations compatible with ASE calculators. - -Additionally, it allows using arbitrary models with prediction targets which are -not just the energy, through the -:py:meth:`ase_calculator.MetatomicCalculator.run_model` function. - -.. _calculator: https://ase-lib.org/ase/calculators/calculators.html - -.. autoclass:: metatomic.torch.ase_calculator.MetatomicCalculator - :show-inheritance: - :members: - -.. autoclass:: metatomic.torch.ase_calculator.SymmetrizedCalculator - :show-inheritance: - :members: +The integration of metatomic with the Atomic Simulation Environment (ASE) was +moved into it's own package, ``metatomic-ase``, which is available on PyPI. The +documentation for this package can be found in the :ref:`corresponding section +of the documentation `. + +Both calculators classes are re-exported from the +``metatomic.torch.ase_calculator`` module for baclwards compatibility, but users +are encouraged to import them from the ``metatomic_ase`` package instead. The +old import paths will be removed in a future release. diff --git a/metatomic-torch/CHANGELOG.md b/metatomic-torch/CHANGELOG.md index 354feb046..520bc0a18 100644 --- a/metatomic-torch/CHANGELOG.md +++ b/metatomic-torch/CHANGELOG.md @@ -28,6 +28,13 @@ a changelog](https://keepachangelog.com/en/1.1.0/) format. This project follows - 3-argument `unit_conversion_factor(quantity, from_unit, to_unit)` is deprecated; the `quantity` parameter is ignored +- `metatomic.torch.ase_calculator` has been split into a separate + `metatomic-ase` package. The code is temporarily re-exported from the old + path, but all users are encouraged to update to explicitly requiring + `metatomic-ase` as a dependency and changing `from + metatomic.torch.ase_calculator import MetatomicCalculator` to `from + metatomic_ase import MetatomicCalculator` + ## [Version 0.1.11](https://github.com/metatensor/metatomic/releases/tag/metatomic-torch-v0.1.11) - 2026-02-27 ### Added diff --git a/pyproject.toml b/pyproject.toml index f2451aa49..500e11a3b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,7 +78,12 @@ ignore = ["B018", "B904"] [tool.ruff.lint.isort] lines-after-imports = 2 -known-first-party = ["metatomic", "metatomic_torchsim", "metatomic_lj_test"] +known-first-party = [ + "metatomic", + "metatomic_ase", + "metatomic_torchsim", + "metatomic_lj_test", +] known-third-party = ["torch"] [tool.ruff.format] diff --git a/python/examples/2-running-ase-md.py b/python/examples/2-running-ase-md.py index 3fe313c50..fcde979e5 100644 --- a/python/examples/2-running-ase-md.py +++ b/python/examples/2-running-ase-md.py @@ -42,7 +42,7 @@ ) # Integration with ASE for metatomic models -from metatomic.torch.ase_calculator import MetatomicCalculator +from metatomic_ase import MetatomicCalculator # %% diff --git a/python/examples/3-atomistic-model-with-nl.py b/python/examples/3-atomistic-model-with-nl.py index 2e478b140..ed6c2850c 100644 --- a/python/examples/3-atomistic-model-with-nl.py +++ b/python/examples/3-atomistic-model-with-nl.py @@ -59,7 +59,7 @@ NeighborListOptions, System, ) -from metatomic.torch.ase_calculator import MetatomicCalculator +from metatomic_ase import MetatomicCalculator # %% diff --git a/python/examples/4-profiling.py b/python/examples/4-profiling.py index 7934a2118..933f5d185 100644 --- a/python/examples/4-profiling.py +++ b/python/examples/4-profiling.py @@ -29,7 +29,7 @@ ModelOutput, System, ) -from metatomic.torch.ase_calculator import MetatomicCalculator +from metatomic_ase import MetatomicCalculator # %% diff --git a/python/metatomic_ase/AUTHORS b/python/metatomic_ase/AUTHORS new file mode 100644 index 000000000..07bbee152 --- /dev/null +++ b/python/metatomic_ase/AUTHORS @@ -0,0 +1,4 @@ +Guillaume Fraux +Qianjun Xu +Filippo Bigi +Paolo Pegolo diff --git a/python/metatomic_ase/CHANGELOG.md b/python/metatomic_ase/CHANGELOG.md new file mode 100644 index 000000000..6508281f9 --- /dev/null +++ b/python/metatomic_ase/CHANGELOG.md @@ -0,0 +1,20 @@ +# Changelog + +All notable changes to metatomic-ase are documented here, following the +[keep a changelog](https://keepachangelog.com/en/1.1.0/) format. This project +follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + + + +## [Unreleased](https://github.com/metatensor/metatomic/) + +- `metatomic-ase` is now a standalone package, containing the ASE integration + for metatomic models. diff --git a/python/metatomic_ase/LICENSE b/python/metatomic_ase/LICENSE new file mode 120000 index 000000000..30cff7403 --- /dev/null +++ b/python/metatomic_ase/LICENSE @@ -0,0 +1 @@ +../../LICENSE \ No newline at end of file diff --git a/python/metatomic_ase/MANIFEST.in b/python/metatomic_ase/MANIFEST.in new file mode 100644 index 000000000..576b01ef3 --- /dev/null +++ b/python/metatomic_ase/MANIFEST.in @@ -0,0 +1,2 @@ +include AUTHORS +include git_version_info diff --git a/python/metatomic_ase/README.md b/python/metatomic_ase/README.md new file mode 100644 index 000000000..b1cb96e33 --- /dev/null +++ b/python/metatomic_ase/README.md @@ -0,0 +1,33 @@ +# `metatomic-ase` + +[ASE](https://ase-lib.org/) integration for metatomic models. + +This package allows you to use metatomic models as ASE +[`Calculator`](https://ase-lib.org/ase/calculators/calculators.html), integrating with any workflow based on ASE. + +## Installation + +```bash +pip install metatomic-ase +``` + +## Usage + +```python +import ase.io +from metatomic_ase import MetatomicCalculator + +# load atoms +atoms = ase.io.read("...") + +# create a calculator from a saved .pt model +atoms.calc = MetatomicCalculator("model.pt", device="cuda") + +# from here, all the normal ASE functionality is available +print(atoms.get_forces()) +print(atoms.get_potential_energies()) +``` + +For full documentation, see the [ASE engine +page](https://docs.metatensor.org/metatomic/latest/engines/ase.html) in +metatomic documentation. diff --git a/python/metatomic_ase/pyproject.toml b/python/metatomic_ase/pyproject.toml new file mode 100644 index 000000000..6bb19c5ee --- /dev/null +++ b/python/metatomic_ase/pyproject.toml @@ -0,0 +1,52 @@ +[project] +name = "metatomic-ase" +dynamic = ["version", "authors", "dependencies"] +requires-python = ">=3.10" + +readme = "README.md" +license = "BSD-3-Clause" +description = "ASE interface to use metatomic atomistic machine learning models" + +keywords = ["machine learning", "molecular modeling", "ase"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "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://docs.metatensor.org/metatomic/latest/engines/ase.html" +documentation = "https://docs.metatensor.org/metatomic/latest/engines/ase.html" +repository = "https://github.com/metatensor/metatomic/blob/main/python/metatomic_ase" +changelog = "https://github.com/metatensor/metatomic/blob/main/python/metatomic_ase/CHANGELOG.md" + +### ======================================================================== ### +[build-system] +requires = [ + "setuptools >=77", + "packaging >=23", +] +build-backend = "setuptools.build_meta" + +### ======================================================================== ### +[tool.pytest.ini_options] +python_files = ["*.py"] +testpaths = ["tests"] +filterwarnings = [ + "error", + "ignore:`torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`:DeprecationWarning", + "ignore:`torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`:DeprecationWarning", + "ignore:`torch.jit.save` is deprecated. Please switch to `torch.export`:DeprecationWarning", + "ignore:`torch.jit.load` is deprecated. Please switch to `torch.export`:DeprecationWarning", + "ignore:.*vesin.metatomic was only tested with metatomic.torch >=0.1.3,<0.2.*:UserWarning", +] diff --git a/python/metatomic_ase/setup.py b/python/metatomic_ase/setup.py new file mode 100644 index 000000000..3d57b23fe --- /dev/null +++ b/python/metatomic_ase/setup.py @@ -0,0 +1,154 @@ +import os +import subprocess +import sys + +import packaging.version +from setuptools import setup +from setuptools.command.bdist_egg import bdist_egg +from setuptools.command.sdist import sdist + + +ROOT = os.path.realpath(os.path.dirname(__file__)) +METATOMIC_TORCH = os.path.realpath(os.path.join(ROOT, "..", "metatomic_torch")) + +METATOMIC_ASE_VERSION = "0.1.0" + + +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.\nUse `pip install .` or " + "`python -m build --wheel . && pip install dist/metatomic_ase-*.whl` " + "to install from source." + ) + + +class sdist_git_version(sdist): + """ + Create a sdist with an additional generated file containing the extra + version from git. + """ + + 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") + + # run original sdist + super().run() + + os.unlink("git_version_info") + + +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 = "metatomic-ase-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: + # 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__": + install_requires = [ + # No dependency on ASE itself until this package is no longer a direct + # dependency of metatomic-torch + # "ase >=3.22.0", + "vesin >=0.5.2,<0.6", + ] + + # when packaging a sdist for release, we should never use local dependencies + METATOMIC_NO_LOCAL_DEPS = os.environ.get("METATOMIC_NO_LOCAL_DEPS", "0") == "1" + + if not METATOMIC_NO_LOCAL_DEPS and os.path.exists(METATOMIC_TORCH): + # we are building from a git checkout or full repo archive + install_requires.append(f"metatomic-torch @ file://{METATOMIC_TORCH}") + else: + # we are building from a sdist/installing from a wheel + install_requires.append("metatomic-torch >=0.1.11,<0.2.0") + + 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() + + setup( + version=create_version_number(METATOMIC_ASE_VERSION), + author=", ".join(authors), + install_requires=install_requires, + cmdclass={ + "bdist_egg": bdist_egg if "bdist_egg" in sys.argv else bdist_egg_disabled, + "sdist": sdist_git_version, + }, + ) diff --git a/python/metatomic_ase/src/metatomic_ase/__init__.py b/python/metatomic_ase/src/metatomic_ase/__init__.py new file mode 100644 index 000000000..e46cfd5d5 --- /dev/null +++ b/python/metatomic_ase/src/metatomic_ase/__init__.py @@ -0,0 +1,8 @@ +from ._calculator import MetatomicCalculator +from ._symmetry import SymmetrizedCalculator + + +__all__ = [ + "MetatomicCalculator", + "SymmetrizedCalculator", +] diff --git a/python/metatomic_ase/src/metatomic_ase/_calculator.py b/python/metatomic_ase/src/metatomic_ase/_calculator.py new file mode 100644 index 000000000..5d1737372 --- /dev/null +++ b/python/metatomic_ase/src/metatomic_ase/_calculator.py @@ -0,0 +1,963 @@ +import logging +import os +import pathlib +import warnings +from typing import Dict, List, Optional, Union + +import metatensor.torch as mts +import numpy as np +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap +from torch.profiler import record_function + +from metatomic.torch import ( + AtomisticModel, + ModelEvaluationOptions, + ModelMetadata, + ModelOutput, + System, + load_atomistic_model, + pick_device, + pick_output, +) + +from ._neighbors import _compute_requested_neighbors + + +import ase # isort: skip +import ase.calculators.calculator # isort: skip +from ase.calculators.calculator import ( # isort: skip + InputError, + PropertyNotImplementedError, + all_properties as ALL_ASE_PROPERTIES, +) + + +FilePath = Union[str, bytes, pathlib.PurePath] + +LOGGER = logging.getLogger(__name__) + + +STR_TO_DTYPE = { + "float32": torch.float32, + "float64": torch.float64, +} + + +def _get_charges(atoms: ase.Atoms) -> np.ndarray: + try: + return atoms.get_charges() + except Exception: + return atoms.get_initial_charges() + + +ARRAY_QUANTITIES = { + "momenta": { + "quantity": "momentum", + "getter": ase.Atoms.get_momenta, + "unit": "(eV*u)^(1/2)", + }, + "masses": { + "quantity": "mass", + "getter": ase.Atoms.get_masses, + "unit": "u", + }, + "velocities": { + "quantity": "velocity", + "getter": ase.Atoms.get_velocities, + "unit": "(eV/u)^(1/2)", + }, + "charges": { + "quantity": "charge", + "getter": _get_charges, + "unit": "e", + }, + "ase::initial_magmoms": { + "quantity": "magnetic_moment", + "getter": ase.Atoms.get_initial_magnetic_moments, + "unit": "", + }, + "ase::magnetic_moment": { + "quantity": "magnetic_moment", + "getter": ase.Atoms.get_magnetic_moment, + "unit": "", + }, + "ase::magnetic_moments": { + "quantity": "magnetic_moment", + "getter": ase.Atoms.get_magnetic_moments, + "unit": "", + }, + "ase::initial_charges": { + "quantity": "charge", + "getter": ase.Atoms.get_initial_charges, + "unit": "e", + }, + "ase::dipole_moment": { + "quantity": "dipole_moment", + "getter": ase.Atoms.get_dipole_moment, + "unit": "", + }, +} + + +class MetatomicCalculator(ase.calculators.calculator.Calculator): + """ + The :py:class:`MetatomicCalculator` class implements ASE's + :py:class:`ase.calculators.calculator.Calculator` API using metatomic models to + compute energy, forces and any other supported property. + + This class can be initialized with any :py:class:`metatomic.torch.AtomisticModel`, + and used to run simulations using ASE's MD facilities. + + Neighbor lists are computed using the fast `vesin + `_ neighbor list library, either on CPU + or GPU depending on the device of the model. If `nvalchemiops + `_ is installed, full neighbor + lists on GPU will be computed with it instead. + """ + + def __init__( + self, + model: Union[FilePath, AtomisticModel], + *, + additional_outputs: Optional[Dict[str, ModelOutput]] = None, + extensions_directory=None, + check_consistency=False, + device=None, + variants: Optional[Dict[str, Optional[str]]] = None, + non_conservative=False, + do_gradients_with_energy=True, + uncertainty_threshold=0.1, + ): + """ + :param model: model to use for the calculation. This can be a file path, a + Python instance of :py:class:`metatomic.torch.AtomisticModel`, or the output + of :py:func:`torch.jit.script` on + :py:class:`metatomic.torch.AtomisticModel`. + :param additional_outputs: Dictionary of additional outputs to be computed by + the model. These outputs will always be computed whenever the + :py:meth:`calculate` function is called (e.g. by + :py:meth:`ase.Atoms.get_potential_energy`, + :py:meth:`ase.optimize.optimize.Dynamics.run`, *etc.*) and stored in the + :py:attr:`additional_outputs` attribute. If you want more control over when + and how to compute specific outputs, you should use :py:meth:`run_model` + instead. + :param extensions_directory: if the model uses extensions, we will try to load + them from this directory + :param check_consistency: should we check the model for consistency when + running, defaults to False. + :param device: torch device to use for the calculation. If ``None``, we will try + the options in the model's ``supported_device`` in order. + :param variants: dictionary mapping output names to a variant that should be + used for the calculations (e.g. ``{"energy": "PBE"}``). If ``"energy"`` is + set to a variant also the uncertainty and non conservative outputs will be + taken from this variant. This behaviour can be overriden by setting the + corresponding keys explicitly to ``None`` or to another value (e.g. + ``{"energy_uncertainty": "r2scan"}``). + :param non_conservative: if ``True``, the model will be asked to compute + non-conservative forces and stresses. This can afford a speed-up, + potentially at the expense of physical correctness (especially in molecular + dynamics simulations). + :param do_gradients_with_energy: if ``True``, this calculator will always + compute the energy gradients (forces and stress) when the energy is + requested (e.g. through ``atoms.get_potential_energy()``). Because the + results of a calculation are cached by ASE, this means future calls to + ``atom.get_forces()`` will return immediately, without needing to execute + the model again. If you are mainly interested in the energy, you can set + this to ``False`` and enjoy a faster model. Forces will still be calculated + if requested with ``atoms.get_forces()``. + :param uncertainty_threshold: threshold for the atomic energy uncertainty in eV. + This will only be used if the model supports atomic uncertainty estimation + (https://pubs.acs.org/doi/full/10.1021/acs.jctc.3c00704). Set this to + ``None`` to disable uncertainty quantification even if the model supports + it. + """ + super().__init__() + + self.parameters = { + "extensions_directory": extensions_directory, + "check_consistency": bool(check_consistency), + "variants": variants, + "non_conservative": bool(non_conservative), + "do_gradients_with_energy": bool(do_gradients_with_energy), + "additional_outputs": additional_outputs, + "uncertainty_threshold": uncertainty_threshold, + } + + # Load the model + if isinstance(model, (str, bytes, pathlib.PurePath)): + if not os.path.exists(model): + raise InputError(f"given model path '{model}' does not exist") + + # only store the model in self.parameters if is it the path to a file + self.parameters["model"] = str(model) + + model = load_atomistic_model( + model, extensions_directory=extensions_directory + ) + + elif isinstance(model, torch.jit.RecursiveScriptModule): + if model.original_name != "AtomisticModel": + raise InputError( + "torch model must be 'AtomisticModel', " + f"got '{model.original_name}' instead" + ) + elif isinstance(model, AtomisticModel): + # nothing to do + pass + else: + raise TypeError(f"unknown type for model: {type(model)}") + + self.parameters["device"] = str(device) if device is not None else None + # get the best device according what the model supports and what's available on + # the current machine + capabilities = model.capabilities() + self._device = torch.device( + pick_device(capabilities.supported_devices, self.parameters["device"]) + ) + + if capabilities.dtype in STR_TO_DTYPE: + self._dtype = STR_TO_DTYPE[capabilities.dtype] + else: + raise ValueError( + f"found unexpected dtype in model capabilities: {capabilities.dtype}" + ) + + # resolve the output keys to use based on the requested variants + variants = variants or {} + default_variant = variants.get("energy") + + resolved_variants = { + key: variants.get(key, default_variant) + for key in [ + "energy", + "energy_uncertainty", + "non_conservative_forces", + "non_conservative_stress", + ] + } + + outputs = capabilities.outputs + + # Check if the model has an energy output + has_energy = any( + "energy" == key or key.startswith("energy/") for key in outputs.keys() + ) + if has_energy: + self._energy_key = pick_output( + "energy", outputs, resolved_variants["energy"] + ) + else: + self._energy_key = None + + has_energy_uq = any("energy_uncertainty" in key for key in outputs.keys()) + if has_energy_uq and uncertainty_threshold is not None: + self._energy_uq_key = pick_output( + "energy_uncertainty", outputs, resolved_variants["energy_uncertainty"] + ) + else: + self._energy_uq_key = "energy_uncertainty" + + if non_conservative: + if ( + "non_conservative_stress" in variants + and "non_conservative_forces" in variants + and ( + (variants["non_conservative_stress"] is None) + != (variants["non_conservative_forces"] is None) + ) + ): + raise ValueError( + "if both 'non_conservative_stress' and " + "'non_conservative_forces' are present in `variants`, they " + "must either be both `None` or both not `None`." + ) + + self._nc_forces_key = pick_output( + "non_conservative_forces", + outputs, + resolved_variants["non_conservative_forces"], + ) + self._nc_stress_key = pick_output( + "non_conservative_stress", + outputs, + resolved_variants["non_conservative_stress"], + ) + else: + self._nc_forces_key = "non_conservative_forces" + self._nc_stress_key = "non_conservative_stress" + + if additional_outputs is None: + self._additional_output_requests = {} + else: + assert isinstance(additional_outputs, dict) + for name, output in additional_outputs.items(): + assert isinstance(name, str) + assert isinstance(output, torch.ScriptObject) + assert "explicit_gradients_setter" in output._method_names(), ( + "outputs must be ModelOutput instances" + ) + + self._additional_output_requests = additional_outputs + + self._model = model.to(device=self._device) + + self._calculate_uncertainty = ( + self._energy_uq_key in self._model.capabilities().outputs + # we require per-atom uncertainties to capture local effects + and self._model.capabilities().outputs[self._energy_uq_key].per_atom + and uncertainty_threshold is not None + ) + + if self._calculate_uncertainty: + assert uncertainty_threshold is not None + if uncertainty_threshold <= 0.0: + raise ValueError( + f"`uncertainty_threshold` is {uncertainty_threshold} but must " + "be positive" + ) + + # We do our own check to verify if a property is implemented in `calculate()`, + # so we pretend to be able to compute all properties ASE knows about. + self.implemented_properties = ALL_ASE_PROPERTIES + + self.additional_outputs: Dict[str, TensorMap] = {} + """ + Additional outputs computed by :py:meth:`calculate` are stored in this + dictionary. + + The keys will match the keys of the ``additional_outputs`` parameters to the + constructor; and the values will be the corresponding raw + :py:class:`metatensor.torch.TensorMap` produced by the model. + """ + + def todict(self): + if "model" not in self.parameters: + raise RuntimeError( + "can not save metatensor model in ASE `todict`, please initialize " + "`MetatomicCalculator` with a path to a saved model file if you need " + "to use `todict`" + ) + + return self.parameters + + @classmethod + def fromdict(cls, data): + return MetatomicCalculator(**data) + + def metadata(self) -> ModelMetadata: + """Get the metadata of the underlying model""" + return self._model.metadata() + + def run_model( + self, + atoms: Union[ase.Atoms, List[ase.Atoms]], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels] = None, + ) -> Dict[str, TensorMap]: + """ + Run the model on the given ``atoms``, computing the requested ``outputs`` and + only these. + + The output of the model is returned directly, and as such the blocks' ``values`` + will be :py:class:`torch.Tensor`. + + This is intended as an easy way to run metatensor models on + :py:class:`ase.Atoms` when the model can compute outputs not supported by the + standard ASE's calculator interface. + + All the parameters have the same meaning as the corresponding ones in + :py:meth:`metatomic.torch.ModelInterface.forward`. + + :param atoms: :py:class:`ase.Atoms`, or list of :py:class:`ase.Atoms`, on which + to run the model + :param outputs: outputs of the model that should be predicted + :param selected_atoms: subset of atoms on which to run the calculation + """ + if isinstance(atoms, ase.Atoms): + atoms_list = [atoms] + else: + atoms_list = atoms + + systems = [] + for atoms in atoms_list: + types, positions, cell, pbc = _ase_to_torch_data( + atoms=atoms, dtype=self._dtype, device=self._device + ) + system = System(types, positions, cell, pbc) + # Get the additional inputs requested by the model + for name, option in self._model.requested_inputs().items(): + input_tensormap = _get_ase_input( + atoms, name, option, dtype=self._dtype, device=self._device + ) + system.add_data(name, input_tensormap) + systems.append(system) + + # Compute the neighbors lists requested by the model + input_systems = _compute_requested_neighbors( + systems=systems, + requested_options=self._model.requested_neighbor_lists(), + check_consistency=self.parameters["check_consistency"], + ) + + available_outputs = self._model.capabilities().outputs + for key in outputs: + if key not in available_outputs: + raise ValueError(f"this model does not support '{key}' output") + + options = ModelEvaluationOptions( + length_unit="angstrom", + outputs=outputs, + selected_atoms=selected_atoms, + ) + return self._model( + systems=input_systems, + options=options, + check_consistency=self.parameters["check_consistency"], + ) + + def calculate( + self, + atoms: ase.Atoms, + properties: List[str], + system_changes: List[str], + ) -> None: + """ + Compute some ``properties`` with this calculator, and return them in the format + expected by ASE. + + This is not intended to be called directly by users, but to be an implementation + detail of ``atoms.get_energy()`` and related functions. See + :py:meth:`ase.calculators.calculator.Calculator.calculate` for more information. + """ + super().calculate( + atoms=atoms, + properties=properties, + system_changes=system_changes, + ) + + # In the next few lines, we decide which properties to calculate among energy, + # forces and stress. In addition to the requested properties, we calculate the + # energy if any of the three is requested, as it is an intermediate step in the + # calculation of the other two. We also calculate the forces if the stress is + # requested, and vice-versa. The overhead for the latter operation is also + # small, assuming that the majority of the model computes forces and stresses + # by backward propagation as opposed to forward-mode differentiation. + calculate_energy = ( + "energy" in properties + or "energies" in properties + or "forces" in properties + or "stress" in properties + ) + + # Check if energy-related properties are requested but the model doesn't + # support energy + if calculate_energy and self._energy_key is None: + raise PropertyNotImplementedError( + "This calculator does not support energy-related properties " + "(energy, energies, forces, stress) because the underlying model " + "does not have an energy output" + ) + calculate_energies = "energies" in properties + calculate_forces = "forces" in properties or "stress" in properties + calculate_stress = "stress" in properties + if calculate_stress and not atoms.pbc.all(): + warnings.warn( + "stress requested but likely to be wrong, since the system is not " + "periodic in all directions", + stacklevel=2, + ) + if "forces" in properties and atoms.pbc.all(): + # we have PBCs, and, since the user/integrator requested forces, we will run + # backward anyway, so let's do the stress as well for free (this saves + # another forward-backward call later if the stress is requested) + calculate_stress = True + if "stresses" in properties: + raise NotImplementedError("'stresses' are not implemented yet") + + if self.parameters["do_gradients_with_energy"]: + if calculate_energies or calculate_energy: + calculate_forces = True + calculate_stress = True + + with record_function("MetatomicCalculator::prepare_inputs"): + outputs = self._ase_properties_to_metatensor_outputs( + properties, + calculate_forces=calculate_forces, + calculate_stress=calculate_stress, + calculate_stresses=False, + ) + outputs.update(self._additional_output_requests) + if calculate_energy and self._calculate_uncertainty: + outputs[self._energy_uq_key] = ModelOutput( + quantity="energy", + unit="eV", + per_atom=True, + explicit_gradients=[], + ) + + capabilities = self._model.capabilities() + for name in outputs.keys(): + if name not in capabilities.outputs: + raise ValueError( + f"you asked for the calculation of {name}, but this model " + "does not support it" + ) + + types, positions, cell, pbc = _ase_to_torch_data( + atoms=atoms, dtype=self._dtype, device=self._device + ) + + do_backward = False + if calculate_forces and not self.parameters["non_conservative"]: + do_backward = True + positions.requires_grad_(True) + + if calculate_stress and not self.parameters["non_conservative"]: + do_backward = True + + strain = torch.eye( + 3, requires_grad=True, device=self._device, dtype=self._dtype + ) + + positions = positions @ strain + positions.retain_grad() + + cell = cell @ strain + + run_options = ModelEvaluationOptions( + length_unit="angstrom", + outputs=outputs, + selected_atoms=None, + ) + + with record_function("MetatomicCalculator::compute_neighbors"): + # convert from ase.Atoms to metatomic.torch.System + system = System(types, positions, cell, pbc) + input_system = _compute_requested_neighbors( + systems=[system], + requested_options=self._model.requested_neighbor_lists(), + check_consistency=self.parameters["check_consistency"], + )[0] + + with record_function("MetatomicCalculator::get_model_inputs"): + for name, option in self._model.requested_inputs().items(): + input_tensormap = _get_ase_input( + atoms, name, option, dtype=self._dtype, device=self._device + ) + input_system.add_data(name, input_tensormap) + + # no `record_function` here, this will be handled by AtomisticModel + outputs = self._model( + [input_system], + run_options, + check_consistency=self.parameters["check_consistency"], + ) + energy = outputs[self._energy_key] + + with record_function("MetatomicCalculator::sum_energies"): + if run_options.outputs[self._energy_key].per_atom: + assert len(energy) == 1 + assert energy.sample_names == ["system", "atom"] + assert torch.all(energy.block().samples["system"] == 0) + energies = energy + assert energies.block().values.shape == (len(atoms), 1) + + energy = mts.sum_over_samples(energy, sample_names=["atom"]) + + assert len(energy.block().gradients_list()) == 0 + assert energy.block().values.shape == (1, 1) + + with record_function("ASECalculator::uncertainty_warning"): + if calculate_energy and self._calculate_uncertainty: + uncertainty = outputs[self._energy_uq_key].block().values + assert uncertainty.shape == (len(atoms), 1) + uncertainty = uncertainty.detach().cpu().numpy() + + threshold = self.parameters["uncertainty_threshold"] + if np.any(uncertainty > threshold): + warnings.warn( + "Some of the atomic energy uncertainties are larger than the " + f"threshold of {threshold} eV. The prediction is above the " + f"threshold for atoms {np.where(uncertainty > threshold)[0]}.", + stacklevel=2, + ) + + if do_backward: + if energy.block().values.grad_fn is None: + # did the user actually request a gradient, or are we trying to + # compute one just for efficiency? + if "forces" in properties or "stress" in properties: + # the user asked for it, let it fail below + pass + else: + # we added the calculation, let's remove it + do_backward = False + calculate_forces = False + calculate_stress = False + + with record_function("MetatomicCalculator::run_backward"): + if do_backward: + energy.block().values.backward() + + with record_function("MetatomicCalculator::convert_outputs"): + self.results = {} + + if calculate_energies: + energies_values = energies.block().values.detach().reshape(-1) + atom_indexes = energies.block().samples.column("atom") + result = torch.zeros_like(energies_values) + result.index_add_(0, atom_indexes, energies_values) + self.results["energies"] = result.cpu().double().numpy() + + if calculate_energy: + energy_values = energy.block().values.detach() + energy_values = energy_values.cpu().double() + self.results["energy"] = energy_values.numpy()[0, 0] + + if calculate_forces: + if self.parameters["non_conservative"]: + forces_values = outputs[self._nc_forces_key].block().values.detach() + # remove any spurious net force + forces_values = forces_values - forces_values.mean( + dim=0, keepdim=True + ) + else: + forces_values = -system.positions.grad + forces_values = forces_values.reshape(-1, 3) + forces_values = forces_values.cpu().double() + self.results["forces"] = forces_values.numpy() + + if calculate_stress: + if self.parameters["non_conservative"]: + stress_values = outputs[self._nc_stress_key].block().values.detach() + else: + stress_values = strain.grad / atoms.cell.volume + stress_values = stress_values.reshape(3, 3) + stress_values = stress_values.cpu().double() + self.results["stress"] = _full_3x3_to_voigt_6_stress( + stress_values.numpy() + ) + + self.additional_outputs = {} + for name in self._additional_output_requests: + self.additional_outputs[name] = outputs[name] + + def compute_energy( + self, + atoms: Union[ase.Atoms, List[ase.Atoms]], + compute_forces_and_stresses: bool = False, + *, + per_atom: bool = False, + ) -> Dict[str, Union[Union[float, np.ndarray], List[Union[float, np.ndarray]]]]: + """ + Compute the energy of the given ``atoms``. + + Energies are computed in eV, forces in eV/Å, and stresses in 3x3 tensor format + and in units of eV/Å^3. + + :param atoms: :py:class:`ase.Atoms`, or list of :py:class:`ase.Atoms`, on which + to run the model + :param compute_forces_and_stresses: if ``True``, the model will also compute + forces and stresses. IMPORTANT: stresses will only be computed if all + provided systems have periodic boundary conditions in all directions. + :param per_atom: if ``True``, the per-atom energies will also be + computed. + :return: A dictionary with the computed properties. The dictionary will contain + the ``energy`` as a float. If ``compute_forces_and_stresses`` is True, + the ``forces`` and ``stress`` will also be included as numpy arrays. + If ``per_atom`` is True, the ``energies`` key will also be present, + containing the per-atom energies as a numpy array. + In case of a list of :py:class:`ase.Atoms`, the dictionary values will + instead be lists of the corresponding properties, in the same format. + """ + if self._energy_key is None: + raise ValueError( + "This calculator does not support energy computation because " + "the underlying model does not have an energy output" + ) + + if isinstance(atoms, ase.Atoms): + atoms_list = [atoms] + was_single = True + else: + atoms_list = atoms + was_single = False + + properties = ["energy"] + energy_per_atom = False + if per_atom: + energy_per_atom = True + properties.append("energies") + + outputs = self._ase_properties_to_metatensor_outputs( + properties=properties, + calculate_forces=compute_forces_and_stresses, + calculate_stress=compute_forces_and_stresses, + calculate_stresses=False, + ) + + systems = [] + if compute_forces_and_stresses: + strains = [] + for atoms in atoms_list: + types, positions, cell, pbc = _ase_to_torch_data( + atoms=atoms, dtype=self._dtype, device=self._device + ) + if compute_forces_and_stresses and not self.parameters["non_conservative"]: + positions.requires_grad_(True) + strain = torch.eye( + 3, requires_grad=True, device=self._device, dtype=self._dtype + ) + positions = positions @ strain + positions.retain_grad() + cell = cell @ strain + strains.append(strain) + system = System(types, positions, cell, pbc) + systems.append(system) + + # Compute the neighbors lists requested by the model + input_systems = _compute_requested_neighbors( + systems=systems, + requested_options=self._model.requested_neighbor_lists(), + check_consistency=self.parameters["check_consistency"], + ) + + predictions = self._model( + systems=input_systems, + options=ModelEvaluationOptions(length_unit="angstrom", outputs=outputs), + check_consistency=self.parameters["check_consistency"], + ) + energies = predictions[self._energy_key] + + if energy_per_atom: + # Get per-atom energies + sorted_block = mts.sort_block(energies.block(), axes="samples") + energies_values = sorted_block.values.detach().reshape(-1) + + split_sizes = [len(system) for system in systems] + atom_indices = sorted_block.samples.column("atom") + energies_values = torch.split(energies_values, split_sizes, dim=0) + split_atom_indices = torch.split(atom_indices, split_sizes, dim=0) + split_energies = [] + for atom_indices, values in zip( + split_atom_indices, energies_values, strict=True + ): + split_energy = torch.zeros( + len(atom_indices), dtype=values.dtype, device=values.device + ) + split_energy.index_add_(0, atom_indices, values) + split_energies.append(split_energy) + + total_energy = ( + mts.sum_over_samples(energies, ["atom"]) + .block() + .values.detach() + .cpu() + .double() + .numpy() + .flatten() + .tolist() + ) + results_as_numpy_arrays = { + "energy": total_energy, + "energies": [e.cpu().double().numpy() for e in split_energies], + } + else: + results_as_numpy_arrays = { + "energy": energies.block() + .values.squeeze(-1) + .detach() + .cpu() + .double() + .numpy(), + } + + if compute_forces_and_stresses: + if self.parameters["non_conservative"]: + results_as_numpy_arrays["forces"] = ( + predictions[self._nc_forces_key] + .block() + .values.squeeze(-1) + .detach() + .cpu() + .numpy() + ) + # all the forces are concatenated in a single array, so we need to + # split them into the original systems + split_sizes = [len(system) for system in systems] + split_indices = np.cumsum(split_sizes[:-1]) + results_as_numpy_arrays["forces"] = np.split( + results_as_numpy_arrays["forces"], split_indices, axis=0 + ) + + # remove net forces + results_as_numpy_arrays["forces"] = [ + f - f.mean(axis=0, keepdims=True) + for f in results_as_numpy_arrays["forces"] + ] + + if all(atoms.pbc.all() for atoms in atoms_list): + results_as_numpy_arrays["stress"] = [ + s + for s in predictions[self._nc_stress_key] + .block() + .values.squeeze(-1) + .detach() + .cpu() + .numpy() + ] + else: + energy_tensor = energies.block().values + energy_tensor.backward(torch.ones_like(energy_tensor)) + results_as_numpy_arrays["forces"] = [ + -system.positions.grad.cpu().numpy() for system in systems + ] + if all(atoms.pbc.all() for atoms in atoms_list): + results_as_numpy_arrays["stress"] = [ + strain.grad.cpu().numpy() / atoms.cell.volume + for strain, atoms in zip(strains, atoms_list, strict=False) + ] + if was_single: + for key, value in results_as_numpy_arrays.items(): + results_as_numpy_arrays[key] = value[0] + return results_as_numpy_arrays + + def _ase_properties_to_metatensor_outputs( + self, + properties, + *, + calculate_forces, + calculate_stress, + calculate_stresses, + ): + energy_properties = [] + for p in properties: + if p in ["energy", "energies", "forces", "stress", "stresses"]: + energy_properties.append(p) + else: + raise PropertyNotImplementedError( + f"property '{p}' it not yet supported by this calculator, " + "even if it might be supported by the model" + ) + + metatensor_outputs = {} + + # Only add energy output if the model supports it + if self._energy_key is not None: + output = ModelOutput( + quantity="energy", + unit="ev", + explicit_gradients=[], + ) + + if "energies" in properties or "stresses" in properties: + output.per_atom = True + else: + output.per_atom = False + + metatensor_outputs[self._energy_key] = output + if calculate_forces and self.parameters["non_conservative"]: + metatensor_outputs[self._nc_forces_key] = ModelOutput( + quantity="force", + unit="eV/Angstrom", + per_atom=True, + ) + + if calculate_stress and self.parameters["non_conservative"]: + metatensor_outputs[self._nc_stress_key] = ModelOutput( + quantity="pressure", + unit="eV/Angstrom^3", + per_atom=False, + ) + + if calculate_stresses and self.parameters["non_conservative"]: + raise NotImplementedError( + "non conservative, per-atom stress is not yet implemented" + ) + + available_outputs = self._model.capabilities().outputs + for key in metatensor_outputs: + if key not in available_outputs: + raise ValueError(f"this model does not support '{key}' output") + + return metatensor_outputs + + +def _get_ase_input( + atoms: ase.Atoms, + name: str, + option: ModelOutput, + dtype: torch.dtype, + device: torch.device, +) -> "TensorMap": + if name not in ARRAY_QUANTITIES: + raise ValueError( + f"The model requested '{name}', which is not available in `ase`." + ) + + infos = ARRAY_QUANTITIES[name] + + values = infos["getter"](atoms) + if values.shape[0] != len(atoms): + raise NotImplementedError( + f"The model requested the '{name}' input, " + f"but the data is not per-atom (shape {values.shape}). " + ) + # Shape: (n_atoms, n_components) -> (n_atoms, n_components, /* n_properties */ 1) + # for metatensor + values = torch.tensor(values[..., None]) + + components = [] + if values.shape[1] != 1: + components.append(Labels(["xyz"], torch.arange(values.shape[1]).reshape(-1, 1))) + + block = TensorBlock( + values, + samples=Labels( + ["system", "atom"], + torch.vstack( + [torch.full((values.shape[0],), 0), torch.arange(values.shape[0])] + ).T, + ), + components=components, + properties=Labels([infos["quantity"]], torch.tensor([[0]])), + ) + + tensor = TensorMap(Labels(["_"], torch.tensor([[0]])), [block]) + + tensor.set_info("quantity", infos["quantity"]) + tensor.set_info("unit", infos["unit"]) + + tensor = tensor.to(dtype=dtype, device=device) + return tensor + + +def _ase_to_torch_data(atoms, dtype, device): + """Get the positions, cell and pbc from ASE atoms as torch tensors""" + + types = torch.from_numpy(atoms.numbers).to(dtype=torch.int32, device=device) + positions = torch.from_numpy(atoms.positions).to(dtype=dtype).to(device=device) + cell = torch.zeros((3, 3), dtype=dtype, device=device) + pbc = torch.tensor(atoms.pbc, dtype=torch.bool, device=device) + + cell[pbc] = torch.tensor(atoms.cell[atoms.pbc], dtype=dtype, device=device) + + return types, positions, cell, pbc + + +def _full_3x3_to_voigt_6_stress(stress): + """ + Re-implementation of ``ase.stress.full_3x3_to_voigt_6_stress`` which does not do the + stress symmetrization correctly (they do ``(stress[1, 2] + stress[1, 2]) / 2.0``) + """ + return np.array( + [ + stress[0, 0], + stress[1, 1], + stress[2, 2], + (stress[1, 2] + stress[2, 1]) / 2.0, + (stress[0, 2] + stress[2, 0]) / 2.0, + (stress[0, 1] + stress[1, 0]) / 2.0, + ] + ) diff --git a/python/metatomic_ase/src/metatomic_ase/_neighbors.py b/python/metatomic_ase/src/metatomic_ase/_neighbors.py new file mode 100644 index 000000000..9073f62b9 --- /dev/null +++ b/python/metatomic_ase/src/metatomic_ase/_neighbors.py @@ -0,0 +1,139 @@ +from typing import List + +import torch +import vesin.metatomic +from metatensor.torch import Labels, TensorBlock + +from metatomic.torch import NeighborListOptions, System + + +try: + from nvalchemiops.torch.neighbors import neighbor_list as nvalchemi_neighbor_list + + HAS_NVALCHEMIOPS = True +except ImportError: + HAS_NVALCHEMIOPS = False + + +def _compute_requested_neighbors( + systems: List[System], + requested_options: List[NeighborListOptions], + check_consistency=False, +) -> List[System]: + """ + Compute all neighbor lists requested by ``model`` and store them inside the systems. + """ + can_use_nvalchemi = HAS_NVALCHEMIOPS and all( + system.device.type == "cuda" for system in systems + ) + + if can_use_nvalchemi: + full_nl_options = [] + half_nl_options = [] + for options in requested_options: + if options.full_list: + full_nl_options.append(options) + else: + half_nl_options.append(options) + + # Do the full neighbor lists with nvalchemi, and the rest with vesin + systems = _compute_requested_neighbors_nvalchemi( + systems=systems, + requested_options=full_nl_options, + ) + systems = _compute_requested_neighbors_vesin( + systems=systems, + requested_options=half_nl_options, + check_consistency=check_consistency, + ) + else: + systems = _compute_requested_neighbors_vesin( + systems=systems, + requested_options=requested_options, + check_consistency=check_consistency, + ) + + return systems + + +def _compute_requested_neighbors_vesin( + systems: List[System], + requested_options: List[NeighborListOptions], + check_consistency=False, +) -> List[System]: + """ + Compute all neighbor lists requested by ``model`` and store them inside the systems, + using vesin. + """ + + system_devices = [] + moved_systems = [] + for system in systems: + system_devices.append(system.device) + if system.device.type not in ["cpu", "cuda"]: + moved_systems.append(system.to(device="cpu")) + else: + moved_systems.append(system) + + vesin.metatomic.compute_requested_neighbors_from_options( + systems=moved_systems, + system_length_unit="angstrom", + options=requested_options, + check_consistency=check_consistency, + ) + + systems = [] + for system, device in zip(moved_systems, system_devices, strict=True): + systems.append(system.to(device=device)) + + return systems + + +def _compute_requested_neighbors_nvalchemi(systems, requested_options): + """ + Compute all neighbor lists requested by ``model`` and store them inside the systems, + using nvalchemiops. This function should only be called if all systems are on CUDA + and all neighbor list options require a full neighbor list. + """ + + for options in requested_options: + assert options.full_list + for system in systems: + assert system.device.type == "cuda" + + edge_index, _, S = nvalchemi_neighbor_list( + system.positions, + options.engine_cutoff("angstrom"), + cell=system.cell, + pbc=system.pbc, + return_neighbor_list=True, + ) + D = ( + system.positions[edge_index[1]] + - system.positions[edge_index[0]] + + S.to(system.cell.dtype) @ system.cell + ) + P = edge_index.T + + neighbors = TensorBlock( + D.reshape(-1, 3, 1), + samples=Labels( + names=[ + "first_atom", + "second_atom", + "cell_shift_a", + "cell_shift_b", + "cell_shift_c", + ], + values=torch.hstack([P, S]), + ), + components=[ + Labels("xyz", torch.tensor([[0], [1], [2]], device=system.device)) + ], + properties=Labels( + "distance", torch.tensor([[0]], device=system.device) + ), + ) + system.add_neighbor_list(options, neighbors) + + return systems diff --git a/python/metatomic_ase/src/metatomic_ase/_symmetry.py b/python/metatomic_ase/src/metatomic_ase/_symmetry.py new file mode 100644 index 000000000..5f91325b0 --- /dev/null +++ b/python/metatomic_ase/src/metatomic_ase/_symmetry.py @@ -0,0 +1,460 @@ +from typing import Dict, List, Optional, Tuple + +import numpy as np +import torch + +from ._calculator import MetatomicCalculator + + +import ase # isort: skip +import ase.calculators.calculator # isort: skip + + +class SymmetrizedCalculator(ase.calculators.calculator.Calculator): + r""" + Take a MetatomicCalculator and average its predictions to make it (approximately) + equivariant. Only predictions for energy, forces and stress are supported. + + The default is to average over a quadrature of the orthogonal group O(3) composed + this way: + + - Lebedev quadrature of the unit sphere (S^2) + - Equispaced sampling of the unit circle (S^1) + - Both proper and improper rotations are taken into account by including the + inversion operation (if ``include_inversion=True``) + + :param base_calculator: the MetatomicCalculator to be symmetrized + :param l_max: the maximum spherical harmonic degree that the model is expected to + be able to represent. This is used to choose the quadrature order. If ``0``, + no rotational averaging will be performed (it can be useful to average only over + the space group, see ``apply_group_symmetry``). + :param batch_size: number of rotated systems to evaluate at once. If ``None``, all + systems will be evaluated at once (this can lead to high memory usage). + :param include_inversion: if ``True``, the inversion operation will be included in + the averaging. This is required to average over the full orthogonal group O(3). + :param apply_space_group_symmetry: if ``True``, the results will be averaged over + discrete space group of rotations for the input system. The group operations are + computed with `spglib `_, and the average is + performed after the O(3) averaging (if any). This has no effect for non-periodic + systems. + :param store_rotational_std: if ``True``, the results will contain the standard + deviation over the different rotations for each property (e.g., ``energy_std``). + """ + + implemented_properties = ["energy", "energies", "forces", "stress", "stresses"] + + def __init__( + self, + base_calculator: MetatomicCalculator, + *, + l_max: int = 3, + batch_size: Optional[int] = None, + include_inversion: bool = True, + apply_space_group_symmetry: bool = False, + store_rotational_std: bool = False, + ) -> None: + try: + from scipy.integrate import lebedev_rule # noqa: F401 + except ImportError as e: + raise ImportError( + "scipy is required to use the `SymmetrizedCalculator`, please install " + "it with `pip install scipy` or `conda install scipy`" + ) from e + + super().__init__() + + self.base_calculator = base_calculator + if l_max > 131: + raise ValueError( + f"l_max={l_max} is too large, the maximum supported value is 131" + ) + self.l_max = l_max + self.include_inversion = include_inversion + + if l_max > 0: + lebedev_order, n_inplane_rotations = _choose_quadrature(l_max) + self.quadrature_rotations, self.quadrature_weights = _get_quadrature( + lebedev_order, n_inplane_rotations, include_inversion + ) + else: # no quadrature + if include_inversion: # identity and inversion + self.quadrature_rotations = np.array([np.eye(3), -np.eye(3)]) + self.quadrature_weights = np.array([0.5, 0.5]) + else: # only the identity + self.quadrature_rotations = np.array([np.eye(3)]) + self.quadrature_weights = np.array([1.0]) + + self.batch_size = ( + batch_size if batch_size is not None else len(self.quadrature_rotations) + ) + + self.store_rotational_std = store_rotational_std + self.apply_space_group_symmetry = apply_space_group_symmetry + + def calculate( + self, atoms: ase.Atoms, properties: List[str], system_changes: List[str] + ) -> None: + """ + Perform the calculation for the given atoms and properties. + + :param atoms: the :py:class:`ase.Atoms` on which to perform the calculation + :param properties: list of properties to compute, among ``energy``, ``forces``, + and ``stress`` + :param system_changes: list of changes to the system since the last call to + ``calculate`` + """ + super().calculate(atoms, properties, system_changes) + self.base_calculator.calculate(atoms, properties, system_changes) + + compute_forces_and_stresses = "forces" in properties or "stress" in properties + per_atom = "energies" in properties + + if len(self.quadrature_rotations) > 0: + rotated_atoms_list = _rotate_atoms(atoms, self.quadrature_rotations) + batches = [ + rotated_atoms_list[i : i + self.batch_size] + for i in range(0, len(rotated_atoms_list), self.batch_size) + ] + results: Dict[str, np.ndarray] = {} + for batch in batches: + try: + batch_results = self.base_calculator.compute_energy( + batch, + compute_forces_and_stresses, + per_atom=per_atom, + ) + for key, value in batch_results.items(): + results.setdefault(key, []) + results[key].extend( + [value] if isinstance(value, float) else value + ) + except torch.cuda.OutOfMemoryError as e: + raise RuntimeError( + "Out of memory error encountered during rotational averaging. " + "Please reduce the batch size or use lower rotational " + "averaging parameters. This can be done by setting the " + "`batch_size` and `l_max` parameters while initializing the " + "calculator." + ) from e + + self.results.update( + _compute_rotational_average( + results, + self.quadrature_rotations, + self.quadrature_weights, + self.store_rotational_std, + ) + ) + + if self.apply_space_group_symmetry: + # Apply the discrete space group of the system a posteriori + Q_list, P_list = _get_group_operations(atoms) + self.results.update(_average_over_group(self.results, Q_list, P_list)) + + +def _choose_quadrature(L_max: int) -> Tuple[int, int]: + """ + Choose a Lebedev quadrature order and number of in-plane rotations to integrate + spherical harmonics up to degree ``L_max``. + + :param L_max: maximum spherical harmonic degree + :return: (lebedev_order, n_inplane_rotations) + """ + # fmt: off + available = [ + 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31, 35, 41, + 47, 53, 59, 65, 71, 77, 83, 89, 95, 101, 107, 113, 119, 125, 131, + ] + # fmt: on + + # pick smallest order >= L_max + n = min(o for o in available if o >= L_max) + # minimal gamma count + K = 2 * L_max + 1 + return n, K + + +def _rotate_atoms(atoms: ase.Atoms, rotations: List[np.ndarray]) -> List[ase.Atoms]: + """ + Create a list of copies of ``atoms``, rotated by each of the given ``rotations``. + + :param atoms: the :py:class:`ase.Atoms` to be rotated + :param rotations: (N, 3, 3) array of orthogonal matrices + :return: list of N :py:class:`ase.Atoms`, each rotated by the corresponding matrix + """ + rotated_atoms_list = [] + has_cell = atoms.cell is not None and atoms.cell.rank > 0 + for rot in rotations: + new_atoms = atoms.copy() + new_atoms.positions = new_atoms.positions @ rot.T + if has_cell: + new_atoms.set_cell( + new_atoms.cell.array @ rot.T, scale_atoms=False, apply_constraint=False + ) + new_atoms.wrap() + rotated_atoms_list.append(new_atoms) + return rotated_atoms_list + + +def _get_quadrature(lebedev_order: int, n_rotations: int, include_inversion: bool): + """ + Lebedev(S^2) x uniform angle quadrature on SO(3). + If include_inversion=True, extend to O(3) by adding inversion * R. + + :param lebedev_order: order of the Lebedev quadrature on the unit sphere + :param n_rotations: number of in-plane rotations per Lebedev node + :param include_inversion: if ``True``, include the inversion operation in the + quadrature + :return: (N, 3, 3) array of orthogonal matrices, and (N,) array of weights + associated to each matrix + """ + from scipy.integrate import lebedev_rule + + # Lebedev nodes (X: (3, M)) + X, w = lebedev_rule(lebedev_order) # w sums to 4*pi + x, y, z = X + alpha = np.arctan2(y, x) # (M,) + beta = np.arccos(z) # (M,) + # beta = np.arccos(np.clip(z, -1.0, 1.0)) # (M,) + + K = int(n_rotations) + gamma = np.linspace(0.0, 2 * np.pi, K, endpoint=False) # (K,) + + Rot = _rotations_from_angles(alpha, beta, gamma) + R_so3 = Rot.as_matrix() # (N, 3, 3) + + # SO(3) Haar–probability weights: w_i/(4*pi*K), repeated over gamma + w_so3 = np.repeat(w / (4 * np.pi * K), repeats=gamma.size) # (N,) + + if not include_inversion: + return R_so3, w_so3 + + # Extend to O(3) by appending inversion * R + P = -np.eye(3) + R_o3 = np.concatenate([R_so3, P @ R_so3], axis=0) # (2N, 3, 3) + w_o3 = np.concatenate([0.5 * w_so3, 0.5 * w_so3], axis=0) + + return R_o3, w_o3 + + +def _rotations_from_angles(alpha, beta, gamma): + from scipy.spatial.transform import Rotation + + # Build all combinations (alpha_i, beta_i, gamma_j) + A = np.repeat(alpha, gamma.size).reshape(-1, 1) # (N, 1) + B = np.repeat(beta, gamma.size).reshape(-1, 1) # (N, 1) + G = np.tile(gamma, alpha.size).reshape(-1, 1) # (N, 1) + + # Compose ZYZ rotations in SO(3) + Rot = ( + Rotation.from_euler("z", A) + * Rotation.from_euler("y", B) + * Rotation.from_euler("z", G) + ) + + return Rot + + +def _compute_rotational_average(results, rotations, weights, store_std): + R = rotations + B = R.shape[0] + w = weights + w = w / w.sum() + + def _wreshape(x): + return w.reshape((B,) + (1,) * (x.ndim - 1)) + + def _wmean(x): + return np.sum(_wreshape(x) * x, axis=0) + + def _wstd(x): + mu = _wmean(x) + return np.sqrt(np.sum(_wreshape(x) * (x - mu) ** 2, axis=0)) + + out = {} + + # Energy (B,) + if "energy" in results: + E = np.asarray(results["energy"], dtype=float) # (B,) + out["energy"] = _wmean(E) # () + if store_std: + out["energy_rot_std"] = _wstd(E) # () + + if "energies" in results: + E = np.asarray(results["energies"], dtype=float) # (B,N) + out["energies"] = _wmean(E) # (N,) + if store_std: + out["energies_rot_std"] = _wstd(E) # (N,) + + # Forces (B,N,3) from rotated structures: back-rotate with F' R + if "forces" in results: + F = np.asarray(results["forces"], dtype=float) # (B,N,3) + F_back = F @ R # F' R + out["forces"] = _wmean(F_back) # (N,3) + if store_std: + out["forces_rot_std"] = _wstd(F_back) # (N,3) + + # Stress (B,3,3) from rotated structures: back-rotate with R^T S' R + if "stress" in results: + S = np.asarray(results["stress"], dtype=float) # (B,3,3) + RT = np.swapaxes(R, 1, 2) + S_back = RT @ S @ R # R^T S' R + out["stress"] = _wmean(S_back) # (3,3) + if store_std: + out["stress_rot_std"] = _wstd(S_back) # (3,3) + + if "stresses" in results: + S = np.asarray(results["stresses"], dtype=float) # (B,N,3,3) + RT = np.swapaxes(R, 1, 2) + S_back = RT[:, None, :, :] @ S @ R[:, None, :, :] # R^T S' R + out["stresses"] = _wmean(S_back) # (N,3,3) + if store_std: + out["stresses_rot_std"] = _wstd(S_back) # (N,3,3) + + return out + + +def _get_group_operations( + atoms: ase.Atoms, symprec: float = 1e-6, angle_tolerance: float = -1.0 +) -> Tuple[List[np.ndarray], List[np.ndarray]]: + """ + Extract point-group rotations Q_g (Cartesian, 3x3) and the corresponding + atom-index permutations P_g (N x N) induced by the space-group operations. + Returns Q_list, Cartesian rotation matrices of the point group, + and P_list, permutation matrices mapping original indexing -> indexing after (R,t), + + :param atoms: input structure + :param symprec: tolerance for symmetry finding + :param angle_tolerance: tolerance for symmetry finding (in degrees). If less than 0, + a value depending on ``symprec`` will be chosen automatically by spglib. + :return: List of rotation matrices and permutation matrices. + + """ + try: + import spglib + except ImportError as e: + raise ImportError( + "spglib is required to use the SymmetrizedCalculator with " + "`apply_group_symmetry=True`. Please install it with " + "`pip install spglib` or `conda install -c conda-forge spglib`" + ) from e + + if not (atoms.pbc.all()): + # No periodic boundary conditions: no symmetry + return [], [] + + # Lattice with column vectors a1,a2,a3 (spglib expects (cell, frac, Z)) + A = atoms.cell.array.T # (3,3) + frac = atoms.get_scaled_positions() # (N,3) in [0,1) + numbers = atoms.numbers + N = len(atoms) + + data = spglib.get_symmetry_dataset( + (atoms.cell.array, frac, numbers), + symprec=symprec, + angle_tolerance=angle_tolerance, + _throw=True, + ) + + if data is None: + # No symmetry found + return [], [] + R_frac = data.rotations # (n_ops, 3,3), integer + t_frac = data.translations # (n_ops, 3) + Z = numbers + + # Match fractional coords modulo 1 within a tolerance, respecting chemical species + def _match_index(x_new, frac_ref, Z_ref, Z_i, tol=1e-6): + d = np.abs(frac_ref - x_new) # (N,3) + d = np.minimum(d, 1.0 - d) # periodic distance + # Mask by identical species + mask = Z_ref == Z_i + if not np.any(mask): + raise RuntimeError("No matching species found while building permutation.") + # Choose argmin over max-norm within species + idx = np.where(mask)[0] + j = idx[np.argmin(np.max(d[idx], axis=1))] + + # Sanity check + if np.max(d[j]) > tol: + raise RuntimeError( + ( + f"Sanity check failed in _match_index: max distance {np.max(d[j])} " + f"exceeds tolerance {tol}." + ) + ) + return j + + Q_list, P_list = [], [] + seen = set() + Ainv = np.linalg.inv(A) + + for Rf, tf in zip(R_frac, t_frac, strict=False): + # Cartesian rotation: Q = A Rf A^{-1} + Q = A @ Rf @ Ainv + # Deduplicate rotations (point group) by rounding + key = tuple(np.round(Q.flatten(), 12)) + if key in seen: + continue + seen.add(key) + + # Build the permutation P from i to j + P = np.zeros((N, N), dtype=int) + new_frac = (frac @ Rf.T + tf) % 1.0 # images after (Rf,tf) + for i in range(N): + j = _match_index(new_frac[i], frac, Z, Z[i]) + P[j, i] = 1 # column i maps to row j + + Q_list.append(Q.astype(float)) + P_list.append(P) + + return Q_list, P_list + + +def _average_over_group( + results: dict, Q_list: List[np.ndarray], P_list: List[np.ndarray] +) -> dict: + """ + Apply the point-group projector in output space. + + :param results: Must contain 'energy' (scalar), and/or 'forces' (N,3), and/or + 'stress' (3,3). These are predictions for the current structure in the reference + frame. + :param Q_list: Rotation matrices of the point group, from + :py:func:`_get_group_operations` + :param P_list: Permutation matrices of the point group, from + :py:func:`_get_group_operations` + :return out: Projected quantities. + """ + m = len(Q_list) + if m == 0: + return results # nothing to do + + out = {} + # Energy: unchanged by the projector (scalar) + if "energy" in results: + out["energy"] = float(results["energy"]) + + # Forces: (N,3) row-vectors; projector: (1/|G|) \sum_g P_g^T F Q_g + if "forces" in results: + F = np.asarray(results["forces"], float) + if F.ndim != 2 or F.shape[1] != 3: + raise ValueError(f"'forces' must be (N,3), got {F.shape}") + acc = np.zeros_like(F) + for Q, P in zip(Q_list, P_list, strict=False): + acc += P.T @ (F @ Q) + out["forces"] = acc / m + + # Stress: (3,3); projector: (1/|G|) \sum_g Q_g^T S Q_g + if "stress" in results: + S = np.asarray(results["stress"], float) + if S.shape != (3, 3): + raise ValueError(f"'stress' must be (3,3), got {S.shape}") + # S = 0.5 * (S + S.T) # symmetrize just in case + acc = np.zeros_like(S) + for Q in Q_list: + acc += Q.T @ S @ Q + S_pg = acc / m + out["stress"] = S_pg + + return out diff --git a/python/metatomic_ase/tests/__init__.py b/python/metatomic_ase/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/metatomic_ase/tests/_tests_utils.py b/python/metatomic_ase/tests/_tests_utils.py new file mode 100644 index 000000000..7ea6f2b21 --- /dev/null +++ b/python/metatomic_ase/tests/_tests_utils.py @@ -0,0 +1,28 @@ +import os + +import torch + + +def _can_use_mps_backend(): + return ( + # Github Actions M1 runners don't have a GPU accessible + os.environ.get("GITHUB_ACTIONS") is None + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_built() + and torch.backends.mps.is_available() + ) + + +ALL_DEVICE_DTYPE = [("cpu", "float64"), ("cpu", "float32")] +if torch.cuda.is_available(): + ALL_DEVICE_DTYPE.append(("cuda", "float64")) + ALL_DEVICE_DTYPE.append(("cuda", "float32")) + +if _can_use_mps_backend(): + ALL_DEVICE_DTYPE.append(("mps", "float32")) + + +STR_TO_DTYPE = { + "float32": torch.float32, + "float64": torch.float64, +} diff --git a/python/metatomic_ase/tests/calculator.py b/python/metatomic_ase/tests/calculator.py new file mode 100644 index 000000000..acef0ed60 --- /dev/null +++ b/python/metatomic_ase/tests/calculator.py @@ -0,0 +1,847 @@ +import os +import subprocess +import sys +from typing import Dict, List, Optional + +import ase.build +import ase.calculators.lj +import ase.md +import ase.units +import numpy as np +import pytest +import torch +from ase.calculators.calculator import PropertyNotImplementedError +from ase.md.velocitydistribution import MaxwellBoltzmannDistribution +from metatensor.torch import Labels, TensorBlock, TensorMap + +import metatomic_lj_test +from metatomic.torch import ( + AtomisticModel, + ModelCapabilities, + ModelMetadata, + ModelOutput, + System, +) +from metatomic_ase import MetatomicCalculator +from metatomic_ase._calculator import ( + ARRAY_QUANTITIES, + _full_3x3_to_voigt_6_stress, +) + +from ._tests_utils import ALL_DEVICE_DTYPE, STR_TO_DTYPE + + +CUTOFF = 5.0 +SIGMA = 1.5808 +EPSILON = 0.1729 + + +@pytest.fixture +def model(): + return metatomic_lj_test.lennard_jones_model( + atomic_type=28, + cutoff=CUTOFF, + sigma=SIGMA, + epsilon=EPSILON, + length_unit="Angstrom", + energy_unit="eV", + with_extension=False, + ) + + +@pytest.fixture +def model_different_units(): + return metatomic_lj_test.lennard_jones_model( + atomic_type=28, + cutoff=CUTOFF / ase.units.Bohr, + sigma=SIGMA / ase.units.Bohr, + epsilon=EPSILON / ase.units.kJ * ase.units.mol, + length_unit="Bohr", + energy_unit="kJ/mol", + with_extension=False, + ) + + +@pytest.fixture +def atoms(): + np.random.seed(0xDEADBEEF) + + atoms = ase.build.make_supercell( + ase.build.bulk("Ni", "fcc", a=3.6, cubic=True), 2 * np.eye(3) + ) + atoms.positions += 0.2 * np.random.rand(*atoms.positions.shape) + + return atoms + + +def check_against_ase_lj(atoms, calculator, dtype): + ref = atoms.copy() + ref.calc = ase.calculators.lj.LennardJones( + sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False + ) + + atoms.calc = calculator + + if dtype == "float32": + rtol = 1e-5 + atol = 1e-5 + elif dtype == "float64": + rtol = 1e-5 + atol = 1e-8 + else: + raise ValueError(f"unsupported dtype: {dtype}") + + assert np.allclose( + ref.get_potential_energy(), atoms.get_potential_energy(), atol=atol, rtol=rtol + ) + assert np.allclose( + ref.get_potential_energies(), + atoms.get_potential_energies(), + atol=atol, + rtol=rtol, + ) + assert np.allclose(ref.get_forces(), atoms.get_forces(), atol=atol, rtol=rtol) + assert np.allclose(ref.get_stress(), atoms.get_stress(), atol=atol, rtol=rtol) + + +def _set_model_dtype(model, dtype): + model._capabilities.dtype = dtype + model._model_dtype = STR_TO_DTYPE[dtype] + return model + + +@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) +def test_python_model(model, model_different_units, atoms, device, dtype): + model = _set_model_dtype(model, dtype) + model_different_units = _set_model_dtype(model_different_units, dtype) + + check_against_ase_lj( + atoms, + MetatomicCalculator( + model, + check_consistency=True, + uncertainty_threshold=None, + device=device, + ), + dtype=dtype, + ) + check_against_ase_lj( + atoms, + MetatomicCalculator( + model_different_units, + check_consistency=True, + uncertainty_threshold=None, + device=device, + ), + dtype=dtype, + ) + + +@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) +def test_torch_script_model(model, model_different_units, atoms, device, dtype): + model = _set_model_dtype(model, dtype) + model_different_units = _set_model_dtype(model_different_units, dtype) + + model = torch.jit.script(model) + check_against_ase_lj( + atoms, + MetatomicCalculator( + model, + check_consistency=True, + uncertainty_threshold=None, + device=device, + ), + dtype=dtype, + ) + + model_different_units = torch.jit.script(model_different_units) + check_against_ase_lj( + atoms, + MetatomicCalculator( + model_different_units, + check_consistency=True, + uncertainty_threshold=None, + device=device, + ), + dtype=dtype, + ) + + +@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) +def test_exported_model(tmpdir, model, model_different_units, atoms, device, dtype): + model = _set_model_dtype(model, dtype) + model_different_units = _set_model_dtype(model_different_units, dtype) + + path = os.path.join(tmpdir, "exported-model.pt") + model.save(path) + check_against_ase_lj( + atoms, + MetatomicCalculator( + path, + check_consistency=True, + uncertainty_threshold=None, + device=device, + ), + dtype=dtype, + ) + + model_different_units.save(path) + check_against_ase_lj( + atoms, + MetatomicCalculator( + path, + check_consistency=True, + uncertainty_threshold=None, + device=device, + ), + dtype=dtype, + ) + + +@pytest.mark.parametrize("non_conservative", [True, False]) +def test_get_properties(model, atoms, non_conservative): + atoms.calc = MetatomicCalculator( + model, + check_consistency=True, + non_conservative=non_conservative, + uncertainty_threshold=None, + ) + + properties = atoms.get_properties(["energy", "energies", "forces", "stress"]) + + assert np.all(properties["energies"] == atoms.get_potential_energies()) + assert np.all(properties["energy"] == atoms.get_potential_energy()) + assert np.all(properties["forces"] == atoms.get_forces()) + assert np.all(properties["stress"] == atoms.get_stress()) + + # check that we can use all of the `.get_xxx` functions independantly + atoms.calc = MetatomicCalculator( + model, non_conservative=non_conservative, uncertainty_threshold=None + ) + atoms.get_potential_energy() + + atoms.calc = MetatomicCalculator( + model, non_conservative=non_conservative, uncertainty_threshold=None + ) + atoms.get_potential_energies() + + atoms.calc = MetatomicCalculator( + model, non_conservative=non_conservative, uncertainty_threshold=None + ) + atoms.get_forces() + + atoms.calc = MetatomicCalculator( + model, non_conservative=non_conservative, uncertainty_threshold=None + ) + atoms.get_stress() + + +def test_accuracy_warning(model, atoms): + # our dummy model artificially gives a high uncertainty for large structures + big_atoms = atoms * (2, 2, 2) + big_atoms.calc = MetatomicCalculator(model, check_consistency=True) + + with pytest.warns( + UserWarning, + match="Some of the atomic energy uncertainties are large", + ): + big_atoms.get_forces() + + +def accuracy_is_zero_error(atoms): + match = "`uncertainty_threshold` is 0.0 but must be positive" + with pytest.raises(ValueError, match=match): + atoms.calc = MetatomicCalculator(model, uncertainty_threshold=0.0) + + +def test_run_model(tmpdir, model, atoms): + ref = atoms.copy() + ref.calc = ase.calculators.lj.LennardJones( + sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False + ) + + path = os.path.join(tmpdir, "exported-model.pt") + model.save(path) + calculator = MetatomicCalculator( + path, check_consistency=True, uncertainty_threshold=None + ) + + first_mask = [a % 2 == 0 for a in range(len(atoms))] + first_half = Labels( + ["system", "atom"], + torch.tensor([[0, a] for a in range(len(atoms)) if a % 2 == 0]), + ) + + second_mask = [a % 2 == 1 for a in range(len(atoms))] + second_half = Labels( + ["system", "atom"], + torch.tensor([[0, a] for a in range(len(atoms)) if a % 2 == 1]), + ) + + # check overall prediction + requested = {"energy": ModelOutput(per_atom=False)} + outputs = calculator.run_model(atoms, outputs=requested) + assert np.allclose( + ref.get_potential_energy(), outputs["energy"].block().values.item() + ) + + # check per atom energy + requested = {"energy": ModelOutput(per_atom=True)} + outputs = calculator.run_model(atoms, outputs=requested, selected_atoms=first_half) + first_energies = outputs["energy"].block().values.numpy().reshape(-1) + + outputs = calculator.run_model(atoms, outputs=requested, selected_atoms=second_half) + second_energies = outputs["energy"].block().values.numpy().reshape(-1) + + expected = ref.get_potential_energies() + assert np.allclose(expected[first_mask], first_energies) + assert np.allclose(expected[second_mask], second_energies) + + # check total energy + requested = {"energy": ModelOutput(per_atom=False)} + outputs = calculator.run_model(atoms, outputs=requested, selected_atoms=first_half) + first_energies = outputs["energy"].block().values.numpy().reshape(-1) + + outputs = calculator.run_model(atoms, outputs=requested, selected_atoms=second_half) + second_energies = outputs["energy"].block().values.numpy().reshape(-1) + + expected = ref.get_potential_energy() + assert np.allclose(expected, first_energies + second_energies) + + # check batched prediction + requested = {"energy": ModelOutput(per_atom=False)} + outputs = calculator.run_model([atoms, atoms], outputs=requested) + assert np.allclose( + ref.get_potential_energy(), outputs["energy"].block().values[[0]] + ) + + # check non-conservative forces and stresses + requested = { + "energy": ModelOutput(per_atom=False), + "non_conservative_forces": ModelOutput(per_atom=True), + "non_conservative_stress": ModelOutput(per_atom=False), + } + outputs = calculator.run_model([atoms, atoms], outputs=requested) + assert np.allclose( + ref.get_potential_energy(), outputs["energy"].block().values[[0]] + ) + assert np.allclose( + ref.get_potential_energy(), outputs["energy"].block().values[[1]] + ) + assert "non_conservative_forces" in outputs + assert outputs["non_conservative_forces"].block().values.shape == ( + 2 * len(atoms), + 3, + 1, + ) + assert "non_conservative_stress" in outputs + assert outputs["non_conservative_stress"].block().values.shape == (2, 3, 3, 1) + + +@pytest.mark.parametrize("non_conservative", [True, False]) +@pytest.mark.parametrize("per_atom", [True, False]) +def test_compute_energy(tmpdir, model, atoms, non_conservative, per_atom): + ref = atoms.copy() + ref.calc = ase.calculators.lj.LennardJones( + sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False + ) + + path = os.path.join(tmpdir, "exported-model.pt") + model.save(path) + calculator = MetatomicCalculator( + path, + check_consistency=True, + non_conservative=non_conservative, + ) + + results = calculator.compute_energy(atoms, per_atom=per_atom) + if per_atom: + energies = results["energies"] + assert np.allclose(ref.get_potential_energies(), energies) + assert np.allclose(ref.get_potential_energy(), results["energy"]) + + results = calculator.compute_energy( + atoms, compute_forces_and_stresses=True, per_atom=per_atom + ) + assert np.allclose(ref.get_potential_energy(), results["energy"]) + if not non_conservative: + assert np.allclose(ref.get_forces(), results["forces"]) + assert np.allclose( + ref.get_stress(), _full_3x3_to_voigt_6_stress(results["stress"]) + ) + if per_atom: + assert np.allclose(ref.get_potential_energies(), results["energies"]) + + results = calculator.compute_energy([atoms, atoms], per_atom=per_atom) + assert np.allclose(ref.get_potential_energy(), results["energy"][0]) + assert np.allclose(ref.get_potential_energy(), results["energy"][1]) + if per_atom: + assert np.allclose(ref.get_potential_energies(), results["energies"][0]) + assert np.allclose(ref.get_potential_energies(), results["energies"][1]) + + results = calculator.compute_energy( + [atoms, atoms], + compute_forces_and_stresses=True, + per_atom=per_atom, + ) + assert np.allclose(ref.get_potential_energy(), results["energy"][0]) + assert np.allclose(ref.get_potential_energy(), results["energy"][1]) + if not non_conservative: + assert np.allclose(ref.get_forces(), results["forces"][0]) + assert np.allclose(ref.get_forces(), results["forces"][1]) + assert np.allclose( + ref.get_stress(), _full_3x3_to_voigt_6_stress(results["stress"][0]) + ) + assert np.allclose( + ref.get_stress(), _full_3x3_to_voigt_6_stress(results["stress"][1]) + ) + if per_atom: + assert np.allclose(ref.get_potential_energies(), results["energies"][0]) + assert np.allclose(ref.get_potential_energies(), results["energies"][1]) + + atoms_no_pbc = atoms.copy() + atoms_no_pbc.pbc = [False, False, False] + assert "stress" not in calculator.compute_energy(atoms_no_pbc) + + +def test_serialize_ase(tmpdir, model, atoms): + calculator = MetatomicCalculator(model, uncertainty_threshold=None) + + message = ( + "can not save metatensor model in ASE `todict`, please initialize " + "`MetatomicCalculator` with a path to a saved model file if you need to use " + "`todict" + ) + with pytest.raises(RuntimeError, match=message): + calculator.todict() + + # save with exported model + path = os.path.join(tmpdir, "exported-model.pt") + model.save(path) + + calculator = MetatomicCalculator(path, uncertainty_threshold=None) + data = calculator.todict() + _ = MetatomicCalculator.fromdict(data) + + # check the standard trajectory format of ASE, which uses `todict`/`fromdict` + atoms.calc = MetatomicCalculator(path, uncertainty_threshold=None) + with tmpdir.as_cwd(): + dyn = ase.md.VelocityVerlet( + atoms, + timestep=2 * ase.units.fs, + trajectory="file.traj", + ) + dyn.run(10) + dyn.close() + + atoms = ase.io.read("file.traj", "-1") + assert atoms.calc is not None + + +@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) +def test_dtype_device(tmpdir, model, atoms, device, dtype): + ref = atoms.copy() + ref.calc = ase.calculators.lj.LennardJones( + sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False + ) + expected = ref.get_potential_energy() + path = os.path.join(tmpdir, "exported-model.pt") + + capabilities = model.capabilities() + capabilities.dtype = dtype + + # re-create the model with a different dtype + dtype_model = AtomisticModel( + model.module.to(STR_TO_DTYPE[dtype]), + model.metadata(), + capabilities, + ) + + dtype_model.save(path) + atoms.calc = MetatomicCalculator( + path, check_consistency=True, device=device, uncertainty_threshold=None + ) + assert np.allclose(atoms.get_potential_energy(), expected) + + +@pytest.mark.parametrize("device", ["cpu"]) +def test_model_with_extensions(tmpdir, atoms, device): + ref = atoms.copy() + ref.calc = ase.calculators.lj.LennardJones( + sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False + ) + + model_path = os.path.join(tmpdir, "model.pt") + extensions_directory = os.path.join(tmpdir, "extensions") + + # use forward slash as path separator even on windows. This removes the need to + # escape the path below + model_path = model_path.replace("\\", "/") + extensions_directory = extensions_directory.replace("\\", "/") + + # export the model in a sub-process, to prevent loading of the extension in the + # current interpreter (until we try to execute the code) + script = f""" +import metatomic_lj_test + +model = metatomic_lj_test.lennard_jones_model( + atomic_type=28, + cutoff={CUTOFF}, + sigma={SIGMA}, + epsilon={EPSILON}, + length_unit="Angstrom", + energy_unit="eV", + with_extension=True, +) + +model.save("{model_path}", collect_extensions="{extensions_directory}") + """ + + subprocess.run([sys.executable, "-c", script], check=True, cwd=tmpdir) + + message = ( + "This is likely due to missing TorchScript extensions.\nMake sure to provide " + "the `extensions_directory` argument if your extensions are not installed " + "system-wide" + ) + with pytest.raises(RuntimeError, match=message): + MetatomicCalculator(model_path, check_consistency=True) + + # Now actually loading the extensions + atoms.calc = MetatomicCalculator( + model_path, + extensions_directory=extensions_directory, + check_consistency=True, + uncertainty_threshold=None, + ) + + assert np.allclose(ref.get_potential_energy(), atoms.get_potential_energy()) + + +class MultipleOutputModel(torch.nn.Module): + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels] = None, + ) -> Dict[str, TensorMap]: + results = {} + device = systems[0].positions.device + for name, requested in outputs.items(): + assert not requested.per_atom + + block = TensorBlock( + values=torch.tensor([[0.0]], dtype=torch.float64, device=device), + samples=Labels("system", torch.tensor([[0]], device=device)), + components=torch.jit.annotate(List[Labels], []), + properties=Labels( + name.split(":")[0], torch.tensor([[0]], device=device) + ), + ) + tensor = TensorMap(Labels("_", torch.tensor([[0]], device=device)), [block]) + results[name] = tensor + + return results + + +def test_additional_outputs(atoms): + capabilities = ModelCapabilities( + outputs={ + "energy": ModelOutput(per_atom=False), + "test::test": ModelOutput(per_atom=False), + "another::one": ModelOutput(per_atom=False), + }, + atomic_types=[28], + interaction_range=0.0, + supported_devices=["cpu"], + dtype="float64", + ) + model = AtomisticModel(MultipleOutputModel().eval(), ModelMetadata(), capabilities) + + atoms.calc = MetatomicCalculator( + model, + check_consistency=True, + uncertainty_threshold=None, + ) + + assert atoms.get_potential_energy() == 0.0 + assert atoms.calc.additional_outputs == {} + + atoms.calc = MetatomicCalculator( + model, + additional_outputs={ + "test::test": ModelOutput(per_atom=False), + "another::one": ModelOutput(per_atom=False), + }, + check_consistency=True, + uncertainty_threshold=None, + ) + assert atoms.get_potential_energy() == 0.0 + + test = atoms.calc.additional_outputs["test::test"] + assert test.block().properties.names == ["test"] + + another = atoms.calc.additional_outputs["another::one"] + assert another.block().properties.names == ["another"] + + +@pytest.mark.parametrize("non_conservative", [True, False]) +def test_variants(atoms, model, non_conservative): + atoms.calc = MetatomicCalculator( + model, + check_consistency=True, + non_conservative=non_conservative, + uncertainty_threshold=None, + ) + + atoms_variant = atoms.copy() + atoms_variant.calc = MetatomicCalculator( + model, + check_consistency=True, + non_conservative=non_conservative, + variants={"energy": "doubled"}, + uncertainty_threshold=None, + ) + + np.allclose( + 2.0 * atoms.get_potential_energy(), atoms_variant.get_potential_energy() + ) + np.allclose(2.0 * atoms.get_forces(), atoms_variant.get_forces()) + np.allclose(2.0 * atoms.get_stress(), atoms_variant.get_stress()) + + +@pytest.mark.parametrize( + "default_output", + [ + "energy", + "energy_uncertainty", + "non_conservative_forces", + "non_conservative_stress", + ], +) +def test_variant_default(atoms, model, default_output): + """Allow setting a variant explicitly to None to use the default output.""" + + atoms.calc = MetatomicCalculator( + model, + check_consistency=True, + non_conservative=True, + uncertainty_threshold=None, + ) + + variants = { + v: "doubled" + for v in [ + "energy", + "energy_uncertainty", + "non_conservative_forces", + "non_conservative_stress", + ] + } + variants[default_output] = None + + atoms_variant = atoms.copy() + atoms_variant.calc = MetatomicCalculator( + model, + check_consistency=True, + non_conservative=True, + variants={"energy": "doubled"}, + uncertainty_threshold=None, + ) + + if default_output == "energy": + np.allclose(atoms.get_potential_energy(), atoms_variant.get_potential_energy()) + np.allclose(2.0 * atoms.get_forces(), atoms_variant.get_forces()) + np.allclose(2.0 * atoms.get_stress(), atoms_variant.get_stress()) + elif default_output == "non_conservative_forces": + np.allclose( + 2.0 * atoms.get_potential_energy(), atoms_variant.get_potential_energy() + ) + np.allclose(atoms.get_forces(), atoms_variant.get_forces()) + np.allclose(2.0 * atoms.get_stress(), atoms_variant.get_stress()) + elif default_output == "non_conservative_stress": + np.allclose( + 2.0 * atoms.get_potential_energy(), atoms_variant.get_potential_energy() + ) + np.allclose(2.0 * atoms.get_forces(), atoms_variant.get_forces()) + np.allclose(atoms.get_stress(), atoms_variant.get_stress()) + + +@pytest.mark.parametrize("force_is_None", [True, False]) +def test_variant_non_conservative_error(atoms, model, force_is_None): + variants = { + "energy": "doubled", + "non_conservative_forces": "doubled", + "non_conservative_stress": "doubled", + } + + if force_is_None: + variants["non_conservative_forces"] = None + else: + variants["non_conservative_stress"] = None + + match = "must either be both `None` or both not `None`." + with pytest.raises(ValueError, match=match): + MetatomicCalculator( + model, + check_consistency=True, + non_conservative=True, + variants=variants, + uncertainty_threshold=None, + ) + + +def test_model_without_energy(atoms): + """ + Test that a MetatomicCalculator can be created with a model without energy + output. + """ + + # Create a model that only outputs a custom property, no energy + class NoEnergyModel(torch.nn.Module): + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels] = None, + ) -> Dict[str, TensorMap]: + results = {} + for name in outputs: + # Return dummy data for each requested output + block = TensorBlock( + values=torch.tensor([[1.0]], dtype=torch.float64), + samples=Labels("system", torch.tensor([[0]])), + components=torch.jit.annotate(List[Labels], []), + properties=Labels(name.split(":")[0], torch.tensor([[0]])), + ) + tensor = TensorMap(Labels("_", torch.tensor([[0]])), [block]) + results[name] = tensor + return results + + # Create model capabilities without energy output + capabilities = ModelCapabilities( + outputs={ + "features": ModelOutput(per_atom=False), + "custom::output": ModelOutput(per_atom=False), + }, + atomic_types=[28], + interaction_range=0.0, + supported_devices=["cpu"], + dtype="float64", + ) + model = AtomisticModel(NoEnergyModel().eval(), ModelMetadata(), capabilities) + + # Should be able to create calculator without error + calc = MetatomicCalculator( + model, + check_consistency=True, + uncertainty_threshold=None, + ) + + # The calculator should work for additional outputs + atoms.calc = MetatomicCalculator( + model, + additional_outputs={ + "features": ModelOutput(per_atom=False), + }, + check_consistency=True, + uncertainty_threshold=None, + ) + + # Should be able to call run_model directly with custom outputs + outputs = atoms.calc.run_model( + atoms, + outputs={"features": ModelOutput(per_atom=False)}, + ) + assert "features" in outputs + + # But trying to get energy should fail with a clear error + match = "does not support energy-related properties" + with pytest.raises(PropertyNotImplementedError, match=match): + atoms.get_potential_energy() + + with pytest.raises(PropertyNotImplementedError, match=match): + atoms.get_forces() + + # compute_energy should also fail + match = "does not support energy computation" + with pytest.raises(ValueError, match=match): + calc.compute_energy(atoms) + + +class AdditionalInputModel(torch.nn.Module): + def __init__(self, inputs): + super().__init__() + self._requested_inputs = inputs + + def requested_inputs(self) -> Dict[str, ModelOutput]: + return self._requested_inputs + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels] = None, + ) -> Dict[str, TensorMap]: + return { + ("extra::" + input): systems[0].get_data(input) + for input in self._requested_inputs + } + + +def test_additional_input(atoms): + inputs = { + "masses": ModelOutput(quantity="mass", unit="u", per_atom=True), + "velocities": ModelOutput(quantity="velocity", unit="A/fs", per_atom=True), + "charges": ModelOutput(quantity="charge", unit="e", per_atom=True), + "ase::initial_charges": ModelOutput(quantity="charge", unit="e", per_atom=True), + } + outputs = {("extra::" + n): inputs[n] for n in inputs} + capabilities = ModelCapabilities( + outputs=outputs, + atomic_types=[28], + interaction_range=0.0, + supported_devices=["cpu"], + dtype="float64", + ) + + model = AtomisticModel( + AdditionalInputModel(inputs).eval(), ModelMetadata(), capabilities + ) + MaxwellBoltzmannDistribution(atoms, temperature_K=300.0) + atoms.set_initial_charges([0.0] * len(atoms)) + calculator = MetatomicCalculator(model, check_consistency=True) + results = calculator.run_model(atoms, outputs) + for name, tensor in results.items(): + head, name = name.split("::", maxsplit=1) + assert head == "extra" + assert name in inputs + + assert tensor.get_info("quantity") == inputs[name].quantity + values = tensor[0].values.numpy() + + expected = ARRAY_QUANTITIES[name]["getter"](atoms).reshape(values.shape) + if name == "velocities": + expected /= ( + ase.units.Angstrom / ase.units.fs + ) # ase velocity is in (eV/u)^(1/2) and we want A/fs + + assert np.allclose(values, expected) + + +@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) +def test_mixed_pbc(model, device, dtype): + """Test that the calculator works on a mixed-PBC system""" + atoms = ase.build.fcc111("Ni", size=(2, 2, 3), vacuum=10.0) + atoms.set_pbc((True, True, False)) + + model = _set_model_dtype(model, dtype) + + atoms.calc = MetatomicCalculator( + model, + check_consistency=True, + uncertainty_threshold=None, + device=device, + ) + atoms.get_potential_energy() + atoms.get_forces() diff --git a/python/metatomic_ase/tests/neighbors.py b/python/metatomic_ase/tests/neighbors.py new file mode 100644 index 000000000..89bd65686 --- /dev/null +++ b/python/metatomic_ase/tests/neighbors.py @@ -0,0 +1,111 @@ +import glob +import json +import os + +import pytest +import torch +from metatensor.torch import Labels, TensorBlock + +from metatomic.torch import ( + NeighborListOptions, + System, +) +from metatomic_ase._neighbors import ( + _compute_requested_neighbors_nvalchemi, + _compute_requested_neighbors_vesin, +) + +from ._tests_utils import ALL_DEVICE_DTYPE + + +def _read_neighbor_check(path): + with open(path) as fd: + data = json.load(fd) + + dtype = torch.float64 + + positions = torch.tensor(data["system"]["positions"], dtype=dtype).reshape(-1, 3) + system = System( + types=torch.tensor([1] * positions.shape[0], dtype=torch.int32), + positions=positions, + cell=torch.tensor(data["system"]["cell"], dtype=dtype), + pbc=torch.tensor([True, True, True]), + ) + + options = NeighborListOptions( + cutoff=data["options"]["cutoff"], + full_list=data["options"]["full_list"], + # ASE can only compute strict NL + strict=True, + ) + + samples = torch.tensor( + data["expected-neighbors"]["samples"], dtype=torch.int32 + ).reshape(-1, 5) + distances = torch.tensor( + data["expected-neighbors"]["distances"], dtype=dtype + ).reshape(-1, 3, 1) + + neighbors = TensorBlock( + values=distances, + samples=Labels( + [ + "first_atom", + "second_atom", + "cell_shift_a", + "cell_shift_b", + "cell_shift_c", + ], + samples, + ), + components=[Labels.range("xyz", 3)], + properties=Labels.range("distance", 1), + ) + + return system, options, neighbors + + +def _check_same_set_of_neighbors(expected, actual, full_list): + assert expected.samples.names == actual.samples.names + assert len(expected.samples) == len(actual.samples) + + for sample_i, sample in enumerate(expected.samples): + sign = 1.0 + position = actual.samples.position(sample) + + if position is None and not full_list: + # try looking for the inverse pair + sign = -1.0 + position = actual.samples.position( + [sample[1], sample[0], -sample[2], -sample[3], -sample[4]] + ) + + if position is None: + raise AssertionError(f"missing expected neighbors sample: {sample}") + + assert torch.allclose(expected.values[sample_i], sign * actual.values[position]) + + +@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) +@pytest.mark.parametrize( + "neighbor_fn", + [_compute_requested_neighbors_vesin, _compute_requested_neighbors_nvalchemi], +) +def test_neighbor_list_adapter(device, dtype, neighbor_fn): + if neighbor_fn == _compute_requested_neighbors_nvalchemi and device != "cuda": + pytest.skip("nvalchemiops neighbor list is only implemented for CUDA") + + HERE = os.path.realpath(os.path.dirname(__file__)) + test_files = os.path.join( + HERE, "..", "..", "..", "..", "metatensor-torch", "tests", "neighbor-checks" + ) + + for path in glob.glob(os.path.join(test_files, "*.json")): + system, options, expected_neighbors = _read_neighbor_check(path) + system = system.to(device=device, dtype=dtype) + expected_neighbors = expected_neighbors.to(device=device, dtype=dtype) + + neighbor_fn([system], [options], check_consistency=True) + _check_same_set_of_neighbors( + expected_neighbors, system.get_neighbor_list(options), options.full_list + ) diff --git a/python/metatomic_torch/tests/symmetrized_ase_calculator.py b/python/metatomic_ase/tests/symmetrized.py similarity index 93% rename from python/metatomic_torch/tests/symmetrized_ase_calculator.py rename to python/metatomic_ase/tests/symmetrized.py index 3e1e83a96..20650b22b 100644 --- a/python/metatomic_torch/tests/symmetrized_ase_calculator.py +++ b/python/metatomic_ase/tests/symmetrized.py @@ -7,13 +7,23 @@ from ase.build import bulk, molecule from metatensor.torch import Labels, TensorBlock, TensorMap -import metatomic.torch as mta from metatomic.torch import ( + AtomisticModel, + ModelCapabilities, + ModelMetadata, ModelOutput, NeighborListOptions, System, ) -from metatomic.torch.ase_calculator import SymmetrizedCalculator, _get_quadrature +from metatomic_ase import MetatomicCalculator, SymmetrizedCalculator +from metatomic_ase._symmetry import ( + _average_over_group, + _choose_quadrature, + _compute_rotational_average, + _get_group_operations, + _get_quadrature, + _rotate_atoms, +) def _body_axis_from_system(system: System) -> torch.Tensor: @@ -270,7 +280,7 @@ def mock_calculator( p_iso: float = 1.0, tensor_forces: bool = False, tensor_amp: float = 0.5, -) -> mta.ase_calculator.MetatomicCalculator: +) -> MetatomicCalculator: model = MockAnisoModel( a=a, b=b, @@ -281,14 +291,14 @@ def mock_calculator( ) model.eval() - atomistic_model = mta.AtomisticModel( + atomistic_model = AtomisticModel( model, - mta.ModelMetadata("mock_aniso", "Mock anisotropic model for testing"), - mta.ModelCapabilities( + ModelMetadata("mock_aniso", "Mock anisotropic model for testing"), + ModelCapabilities( { - "energy": mta.ModelOutput(per_atom=False), - "non_conservative_forces": mta.ModelOutput(per_atom=True), - "non_conservative_stress": mta.ModelOutput(per_atom=False), + "energy": ModelOutput(per_atom=False), + "non_conservative_forces": ModelOutput(per_atom=True), + "non_conservative_stress": ModelOutput(per_atom=False), }, list(range(1, 102)), 100, @@ -297,14 +307,14 @@ def mock_calculator( "float64", ), ) - return mta.ase_calculator.MetatomicCalculator( + return MetatomicCalculator( atomistic_model, non_conservative=True, do_gradients_with_energy=False, additional_outputs={ - "energy": mta.ModelOutput(per_atom=False), - "non_conservative_forces": mta.ModelOutput(per_atom=True), - "non_conservative_stress": mta.ModelOutput(per_atom=False), + "energy": ModelOutput(per_atom=False), + "non_conservative_forces": ModelOutput(per_atom=True), + "non_conservative_stress": ModelOutput(per_atom=False), }, ) @@ -456,8 +466,6 @@ def test_rotate_atoms_preserves_geometry(tmp_path): """Check that _rotate_atoms applies rotations correctly and preserves distances.""" from scipy.spatial.transform import Rotation - from metatomic.torch.ase_calculator import _rotate_atoms - # Build simple cubic cell with 2 atoms along x atoms = Atoms("H2", positions=[[0, 0, 0], [1, 0, 0]], cell=np.eye(3)) R = Rotation.from_euler("z", 90, degrees=True).as_matrix()[None, ...] # 90° about z @@ -477,8 +485,6 @@ def test_rotate_atoms_preserves_geometry(tmp_path): def test_choose_quadrature_rules(): """Check that _choose_quadrature selects appropriate rules.""" - from metatomic.torch.ase_calculator import _choose_quadrature - for L in [0, 5, 17, 50]: lebedev_order, n_gamma = _choose_quadrature(L) assert lebedev_order >= L @@ -487,8 +493,6 @@ def test_choose_quadrature_rules(): def test_get_quadrature_properties(): """Check properties of the quadrature returned by _get_quadrature.""" - from metatomic.torch.ase_calculator import _get_quadrature - R, w = _get_quadrature(lebedev_order=11, n_rotations=5, include_inversion=False) assert np.isclose(np.sum(w), 1.0) assert np.allclose([np.dot(r.T, r) for r in R], np.eye(3), atol=1e-12) @@ -505,7 +509,6 @@ def test_get_quadrature_properties(): def test_compute_rotational_average_identity(): """Check that _compute_rotational_average produces correct averages.""" - from metatomic.torch.ase_calculator import _compute_rotational_average R = np.repeat(np.eye(3)[None, :, :], 3, axis=0) w = np.ones(3) / 3 @@ -530,10 +533,6 @@ def test_average_over_fcc_group(fcc_bulk: Atoms): Check that averaging over the space group of an FCC crystal produces an isotropic (scalar) stress tensor. """ - from metatomic.torch.ase_calculator import ( - _average_over_group, - _get_group_operations, - ) # FCC conventional cubic cell (4 atoms) atoms = fcc_bulk @@ -568,10 +567,6 @@ def test_space_group_average_non_periodic(): Check that averaging over the space group of a non-periodic system leaves the results unchanged. """ - from metatomic.torch.ase_calculator import ( - _average_over_group, - _get_group_operations, - ) # Methane molecule (Td symmetry) atoms = molecule("CH4") diff --git a/python/metatomic_torch/metatomic/torch/__init__.py b/python/metatomic_torch/metatomic/torch/__init__.py index 30cdce149..abee54a80 100644 --- a/python/metatomic_torch/metatomic/torch/__init__.py +++ b/python/metatomic_torch/metatomic/torch/__init__.py @@ -50,6 +50,7 @@ pick_device = torch.ops.metatomic.pick_device pick_output = torch.ops.metatomic.pick_output +from . import ase_calculator # noqa: F401 from .model import ( # noqa: F401 AtomisticModel, ModelInterface, @@ -63,15 +64,3 @@ save_buffer, ) from .systems_to_torch import systems_to_torch # noqa: F401 - - -def __getattr__(name): - # lazy import for ase_calculator, making it accessible as - # ``metatomic.torch.ase_calculator`` without requiring a separate import from - # ``metatomic.torch``, but only importing the code when actually required. - if name == "ase_calculator": - import metatomic.torch.ase_calculator - - return metatomic.torch.ase_calculator - else: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/metatomic_torch/metatomic/torch/ase_calculator.py b/python/metatomic_torch/metatomic/torch/ase_calculator.py index a4919d654..b31f8230a 100644 --- a/python/metatomic_torch/metatomic/torch/ase_calculator.py +++ b/python/metatomic_torch/metatomic/torch/ase_calculator.py @@ -1,1543 +1,30 @@ -import logging -import os -import pathlib import warnings -from typing import Dict, List, Optional, Tuple, Union -import metatensor.torch as mts -import numpy as np -import torch -import vesin.metatomic -from metatensor.torch import Labels, TensorBlock, TensorMap -from torch.profiler import record_function -from . import ( - AtomisticModel, - ModelEvaluationOptions, - ModelMetadata, - ModelOutput, - NeighborListOptions, - System, - load_atomistic_model, - pick_device, - pick_output, -) +def __getattr__(name: str): + if name == "MetatomicCalculator": + from metatomic_ase import MetatomicCalculator - -import ase # isort: skip -import ase.calculators.calculator # isort: skip -from ase.calculators.calculator import ( # isort: skip - InputError, - PropertyNotImplementedError, - all_properties as ALL_ASE_PROPERTIES, -) - - -try: - from nvalchemiops.neighborlist import neighbor_list as nvalchemi_neighbor_list - - HAS_NVALCHEMIOPS = True -except ImportError: - HAS_NVALCHEMIOPS = False - -FilePath = Union[str, bytes, pathlib.PurePath] - -LOGGER = logging.getLogger(__name__) - - -STR_TO_DTYPE = { - "float32": torch.float32, - "float64": torch.float64, -} - - -def _get_charges(atoms: ase.Atoms) -> np.ndarray: - try: - return atoms.get_charges() - except Exception: - return atoms.get_initial_charges() - - -ARRAY_QUANTITIES = { - "momenta": { - "quantity": "momentum", - "getter": ase.Atoms.get_momenta, - "unit": "(eV*u)^(1/2)", - }, - "masses": { - "quantity": "mass", - "getter": ase.Atoms.get_masses, - "unit": "u", - }, - "velocities": { - "quantity": "velocity", - "getter": ase.Atoms.get_velocities, - "unit": "(eV/u)^(1/2)", - }, - "charges": { - "quantity": "charge", - "getter": _get_charges, - "unit": "e", - }, - "ase::initial_magmoms": { - "quantity": "magnetic_moment", - "getter": ase.Atoms.get_initial_magnetic_moments, - "unit": "", - }, - "ase::magnetic_moment": { - "quantity": "magnetic_moment", - "getter": ase.Atoms.get_magnetic_moment, - "unit": "", - }, - "ase::magnetic_moments": { - "quantity": "magnetic_moment", - "getter": ase.Atoms.get_magnetic_moments, - "unit": "", - }, - "ase::initial_charges": { - "quantity": "charge", - "getter": ase.Atoms.get_initial_charges, - "unit": "e", - }, - "ase::dipole_moment": { - "quantity": "dipole_moment", - "getter": ase.Atoms.get_dipole_moment, - "unit": "", - }, -} - - -class MetatomicCalculator(ase.calculators.calculator.Calculator): - """ - The :py:class:`MetatomicCalculator` class implements ASE's - :py:class:`ase.calculators.calculator.Calculator` API using metatomic - models to compute energy, forces and any other supported property. - - This class can be initialized with any :py:class:`AtomisticModel`, and - used to run simulations using ASE's MD facilities. - - Neighbor lists are computed using the fast - `vesin `_ neighbor list library, - either on CPU or GPU depending on the device of the model. If - `nvalchemiops `_ - is installed, full neighbor lists on GPU will be computed with it instead. - """ - - def __init__( - self, - model: Union[FilePath, AtomisticModel], - *, - additional_outputs: Optional[Dict[str, ModelOutput]] = None, - extensions_directory=None, - check_consistency=False, - device=None, - variants: Optional[Dict[str, Optional[str]]] = None, - non_conservative=False, - do_gradients_with_energy=True, - uncertainty_threshold=0.1, - ): - """ - :param model: model to use for the calculation. This can be a file path, a - Python instance of :py:class:`AtomisticModel`, or the output of - :py:func:`torch.jit.script` on :py:class:`AtomisticModel`. - :param additional_outputs: Dictionary of additional outputs to be computed by - the model. These outputs will always be computed whenever the - :py:meth:`calculate` function is called (e.g. by - :py:meth:`ase.Atoms.get_potential_energy`, - :py:meth:`ase.optimize.optimize.Dynamics.run`, *etc.*) and stored in the - :py:attr:`additional_outputs` attribute. If you want more control over when - and how to compute specific outputs, you should use :py:meth:`run_model` - instead. - :param extensions_directory: if the model uses extensions, we will try to load - them from this directory - :param check_consistency: should we check the model for consistency when - running, defaults to False. - :param device: torch device to use for the calculation. If ``None``, we will try - the options in the model's ``supported_device`` in order. - :param variants: dictionary mapping output names to a variant that should be - used for the calculations (e.g. ``{"energy": "PBE"}``). If ``"energy"`` is - set to a variant also the uncertainty and non conservative outputs will be - taken from this variant. This behaviour can be overriden by setting the - corresponding keys explicitly to ``None`` or to another value (e.g. - ``{"energy_uncertainty": "r2scan"}``). - :param non_conservative: if ``True``, the model will be asked to compute - non-conservative forces and stresses. This can afford a speed-up, - potentially at the expense of physical correctness (especially in molecular - dynamics simulations). - :param do_gradients_with_energy: if ``True``, this calculator will always - compute the energy gradients (forces and stress) when the energy is - requested (e.g. through ``atoms.get_potential_energy()``). Because the - results of a calculation are cached by ASE, this means future calls to - ``atom.get_forces()`` will return immediately, without needing to execute - the model again. If you are mainly interested in the energy, you can set - this to ``False`` and enjoy a faster model. Forces will still be calculated - if requested with ``atoms.get_forces()``. - :param uncertainty_threshold: threshold for the atomic energy uncertainty in eV. - This will only be used if the model supports atomic uncertainty estimation - (https://pubs.acs.org/doi/full/10.1021/acs.jctc.3c00704). Set this to - ``None`` to disable uncertainty quantification even if the model supports - it. - """ - super().__init__() - - self.parameters = { - "extensions_directory": extensions_directory, - "check_consistency": bool(check_consistency), - "variants": variants, - "non_conservative": bool(non_conservative), - "do_gradients_with_energy": bool(do_gradients_with_energy), - "additional_outputs": additional_outputs, - "uncertainty_threshold": uncertainty_threshold, - } - - # Load the model - if isinstance(model, (str, bytes, pathlib.PurePath)): - if not os.path.exists(model): - raise InputError(f"given model path '{model}' does not exist") - - # only store the model in self.parameters if is it the path to a file - self.parameters["model"] = str(model) - - model = load_atomistic_model( - model, extensions_directory=extensions_directory - ) - - elif isinstance(model, torch.jit.RecursiveScriptModule): - if model.original_name != "AtomisticModel": - raise InputError( - "torch model must be 'AtomisticModel', " - f"got '{model.original_name}' instead" - ) - elif isinstance(model, AtomisticModel): - # nothing to do - pass - else: - raise TypeError(f"unknown type for model: {type(model)}") - - self.parameters["device"] = str(device) if device is not None else None - # get the best device according what the model supports and what's available on - # the current machine - capabilities = model.capabilities() - self._device = torch.device( - pick_device(capabilities.supported_devices, self.parameters["device"]) - ) - - if capabilities.dtype in STR_TO_DTYPE: - self._dtype = STR_TO_DTYPE[capabilities.dtype] - else: - raise ValueError( - f"found unexpected dtype in model capabilities: {capabilities.dtype}" - ) - - # resolve the output keys to use based on the requested variants - variants = variants or {} - default_variant = variants.get("energy") - - resolved_variants = { - key: variants.get(key, default_variant) - for key in [ - "energy", - "energy_uncertainty", - "non_conservative_forces", - "non_conservative_stress", - ] - } - - outputs = capabilities.outputs - - # Check if the model has an energy output - has_energy = any( - "energy" == key or key.startswith("energy/") for key in outputs.keys() - ) - if has_energy: - self._energy_key = pick_output( - "energy", outputs, resolved_variants["energy"] - ) - else: - self._energy_key = None - - has_energy_uq = any("energy_uncertainty" in key for key in outputs.keys()) - if has_energy_uq and uncertainty_threshold is not None: - self._energy_uq_key = pick_output( - "energy_uncertainty", outputs, resolved_variants["energy_uncertainty"] - ) - else: - self._energy_uq_key = "energy_uncertainty" - - if non_conservative: - if ( - "non_conservative_stress" in variants - and "non_conservative_forces" in variants - and ( - (variants["non_conservative_stress"] is None) - != (variants["non_conservative_forces"] is None) - ) - ): - raise ValueError( - "if both 'non_conservative_stress' and " - "'non_conservative_forces' are present in `variants`, they " - "must either be both `None` or both not `None`." - ) - - self._nc_forces_key = pick_output( - "non_conservative_forces", - outputs, - resolved_variants["non_conservative_forces"], - ) - self._nc_stress_key = pick_output( - "non_conservative_stress", - outputs, - resolved_variants["non_conservative_stress"], - ) - else: - self._nc_forces_key = "non_conservative_forces" - self._nc_stress_key = "non_conservative_stress" - - if additional_outputs is None: - self._additional_output_requests = {} - else: - assert isinstance(additional_outputs, dict) - for name, output in additional_outputs.items(): - assert isinstance(name, str) - assert isinstance(output, torch.ScriptObject) - assert "explicit_gradients_setter" in output._method_names(), ( - "outputs must be ModelOutput instances" - ) - - self._additional_output_requests = additional_outputs - - self._model = model.to(device=self._device) - - self._calculate_uncertainty = ( - self._energy_uq_key in self._model.capabilities().outputs - # we require per-atom uncertainties to capture local effects - and self._model.capabilities().outputs[self._energy_uq_key].per_atom - and uncertainty_threshold is not None - ) - - if self._calculate_uncertainty: - assert uncertainty_threshold is not None - if uncertainty_threshold <= 0.0: - raise ValueError( - f"`uncertainty_threshold` is {uncertainty_threshold} but must " - "be positive" - ) - - # We do our own check to verify if a property is implemented in `calculate()`, - # so we pretend to be able to compute all properties ASE knows about. - self.implemented_properties = ALL_ASE_PROPERTIES - - self.additional_outputs: Dict[str, TensorMap] = {} - """ - Additional outputs computed by :py:meth:`calculate` are stored in this - dictionary. - - The keys will match the keys of the ``additional_outputs`` parameters to the - constructor; and the values will be the corresponding raw - :py:class:`metatensor.torch.TensorMap` produced by the model. - """ - - def todict(self): - if "model" not in self.parameters: - raise RuntimeError( - "can not save metatensor model in ASE `todict`, please initialize " - "`MetatomicCalculator` with a path to a saved model file if you need " - "to use `todict`" - ) - - return self.parameters - - @classmethod - def fromdict(cls, data): - return MetatomicCalculator(**data) - - def metadata(self) -> ModelMetadata: - """Get the metadata of the underlying model""" - return self._model.metadata() - - def run_model( - self, - atoms: Union[ase.Atoms, List[ase.Atoms]], - outputs: Dict[str, ModelOutput], - selected_atoms: Optional[Labels] = None, - ) -> Dict[str, TensorMap]: - """ - Run the model on the given ``atoms``, computing the requested ``outputs`` and - only these. - - The output of the model is returned directly, and as such the blocks' ``values`` - will be :py:class:`torch.Tensor`. - - This is intended as an easy way to run metatensor models on - :py:class:`ase.Atoms` when the model can compute outputs not supported by the - standard ASE's calculator interface. - - All the parameters have the same meaning as the corresponding ones in - :py:meth:`metatomic.torch.ModelInterface.forward`. - - :param atoms: :py:class:`ase.Atoms`, or list of :py:class:`ase.Atoms`, on which - to run the model - :param outputs: outputs of the model that should be predicted - :param selected_atoms: subset of atoms on which to run the calculation - """ - if isinstance(atoms, ase.Atoms): - atoms_list = [atoms] - else: - atoms_list = atoms - - systems = [] - for atoms in atoms_list: - types, positions, cell, pbc = _ase_to_torch_data( - atoms=atoms, dtype=self._dtype, device=self._device - ) - system = System(types, positions, cell, pbc) - # Get the additional inputs requested by the model - for name, option in self._model.requested_inputs().items(): - input_tensormap = _get_ase_input( - atoms, name, option, dtype=self._dtype, device=self._device - ) - system.add_data(name, input_tensormap) - systems.append(system) - - # Compute the neighbors lists requested by the model - input_systems = _compute_requested_neighbors( - systems=systems, - requested_options=self._model.requested_neighbor_lists(), - check_consistency=self.parameters["check_consistency"], - ) - - available_outputs = self._model.capabilities().outputs - for key in outputs: - if key not in available_outputs: - raise ValueError(f"this model does not support '{key}' output") - - options = ModelEvaluationOptions( - length_unit="angstrom", - outputs=outputs, - selected_atoms=selected_atoms, - ) - return self._model( - systems=input_systems, - options=options, - check_consistency=self.parameters["check_consistency"], - ) - - def calculate( - self, - atoms: ase.Atoms, - properties: List[str], - system_changes: List[str], - ) -> None: - """ - Compute some ``properties`` with this calculator, and return them in the format - expected by ASE. - - This is not intended to be called directly by users, but to be an implementation - detail of ``atoms.get_energy()`` and related functions. See - :py:meth:`ase.calculators.calculator.Calculator.calculate` for more information. - """ - super().calculate( - atoms=atoms, - properties=properties, - system_changes=system_changes, - ) - - # In the next few lines, we decide which properties to calculate among energy, - # forces and stress. In addition to the requested properties, we calculate the - # energy if any of the three is requested, as it is an intermediate step in the - # calculation of the other two. We also calculate the forces if the stress is - # requested, and vice-versa. The overhead for the latter operation is also - # small, assuming that the majority of the model computes forces and stresses - # by backward propagation as opposed to forward-mode differentiation. - calculate_energy = ( - "energy" in properties - or "energies" in properties - or "forces" in properties - or "stress" in properties - ) - - # Check if energy-related properties are requested but the model doesn't - # support energy - if calculate_energy and self._energy_key is None: - raise PropertyNotImplementedError( - "This calculator does not support energy-related properties " - "(energy, energies, forces, stress) because the underlying model " - "does not have an energy output" - ) - calculate_energies = "energies" in properties - calculate_forces = "forces" in properties or "stress" in properties - calculate_stress = "stress" in properties - if calculate_stress and not atoms.pbc.all(): - warnings.warn( - "stress requested but likely to be wrong, since the system is not " - "periodic in all directions", - stacklevel=2, - ) - if "forces" in properties and atoms.pbc.all(): - # we have PBCs, and, since the user/integrator requested forces, we will run - # backward anyway, so let's do the stress as well for free (this saves - # another forward-backward call later if the stress is requested) - calculate_stress = True - if "stresses" in properties: - raise NotImplementedError("'stresses' are not implemented yet") - - if self.parameters["do_gradients_with_energy"]: - if calculate_energies or calculate_energy: - calculate_forces = True - calculate_stress = True - - with record_function("MetatomicCalculator::prepare_inputs"): - outputs = self._ase_properties_to_metatensor_outputs( - properties, - calculate_forces=calculate_forces, - calculate_stress=calculate_stress, - calculate_stresses=False, - ) - outputs.update(self._additional_output_requests) - if calculate_energy and self._calculate_uncertainty: - outputs[self._energy_uq_key] = ModelOutput( - quantity="energy", - unit="eV", - per_atom=True, - explicit_gradients=[], - ) - - capabilities = self._model.capabilities() - for name in outputs.keys(): - if name not in capabilities.outputs: - raise ValueError( - f"you asked for the calculation of {name}, but this model " - "does not support it" - ) - - types, positions, cell, pbc = _ase_to_torch_data( - atoms=atoms, dtype=self._dtype, device=self._device - ) - - do_backward = False - if calculate_forces and not self.parameters["non_conservative"]: - do_backward = True - positions.requires_grad_(True) - - if calculate_stress and not self.parameters["non_conservative"]: - do_backward = True - - strain = torch.eye( - 3, requires_grad=True, device=self._device, dtype=self._dtype - ) - - positions = positions @ strain - positions.retain_grad() - - cell = cell @ strain - - run_options = ModelEvaluationOptions( - length_unit="angstrom", - outputs=outputs, - selected_atoms=None, - ) - - with record_function("MetatomicCalculator::compute_neighbors"): - # convert from ase.Atoms to metatomic.torch.System - system = System(types, positions, cell, pbc) - input_system = _compute_requested_neighbors( - systems=[system], - requested_options=self._model.requested_neighbor_lists(), - check_consistency=self.parameters["check_consistency"], - )[0] - - with record_function("MetatomicCalculator::get_model_inputs"): - for name, option in self._model.requested_inputs().items(): - input_tensormap = _get_ase_input( - atoms, name, option, dtype=self._dtype, device=self._device - ) - input_system.add_data(name, input_tensormap) - - # no `record_function` here, this will be handled by AtomisticModel - outputs = self._model( - [input_system], - run_options, - check_consistency=self.parameters["check_consistency"], - ) - energy = outputs[self._energy_key] - - with record_function("MetatomicCalculator::sum_energies"): - if run_options.outputs[self._energy_key].per_atom: - assert len(energy) == 1 - assert energy.sample_names == ["system", "atom"] - assert torch.all(energy.block().samples["system"] == 0) - energies = energy - assert energies.block().values.shape == (len(atoms), 1) - - energy = mts.sum_over_samples(energy, sample_names=["atom"]) - - assert len(energy.block().gradients_list()) == 0 - assert energy.block().values.shape == (1, 1) - - with record_function("ASECalculator::uncertainty_warning"): - if calculate_energy and self._calculate_uncertainty: - uncertainty = outputs[self._energy_uq_key].block().values - assert uncertainty.shape == (len(atoms), 1) - uncertainty = uncertainty.detach().cpu().numpy() - - threshold = self.parameters["uncertainty_threshold"] - if np.any(uncertainty > threshold): - warnings.warn( - "Some of the atomic energy uncertainties are larger than the " - f"threshold of {threshold} eV. The prediction is above the " - f"threshold for atoms {np.where(uncertainty > threshold)[0]}.", - stacklevel=2, - ) - - if do_backward: - if energy.block().values.grad_fn is None: - # did the user actually request a gradient, or are we trying to - # compute one just for efficiency? - if "forces" in properties or "stress" in properties: - # the user asked for it, let it fail below - pass - else: - # we added the calculation, let's remove it - do_backward = False - calculate_forces = False - calculate_stress = False - - with record_function("MetatomicCalculator::run_backward"): - if do_backward: - energy.block().values.backward() - - with record_function("MetatomicCalculator::convert_outputs"): - self.results = {} - - if calculate_energies: - energies_values = energies.block().values.detach().reshape(-1) - atom_indexes = energies.block().samples.column("atom") - result = torch.zeros_like(energies_values) - result.index_add_(0, atom_indexes, energies_values) - self.results["energies"] = result.cpu().double().numpy() - - if calculate_energy: - energy_values = energy.block().values.detach() - energy_values = energy_values.cpu().double() - self.results["energy"] = energy_values.numpy()[0, 0] - - if calculate_forces: - if self.parameters["non_conservative"]: - forces_values = outputs[self._nc_forces_key].block().values.detach() - # remove any spurious net force - forces_values = forces_values - forces_values.mean( - dim=0, keepdim=True - ) - else: - forces_values = -system.positions.grad - forces_values = forces_values.reshape(-1, 3) - forces_values = forces_values.cpu().double() - self.results["forces"] = forces_values.numpy() - - if calculate_stress: - if self.parameters["non_conservative"]: - stress_values = outputs[self._nc_stress_key].block().values.detach() - else: - stress_values = strain.grad / atoms.cell.volume - stress_values = stress_values.reshape(3, 3) - stress_values = stress_values.cpu().double() - self.results["stress"] = _full_3x3_to_voigt_6_stress( - stress_values.numpy() - ) - - self.additional_outputs = {} - for name in self._additional_output_requests: - self.additional_outputs[name] = outputs[name] - - def compute_energy( - self, - atoms: Union[ase.Atoms, List[ase.Atoms]], - compute_forces_and_stresses: bool = False, - *, - per_atom: bool = False, - ) -> Dict[str, Union[Union[float, np.ndarray], List[Union[float, np.ndarray]]]]: - """ - Compute the energy of the given ``atoms``. - - Energies are computed in eV, forces in eV/Å, and stresses in 3x3 tensor format - and in units of eV/Å^3. - - :param atoms: :py:class:`ase.Atoms`, or list of :py:class:`ase.Atoms`, on which - to run the model - :param compute_forces_and_stresses: if ``True``, the model will also compute - forces and stresses. IMPORTANT: stresses will only be computed if all - provided systems have periodic boundary conditions in all directions. - :param per_atom: if ``True``, the per-atom energies will also be - computed. - :return: A dictionary with the computed properties. The dictionary will contain - the ``energy`` as a float. If ``compute_forces_and_stresses`` is True, - the ``forces`` and ``stress`` will also be included as numpy arrays. - If ``per_atom`` is True, the ``energies`` key will also be present, - containing the per-atom energies as a numpy array. - In case of a list of :py:class:`ase.Atoms`, the dictionary values will - instead be lists of the corresponding properties, in the same format. - """ - if self._energy_key is None: - raise ValueError( - "This calculator does not support energy computation because " - "the underlying model does not have an energy output" - ) - - if isinstance(atoms, ase.Atoms): - atoms_list = [atoms] - was_single = True - else: - atoms_list = atoms - was_single = False - - properties = ["energy"] - energy_per_atom = False - if per_atom: - energy_per_atom = True - properties.append("energies") - - outputs = self._ase_properties_to_metatensor_outputs( - properties=properties, - calculate_forces=compute_forces_and_stresses, - calculate_stress=compute_forces_and_stresses, - calculate_stresses=False, - ) - - systems = [] - if compute_forces_and_stresses: - strains = [] - for atoms in atoms_list: - types, positions, cell, pbc = _ase_to_torch_data( - atoms=atoms, dtype=self._dtype, device=self._device - ) - if compute_forces_and_stresses and not self.parameters["non_conservative"]: - positions.requires_grad_(True) - strain = torch.eye( - 3, requires_grad=True, device=self._device, dtype=self._dtype - ) - positions = positions @ strain - positions.retain_grad() - cell = cell @ strain - strains.append(strain) - system = System(types, positions, cell, pbc) - systems.append(system) - - # Compute the neighbors lists requested by the model - input_systems = _compute_requested_neighbors( - systems=systems, - requested_options=self._model.requested_neighbor_lists(), - check_consistency=self.parameters["check_consistency"], - ) - - predictions = self._model( - systems=input_systems, - options=ModelEvaluationOptions(length_unit="angstrom", outputs=outputs), - check_consistency=self.parameters["check_consistency"], - ) - energies = predictions[self._energy_key] - - if energy_per_atom: - # Get per-atom energies - sorted_block = mts.sort_block(energies.block(), axes="samples") - energies_values = sorted_block.values.detach().reshape(-1) - - split_sizes = [len(system) for system in systems] - atom_indices = sorted_block.samples.column("atom") - energies_values = torch.split(energies_values, split_sizes, dim=0) - split_atom_indices = torch.split(atom_indices, split_sizes, dim=0) - split_energies = [] - for atom_indices, values in zip( - split_atom_indices, energies_values, strict=True - ): - split_energy = torch.zeros( - len(atom_indices), dtype=values.dtype, device=values.device - ) - split_energy.index_add_(0, atom_indices, values) - split_energies.append(split_energy) - - total_energy = ( - mts.sum_over_samples(energies, ["atom"]) - .block() - .values.detach() - .cpu() - .double() - .numpy() - .flatten() - .tolist() - ) - results_as_numpy_arrays = { - "energy": total_energy, - "energies": [e.cpu().double().numpy() for e in split_energies], - } - else: - results_as_numpy_arrays = { - "energy": energies.block() - .values.squeeze(-1) - .detach() - .cpu() - .double() - .numpy(), - } - - if compute_forces_and_stresses: - if self.parameters["non_conservative"]: - results_as_numpy_arrays["forces"] = ( - predictions[self._nc_forces_key] - .block() - .values.squeeze(-1) - .detach() - .cpu() - .numpy() - ) - # all the forces are concatenated in a single array, so we need to - # split them into the original systems - split_sizes = [len(system) for system in systems] - split_indices = np.cumsum(split_sizes[:-1]) - results_as_numpy_arrays["forces"] = np.split( - results_as_numpy_arrays["forces"], split_indices, axis=0 - ) - - # remove net forces - results_as_numpy_arrays["forces"] = [ - f - f.mean(axis=0, keepdims=True) - for f in results_as_numpy_arrays["forces"] - ] - - if all(atoms.pbc.all() for atoms in atoms_list): - results_as_numpy_arrays["stress"] = [ - s - for s in predictions[self._nc_stress_key] - .block() - .values.squeeze(-1) - .detach() - .cpu() - .numpy() - ] - else: - energy_tensor = energies.block().values - energy_tensor.backward(torch.ones_like(energy_tensor)) - results_as_numpy_arrays["forces"] = [ - -system.positions.grad.cpu().numpy() for system in systems - ] - if all(atoms.pbc.all() for atoms in atoms_list): - results_as_numpy_arrays["stress"] = [ - strain.grad.cpu().numpy() / atoms.cell.volume - for strain, atoms in zip(strains, atoms_list, strict=False) - ] - if was_single: - for key, value in results_as_numpy_arrays.items(): - results_as_numpy_arrays[key] = value[0] - return results_as_numpy_arrays - - def _ase_properties_to_metatensor_outputs( - self, - properties, - *, - calculate_forces, - calculate_stress, - calculate_stresses, - ): - energy_properties = [] - for p in properties: - if p in ["energy", "energies", "forces", "stress", "stresses"]: - energy_properties.append(p) - else: - raise PropertyNotImplementedError( - f"property '{p}' it not yet supported by this calculator, " - "even if it might be supported by the model" - ) - - metatensor_outputs = {} - - # Only add energy output if the model supports it - if self._energy_key is not None: - output = ModelOutput( - quantity="energy", - unit="ev", - explicit_gradients=[], - ) - - if "energies" in properties or "stresses" in properties: - output.per_atom = True - else: - output.per_atom = False - - metatensor_outputs[self._energy_key] = output - if calculate_forces and self.parameters["non_conservative"]: - metatensor_outputs[self._nc_forces_key] = ModelOutput( - quantity="force", - unit="eV/Angstrom", - per_atom=True, - ) - - if calculate_stress and self.parameters["non_conservative"]: - metatensor_outputs[self._nc_stress_key] = ModelOutput( - quantity="pressure", - unit="eV/Angstrom^3", - per_atom=False, - ) - - if calculate_stresses and self.parameters["non_conservative"]: - raise NotImplementedError( - "non conservative, per-atom stress is not yet implemented" - ) - - available_outputs = self._model.capabilities().outputs - for key in metatensor_outputs: - if key not in available_outputs: - raise ValueError(f"this model does not support '{key}' output") - - return metatensor_outputs - - -def _get_ase_input( - atoms: ase.Atoms, - name: str, - option: ModelOutput, - dtype: torch.dtype, - device: torch.device, -) -> "TensorMap": - if name not in ARRAY_QUANTITIES: - raise ValueError( - f"The model requested '{name}', which is not available in `ase`." + warnings.warn( + "Importing MetatomicCalculator from metatomic.torch.ase_calculator is " + "deprecated and will be removed in a future release. Please import " + "from metatomic_ase instead.", + DeprecationWarning, + stacklevel=2, ) - infos = ARRAY_QUANTITIES[name] + return MetatomicCalculator + elif name == "SymmetrizedCalculator": + from metatomic_ase import SymmetrizedCalculator - values = infos["getter"](atoms) - if values.shape[0] != len(atoms): - raise NotImplementedError( - f"The model requested the '{name}' input, " - f"but the data is not per-atom (shape {values.shape}). " + warnings.warn( + "Importing SymmetrizedCalculator from metatomic.torch.ase_calculator is " + "deprecated and will be removed in a future release. Please import " + "from metatomic_ase instead.", + DeprecationWarning, + stacklevel=2, ) - # Shape: (n_atoms, n_components) -> (n_atoms, n_components, /* n_properties */ 1) - # for metatensor - values = torch.tensor(values[..., None]) - - components = [] - if values.shape[1] != 1: - components.append(Labels(["xyz"], torch.arange(values.shape[1]).reshape(-1, 1))) - - block = TensorBlock( - values, - samples=Labels( - ["system", "atom"], - torch.vstack( - [torch.full((values.shape[0],), 0), torch.arange(values.shape[0])] - ).T, - ), - components=components, - properties=Labels([infos["quantity"]], torch.tensor([[0]])), - ) - - tensor = TensorMap(Labels(["_"], torch.tensor([[0]])), [block]) - - tensor.set_info("quantity", infos["quantity"]) - tensor.set_info("unit", infos["unit"]) - - tensor = tensor.to(dtype=dtype, device=device) - return tensor - - -def _ase_to_torch_data(atoms, dtype, device): - """Get the positions, cell and pbc from ASE atoms as torch tensors""" - - types = torch.from_numpy(atoms.numbers).to(dtype=torch.int32, device=device) - positions = torch.from_numpy(atoms.positions).to(dtype=dtype).to(device=device) - cell = torch.zeros((3, 3), dtype=dtype, device=device) - pbc = torch.tensor(atoms.pbc, dtype=torch.bool, device=device) - - cell[pbc] = torch.tensor(atoms.cell[atoms.pbc], dtype=dtype, device=device) - - return types, positions, cell, pbc - - -def _full_3x3_to_voigt_6_stress(stress): - """ - Re-implementation of ``ase.stress.full_3x3_to_voigt_6_stress`` which does not do the - stress symmetrization correctly (they do ``(stress[1, 2] + stress[1, 2]) / 2.0``) - """ - return np.array( - [ - stress[0, 0], - stress[1, 1], - stress[2, 2], - (stress[1, 2] + stress[2, 1]) / 2.0, - (stress[0, 2] + stress[2, 0]) / 2.0, - (stress[0, 1] + stress[1, 0]) / 2.0, - ] - ) - - -class SymmetrizedCalculator(ase.calculators.calculator.Calculator): - r""" - Take a MetatomicCalculator and average its predictions to make it (approximately) - equivariant. Only predictions for energy, forces and stress are supported. - - The default is to average over a quadrature of the orthogonal group O(3) composed - this way: - - - Lebedev quadrature of the unit sphere (S^2) - - Equispaced sampling of the unit circle (S^1) - - Both proper and improper rotations are taken into account by including the - inversion operation (if ``include_inversion=True``) - - :param base_calculator: the MetatomicCalculator to be symmetrized - :param l_max: the maximum spherical harmonic degree that the model is expected to - be able to represent. This is used to choose the quadrature order. If ``0``, - no rotational averaging will be performed (it can be useful to average only over - the space group, see ``apply_group_symmetry``). - :param batch_size: number of rotated systems to evaluate at once. If ``None``, all - systems will be evaluated at once (this can lead to high memory usage). - :param include_inversion: if ``True``, the inversion operation will be included in - the averaging. This is required to average over the full orthogonal group O(3). - :param apply_space_group_symmetry: if ``True``, the results will be averaged over - discrete space group of rotations for the input system. The group operations are - computed with `spglib `_, and the average is - performed after the O(3) averaging (if any). This has no effect for non-periodic - systems. - :param store_rotational_std: if ``True``, the results will contain the standard - deviation over the different rotations for each property (e.g., ``energy_std``). - """ - - implemented_properties = ["energy", "energies", "forces", "stress", "stresses"] - - def __init__( - self, - base_calculator: MetatomicCalculator, - *, - l_max: int = 3, - batch_size: Optional[int] = None, - include_inversion: bool = True, - apply_space_group_symmetry: bool = False, - store_rotational_std: bool = False, - ) -> None: - try: - from scipy.integrate import lebedev_rule # noqa: F401 - except ImportError as e: - raise ImportError( - "scipy is required to use the `SymmetrizedCalculator`, please install " - "it with `pip install scipy` or `conda install scipy`" - ) from e - - super().__init__() - - self.base_calculator = base_calculator - if l_max > 131: - raise ValueError( - f"l_max={l_max} is too large, the maximum supported value is 131" - ) - self.l_max = l_max - self.include_inversion = include_inversion - - if l_max > 0: - lebedev_order, n_inplane_rotations = _choose_quadrature(l_max) - self.quadrature_rotations, self.quadrature_weights = _get_quadrature( - lebedev_order, n_inplane_rotations, include_inversion - ) - else: # no quadrature - if include_inversion: # identity and inversion - self.quadrature_rotations = np.array([np.eye(3), -np.eye(3)]) - self.quadrature_weights = np.array([0.5, 0.5]) - else: # only the identity - self.quadrature_rotations = np.array([np.eye(3)]) - self.quadrature_weights = np.array([1.0]) - - self.batch_size = ( - batch_size if batch_size is not None else len(self.quadrature_rotations) - ) - - self.store_rotational_std = store_rotational_std - self.apply_space_group_symmetry = apply_space_group_symmetry - - def calculate( - self, atoms: ase.Atoms, properties: List[str], system_changes: List[str] - ) -> None: - """ - Perform the calculation for the given atoms and properties. - - :param atoms: the :py:class:`ase.Atoms` on which to perform the calculation - :param properties: list of properties to compute, among ``energy``, ``forces``, - and ``stress`` - :param system_changes: list of changes to the system since the last call to - ``calculate`` - """ - super().calculate(atoms, properties, system_changes) - self.base_calculator.calculate(atoms, properties, system_changes) - - compute_forces_and_stresses = "forces" in properties or "stress" in properties - per_atom = "energies" in properties - - if len(self.quadrature_rotations) > 0: - rotated_atoms_list = _rotate_atoms(atoms, self.quadrature_rotations) - batches = [ - rotated_atoms_list[i : i + self.batch_size] - for i in range(0, len(rotated_atoms_list), self.batch_size) - ] - results: Dict[str, np.ndarray] = {} - for batch in batches: - try: - batch_results = self.base_calculator.compute_energy( - batch, - compute_forces_and_stresses, - per_atom=per_atom, - ) - for key, value in batch_results.items(): - results.setdefault(key, []) - results[key].extend( - [value] if isinstance(value, float) else value - ) - except torch.cuda.OutOfMemoryError as e: - raise RuntimeError( - "Out of memory error encountered during rotational averaging. " - "Please reduce the batch size or use lower rotational " - "averaging parameters. This can be done by setting the " - "`batch_size` and `l_max` parameters while initializing the " - "calculator." - ) from e - - self.results.update( - _compute_rotational_average( - results, - self.quadrature_rotations, - self.quadrature_weights, - self.store_rotational_std, - ) - ) - - if self.apply_space_group_symmetry: - # Apply the discrete space group of the system a posteriori - Q_list, P_list = _get_group_operations(atoms) - self.results.update(_average_over_group(self.results, Q_list, P_list)) - - -def _choose_quadrature(L_max: int) -> Tuple[int, int]: - """ - Choose a Lebedev quadrature order and number of in-plane rotations to integrate - spherical harmonics up to degree ``L_max``. - - :param L_max: maximum spherical harmonic degree - :return: (lebedev_order, n_inplane_rotations) - """ - # fmt: off - available = [ - 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31, 35, 41, - 47, 53, 59, 65, 71, 77, 83, 89, 95, 101, 107, 113, 119, 125, 131, - ] - # fmt: on - - # pick smallest order >= L_max - n = min(o for o in available if o >= L_max) - # minimal gamma count - K = 2 * L_max + 1 - return n, K - - -def _rotate_atoms(atoms: ase.Atoms, rotations: List[np.ndarray]) -> List[ase.Atoms]: - """ - Create a list of copies of ``atoms``, rotated by each of the given ``rotations``. - - :param atoms: the :py:class:`ase.Atoms` to be rotated - :param rotations: (N, 3, 3) array of orthogonal matrices - :return: list of N :py:class:`ase.Atoms`, each rotated by the corresponding matrix - """ - rotated_atoms_list = [] - has_cell = atoms.cell is not None and atoms.cell.rank > 0 - for rot in rotations: - new_atoms = atoms.copy() - new_atoms.positions = new_atoms.positions @ rot.T - if has_cell: - new_atoms.set_cell( - new_atoms.cell.array @ rot.T, scale_atoms=False, apply_constraint=False - ) - new_atoms.wrap() - rotated_atoms_list.append(new_atoms) - return rotated_atoms_list - - -def _get_quadrature(lebedev_order: int, n_rotations: int, include_inversion: bool): - """ - Lebedev(S^2) x uniform angle quadrature on SO(3). - If include_inversion=True, extend to O(3) by adding inversion * R. - - :param lebedev_order: order of the Lebedev quadrature on the unit sphere - :param n_rotations: number of in-plane rotations per Lebedev node - :param include_inversion: if ``True``, include the inversion operation in the - quadrature - :return: (N, 3, 3) array of orthogonal matrices, and (N,) array of weights - associated to each matrix - """ - from scipy.integrate import lebedev_rule - - # Lebedev nodes (X: (3, M)) - X, w = lebedev_rule(lebedev_order) # w sums to 4*pi - x, y, z = X - alpha = np.arctan2(y, x) # (M,) - beta = np.arccos(z) # (M,) - # beta = np.arccos(np.clip(z, -1.0, 1.0)) # (M,) - - K = int(n_rotations) - gamma = np.linspace(0.0, 2 * np.pi, K, endpoint=False) # (K,) - - Rot = _rotations_from_angles(alpha, beta, gamma) - R_so3 = Rot.as_matrix() # (N, 3, 3) - - # SO(3) Haar–probability weights: w_i/(4*pi*K), repeated over gamma - w_so3 = np.repeat(w / (4 * np.pi * K), repeats=gamma.size) # (N,) - - if not include_inversion: - return R_so3, w_so3 - - # Extend to O(3) by appending inversion * R - P = -np.eye(3) - R_o3 = np.concatenate([R_so3, P @ R_so3], axis=0) # (2N, 3, 3) - w_o3 = np.concatenate([0.5 * w_so3, 0.5 * w_so3], axis=0) - - return R_o3, w_o3 - - -def _rotations_from_angles(alpha, beta, gamma): - from scipy.spatial.transform import Rotation - # Build all combinations (alpha_i, beta_i, gamma_j) - A = np.repeat(alpha, gamma.size).reshape(-1, 1) # (N, 1) - B = np.repeat(beta, gamma.size).reshape(-1, 1) # (N, 1) - G = np.tile(gamma, alpha.size).reshape(-1, 1) # (N, 1) - - # Compose ZYZ rotations in SO(3) - Rot = ( - Rotation.from_euler("z", A) - * Rotation.from_euler("y", B) - * Rotation.from_euler("z", G) - ) - - return Rot - - -def _compute_rotational_average(results, rotations, weights, store_std): - R = rotations - B = R.shape[0] - w = weights - w = w / w.sum() - - def _wreshape(x): - return w.reshape((B,) + (1,) * (x.ndim - 1)) - - def _wmean(x): - return np.sum(_wreshape(x) * x, axis=0) - - def _wstd(x): - mu = _wmean(x) - return np.sqrt(np.sum(_wreshape(x) * (x - mu) ** 2, axis=0)) - - out = {} - - # Energy (B,) - if "energy" in results: - E = np.asarray(results["energy"], dtype=float) # (B,) - out["energy"] = _wmean(E) # () - if store_std: - out["energy_rot_std"] = _wstd(E) # () - - if "energies" in results: - E = np.asarray(results["energies"], dtype=float) # (B,N) - out["energies"] = _wmean(E) # (N,) - if store_std: - out["energies_rot_std"] = _wstd(E) # (N,) - - # Forces (B,N,3) from rotated structures: back-rotate with F' R - if "forces" in results: - F = np.asarray(results["forces"], dtype=float) # (B,N,3) - F_back = F @ R # F' R - out["forces"] = _wmean(F_back) # (N,3) - if store_std: - out["forces_rot_std"] = _wstd(F_back) # (N,3) - - # Stress (B,3,3) from rotated structures: back-rotate with R^T S' R - if "stress" in results: - S = np.asarray(results["stress"], dtype=float) # (B,3,3) - RT = np.swapaxes(R, 1, 2) - S_back = RT @ S @ R # R^T S' R - out["stress"] = _wmean(S_back) # (3,3) - if store_std: - out["stress_rot_std"] = _wstd(S_back) # (3,3) - - if "stresses" in results: - S = np.asarray(results["stresses"], dtype=float) # (B,N,3,3) - RT = np.swapaxes(R, 1, 2) - S_back = RT[:, None, :, :] @ S @ R[:, None, :, :] # R^T S' R - out["stresses"] = _wmean(S_back) # (N,3,3) - if store_std: - out["stresses_rot_std"] = _wstd(S_back) # (N,3,3) - - return out - - -def _get_group_operations( - atoms: ase.Atoms, symprec: float = 1e-6, angle_tolerance: float = -1.0 -) -> Tuple[List[np.ndarray], List[np.ndarray]]: - """ - Extract point-group rotations Q_g (Cartesian, 3x3) and the corresponding - atom-index permutations P_g (N x N) induced by the space-group operations. - Returns Q_list, Cartesian rotation matrices of the point group, - and P_list, permutation matrices mapping original indexing -> indexing after (R,t), - - :param atoms: input structure - :param symprec: tolerance for symmetry finding - :param angle_tolerance: tolerance for symmetry finding (in degrees). If less than 0, - a value depending on ``symprec`` will be chosen automatically by spglib. - :return: List of rotation matrices and permutation matrices. - - """ - try: - import spglib - except ImportError as e: - raise ImportError( - "spglib is required to use the SymmetrizedCalculator with " - "`apply_group_symmetry=True`. Please install it with " - "`pip install spglib` or `conda install -c conda-forge spglib`" - ) from e - - if not (atoms.pbc.all()): - # No periodic boundary conditions: no symmetry - return [], [] - - # Lattice with column vectors a1,a2,a3 (spglib expects (cell, frac, Z)) - A = atoms.cell.array.T # (3,3) - frac = atoms.get_scaled_positions() # (N,3) in [0,1) - numbers = atoms.numbers - N = len(atoms) - - data = spglib.get_symmetry_dataset( - (atoms.cell.array, frac, numbers), - symprec=symprec, - angle_tolerance=angle_tolerance, - _throw=True, - ) - - if data is None: - # No symmetry found - return [], [] - R_frac = data.rotations # (n_ops, 3,3), integer - t_frac = data.translations # (n_ops, 3) - Z = numbers - - # Match fractional coords modulo 1 within a tolerance, respecting chemical species - def _match_index(x_new, frac_ref, Z_ref, Z_i, tol=1e-6): - d = np.abs(frac_ref - x_new) # (N,3) - d = np.minimum(d, 1.0 - d) # periodic distance - # Mask by identical species - mask = Z_ref == Z_i - if not np.any(mask): - raise RuntimeError("No matching species found while building permutation.") - # Choose argmin over max-norm within species - idx = np.where(mask)[0] - j = idx[np.argmin(np.max(d[idx], axis=1))] - - # Sanity check - if np.max(d[j]) > tol: - raise RuntimeError( - ( - f"Sanity check failed in _match_index: max distance {np.max(d[j])} " - f"exceeds tolerance {tol}." - ) - ) - return j - - Q_list, P_list = [], [] - seen = set() - Ainv = np.linalg.inv(A) - - for Rf, tf in zip(R_frac, t_frac, strict=False): - # Cartesian rotation: Q = A Rf A^{-1} - Q = A @ Rf @ Ainv - # Deduplicate rotations (point group) by rounding - key = tuple(np.round(Q.flatten(), 12)) - if key in seen: - continue - seen.add(key) - - # Build the permutation P from i to j - P = np.zeros((N, N), dtype=int) - new_frac = (frac @ Rf.T + tf) % 1.0 # images after (Rf,tf) - for i in range(N): - j = _match_index(new_frac[i], frac, Z, Z[i]) - P[j, i] = 1 # column i maps to row j - - Q_list.append(Q.astype(float)) - P_list.append(P) - - return Q_list, P_list - - -def _average_over_group( - results: dict, Q_list: List[np.ndarray], P_list: List[np.ndarray] -) -> dict: - """ - Apply the point-group projector in output space. - - :param results: Must contain 'energy' (scalar), and/or 'forces' (N,3), and/or - 'stress' (3,3). These are predictions for the current structure in the reference - frame. - :param Q_list: Rotation matrices of the point group, from - :py:func:`_get_group_operations` - :param P_list: Permutation matrices of the point group, from - :py:func:`_get_group_operations` - :return out: Projected quantities. - """ - m = len(Q_list) - if m == 0: - return results # nothing to do - - out = {} - # Energy: unchanged by the projector (scalar) - if "energy" in results: - out["energy"] = float(results["energy"]) - - # Forces: (N,3) row-vectors; projector: (1/|G|) \sum_g P_g^T F Q_g - if "forces" in results: - F = np.asarray(results["forces"], float) - if F.ndim != 2 or F.shape[1] != 3: - raise ValueError(f"'forces' must be (N,3), got {F.shape}") - acc = np.zeros_like(F) - for Q, P in zip(Q_list, P_list, strict=False): - acc += P.T @ (F @ Q) - out["forces"] = acc / m - - # Stress: (3,3); projector: (1/|G|) \sum_g Q_g^T S Q_g - if "stress" in results: - S = np.asarray(results["stress"], float) - if S.shape != (3, 3): - raise ValueError(f"'stress' must be (3,3), got {S.shape}") - # S = 0.5 * (S + S.T) # symmetrize just in case - acc = np.zeros_like(S) - for Q in Q_list: - acc += Q.T @ S @ Q - S_pg = acc / m - out["stress"] = S_pg - - return out - - -def _compute_requested_neighbors( - systems: List[System], - requested_options: List[NeighborListOptions], - check_consistency=False, -) -> List[System]: - """ - Compute all neighbor lists requested by ``model`` and store them inside the systems. - """ - can_use_nvalchemi = HAS_NVALCHEMIOPS and all( - system.device.type == "cuda" for system in systems - ) - - if can_use_nvalchemi: - full_nl_options = [] - half_nl_options = [] - for options in requested_options: - if options.full_list: - full_nl_options.append(options) - else: - half_nl_options.append(options) - - # Do the full neighbor lists with nvalchemi, and the rest with vesin - systems = _compute_requested_neighbors_nvalchemi( - systems=systems, - requested_options=full_nl_options, - ) - systems = _compute_requested_neighbors_vesin( - systems=systems, - requested_options=half_nl_options, - check_consistency=check_consistency, - ) + return SymmetrizedCalculator else: - systems = _compute_requested_neighbors_vesin( - systems=systems, - requested_options=requested_options, - check_consistency=check_consistency, - ) - - return systems - - -def _compute_requested_neighbors_vesin( - systems: List[System], - requested_options: List[NeighborListOptions], - check_consistency=False, -) -> List[System]: - """ - Compute all neighbor lists requested by ``model`` and store them inside the systems, - using vesin. - """ - - system_devices = [] - moved_systems = [] - for system in systems: - system_devices.append(system.device) - if system.device.type not in ["cpu", "cuda"]: - moved_systems.append(system.to(device="cpu")) - else: - moved_systems.append(system) - - vesin.metatomic.compute_requested_neighbors_from_options( - systems=moved_systems, - system_length_unit="angstrom", - options=requested_options, - check_consistency=check_consistency, - ) - - systems = [] - for system, device in zip(moved_systems, system_devices, strict=True): - systems.append(system.to(device=device)) - - return systems - - -def _compute_requested_neighbors_nvalchemi(systems, requested_options): - """ - Compute all neighbor lists requested by ``model`` and store them inside the systems, - using nvalchemiops. This function should only be called if all systems are on CUDA - and all neighbor list options require a full neighbor list. - """ - - for options in requested_options: - assert options.full_list - for system in systems: - assert system.device.type == "cuda" - - edge_index, _, S = nvalchemi_neighbor_list( - system.positions, - options.engine_cutoff("angstrom"), - cell=system.cell, - pbc=system.pbc, - return_neighbor_list=True, - ) - D = ( - system.positions[edge_index[1]] - - system.positions[edge_index[0]] - + S.to(system.cell.dtype) @ system.cell - ) - P = edge_index.T - - neighbors = TensorBlock( - D.reshape(-1, 3, 1), - samples=Labels( - names=[ - "first_atom", - "second_atom", - "cell_shift_a", - "cell_shift_b", - "cell_shift_c", - ], - values=torch.hstack([P, S]), - ), - components=[ - Labels("xyz", torch.tensor([[0], [1], [2]], device=system.device)) - ], - properties=Labels( - "distance", torch.tensor([[0]], device=system.device) - ), - ) - system.add_neighbor_list(options, neighbors) - - return systems + raise AttributeError(f"module {__name__} has no attribute {name}") diff --git a/python/metatomic_torch/setup.py b/python/metatomic_torch/setup.py index 6ab76d98b..9a3f3de05 100644 --- a/python/metatomic_torch/setup.py +++ b/python/metatomic_torch/setup.py @@ -21,7 +21,10 @@ "expected 'debug' or 'release'" ) -METATOMIC_TORCH_SRC = os.path.join(ROOT, "..", "..", "metatomic-torch") +METATOMIC_TORCH_SRC = os.path.realpath( + os.path.join(ROOT, "..", "..", "metatomic-torch") +) +METATOMIC_ASE = os.path.realpath(os.path.join(ROOT, "..", "metatomic_ase")) class universal_wheel(bdist_wheel): @@ -315,11 +318,20 @@ def create_version_number(version): install_requires = [ f"torch {torch_version}", - "vesin >=0.5.1", "metatensor-torch >=0.8.0,<0.9", "metatensor-operations >=0.4.0,<0.5", ] + # when packaging a sdist for release, we should never use local dependencies + METATOMIC_NO_LOCAL_DEPS = os.environ.get("METATOMIC_NO_LOCAL_DEPS", "0") == "1" + + if not METATOMIC_NO_LOCAL_DEPS and os.path.exists(METATOMIC_ASE): + # we are building from a git checkout or full repo archive + install_requires.append(f"metatomic-ase @ file://{METATOMIC_ASE}") + else: + # we are building from a sdist/installing from a wheel + install_requires.append("metatomic-ase >=0.1.0,<0.2.0") + setup( version=create_version_number(METATOMIC_TORCH_VERSION), author=", ".join(authors), diff --git a/python/metatomic_torch/tests/ase_calculator.py b/python/metatomic_torch/tests/ase_calculator.py index 0b84aa565..e4a5be12c 100644 --- a/python/metatomic_torch/tests/ase_calculator.py +++ b/python/metatomic_torch/tests/ase_calculator.py @@ -1,969 +1,19 @@ -import glob -import json -import os -import subprocess -import sys -from typing import Dict, List, Optional - -import ase.build -import ase.calculators.lj -import ase.md -import ase.units -import numpy as np import pytest -import torch -from ase.calculators.calculator import PropertyNotImplementedError -from ase.md.velocitydistribution import MaxwellBoltzmannDistribution -from metatensor.torch import Labels, TensorBlock, TensorMap - -import metatomic_lj_test -from metatomic.torch import ( - AtomisticModel, - ModelCapabilities, - ModelMetadata, - ModelOutput, - NeighborListOptions, - System, -) -from metatomic.torch.ase_calculator import ( - ARRAY_QUANTITIES, - MetatomicCalculator, - _compute_requested_neighbors_nvalchemi, - _compute_requested_neighbors_vesin, - _full_3x3_to_voigt_6_stress, -) - -from . import _tests_utils - - -STR_TO_DTYPE = { - "float32": torch.float32, - "float64": torch.float64, -} - -ALL_DEVICE_DTYPE = [("cpu", torch.float64), ("cpu", torch.float32)] -if torch.cuda.is_available(): - ALL_DEVICE_DTYPE.append(("cuda", torch.float64)) - ALL_DEVICE_DTYPE.append(("cuda", torch.float32)) -if torch.backends.mps.is_available(): - ALL_DEVICE_DTYPE.append(("mps", torch.float32)) - -CUTOFF = 5.0 -SIGMA = 1.5808 -EPSILON = 0.1729 - - -@pytest.fixture -def model(): - return metatomic_lj_test.lennard_jones_model( - atomic_type=28, - cutoff=CUTOFF, - sigma=SIGMA, - epsilon=EPSILON, - length_unit="Angstrom", - energy_unit="eV", - with_extension=False, - ) - - -@pytest.fixture -def model_different_units(): - return metatomic_lj_test.lennard_jones_model( - atomic_type=28, - cutoff=CUTOFF / ase.units.Bohr, - sigma=SIGMA / ase.units.Bohr, - epsilon=EPSILON / ase.units.kJ * ase.units.mol, - length_unit="Bohr", - energy_unit="kJ/mol", - with_extension=False, - ) - - -@pytest.fixture -def atoms(): - np.random.seed(0xDEADBEEF) - - atoms = ase.build.make_supercell( - ase.build.bulk("Ni", "fcc", a=3.6, cubic=True), 2 * np.eye(3) - ) - atoms.positions += 0.2 * np.random.rand(*atoms.positions.shape) - - return atoms - - -def check_against_ase_lj(atoms, calculator, dtype): - ref = atoms.copy() - ref.calc = ase.calculators.lj.LennardJones( - sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False - ) - - atoms.calc = calculator - - if dtype == torch.float32: - rtol = 1e-5 - atol = 1e-5 - elif dtype == torch.float64: - rtol = 1e-5 - atol = 1e-8 - else: - raise ValueError(f"unsupported dtype: {dtype}") - - assert np.allclose( - ref.get_potential_energy(), atoms.get_potential_energy(), atol=atol, rtol=rtol - ) - assert np.allclose( - ref.get_potential_energies(), - atoms.get_potential_energies(), - atol=atol, - rtol=rtol, - ) - assert np.allclose(ref.get_forces(), atoms.get_forces(), atol=atol, rtol=rtol) - assert np.allclose(ref.get_stress(), atoms.get_stress(), atol=atol, rtol=rtol) - - -def _set_model_dtype(model, dtype): - model._capabilities.dtype = str(dtype)[6:] - model._model_dtype = dtype - return model - - -@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) -def test_python_model(model, model_different_units, atoms, device, dtype): - model = _set_model_dtype(model, dtype) - model_different_units = _set_model_dtype(model_different_units, dtype) - - check_against_ase_lj( - atoms, - MetatomicCalculator( - model, - check_consistency=True, - uncertainty_threshold=None, - device=device, - ), - dtype=dtype, - ) - check_against_ase_lj( - atoms, - MetatomicCalculator( - model_different_units, - check_consistency=True, - uncertainty_threshold=None, - device=device, - ), - dtype=dtype, - ) - - -@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) -def test_torch_script_model(model, model_different_units, atoms, device, dtype): - model = _set_model_dtype(model, dtype) - model_different_units = _set_model_dtype(model_different_units, dtype) - - model = torch.jit.script(model) - check_against_ase_lj( - atoms, - MetatomicCalculator( - model, - check_consistency=True, - uncertainty_threshold=None, - device=device, - ), - dtype=dtype, - ) - - model_different_units = torch.jit.script(model_different_units) - check_against_ase_lj( - atoms, - MetatomicCalculator( - model_different_units, - check_consistency=True, - uncertainty_threshold=None, - device=device, - ), - dtype=dtype, - ) - - -@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) -def test_exported_model(tmpdir, model, model_different_units, atoms, device, dtype): - model = _set_model_dtype(model, dtype) - model_different_units = _set_model_dtype(model_different_units, dtype) - - path = os.path.join(tmpdir, "exported-model.pt") - model.save(path) - check_against_ase_lj( - atoms, - MetatomicCalculator( - path, - check_consistency=True, - uncertainty_threshold=None, - device=device, - ), - dtype=dtype, - ) - - model_different_units.save(path) - check_against_ase_lj( - atoms, - MetatomicCalculator( - path, - check_consistency=True, - uncertainty_threshold=None, - device=device, - ), - dtype=dtype, - ) - - -@pytest.mark.parametrize("non_conservative", [True, False]) -def test_get_properties(model, atoms, non_conservative): - atoms.calc = MetatomicCalculator( - model, - check_consistency=True, - non_conservative=non_conservative, - uncertainty_threshold=None, - ) - - properties = atoms.get_properties(["energy", "energies", "forces", "stress"]) - - assert np.all(properties["energies"] == atoms.get_potential_energies()) - assert np.all(properties["energy"] == atoms.get_potential_energy()) - assert np.all(properties["forces"] == atoms.get_forces()) - assert np.all(properties["stress"] == atoms.get_stress()) - - # check that we can use all of the `.get_xxx` functions independantly - atoms.calc = MetatomicCalculator( - model, non_conservative=non_conservative, uncertainty_threshold=None - ) - atoms.get_potential_energy() - - atoms.calc = MetatomicCalculator( - model, non_conservative=non_conservative, uncertainty_threshold=None - ) - atoms.get_potential_energies() - - atoms.calc = MetatomicCalculator( - model, non_conservative=non_conservative, uncertainty_threshold=None - ) - atoms.get_forces() - - atoms.calc = MetatomicCalculator( - model, non_conservative=non_conservative, uncertainty_threshold=None - ) - atoms.get_stress() - - -def test_accuracy_warning(model, atoms): - # our dummy model artificially gives a high uncertainty for large structures - big_atoms = atoms * (2, 2, 2) - big_atoms.calc = MetatomicCalculator(model, check_consistency=True) - with pytest.warns( - UserWarning, - match="Some of the atomic energy uncertainties are large", - ): - big_atoms.get_forces() - - -def accuracy_is_zero_error(atoms): - match = "`uncertainty_threshold` is 0.0 but must be positive" - with pytest.raises(ValueError, match=match): - atoms.calc = MetatomicCalculator(model, uncertainty_threshold=0.0) - - -def test_run_model(tmpdir, model, atoms): - ref = atoms.copy() - ref.calc = ase.calculators.lj.LennardJones( - sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False - ) - - path = os.path.join(tmpdir, "exported-model.pt") - model.save(path) - calculator = MetatomicCalculator( - path, check_consistency=True, uncertainty_threshold=None - ) - - first_mask = [a % 2 == 0 for a in range(len(atoms))] - first_half = Labels( - ["system", "atom"], - torch.tensor([[0, a] for a in range(len(atoms)) if a % 2 == 0]), - ) - - second_mask = [a % 2 == 1 for a in range(len(atoms))] - second_half = Labels( - ["system", "atom"], - torch.tensor([[0, a] for a in range(len(atoms)) if a % 2 == 1]), - ) - - # check overall prediction - requested = {"energy": ModelOutput(per_atom=False)} - outputs = calculator.run_model(atoms, outputs=requested) - assert np.allclose( - ref.get_potential_energy(), outputs["energy"].block().values.item() - ) - - # check per atom energy - requested = {"energy": ModelOutput(per_atom=True)} - outputs = calculator.run_model(atoms, outputs=requested, selected_atoms=first_half) - first_energies = outputs["energy"].block().values.numpy().reshape(-1) - - outputs = calculator.run_model(atoms, outputs=requested, selected_atoms=second_half) - second_energies = outputs["energy"].block().values.numpy().reshape(-1) - - expected = ref.get_potential_energies() - assert np.allclose(expected[first_mask], first_energies) - assert np.allclose(expected[second_mask], second_energies) - - # check total energy - requested = {"energy": ModelOutput(per_atom=False)} - outputs = calculator.run_model(atoms, outputs=requested, selected_atoms=first_half) - first_energies = outputs["energy"].block().values.numpy().reshape(-1) - - outputs = calculator.run_model(atoms, outputs=requested, selected_atoms=second_half) - second_energies = outputs["energy"].block().values.numpy().reshape(-1) - - expected = ref.get_potential_energy() - assert np.allclose(expected, first_energies + second_energies) - - # check batched prediction - requested = {"energy": ModelOutput(per_atom=False)} - outputs = calculator.run_model([atoms, atoms], outputs=requested) - assert np.allclose( - ref.get_potential_energy(), outputs["energy"].block().values[[0]] - ) - - # check non-conservative forces and stresses - requested = { - "energy": ModelOutput(per_atom=False), - "non_conservative_forces": ModelOutput(per_atom=True), - "non_conservative_stress": ModelOutput(per_atom=False), - } - outputs = calculator.run_model([atoms, atoms], outputs=requested) - assert np.allclose( - ref.get_potential_energy(), outputs["energy"].block().values[[0]] - ) - assert np.allclose( - ref.get_potential_energy(), outputs["energy"].block().values[[1]] - ) - assert "non_conservative_forces" in outputs - assert outputs["non_conservative_forces"].block().values.shape == ( - 2 * len(atoms), - 3, - 1, - ) - assert "non_conservative_stress" in outputs - assert outputs["non_conservative_stress"].block().values.shape == (2, 3, 3, 1) - - -@pytest.mark.parametrize("non_conservative", [True, False]) -@pytest.mark.parametrize("per_atom", [True, False]) -def test_compute_energy(tmpdir, model, atoms, non_conservative, per_atom): - ref = atoms.copy() - ref.calc = ase.calculators.lj.LennardJones( - sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False - ) - - path = os.path.join(tmpdir, "exported-model.pt") - model.save(path) - calculator = MetatomicCalculator( - path, - check_consistency=True, - non_conservative=non_conservative, - ) - - results = calculator.compute_energy(atoms, per_atom=per_atom) - if per_atom: - energies = results["energies"] - assert np.allclose(ref.get_potential_energies(), energies) - assert np.allclose(ref.get_potential_energy(), results["energy"]) - - results = calculator.compute_energy( - atoms, compute_forces_and_stresses=True, per_atom=per_atom - ) - assert np.allclose(ref.get_potential_energy(), results["energy"]) - if not non_conservative: - assert np.allclose(ref.get_forces(), results["forces"]) - assert np.allclose( - ref.get_stress(), _full_3x3_to_voigt_6_stress(results["stress"]) - ) - if per_atom: - assert np.allclose(ref.get_potential_energies(), results["energies"]) - - results = calculator.compute_energy([atoms, atoms], per_atom=per_atom) - assert np.allclose(ref.get_potential_energy(), results["energy"][0]) - assert np.allclose(ref.get_potential_energy(), results["energy"][1]) - if per_atom: - assert np.allclose(ref.get_potential_energies(), results["energies"][0]) - assert np.allclose(ref.get_potential_energies(), results["energies"][1]) - - results = calculator.compute_energy( - [atoms, atoms], - compute_forces_and_stresses=True, - per_atom=per_atom, - ) - assert np.allclose(ref.get_potential_energy(), results["energy"][0]) - assert np.allclose(ref.get_potential_energy(), results["energy"][1]) - if not non_conservative: - assert np.allclose(ref.get_forces(), results["forces"][0]) - assert np.allclose(ref.get_forces(), results["forces"][1]) - assert np.allclose( - ref.get_stress(), _full_3x3_to_voigt_6_stress(results["stress"][0]) - ) - assert np.allclose( - ref.get_stress(), _full_3x3_to_voigt_6_stress(results["stress"][1]) - ) - if per_atom: - assert np.allclose(ref.get_potential_energies(), results["energies"][0]) - assert np.allclose(ref.get_potential_energies(), results["energies"][1]) - - atoms_no_pbc = atoms.copy() - atoms_no_pbc.pbc = [False, False, False] - assert "stress" not in calculator.compute_energy(atoms_no_pbc) - - -def test_serialize_ase(tmpdir, model, atoms): - calculator = MetatomicCalculator(model, uncertainty_threshold=None) +def test_import(): message = ( - "can not save metatensor model in ASE `todict`, please initialize " - "`MetatomicCalculator` with a path to a saved model file if you need to use " - "`todict" - ) - with pytest.raises(RuntimeError, match=message): - calculator.todict() - - # save with exported model - path = os.path.join(tmpdir, "exported-model.pt") - model.save(path) - - calculator = MetatomicCalculator(path, uncertainty_threshold=None) - data = calculator.todict() - _ = MetatomicCalculator.fromdict(data) - - # check the standard trajectory format of ASE, which uses `todict`/`fromdict` - atoms.calc = MetatomicCalculator(path, uncertainty_threshold=None) - with tmpdir.as_cwd(): - dyn = ase.md.VelocityVerlet( - atoms, - timestep=2 * ase.units.fs, - trajectory="file.traj", - ) - dyn.run(10) - dyn.close() - - atoms = ase.io.read("file.traj", "-1") - assert atoms.calc is not None - - -def test_dtype_device(tmpdir, model, atoms): - ref = atoms.copy() - ref.calc = ase.calculators.lj.LennardJones( - sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False + "Importing MetatomicCalculator from metatomic.torch.ase_calculator is " + "deprecated and will be removed in a future release. Please import from " + "metatomic_ase instead." ) - expected = ref.get_potential_energy() - path = os.path.join(tmpdir, "exported-model.pt") - - dtype_device = [ - ("float64", "cpu"), - ("float32", "cpu"), - ] - - if _tests_utils.can_use_mps_backend(): - dtype_device.append(("float32", "mps")) - - if torch.cuda.is_available(): - dtype_device.append(("float32", "cuda")) - dtype_device.append(("float64", "cuda")) - - for dtype, device in dtype_device: - capabilities = model.capabilities() - capabilities.dtype = dtype - - # re-create the model with a different dtype - dtype_model = AtomisticModel( - model.module.to(STR_TO_DTYPE[dtype]), - model.metadata(), - capabilities, - ) - - dtype_model.save(path) - atoms.calc = MetatomicCalculator( - path, check_consistency=True, device=device, uncertainty_threshold=None - ) - assert np.allclose(atoms.get_potential_energy(), expected) - - -@pytest.mark.parametrize("device", ["cpu"]) -def test_model_with_extensions(tmpdir, atoms, device): - ref = atoms.copy() - ref.calc = ase.calculators.lj.LennardJones( - sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False - ) - - model_path = os.path.join(tmpdir, "model.pt") - extensions_directory = os.path.join(tmpdir, "extensions") - - # use forward slash as path separator even on windows. This removes the need to - # escape the path below - model_path = model_path.replace("\\", "/") - extensions_directory = extensions_directory.replace("\\", "/") - - # export the model in a sub-process, to prevent loading of the extension in the - # current interpreter (until we try to execute the code) - script = f""" -import metatomic_lj_test - -model = metatomic_lj_test.lennard_jones_model( - atomic_type=28, - cutoff={CUTOFF}, - sigma={SIGMA}, - epsilon={EPSILON}, - length_unit="Angstrom", - energy_unit="eV", - with_extension=True, -) - -model.save("{model_path}", collect_extensions="{extensions_directory}") - """ - - subprocess.run([sys.executable, "-c", script], check=True, cwd=tmpdir) + with pytest.warns(DeprecationWarning, match=message): + from metatomic.torch.ase_calculator import MetatomicCalculator # noqa: F401 message = ( - "This is likely due to missing TorchScript extensions.\nMake sure to provide " - "the `extensions_directory` argument if your extensions are not installed " - "system-wide" - ) - with pytest.raises(RuntimeError, match=message): - MetatomicCalculator(model_path, check_consistency=True) - - # Now actually loading the extensions - atoms.calc = MetatomicCalculator( - model_path, - extensions_directory=extensions_directory, - check_consistency=True, - uncertainty_threshold=None, - ) - - assert np.allclose(ref.get_potential_energy(), atoms.get_potential_energy()) - - -def _read_neighbor_check(path): - with open(path) as fd: - data = json.load(fd) - - dtype = torch.float64 - - positions = torch.tensor(data["system"]["positions"], dtype=dtype).reshape(-1, 3) - system = System( - types=torch.tensor([1] * positions.shape[0], dtype=torch.int32), - positions=positions, - cell=torch.tensor(data["system"]["cell"], dtype=dtype), - pbc=torch.tensor([True, True, True]), - ) - - options = NeighborListOptions( - cutoff=data["options"]["cutoff"], - full_list=data["options"]["full_list"], - # ASE can only compute strict NL - strict=True, - ) - - samples = torch.tensor( - data["expected-neighbors"]["samples"], dtype=torch.int32 - ).reshape(-1, 5) - distances = torch.tensor( - data["expected-neighbors"]["distances"], dtype=dtype - ).reshape(-1, 3, 1) - - neighbors = TensorBlock( - values=distances, - samples=Labels( - [ - "first_atom", - "second_atom", - "cell_shift_a", - "cell_shift_b", - "cell_shift_c", - ], - samples, - ), - components=[Labels.range("xyz", 3)], - properties=Labels.range("distance", 1), - ) - - return system, options, neighbors - - -def _check_same_set_of_neighbors(expected, actual, full_list): - assert expected.samples.names == actual.samples.names - assert len(expected.samples) == len(actual.samples) - - for sample_i, sample in enumerate(expected.samples): - sign = 1.0 - position = actual.samples.position(sample) - - if position is None and not full_list: - # try looking for the inverse pair - sign = -1.0 - position = actual.samples.position( - [sample[1], sample[0], -sample[2], -sample[3], -sample[4]] - ) - - if position is None: - raise AssertionError(f"missing expected neighbors sample: {sample}") - - assert torch.allclose(expected.values[sample_i], sign * actual.values[position]) - - -@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) -@pytest.mark.parametrize( - "neighbor_fn", - [_compute_requested_neighbors_vesin, _compute_requested_neighbors_nvalchemi], -) -def test_neighbor_list_adapter(device, dtype, neighbor_fn): - if neighbor_fn == _compute_requested_neighbors_nvalchemi and device != "cuda": - pytest.skip("nvalchemiops neighbor list is only implemented for CUDA") - - HERE = os.path.realpath(os.path.dirname(__file__)) - test_files = os.path.join( - HERE, "..", "..", "..", "..", "metatensor-torch", "tests", "neighbor-checks" - ) - - for path in glob.glob(os.path.join(test_files, "*.json")): - system, options, expected_neighbors = _read_neighbor_check(path) - system = system.to(device=device, dtype=dtype) - expected_neighbors = expected_neighbors.to(device=device, dtype=dtype) - - neighbor_fn([system], [options], check_consistency=True) - _check_same_set_of_neighbors( - expected_neighbors, system.get_neighbor_list(options), options.full_list - ) - - -class MultipleOutputModel(torch.nn.Module): - def forward( - self, - systems: List[System], - outputs: Dict[str, ModelOutput], - selected_atoms: Optional[Labels] = None, - ) -> Dict[str, TensorMap]: - results = {} - device = systems[0].positions.device - for name, requested in outputs.items(): - assert not requested.per_atom - - block = TensorBlock( - values=torch.tensor([[0.0]], dtype=torch.float64, device=device), - samples=Labels("system", torch.tensor([[0]], device=device)), - components=torch.jit.annotate(List[Labels], []), - properties=Labels( - name.split(":")[0], torch.tensor([[0]], device=device) - ), - ) - tensor = TensorMap(Labels("_", torch.tensor([[0]], device=device)), [block]) - results[name] = tensor - - return results - - -def test_additional_outputs(atoms): - capabilities = ModelCapabilities( - outputs={ - "energy": ModelOutput(per_atom=False), - "test::test": ModelOutput(per_atom=False), - "another::one": ModelOutput(per_atom=False), - }, - atomic_types=[28], - interaction_range=0.0, - supported_devices=["cpu"], - dtype="float64", - ) - model = AtomisticModel(MultipleOutputModel().eval(), ModelMetadata(), capabilities) - - atoms.calc = MetatomicCalculator( - model, - check_consistency=True, - uncertainty_threshold=None, - ) - - assert atoms.get_potential_energy() == 0.0 - assert atoms.calc.additional_outputs == {} - - atoms.calc = MetatomicCalculator( - model, - additional_outputs={ - "test::test": ModelOutput(per_atom=False), - "another::one": ModelOutput(per_atom=False), - }, - check_consistency=True, - uncertainty_threshold=None, - ) - assert atoms.get_potential_energy() == 0.0 - - test = atoms.calc.additional_outputs["test::test"] - assert test.block().properties.names == ["test"] - - another = atoms.calc.additional_outputs["another::one"] - assert another.block().properties.names == ["another"] - - -@pytest.mark.parametrize("non_conservative", [True, False]) -def test_variants(atoms, model, non_conservative): - atoms.calc = MetatomicCalculator( - model, - check_consistency=True, - non_conservative=non_conservative, - uncertainty_threshold=None, - ) - - atoms_variant = atoms.copy() - atoms_variant.calc = MetatomicCalculator( - model, - check_consistency=True, - non_conservative=non_conservative, - variants={"energy": "doubled"}, - uncertainty_threshold=None, - ) - - np.allclose( - 2.0 * atoms.get_potential_energy(), atoms_variant.get_potential_energy() - ) - np.allclose(2.0 * atoms.get_forces(), atoms_variant.get_forces()) - np.allclose(2.0 * atoms.get_stress(), atoms_variant.get_stress()) - - -@pytest.mark.parametrize( - "default_output", - [ - "energy", - "energy_uncertainty", - "non_conservative_forces", - "non_conservative_stress", - ], -) -def test_variant_default(atoms, model, default_output): - """Allow setting a variant explicitly to None to use the default output.""" - - atoms.calc = MetatomicCalculator( - model, - check_consistency=True, - non_conservative=True, - uncertainty_threshold=None, - ) - - variants = { - v: "doubled" - for v in [ - "energy", - "energy_uncertainty", - "non_conservative_forces", - "non_conservative_stress", - ] - } - variants[default_output] = None - - atoms_variant = atoms.copy() - atoms_variant.calc = MetatomicCalculator( - model, - check_consistency=True, - non_conservative=True, - variants={"energy": "doubled"}, - uncertainty_threshold=None, - ) - - if default_output == "energy": - np.allclose(atoms.get_potential_energy(), atoms_variant.get_potential_energy()) - np.allclose(2.0 * atoms.get_forces(), atoms_variant.get_forces()) - np.allclose(2.0 * atoms.get_stress(), atoms_variant.get_stress()) - elif default_output == "non_conservative_forces": - np.allclose( - 2.0 * atoms.get_potential_energy(), atoms_variant.get_potential_energy() - ) - np.allclose(atoms.get_forces(), atoms_variant.get_forces()) - np.allclose(2.0 * atoms.get_stress(), atoms_variant.get_stress()) - elif default_output == "non_conservative_stress": - np.allclose( - 2.0 * atoms.get_potential_energy(), atoms_variant.get_potential_energy() - ) - np.allclose(2.0 * atoms.get_forces(), atoms_variant.get_forces()) - np.allclose(atoms.get_stress(), atoms_variant.get_stress()) - - -@pytest.mark.parametrize("force_is_None", [True, False]) -def test_variant_non_conservative_error(atoms, model, force_is_None): - variants = { - "energy": "doubled", - "non_conservative_forces": "doubled", - "non_conservative_stress": "doubled", - } - - if force_is_None: - variants["non_conservative_forces"] = None - else: - variants["non_conservative_stress"] = None - - match = "must either be both `None` or both not `None`." - with pytest.raises(ValueError, match=match): - MetatomicCalculator( - model, - check_consistency=True, - non_conservative=True, - variants=variants, - uncertainty_threshold=None, - ) - - -def test_model_without_energy(atoms): - """ - Test that a MetatomicCalculator can be created with a model without energy - output. - """ - - # Create a model that only outputs a custom property, no energy - class NoEnergyModel(torch.nn.Module): - def forward( - self, - systems: List[System], - outputs: Dict[str, ModelOutput], - selected_atoms: Optional[Labels] = None, - ) -> Dict[str, TensorMap]: - results = {} - for name in outputs: - # Return dummy data for each requested output - block = TensorBlock( - values=torch.tensor([[1.0]], dtype=torch.float64), - samples=Labels("system", torch.tensor([[0]])), - components=torch.jit.annotate(List[Labels], []), - properties=Labels(name.split(":")[0], torch.tensor([[0]])), - ) - tensor = TensorMap(Labels("_", torch.tensor([[0]])), [block]) - results[name] = tensor - return results - - # Create model capabilities without energy output - capabilities = ModelCapabilities( - outputs={ - "features": ModelOutput(per_atom=False), - "custom::output": ModelOutput(per_atom=False), - }, - atomic_types=[28], - interaction_range=0.0, - supported_devices=["cpu"], - dtype="float64", - ) - model = AtomisticModel(NoEnergyModel().eval(), ModelMetadata(), capabilities) - - # Should be able to create calculator without error - calc = MetatomicCalculator( - model, - check_consistency=True, - uncertainty_threshold=None, - ) - - # The calculator should work for additional outputs - atoms.calc = MetatomicCalculator( - model, - additional_outputs={ - "features": ModelOutput(per_atom=False), - }, - check_consistency=True, - uncertainty_threshold=None, - ) - - # Should be able to call run_model directly with custom outputs - outputs = atoms.calc.run_model( - atoms, - outputs={"features": ModelOutput(per_atom=False)}, - ) - assert "features" in outputs - - # But trying to get energy should fail with a clear error - match = "does not support energy-related properties" - with pytest.raises(PropertyNotImplementedError, match=match): - atoms.get_potential_energy() - - with pytest.raises(PropertyNotImplementedError, match=match): - atoms.get_forces() - - # compute_energy should also fail - match = "does not support energy computation" - with pytest.raises(ValueError, match=match): - calc.compute_energy(atoms) - - -class AdditionalInputModel(torch.nn.Module): - def __init__(self, inputs): - super().__init__() - self._requested_inputs = inputs - - def requested_inputs(self) -> Dict[str, ModelOutput]: - return self._requested_inputs - - def forward( - self, - systems: List[System], - outputs: Dict[str, ModelOutput], - selected_atoms: Optional[Labels] = None, - ) -> Dict[str, TensorMap]: - return { - ("extra::" + input): systems[0].get_data(input) - for input in self._requested_inputs - } - - -def test_additional_input(atoms): - inputs = { - "masses": ModelOutput(quantity="mass", unit="u", per_atom=True), - "velocities": ModelOutput(quantity="velocity", unit="A/fs", per_atom=True), - "charges": ModelOutput(quantity="charge", unit="e", per_atom=True), - "ase::initial_charges": ModelOutput(quantity="charge", unit="e", per_atom=True), - } - outputs = {("extra::" + n): inputs[n] for n in inputs} - capabilities = ModelCapabilities( - outputs=outputs, - atomic_types=[28], - interaction_range=0.0, - supported_devices=["cpu"], - dtype="float64", - ) - - model = AtomisticModel( - AdditionalInputModel(inputs).eval(), ModelMetadata(), capabilities - ) - MaxwellBoltzmannDistribution(atoms, temperature_K=300.0) - atoms.set_initial_charges([0.0] * len(atoms)) - calculator = MetatomicCalculator(model, check_consistency=True) - results = calculator.run_model(atoms, outputs) - for name, tensor in results.items(): - head, name = name.split("::", maxsplit=1) - assert head == "extra" - assert name in inputs - - assert tensor.get_info("quantity") == inputs[name].quantity - values = tensor[0].values.numpy() - - expected = ARRAY_QUANTITIES[name]["getter"](atoms).reshape(values.shape) - if name == "velocities": - expected /= ( - ase.units.Angstrom / ase.units.fs - ) # ase velocity is in (eV/u)^(1/2) and we want A/fs - - assert np.allclose(values, expected) - - -@pytest.mark.parametrize("device,dtype", ALL_DEVICE_DTYPE) -def test_mixed_pbc(model, device, dtype): - """Test that the calculator works on a mixed-PBC system""" - atoms = ase.build.fcc111("Ni", size=(2, 2, 3), vacuum=10.0) - atoms.set_pbc((True, True, False)) - - model = _set_model_dtype(model, dtype) - - atoms.calc = MetatomicCalculator( - model, - check_consistency=True, - uncertainty_threshold=None, - device=device, + "Importing SymmetrizedCalculator from metatomic.torch.ase_calculator is " + "deprecated and will be removed in a future release. Please import from " + "metatomic_ase instead." ) - atoms.get_potential_energy() - atoms.get_forces() + with pytest.warns(DeprecationWarning, match=message): + from metatomic.torch.ase_calculator import SymmetrizedCalculator # noqa: F401 diff --git a/python/metatomic_torch/tests/neighbors.py b/python/metatomic_torch/tests/neighbors.py index 1b5909c31..926aef07e 100644 --- a/python/metatomic_torch/tests/neighbors.py +++ b/python/metatomic_torch/tests/neighbors.py @@ -1,15 +1,14 @@ -import metatensor.torch as mts +import numpy as np import pytest import torch -from metatensor.torch import TensorBlock +from metatensor.torch import Labels, TensorBlock from metatomic.torch import NeighborListOptions, System, register_autograd_neighbors try: - import ase # noqa: F401 - - from metatomic.torch.ase_calculator import _compute_requested_neighbors + import ase + import ase.neighborlist HAVE_ASE = True except ImportError: @@ -61,6 +60,83 @@ def test_neighbor_list_options(): assert repr(options) == expected +def compute_neighbors_with_ase(system, options): + # options.strict is ignored by this function, since `ase.neighborlist.neighbor_list` + # only computes strict NL, and these are valid even with `strict=False` + + dtype = system.positions.dtype + device = system.positions.device + + atoms = ase.Atoms( + numbers=system.types.cpu().numpy(), + positions=system.positions.cpu().detach().numpy(), + cell=system.cell.cpu().detach().numpy(), + pbc=system.pbc.cpu().numpy(), + ) + + nl_i, nl_j, nl_S, nl_D = ase.neighborlist.neighbor_list( + "ijSD", + atoms, + cutoff=options.engine_cutoff(engine_length_unit="angstrom"), + ) + + if not options.full_list: + # The pair selection code here below avoids a relatively slow loop over + # all pairs to improve performance + reject_condition = ( + # we want a half neighbor list, so drop all duplicated neighbors + (nl_j < nl_i) + | ( + (nl_i == nl_j) + & ( + # only create pairs with the same atom twice if the pair spans more + # than one unit cell + ((nl_S[:, 0] == 0) & (nl_S[:, 1] == 0) & (nl_S[:, 2] == 0)) + # When creating pairs between an atom and one of its periodic + # images, the code generates multiple redundant pairs + # (e.g. with shifts 0 1 1 and 0 -1 -1); and we want to only keep one + # of these. We keep the pair in the positive half plane of shifts. + | ( + (nl_S.sum(axis=1) < 0) + | ( + (nl_S.sum(axis=1) == 0) + & ( + (nl_S[:, 2] < 0) + | ((nl_S[:, 2] == 0) & (nl_S[:, 1] < 0)) + ) + ) + ) + ) + ) + ) + selected = np.logical_not(reject_condition) + nl_i = nl_i[selected] + nl_j = nl_j[selected] + nl_S = nl_S[selected] + nl_D = nl_D[selected] + + samples = np.concatenate([nl_i[:, None], nl_j[:, None], nl_S], axis=1) + + distances = torch.from_numpy(nl_D).to(dtype=dtype, device=device) + + return TensorBlock( + values=distances.reshape(-1, 3, 1), + samples=Labels( + names=[ + "first_atom", + "second_atom", + "cell_shift_a", + "cell_shift_b", + "cell_shift_c", + ], + values=torch.from_numpy(samples).to(dtype=torch.int32, device=device), + assume_unique=True, + ), + components=[Labels.range("xyz", 3).to(device)], + properties=Labels.range("distance", 1).to(device), + ) + + @pytest.mark.skipif(not HAVE_ASE, reason="this tests requires ASE neighbor list") @pytest.mark.parametrize("device", ALL_DEVICES) def test_neighbors_autograd(device): @@ -83,8 +159,8 @@ def compute(positions, cell, options): cell, pbc=torch.tensor([True, True, True], device=positions.device), ) - _compute_requested_neighbors([system], [options], check_consistency=True) - neighbors = system.get_neighbor_list(options) + neighbors = compute_neighbors_with_ase(system, options) + register_autograd_neighbors(system, neighbors, check_consistency=True) return neighbors.values.sum() options = NeighborListOptions(cutoff=2.0, full_list=False, strict=True) @@ -118,8 +194,8 @@ def test_neighbor_autograd_errors(): pbc=torch.tensor([True, True, True]), ) options = NeighborListOptions(cutoff=2.0, full_list=False, strict=True) - _compute_requested_neighbors([system], [options], check_consistency=True) - neighbors = system.get_neighbor_list(options) + neighbors = compute_neighbors_with_ase(system, options) + register_autograd_neighbors(system, neighbors, check_consistency=True) system = System( torch.tensor([6] * len(positions)), @@ -139,8 +215,6 @@ def test_neighbor_autograd_errors(): r"atom \d+ for the \[.*?\] cell shift should have a distance vector " r"of \[.*?\] but has a distance vector of \[.*?\]" ) - _compute_requested_neighbors([system], [options], check_consistency=True) - neighbors = system.get_neighbor_list(options) system = System( torch.tensor([6] * len(positions)), @@ -148,7 +222,7 @@ def test_neighbor_autograd_errors(): cell, pbc=torch.tensor([True, True, True]), ) - neighbors = mts.detach_block(neighbors) + neighbors = compute_neighbors_with_ase(system, options) neighbors.values[:] *= 3 with pytest.raises(ValueError, match=message): register_autograd_neighbors(system, neighbors, check_consistency=True) diff --git a/python/metatomic_torch/tests/utils.py b/python/metatomic_torch/tests/utils.py index bbd31bd37..3e0aa4cde 100644 --- a/python/metatomic_torch/tests/utils.py +++ b/python/metatomic_torch/tests/utils.py @@ -20,10 +20,19 @@ def test_version_compatible(): def test_lazy_ase_calculator_import(): - mod = metatomic.torch.ase_calculator - assert hasattr(mod, "MetatomicCalculator") - - -def test_lazy_import_missing_attribute(): - with pytest.raises(AttributeError, match="has no attribute"): - _ = metatomic.torch.nonexistent_attribute + module = metatomic.torch.ase_calculator + message = ( + "Importing MetatomicCalculator from metatomic.torch.ase_calculator is " + "deprecated and will be removed in a future release. Please import from " + "metatomic_ase instead." + ) + with pytest.warns(DeprecationWarning, match=message): + assert hasattr(module, "MetatomicCalculator") + + message = ( + "Importing SymmetrizedCalculator from metatomic.torch.ase_calculator is " + "deprecated and will be removed in a future release. Please import from " + "metatomic_ase instead." + ) + with pytest.warns(DeprecationWarning, match=message): + assert hasattr(module, "SymmetrizedCalculator") diff --git a/python/metatomic_torchsim/README.md b/python/metatomic_torchsim/README.md index cf474c860..f63e75341 100644 --- a/python/metatomic_torchsim/README.md +++ b/python/metatomic_torchsim/README.md @@ -1,6 +1,7 @@ -# metatomic-torchsim +# `metatomic-torchsim` -TorchSim integration for metatomic models. +[TorchSim](https://torchsim.github.io/torch-sim/) integration for metatomic +models. This package allows you to wrap metatomic models as TorchSim `ModelInterface` instances, enabling their use in TorchSim molecular dynamics and other diff --git a/tox.ini b/tox.ini index 7f0c85c46..b203e5815 100644 --- a/tox.ini +++ b/tox.ini @@ -10,8 +10,10 @@ envlist = torch-tests-cxx torch-install-tests-cxx docs-tests + ase-tests torchsim-tests + [testenv] passenv = * setenv = @@ -20,8 +22,7 @@ setenv = # store code coverage in a per-env file, so different envs don't override each other COVERAGE_FILE={env_dir}/.coverage -package = external -package_env = build-metatomic-torch +package = skip lint_folders = "{toxinidir}/python" "{toxinidir}/setup.py" build_single_wheel = --no-deps --no-build-isolation --check-build-dependencies @@ -38,26 +39,6 @@ metatensor_deps = metatensor-torch >=0.8.0,<0.9 metatensor-operations >=0.4.0,<0.5 -[testenv:build-metatomic-torch] -# note: this is not redundant with the same value in the root [testenv] without this -# one, cmake can not find the MSVC compiler on Windows CI -passenv = * -setenv = - # Do not use the user PYTHONPATH, tox is handling everything - PYTHONPATH= - -description = - Used to only build the wheels which are then re-used by all other environments - requiring metatomic-torch to be installed -deps = - pip - {[testenv]packaging_deps} - {[testenv]metatensor_deps} - torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.10}.* - -commands = - pip wheel python/metatomic_torch {[testenv]build_single_wheel} --wheel-dir {envtmpdir}/dist - ################################################################################ ##### C++ tests setup ##### @@ -65,7 +46,6 @@ commands = [testenv:torch-tests-cxx] description = Run the C++ tests for metatomic-torch -package = skip deps = cmake {[testenv]metatensor_deps} @@ -89,7 +69,6 @@ commands = [testenv:torch-install-tests-cxx] description = Run the C++ tests for metatomic-torch -package = skip deps = cmake {[testenv]metatensor_deps} @@ -138,25 +117,19 @@ commands = description = Run the tests of the metatomic-torch Python package deps = {[testenv]testing_deps} + {[testenv]packaging_deps} {[testenv]metatensor_deps} + torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.10}.* numpy {env:METATOMIC_TESTS_NUMPY_VERSION_PIN} + vesin ase - # for metatensor-lj-test - setuptools-scm - cmake - pip - # for symmetrized calculator - scipy - spglib - # uncomment for testing nvalchemiops integration on GPU (requires Python 3.11+) - # nvalchemi-toolkit-ops changedir = python/metatomic_torch commands = - # use the reference LJ implementation for tests - pip install {[testenv]build_single_wheel} git+https://github.com/metatensor/lj-test@f7401a8 + pip install {[testenv]build_single_wheel} . + pip install {[testenv]build_single_wheel} ../metatomic_ase # Make torch.autograd.gradcheck works with pytest python {toxinidir}/scripts/pytest-dont-rewrite-torch.py @@ -179,46 +152,97 @@ commands = pip wheel python/metatomic_torchsim {[testenv]build_single_wheel} --wheel-dir {envtmpdir}/dist -[testenv:torchsim-tests] -description = Run the tests of the metatomic-torchsim Python package -package = external -package_env = build-metatomic-torch +[testenv:docs-tests] +description = Run the doctests defined in any metatomic package deps = {[testenv]testing_deps} + {[testenv]packaging_deps} {[testenv]metatensor_deps} + torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.10}.* numpy {env:METATOMIC_TESTS_NUMPY_VERSION_PIN} + vesin ase - torch-sim-atomistic + +commands = + pip install {[testenv]build_single_wheel} python/metatomic_torch + pip install {[testenv]build_single_wheel} python/metatomic_ase + + pytest --doctest-modules --pyargs metatomic + + +################################################################################ +##### Engine integration tests ##### +################################################################################ + +[testenv:ase-tests] +description = Run the tests of the metatomic-ase Python package +deps = + {[testenv]testing_deps} + {[testenv]packaging_deps} + {[testenv]metatensor_deps} + + torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.10}.* + numpy {env:METATOMIC_TESTS_NUMPY_VERSION_PIN} + + ase + vesin + nvalchemi-toolkit-ops + # for metatensor-lj-test setuptools-scm - cmake + pip -changedir = python/metatomic_torchsim + # for symmetrized calculator + scipy + spglib + +changedir = python/metatomic_ase commands = - # install metatomic-torchsim - pip install {[testenv]build_single_wheel} {toxinidir}/python/metatomic_torchsim + pip install {[testenv]build_single_wheel} . + pip install {[testenv]build_single_wheel} ../metatomic_torch # use the reference LJ implementation for tests pip install {[testenv]build_single_wheel} git+https://github.com/metatensor/lj-test@f7401a8 - pytest --cov={env_site_packages_dir}/metatomic_torchsim --cov-report= --import-mode=append {posargs} - + pytest --cov={env_site_packages_dir}/metatomic_ase \ + --cov={env_site_packages_dir}/metatomic \ + --cov-report= \ + --import-mode=append \ + {posargs} -[testenv:docs-tests] -description = Run the doctests defined in any metatomic package +[testenv:torchsim-tests] +description = Run the tests of the metatomic-torchsim Python package +package_env = build-metatomic-torch deps = {[testenv]testing_deps} + {[testenv]packaging_deps} + {[testenv]metatensor_deps} torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.10}.* numpy {env:METATOMIC_TESTS_NUMPY_VERSION_PIN} + vesin ase + torch-sim-atomistic + # for metatensor-lj-test + setuptools-scm + cmake + +changedir = python/metatomic_torchsim commands = - pytest --doctest-modules --pyargs metatomic + pip install {[testenv]build_single_wheel} . + pip install {[testenv]build_single_wheel} ../metatomic_torch + + # use the reference LJ implementation for tests + pip install {[testenv]build_single_wheel} git+https://github.com/metatensor/lj-test@f7401a8 + + pytest --cov={env_site_packages_dir}/metatomic_torchsim --cov-report= --import-mode=append {posargs} + + ################################################################################ ##### Linter and formatter setup ##### @@ -226,7 +250,6 @@ commands = [testenv:lint] description = Run linters and type checks -package = skip deps = ruff @@ -237,7 +260,6 @@ commands = [testenv:format] description = Abuse tox to do actual formatting on all files. -package = skip deps = ruff commands = @@ -254,6 +276,7 @@ setenv = # build the docs against the CPU only version of torch PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cpu {env:PIP_EXTRA_INDEX_URL:} deps = + {[testenv]metatensor_deps} {[testenv]packaging_deps} {[testenv]testing_deps} @@ -276,9 +299,12 @@ deps = # required for examples ase + vesin chemiscope commands = - pip install {[testenv]build_single_wheel} {toxinidir}/python/metatomic_torchsim + pip install {[testenv]build_single_wheel} python/metatomic_torch + pip install {[testenv]build_single_wheel} python/metatomic_ase + pip install {[testenv]build_single_wheel} python/metatomic_torchsim sphinx-build -d docs/build/doctrees -W -b html docs/src docs/build/html