diff --git a/.gitattributes b/.gitattributes index 176a458f9..7d0ace6fc 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,3 @@ * text=auto + +regtests/**/*.lock linguist-generated=true diff --git a/.github/workflows/regtests.yml b/.github/workflows/regtests.yml new file mode 100644 index 000000000..41e960e8f --- /dev/null +++ b/.github/workflows/regtests.yml @@ -0,0 +1,46 @@ +name: Regression tests + +on: + push: + branches: [main] + pull_request: + # Check all PR + +concurrency: + group: regtests-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} + +jobs: + regtests: + runs-on: ubuntu-24.04 + name: regression tests + steps: + - uses: actions/checkout@v7 + + - uses: actions/setup-python@v7 + with: + python-version: "3.13" + + - name: setup sccache + uses: mozilla-actions/sccache-action@v0.0.11 + with: + version: "v0.10.0" + + - name: cache the models + uses: actions/cache@v6 + with: + path: regtests/cache + key: metatomic-regtests-${{ hashFiles('regtests/models.lock') }} + restore-keys: metatomic-regtests- + + - name: install python dependencies + run: python -m pip install tox + + - name: run regression tests + run: tox -e regtests + env: + PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu + SCCACHE_GHA_ENABLED: "true" + RUSTC_WRAPPER: sccache + CMAKE_C_COMPILER_LAUNCHER: sccache + CMAKE_CXX_COMPILER_LAUNCHER: sccache diff --git a/pyproject.toml b/pyproject.toml index 88dc392b9..f470ed72c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,6 +85,7 @@ known-first-party = [ "metatomic_ase", "metatomic_torchsim", "metatomic_lj_test", + "metatomic_regtest", ] known-third-party = ["torch"] diff --git a/regtests/README.md b/regtests/README.md new file mode 100644 index 000000000..c769f7795 --- /dev/null +++ b/regtests/README.md @@ -0,0 +1,96 @@ +# Metatomic ecosystem regression tests + +These tests run models from the metatomic ecosystem (uPET/PET-MAD, FlashMD, +ShiftML, …) that were **exported by previous versions of metatomic**, and check +that they still load and still produce the same outputs. + +You can run the tests with: + +```bash +tox -e regtests # all of them +tox -e regtests -- -k pet-mad-s-v1.5.0 # a single case +``` + +### Adding a model + +`models.json` says *what* to build: the packages a model needs, and the code +that writes it out. Replace `` with the name of the model. + +```json +"": { + "dependencies": ["upet"], + "source": "import upet; upet.save_upet(model='pet-mad', size='s', version='1.5.0')" +} +``` + +The `source` runs in an empty directory and must write exactly one `.pt` file; +its name does not matter. + +```bash +python export-models.py --list +python export-models.py +python export-models.py --relock # re-resolve the dependencies +``` + +This creates a virtual environment under `cache/venvs/`, installs the +dependencies in it, runs the ``source``, stores the model in +`build/.pt`, and records in `models.lock` the exact version of *every* +package that was installed. + +Once a model is in `models.lock`, building it again reuses those pinned versions +rather than resolving them afresh, so the same model can be rebuilt later even +as the upstream projects move on. Pass `--relock` to deliberately move to a +newer set of dependencies. + +The next step is to upload the model so it can be used for tests. This requires +write access to the https://huggingface.co/metatensor/metatomic-regtests +repository. The address of the uploaded model is recorded in `models.lock`. + +```bash +python export-models.py --upload +``` + +### Adding a test case + +Create `references//input.json`: + +```json +{ + "model": "pet-mad-s-v1.5.0", + "length_unit": "angstrom", + "systems": [ + { + "types": [14, ...], + "positions": [ + [1.0, 2.0, 3.0], + [...] + ], + "cell": [ + [9.0, 0.0, 0.0], + [0.0, 9.0, 0.0], + [0.0, 0.0, 9.0] + ], + "pbc": [true, false, true] + } + ], + "selected_atoms": null, + "outputs": { + "energy": { + "unit": "eV", + "sample_kind": "system", + "gradients": ["positions", "strain"], + "reference": "energy.mts", + "rtol": 1e-6, + "atol": 1e-6 + } + } +} + +``` + +Then generate the reference outputs: + +```bash +python update-references.py --list +python update-references.py +``` diff --git a/regtests/cache/.gitignore b/regtests/cache/.gitignore new file mode 100644 index 000000000..d6b7ef32c --- /dev/null +++ b/regtests/cache/.gitignore @@ -0,0 +1,2 @@ +* +!.gitignore diff --git a/regtests/export-models.py b/regtests/export-models.py new file mode 100644 index 000000000..405d029f8 --- /dev/null +++ b/regtests/export-models.py @@ -0,0 +1,523 @@ +#!/usr/bin/env python3 +""" +Build the models used by the regression tests, and lock the environment used to build +them. + + python regtests/export-models.py --list + python regtests/export-models.py pet-mad-s-v1.5.0 + python regtests/export-models.py --all + python regtests/export-models.py --all --relock + python regtests/export-models.py --all --upload + python regtests/export-models.py --check + +``--check`` verifies, without building anything, that every model is built from the +current ``models.json`` and can be downloaded. This is what the test suite runs before +the tests themselves. + +``--upload`` publishes the models it built to HuggingFace and records their address in +``models.lock``, so that everyone else can run the tests without building anything. It +needs the ``huggingface_hub`` package and write access to the repository, through +``HF_TOKEN`` or a previous ``hf auth login``. +""" + +import argparse +import hashlib +import json +import os +import shutil +import subprocess +import sys +import tempfile +import urllib.error +import urllib.request + + +ROOT = os.path.dirname(os.path.abspath(__file__)) +MODELS_PATH = os.path.join(ROOT, "models.json") +LOCK_PATH = os.path.join(ROOT, "models.lock") +# Models are built here, and *not* in the `cache/` the tests download into so both CI +# and local tests use the exact same file. +BUILD_PATH = os.path.join(ROOT, "build") +# one subdirectory per kind of virtualenv, so that each script only ever prunes the +# ones it made: update-references.py keeps its own under build/venvs/references +VENVS_PATH = os.path.join(BUILD_PATH, "venvs", "models") + +# torch publishes CPU-only wheels on its own index, and they are much smaller +TORCH_CPU_INDEX = "https://download.pytorch.org/whl/cpu" + +# where `--upload` publishes the models +DEFAULT_HF_REPO = "metatensor/metatomic-regtests" + + +def save_lock(data: dict): + with open(LOCK_PATH, "w") as fd: + json.dump(data, fd, indent=2, sort_keys=True) + fd.write("\n") + + +def sha256(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as fd: + while chunk := fd.read(1024 * 1024): + digest.update(chunk) + + return digest.hexdigest() + + +def venv_python(directory: str) -> str: + if os.name == "nt": + return os.path.join(directory, "Scripts", "python.exe") + + return os.path.join(directory, "bin", "python") + + +def create_venv(name: str, dependencies: list, python: str) -> str: + """ + Create the virtual environment used to build ``name``, and install + ``dependencies`` in it. + + The environment is recreated from scratch every time: a leftover environment from a + previous set of dependencies would silently change what ends up in the model. + """ + directory = os.path.join(VENVS_PATH, name) + if os.path.exists(directory): + shutil.rmtree(directory) + + os.makedirs(VENVS_PATH, exist_ok=True) + print(f"+++++ creating a virtualenv in {directory}") + subprocess.run([python, "-m", "venv", directory], check=True) + + environment = dict(os.environ) + environment.setdefault("PIP_EXTRA_INDEX_URL", TORCH_CPU_INDEX) + + pip = [venv_python(directory), "-m", "pip", "--disable-pip-version-check"] + subprocess.run( + pip + ["--quiet", "install", "--upgrade", "pip"], check=True, env=environment + ) + + print("+++++ installing dependencies") + subprocess.run(pip + ["install"] + list(dependencies), check=True, env=environment) + + return directory + + +def freeze(directory: str) -> list: + """The exact version of everything installed in the environment.""" + output = subprocess.run( + [venv_python(directory), "-m", "pip", "freeze", "--disable-pip-version-check"], + check=True, + capture_output=True, + text=True, + ) + + return sorted(line for line in output.stdout.splitlines() if line.strip()) + + +# Patch the metatomic export to always record the same date, so that the checksum of a +# model does not change every time it is rebuilt. +FROZEN_DATE = "1970-01-01T00:00:00+00:00" + +PRELUDE = f""" +import datetime as _datetime +import types as _types + +import metatomic.torch.model as _model + + +_FROZEN = _datetime.datetime.fromisoformat({FROZEN_DATE!r}) + + +class _FrozenDatetime(_datetime.datetime): + @classmethod + def now(cls, tz=None): + return _FROZEN if tz is None else _FROZEN.astimezone(tz) + + +# `metatomic.torch.model` does `import datetime` and then `datetime.datetime.now(...)`, +# so replacing the module it sees is enough to freeze the date it records. +_model.datetime = _types.SimpleNamespace( + datetime=_FrozenDatetime, timezone=_datetime.timezone +) +""" + + +def build_model(directory: str, source: str) -> str: + """ + Run the ``source`` code of a model inside the given environment, and return the path + of the file it produced. + + The code runs in an empty directory, so that the model it writes can be picked up + without every project having to agree on a file name. + """ + workdir = tempfile.mkdtemp(prefix="metatomic-regtests-") + try: + script = os.path.join(workdir, "build_model.py") + with open(script, "w") as fd: + fd.write(PRELUDE) + fd.write("\n") + fd.write(source) + + print("+++++ running the model export") + environment = dict(os.environ) + # keep TorchScript from emitting class attributes in a random order + environment["PYTHONHASHSEED"] = "0" + # make sure the current virtualenv is available in PATH for packages that spawn + # subprocesses + environment["PATH"] = ( + os.path.dirname(venv_python(directory)) + os.pathsep + environment["PATH"] + ) + subprocess.run( + [venv_python(directory), "build_model.py"], + check=True, + cwd=workdir, + env=environment, + ) + + produced = [f for f in os.listdir(workdir) if f.endswith(".pt")] + if len(produced) != 1: + raise RuntimeError( + f"expected the source to write exactly one '.pt' file, got {produced}. " + f"Make the 'source' in models.json write a single model file." + ) + + # move it out of the temporary directory before it is deleted + path = os.path.join(BUILD_PATH, "building.pt") + os.makedirs(BUILD_PATH, exist_ok=True) + shutil.move(os.path.join(workdir, produced[0]), path) + return path + finally: + shutil.rmtree(workdir, ignore_errors=True) + + +def export( + name: str, models_json: dict, locked: dict, relock: bool, python: str +) -> dict: + print(f"+++++ {name}:") + + if locked and not relock and locked.get("spec") == models_json: + # nothing changed and we can reuse the previous build + built = os.path.join(BUILD_PATH, locked["sha256"] + ".pt") + if os.path.exists(built) and sha256(built) == locked["sha256"]: + print("+++++ already built from the current models.json") + return locked + + if locked and not relock and "spec" in locked: + # rebuild exactly what the lock recorded + dependencies = locked["dependencies"] + spec = locked["spec"] + print(f"+++++ using the {len(dependencies)} dependencies locked in models.lock") + + if spec != models_json: + print("+++++ warning: models.json has changed since this model was built,") + print( + "+++++ pass --relock to build what models.json now describes" + ) + else: + # build the model from the current models.json, and record what dependencies + # were used + spec = models_json + dependencies = spec["dependencies"] + if locked and not relock: + print( + "+++++ models.lock predates the 'spec' field, re-resolving dependencies" + ) + elif locked: + print("+++++ re-resolving the dependencies (--relock)") + + directory = create_venv(name, dependencies, python) + installed = freeze(directory) + path = build_model(directory, "\n".join(spec["source"])) + + digest = sha256(path) + final = os.path.join(BUILD_PATH, f"{digest}.pt") + os.replace(path, final) + + # a URL describes the model that was uploaded to it, so it is only worth keeping + # while the model is unchanged. Keeping it next to a different checksum would make + # the lock describe something that does not exist: downloading from that URL then + # fails the checksum check. + unchanged = bool(locked) and locked["sha256"] == digest + + entry = { + "sha256": digest, + "size": os.path.getsize(final), + "url": locked.get("url") if unchanged else None, + "python": subprocess.run( + [ + venv_python(directory), + "-c", + "import platform; print(platform.python_version())", + ], + check=True, + capture_output=True, + text=True, + ).stdout.strip(), + # the models.json entry this was built from, so that we can tell whether + # models.json has moved on since + "spec": spec, + # what `pip freeze` reported, unlike the requirements in `spec` + "dependencies": installed, + } + + if locked and locked["sha256"] != digest: + print( + f"+++++ the model changed: {locked['sha256'][:12]}… -> {digest[:12]}…", + file=sys.stderr, + ) + if locked["dependencies"] == installed: + print( + "+++++ the dependencies are unchanged, so either the model does\n" + "+++++ not serialize deterministically, or an upstream package\n" + "+++++ changed without its version changing", + file=sys.stderr, + ) + else: + added = set(installed) - set(locked["dependencies"]) + removed = set(locked["dependencies"]) - set(installed) + for dependency in sorted(removed): + print(f"+++++ - {dependency}") + for dependency in sorted(added): + print(f"+++++ + {dependency}") + + if locked.get("url") is not None: + print( + "+++++ dropping the 'url' in models.lock, it pointed at the\n" + "+++++ previous model: re-run with --upload to publish this one", + file=sys.stderr, + ) + + print(f"+++++ wrote {os.path.relpath(final, ROOT)} ({entry['size'] / 1e6:.1f} MB)") + + return entry + + +def upload(name: str, path: str, repo: str) -> str: + """ + Publish a model on HuggingFace, and return the URL to download it from. + + The URL points at the commit this upload creates rather than at a branch, so that it + keeps referring to this exact model even once a newer one is uploaded over it. + """ + try: + from huggingface_hub import upload_file + except ImportError: + raise RuntimeError( + "--upload requires the `huggingface_hub` package, install it with " + "`pip install huggingface_hub`" + ) from None + + path_in_repo = f"{name}.pt" + print(f"+++++ uploading to {repo}/{path_in_repo}") + + commit = upload_file( + path_or_fileobj=path, + path_in_repo=path_in_repo, + repo_id=repo, + repo_type="model", + commit_message=f"{name} ({sha256(path)[:12]})", + ) + + return f"https://huggingface.co/{repo}/resolve/{commit.oid}/{path_in_repo}" + + +def downloadable(url: str, size: int): + """Check that ``url`` can be downloaded, without downloading it.""" + request = urllib.request.Request(url, method="HEAD") + try: + with urllib.request.urlopen(request, timeout=30) as response: + length = response.headers.get("Content-Length") + except (urllib.error.URLError, TimeoutError, OSError) as e: + return f"can not be downloaded from {url}: {e}" + + if length is not None and int(length) != size: + return ( + f"the file at {url} is {length} bytes, but models.lock says it should be " + f"{size}: the URL points at a different model" + ) + + return None + + +def check(name: str, models_json: dict, locked: dict) -> list: + """ + Everything that stops the tests from using ``name``, as ``(fatal, message)`` pairs. + + This only reads ``models.lock`` and asks whether the model can be downloaded; it + never builds anything. + """ + if locked is None: + return [ + ( + True, + f"has never been built: run `python regtests/export-models.py {name}`", + ) + ] + + problems = [] + if "spec" not in locked: + problems.append( + ( + True, + "models.lock predates the 'spec' field, so there is no way to tell " + "what it was built from: " + f"run `python regtests/export-models.py {name} --relock`", + ) + ) + elif locked["spec"] != models_json: + problems.append( + ( + True, + "models.json has changed since it was built: run " + f"`python regtests/export-models.py {name} --relock`", + ) + ) + + if locked.get("url") is None: + problems.append( + ( + True, + "has not been published, so nothing can download it: run " + f"`python regtests/export-models.py {name} --upload`", + ) + ) + else: + error = downloadable(locked["url"], locked["size"]) + if error is not None: + problems.append((True, error)) + + return problems + + +def prune_lock(locked: dict, models: dict): + """remove entries from models.lock that are no in models.json""" + for name in sorted(set(locked) - set(models)): + print(f"+++++ removing {name} from models.lock since it is not in models.json") + del locked[name] + + +def prune_builds(locked: dict): + """ + Delete what the lock no longer refers to: models left behind by a rebuild that + changed one, and the virtualenvs of models that are gone. + """ + keep = {entry["sha256"] + ".pt" for entry in locked.values()} + for name in sorted(os.listdir(BUILD_PATH)) if os.path.exists(BUILD_PATH) else []: + if not name.endswith(".pt") or name in keep: + continue + + os.unlink(os.path.join(BUILD_PATH, name)) + + for name in sorted(os.listdir(VENVS_PATH)) if os.path.exists(VENVS_PATH) else []: + if name not in locked: + shutil.rmtree(os.path.join(VENVS_PATH, name), ignore_errors=True) + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("models", nargs="*", help="models to build") + parser.add_argument("--all", action="store_true", help="build every model") + parser.add_argument("--list", action="store_true", help="list the known models") + parser.add_argument( + "--check", + action="store_true", + help="check that the models are built from the current models.json and can be " + "downloaded, without building anything (every model, unless some are named)", + ) + parser.add_argument( + "--relock", + action="store_true", + help="re-resolve the dependencies instead of using the locked ones", + ) + parser.add_argument( + "--python", + default=sys.executable, + help="python interpreter used to create the virtualenvs", + ) + parser.add_argument( + "--upload", + action="store_true", + help="publish the models on HuggingFace and record their URL in models.lock", + ) + parser.add_argument( + "--repo", + default=DEFAULT_HF_REPO, + help=f"HuggingFace repository to upload to (default: {DEFAULT_HF_REPO})", + ) + args = parser.parse_args() + + # show the messages from this process and any subprocess in the correct order + sys.stdout.reconfigure(line_buffering=True) + + with open(MODELS_PATH) as fd: + models = json.load(fd) + + if not os.path.exists(LOCK_PATH): + locked = {} + else: + with open(LOCK_PATH) as fd: + locked = json.load(fd) + + if args.list: + for name in sorted(models): + print(name) + return 0 + + if args.all or (args.check and not args.models): + names = sorted(models) + elif args.models: + names = args.models + else: + parser.error("give at least one model, or --all") + + for name in names: + if name not in models: + known = "\n".join(f" - {key}" for key in sorted(models)) + parser.error(f"unknown model '{name}', models.json contains:\n{known}") + + if args.check: + failed = False + for name in names: + for fatal, message in check(name, models[name], locked.get(name)): + failed = failed or fatal + label = "ERROR" if fatal else "warning" + print(f"{label}: {name}: {message}", file=sys.stderr) + + if failed: + print( + "\nsome models are not usable, see above", + file=sys.stderr, + ) + return 1 + + print(f"+++++ all {len(names)} models are up to date and available") + return 0 + + for name in names: + entry = locked.get(name) + locked[name] = export(name, models[name], entry, args.relock, args.python) + + if args.upload: + built = locked[name] + already = entry is not None and entry.get("url") is not None + if already and entry["sha256"] == built["sha256"]: + print("+++++ the model is unchanged, keeping the URL it already has") + else: + built["url"] = upload( + name, os.path.join(BUILD_PATH, built["sha256"] + ".pt"), args.repo + ) + + save_lock(locked) + + prune_lock(locked, models) + save_lock(locked) + prune_builds(locked) + + print(f"\n+++++ updated {os.path.relpath(LOCK_PATH, ROOT)}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regtests/metatomic_regtest/__init__.py b/regtests/metatomic_regtest/__init__.py new file mode 100644 index 000000000..66aa4fe86 --- /dev/null +++ b/regtests/metatomic_regtest/__init__.py @@ -0,0 +1,9 @@ +from .engine import load_expected, load_input, load_model, run_model + + +__all__ = [ + "load_expected", + "load_input", + "load_model", + "run_model", +] diff --git a/regtests/metatomic_regtest/engine.py b/regtests/metatomic_regtest/engine.py new file mode 100644 index 000000000..5dbc9dcda --- /dev/null +++ b/regtests/metatomic_regtest/engine.py @@ -0,0 +1,379 @@ +"""A minimal engine, used to run the frozen models of the regression tests""" + +import json +import os +from collections import namedtuple +from typing import Dict, List + +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap +from vesin.metatomic import neighbor_lists_for_model + +from metatomic.torch import ( + AtomisticModel, + ModelEvaluationOptions, + ModelOutput, + System, + load_atomistic_model, +) + +from .models import model_path + + +Input = namedtuple( + "Input", + [ + "path", + "directory", + "model", + "length_unit", + "systems", + "selected_atoms", + "outputs", + ], +) + +STR_TO_DTYPE = {"float32": torch.float32, "float64": torch.float64} + + +def load_input(reference: str) -> Input: + """ + Read the system, requested outputs, and expected values of the outputs for the given + reference. + """ + with open(reference) as fd: + data = json.load(fd) + + selected_atoms = data.get("selected_atoms") + if selected_atoms is not None: + selected_atoms = Labels( + ["system", "atom"], torch.tensor(selected_atoms, dtype=torch.int32) + ) + + return Input( + path=reference, + directory=os.path.dirname(reference), + model=data["model"], + length_unit=data["length_unit"], + systems=data["systems"], + selected_atoms=selected_atoms, + outputs=data["outputs"], + ) + + +def load_expected(input: Input) -> Dict[str, TensorMap]: + """Read the reference outputs stored next to the given input.""" + import metatensor.torch as mts + + expected = {} + for name, output in input.outputs.items(): + expected[name] = mts.load(os.path.join(input.directory, output["reference"])) + + return expected + + +def load_model(reference: str) -> AtomisticModel: + """ + Download & cache the required model for the given reference and then load it. + """ + return load_atomistic_model(model_path(reference)) + + +def _build_systems(input: Input, dtype: torch.dtype): + """ + Create the :py:class:`System` for this input, tracking the tensors we'll later need + to extract gradients from. + + Returns the systems, the per-system positions, and the per-system strain. Each + system gets its own strain tensor: a single shared one would accumulate the + gradients of all the systems into a single value. + """ + gradients = set() + for output in input.outputs.values(): + gradients.update(output.get("gradients", [])) + + systems = [] + all_positions = [] + all_strains = [] + for data in input.systems: + types = torch.tensor(data["types"], dtype=torch.int32) + positions = torch.tensor(data["positions"], dtype=dtype) + cell = torch.tensor(data["cell"], dtype=dtype) + pbc = torch.tensor(data["pbc"], dtype=torch.bool) + + if "positions" in gradients: + positions.requires_grad_(True) + all_positions.append(positions) + + if "strain" in gradients: + strain = torch.eye(3, requires_grad=True, dtype=dtype) + positions = positions @ strain + if "positions" in gradients: + # `positions` is no longer a leaf of the autograd graph + all_positions[-1].retain_grad() + cell = cell @ strain + all_strains.append(strain) + + systems.append(System(types, positions, cell, pbc)) + + return systems, all_positions, all_strains + + +def _add_model_inputs( + input: Input, model: AtomisticModel, systems: List[System], dtype: torch.dtype +): + """ + Attach the extra data the model declares through ``requested_inputs()``, taking it + from the ``inputs`` of each system in ``input.json``. + + The model says whether it wants one value per atom or one per system, and the shape + of the data in ``input.json`` says whether that value is a scalar or a vector. + """ + requested = model.requested_inputs(use_new_names=True) + + for system, data in zip(systems, input.systems, strict=True): + inputs = data.get("inputs", {}) + for name in requested.keys(): + if name not in inputs: + raise ValueError( + f"the model requires the '{name}' input, but a system in " + f'{input.path} does not provide it; add it to the "inputs" of ' + f"that system" + ) + + values = torch.tensor(inputs[name], dtype=dtype) + sample_kind = requested[name].sample_kind + + if sample_kind == "atom": + count = len(system) + samples = Labels( + ["system", "atom"], + torch.tensor([[0, i] for i in range(count)], dtype=torch.int32), + ) + scalar = values.dim() == 1 + vector = values.dim() == 2 and values.shape[1] == 3 + expected = "one value or one vector per atom" + elif sample_kind == "system": + count = 1 + samples = Labels("system", torch.tensor([[0]])) + scalar = values.dim() == 0 + vector = values.dim() == 1 and values.shape[0] == 3 + expected = "one value or one vector for the whole system" + else: + raise ValueError( + f"the model asks for the '{name}' input with sample kind " + f"'{sample_kind}', which this engine does not know how to build" + ) + + if scalar: + components = [] + values = values.reshape(count, 1) + elif vector: + components = [Labels("xyz", torch.tensor([[0], [1], [2]]))] + values = values.reshape(count, 3, 1) + else: + raise ValueError( + f"can not make the '{name}' input from data of shape " + f"{tuple(values.shape)}: expected {expected}" + ) + + block = TensorBlock( + values=values, + samples=samples, + components=components, + properties=Labels(name, torch.tensor([[0]])), + ) + + system.add_data(name, TensorMap(Labels.single(), [block])) + + +def _gradient_block( + block: TensorBlock, + name: str, + values: List[torch.Tensor], + components: List[Labels], +) -> TensorBlock: + """ + Build the gradient :py:class:`TensorBlock` called ``name`` for ``block``, from the + per-sample gradient ``values``. + """ + if name == "positions": + samples = [] + for sample, gradient in enumerate(values): + for atom in range(gradient.shape[0]): + samples.append([sample, atom]) + + samples = Labels(["sample", "atom"], torch.tensor(samples, dtype=torch.int32)) + else: + assert name == "strain" + samples = Labels( + ["sample"], + torch.arange(len(values), dtype=torch.int32).reshape(-1, 1), + ) + + return TensorBlock( + values=torch.concatenate(values), + samples=samples, + components=components, + properties=block.properties, + ) + + +def _compute_gradients( + output: TensorMap, + gradients: List[str], + all_positions: List[torch.Tensor], + all_strains: List[torch.Tensor], +) -> TensorMap: + """ + Compute the requested gradients of ``output`` with autograd, and return a new + :py:class:`TensorMap` with them attached. + """ + if len(output) != 1: + raise ValueError( + "can only compute gradients for a TensorMap with a single block, " + f"got {len(output)} blocks" + ) + + block = output.block() + if block.samples.names != ["system"]: + raise ValueError( + "can only compute gradients for outputs with one sample per system, got " + f"samples named {block.samples.names}" + ) + + n_systems = len(all_positions) + n_properties = len(block.properties) + + xyz = Labels("xyz", torch.tensor([[0], [1], [2]], dtype=torch.int32)) + xyz_1 = Labels("xyz_1", torch.tensor([[0], [1], [2]], dtype=torch.int32)) + xyz_2 = Labels("xyz_2", torch.tensor([[0], [1], [2]], dtype=torch.int32)) + + # one gradient per property, since autograd only differentiates scalars + positions_gradients = [[] for _ in range(n_systems)] + strain_gradients = [[] for _ in range(n_systems)] + for property in range(n_properties): + wrt = [] + if "positions" in gradients: + wrt += all_positions + if "strain" in gradients: + wrt += all_strains + + computed = torch.autograd.grad( + outputs=block.values[:, property].sum(), + inputs=wrt, + retain_graph=property != n_properties - 1, + ) + + computed = list(computed) + if "positions" in gradients: + for system in range(n_systems): + positions_gradients[system].append(computed.pop(0)) + if "strain" in gradients: + for system in range(n_systems): + strain_gradients[system].append(computed.pop(0)) + + new_block = TensorBlock( + values=block.values.detach(), + samples=block.samples, + components=block.components, + properties=block.properties, + ) + + # keep any gradient the model computed itself + for name, gradient in block.gradients(): + new_block.add_gradient( + name, + TensorBlock( + values=gradient.values.detach(), + samples=gradient.samples, + components=gradient.components, + properties=gradient.properties, + ), + ) + + if "positions" in gradients: + # (n_atoms, 3) per property => (n_atoms, 3, n_properties) + values = [ + torch.stack(per_property, dim=-1).detach() + for per_property in positions_gradients + ] + new_block.add_gradient( + "positions", _gradient_block(block, "positions", values, [xyz]) + ) + + if "strain" in gradients: + # (3, 3) per property => (1, 3, 3, n_properties) + values = [ + torch.stack(per_property, dim=-1).reshape(1, 3, 3, n_properties).detach() + for per_property in strain_gradients + ] + new_block.add_gradient( + "strain", _gradient_block(block, "strain", values, [xyz_1, xyz_2]) + ) + + return TensorMap(output.keys, [new_block]) + + +def _detach(tensor: TensorMap) -> TensorMap: + blocks = [] + for block in tensor.blocks(): + new_block = TensorBlock( + values=block.values.detach(), + samples=block.samples, + components=block.components, + properties=block.properties, + ) + for name, gradient in block.gradients(): + new_block.add_gradient( + name, + TensorBlock( + values=gradient.values.detach(), + samples=gradient.samples, + components=gradient.components, + properties=gradient.properties, + ), + ) + blocks.append(new_block) + + return TensorMap(tensor.keys, blocks) + + +def run_model(model: AtomisticModel, input: Input) -> Dict[str, TensorMap]: + capabilities = model.capabilities() + dtype = STR_TO_DTYPE[capabilities.dtype] + + systems, all_positions, all_strains = _build_systems(input, dtype) + + for calculator in neighbor_lists_for_model(input.length_unit, model, skin=0.0): + calculator.add_neighbor_list(systems) + + _add_model_inputs(input, model, systems, dtype) + + options = ModelEvaluationOptions( + length_unit=input.length_unit, + outputs={ + name: ModelOutput( + unit=output.get("unit", ""), + sample_kind=output.get("sample_kind", "system"), + explicit_gradients=[], + ) + for name, output in input.outputs.items() + }, + selected_atoms=input.selected_atoms, + ) + + outputs = model(systems=systems, options=options, check_consistency=True) + + results = {} + for name, output in input.outputs.items(): + with_autograd = [gradient for gradient in output.get("gradients", [])] + + if with_autograd: + results[name] = _compute_gradients( + outputs[name], with_autograd, all_positions, all_strains + ) + else: + results[name] = _detach(outputs[name]) + + return results diff --git a/regtests/metatomic_regtest/models.py b/regtests/metatomic_regtest/models.py new file mode 100644 index 000000000..1c2856700 --- /dev/null +++ b/regtests/metatomic_regtest/models.py @@ -0,0 +1,141 @@ +""" +The models used by the regression tests, and the code to fetch and cache them locally. + +``models.json`` says how to build each model and ``models.lock`` records the exact +environment it was built in; both are handled by ``export-models.py``. This module only +consumes the result: it looks a model up in the lock, and gets its file from the cache, +downloading it if the lock says where it is published. +""" + +import hashlib +import json +import os +import time +import urllib.error +import urllib.request + + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +MODELS_PATH = os.path.join(ROOT, "models.json") +LOCK_PATH = os.path.join(ROOT, "models.lock") +CACHE_PATH = os.path.join(ROOT, "cache") + +# how many times we retry a download before giving up +N_RETRIES = 3 +# read this many bytes at a time when downloading +CHUNK_SIZE = 1024 * 1024 + + +class ModelNotCached(Exception): + """The model is not in the local cache, and there is nowhere to get it from.""" + + +def _sha256(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as fd: + while True: + chunk = fd.read(CHUNK_SIZE) + if not chunk: + break + digest.update(chunk) + + return digest.hexdigest() + + +def _download(url: str, expected_sha256: str, path: str): + """ + Download ``url`` to ``path``, checking that its content matches ``expected_sha256``. + """ + tmp = f"{path}.tmp.{os.getpid()}" + + last_error = None + for attempt in range(N_RETRIES): + if attempt != 0: + # exponential backoff, to absorb transient CI network failures + time.sleep(2**attempt) + + digest = hashlib.sha256() + try: + with urllib.request.urlopen(url, timeout=60) as response: + with open(tmp, "wb") as fd: + while True: + chunk = response.read(CHUNK_SIZE) + if not chunk: + break + digest.update(chunk) + fd.write(chunk) + except (urllib.error.URLError, TimeoutError, OSError) as e: + last_error = e + if os.path.exists(tmp): + os.unlink(tmp) + continue + + actual = digest.hexdigest() + if actual != expected_sha256: + # this is not a transient error, don't retry it + os.unlink(tmp) + raise RuntimeError( + f"checksum mismatch for {url}: expected sha256 {expected_sha256}, " + f"got {actual}. Either the registry is out of date, or the artifact " + f"was modified in place (it should be pinned to an immutable URL)" + ) + + os.replace(tmp, path) + return + + raise RuntimeError(f"failed to download {url} after {N_RETRIES} tries") from ( + last_error + ) + + +def _fetch(url: str, sha256: str, suffix: str) -> str: + """Download the artifact with the given ``sha256`` into the cache.""" + path = os.path.join(CACHE_PATH, sha256 + suffix) + os.makedirs(CACHE_PATH, exist_ok=True) + + _download(url, sha256, path) + return path + + +def model_path(name: str) -> str: + """ + Get the local path of the model called ``name``, downloading it if required. + """ + with open(MODELS_PATH) as fd: + models = json.load(fd) + + if not os.path.exists(LOCK_PATH): + locked = {} + else: + with open(LOCK_PATH) as fd: + locked = json.load(fd) + + if name not in locked: + if name in models: + raise ModelNotCached( + f"'{name}' has not been built yet: run `python export-models.py " + f"{name} --upload` to build and publish it" + ) + + known = "\n".join(f" - {key}" for key in sorted(models)) + raise KeyError(f"unknown model '{name}', models.json contains:\n{known}") + + entry = locked[name] + path = os.path.join(CACHE_PATH, entry["sha256"] + ".pt") + if os.path.exists(path): + actual = _sha256(path) + if actual != entry["sha256"]: + raise RuntimeError( + f"the cached file {path} does not match its own name: it hashes to " + f"{actual}. Delete it and run the tests again." + ) + + return path + + if entry.get("url") is None: + raise ModelNotCached( + f"'{name}' is not in {CACHE_PATH} and models.lock does not say where it is " + f"published: run `python export-models.py {name} --upload`" + ) + + return _fetch(entry["url"], entry["sha256"], suffix=".pt") diff --git a/regtests/models.json b/regtests/models.json new file mode 100644 index 000000000..21fc4a60e --- /dev/null +++ b/regtests/models.json @@ -0,0 +1,50 @@ +{ + "flashmd-pet-omatpes-v2-16fs": { + "dependencies": ["flashmd"], + "source": [ + "import flashmd", + "_, model = flashmd.get_pretrained('pet-omatpes-v2', 16)", + "model.save('flashmd-pet-omatpes-v2-16fs.pt')" + ] + }, + "pet-mad-dos-v1.0": { + "dependencies": ["upet"], + "source": [ + "import upet.ase.dos", + "model = upet.ase.dos.get_pet_mad_dos(version='1.0')", + "model.save('pet-mad-dos-v1.0.pt')" + ] + }, + "pet-mad-featurizer-v1.0.0": { + "dependencies": ["upet"], + "source": [ + "from upet.ase.explore import PETMADFeaturizer", + "featurizer = PETMADFeaturizer(version='1.0.0')", + "featurizer.mad_explorer.save('pet-mad-featurizer-v1.0.0.pt')" + ] + }, + "pet-mad-s-v1.0.2": { + "dependencies": ["upet"], + "source": [ + "import upet", + "upet.save_upet(model='pet-mad', size='s', version='1.0.2')" + ] + }, + "pet-mad-s-v1.5.0": { + "dependencies": ["upet"], + "source": [ + "import upet", + "upet.save_upet(model='pet-mad', size='s', version='1.5.0')" + ] + }, + "shiftml4": { + "dependencies": ["shiftml"], + "source": [ + "import shiftml.ase.calculator as calculator", + "from shiftml.utils.loading import load_model", + "url = calculator.url_resolve['ShiftML41']", + "calculator.download_with_retry(url, 'downloaded')", + "load_model('downloaded').save('shiftml4.pt')" + ] + } +} diff --git a/regtests/models.lock b/regtests/models.lock new file mode 100644 index 000000000..ef87f4493 --- /dev/null +++ b/regtests/models.lock @@ -0,0 +1,445 @@ +{ + "flashmd-pet-omatpes-v2-16fs": { + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "PyYAML==6.0.3", + "annotated-types==0.8.0", + "antlr4-python3-runtime==4.9.3", + "anyio==4.15.1", + "ase==3.29.0", + "attrs==26.1.0", + "click==8.5.0", + "contourpy==1.4.0", + "ctypes-dlpack==0.1.0", + "cycler==0.12.1", + "filelock==4.0.3", + "flashmd==0.3.0", + "fonttools==4.66.0", + "fsspec==2026.9.0", + "h11==0.16.0", + "hf-xet==1.6.0", + "httpcore2==2.13.1", + "httpx2==2.13.1", + "huggingface_hub==2.0.0", + "idna==3.20", + "jsonschema-specifications==2025.9.1", + "jsonschema==4.26.0", + "kiwisolver==1.5.1", + "matplotlib==3.11.2", + "metatensor-core==0.2.5", + "metatensor-learn==0.5.0", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "metatrain==2026.3.1", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "omegaconf==2.3.1", + "packaging==26.3", + "pillow==12.3.0", + "pydantic==2.13.5", + "pydantic_core==2.46.5", + "pyparsing==3.3.3", + "python-dateutil==2.9.0.post0", + "python_hostlist==2.3.0", + "referencing==0.37.0", + "rpds-py==2026.6.3", + "scipy==1.18.1", + "setuptools==84.0.0", + "six==1.17.0", + "sympy==1.14.0", + "torch==2.13.0", + "tqdm==4.70.1", + "truststore==0.10.4", + "typing-inspection==0.4.4", + "typing_extensions==4.16.0", + "vesin==0.6.1", + "wigners==0.4.1" + ], + "python": "3.14.4", + "sha256": "caa9114c76636687a91f703bfbd61e190b398b2bd4931d1bfd5d61c870c5f7de", + "size": 284189133, + "spec": { + "dependencies": [ + "flashmd" + ], + "source": [ + "import flashmd", + "_, model = flashmd.get_pretrained('pet-omatpes-v2', 16)", + "model.save('flashmd-pet-omatpes-v2-16fs.pt')" + ] + }, + "url": "https://huggingface.co/metatensor/metatomic-regtests/resolve/5ba8c150ff7fafd5fe7d6370cdb0d99b971c25a4/flashmd-pet-omatpes-v2-16fs.pt" + }, + "pet-mad-dos-v1.0": { + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "PyYAML==6.0.3", + "annotated-types==0.8.0", + "antlr4-python3-runtime==4.9.3", + "anyio==4.15.1", + "ase==3.29.0", + "attrs==26.1.0", + "click==8.5.0", + "contourpy==1.4.0", + "ctypes-dlpack==0.1.0", + "cycler==0.12.1", + "filelock==4.0.3", + "fonttools==4.66.0", + "fsspec==2026.9.0", + "h11==0.16.0", + "hf-xet==1.6.0", + "httpcore2==2.13.1", + "httpx2==2.13.1", + "huggingface_hub==2.0.0", + "idna==3.20", + "jsonschema-specifications==2025.9.1", + "jsonschema==4.26.0", + "kiwisolver==1.5.1", + "matplotlib==3.11.2", + "metatensor-core==0.2.5", + "metatensor-learn==0.5.0", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "metatrain==2026.4", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "nvalchemi-toolkit-ops==0.4.1", + "omegaconf==2.3.1", + "packaging==26.3", + "pillow==12.3.0", + "platformdirs==4.11.13", + "pydantic==2.13.5", + "pydantic_core==2.46.5", + "pyparsing==3.3.3", + "python-dateutil==2.9.0.post0", + "python_hostlist==2.3.0", + "referencing==0.37.0", + "rpds-py==2026.6.3", + "scipy==1.18.1", + "setuptools==84.0.0", + "six==1.17.0", + "sympy==1.14.0", + "torch==2.13.0", + "tqdm==4.70.1", + "truststore==0.10.4", + "typing-inspection==0.4.4", + "typing_extensions==4.16.0", + "upet==0.3.1", + "vesin==0.6.1", + "warp-lang==1.17.0", + "wigners==0.4.1" + ], + "python": "3.14.4", + "sha256": "95d98632b2905f88630867e3768604ad76320278e1317c600637f21ee60624b4", + "size": 32220203, + "spec": { + "dependencies": [ + "upet" + ], + "source": [ + "import upet.ase.dos", + "model = upet.ase.dos.get_pet_mad_dos(version='1.0')", + "model.save('pet-mad-dos-v1.0.pt')" + ] + }, + "url": "https://huggingface.co/metatensor/metatomic-regtests/resolve/4ce441371bcab0f3d98cf02c8424ada9bfd24a1f/pet-mad-dos-v1.0.pt" + }, + "pet-mad-featurizer-v1.0.0": { + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "PyYAML==6.0.3", + "annotated-types==0.8.0", + "antlr4-python3-runtime==4.9.3", + "anyio==4.15.1", + "ase==3.29.0", + "attrs==26.1.0", + "click==8.5.0", + "contourpy==1.4.0", + "ctypes-dlpack==0.1.0", + "cycler==0.12.1", + "filelock==4.0.3", + "fonttools==4.66.0", + "fsspec==2026.9.0", + "h11==0.16.0", + "hf-xet==1.6.0", + "httpcore2==2.13.1", + "httpx2==2.13.1", + "huggingface_hub==2.0.0", + "idna==3.20", + "jsonschema-specifications==2025.9.1", + "jsonschema==4.26.0", + "kiwisolver==1.5.1", + "matplotlib==3.11.2", + "metatensor-core==0.2.5", + "metatensor-learn==0.5.0", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "metatrain==2026.4", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "nvalchemi-toolkit-ops==0.4.1", + "omegaconf==2.3.1", + "packaging==26.3", + "pillow==12.3.0", + "platformdirs==4.11.13", + "pydantic==2.13.5", + "pydantic_core==2.46.5", + "pyparsing==3.3.3", + "python-dateutil==2.9.0.post0", + "python_hostlist==2.3.0", + "referencing==0.37.0", + "rpds-py==2026.6.3", + "scipy==1.18.1", + "setuptools==84.0.0", + "six==1.17.0", + "sympy==1.14.0", + "torch==2.13.0", + "tqdm==4.70.1", + "truststore==0.10.4", + "typing-inspection==0.4.4", + "typing_extensions==4.16.0", + "upet==0.3.1", + "vesin==0.6.1", + "warp-lang==1.17.0", + "wigners==0.4.1" + ], + "python": "3.14.4", + "sha256": "24c57354b9c92ecce469520ccc6fe2f87035e4116b43607f4808df2b23b960a8", + "size": 16290655, + "spec": { + "dependencies": [ + "upet" + ], + "source": [ + "from upet.ase.explore import PETMADFeaturizer", + "featurizer = PETMADFeaturizer(version='1.0.0')", + "featurizer.mad_explorer.save('pet-mad-featurizer-v1.0.0.pt')" + ] + }, + "url": "https://huggingface.co/metatensor/metatomic-regtests/resolve/2dda0ca820d4e50060fd76606a6bdfaeeb6a5416/pet-mad-featurizer-v1.0.0.pt" + }, + "pet-mad-s-v1.0.2": { + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "PyYAML==6.0.3", + "annotated-types==0.8.0", + "antlr4-python3-runtime==4.9.3", + "anyio==4.15.1", + "ase==3.29.0", + "attrs==26.1.0", + "certifi==2026.7.22", + "click==8.5.0", + "contourpy==1.4.0", + "ctypes-dlpack==0.1.0", + "cycler==0.12.1", + "filelock==4.0.1", + "fonttools==4.65.0", + "fsspec==2026.9.0", + "h11==0.16.0", + "hf-xet==1.6.0", + "httpcore==1.0.9", + "httpx==0.28.1", + "huggingface_hub==1.32.0", + "idna==3.20", + "jsonschema-specifications==2025.9.1", + "jsonschema==4.26.0", + "kiwisolver==1.5.1", + "matplotlib==3.11.2", + "metatensor-core==0.2.5", + "metatensor-learn==0.5.0", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "metatrain==2026.4", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "nvalchemi-toolkit-ops==0.4.1", + "omegaconf==2.3.1", + "packaging==26.3", + "pillow==12.3.0", + "platformdirs==4.11.12", + "pydantic==2.13.5", + "pydantic_core==2.46.5", + "pyparsing==3.3.3", + "python-dateutil==2.9.0.post0", + "python_hostlist==2.3.0", + "referencing==0.37.0", + "rpds-py==2026.6.3", + "scipy==1.18.1", + "setuptools==84.0.0", + "six==1.17.0", + "sympy==1.14.0", + "torch==2.13.0", + "tqdm==4.70.1", + "typing-inspection==0.4.4", + "typing_extensions==4.16.0", + "upet==0.3.0", + "vesin==0.6.1", + "warp-lang==1.17.0", + "wigners==0.4.1" + ], + "python": "3.13.15", + "sha256": "0289d01a3be9af0c95ba3f2dae85f0a2162a7b3f18c4ea33ac664c0c4ff94722", + "size": 14226761, + "spec": { + "dependencies": [ + "upet" + ], + "source": [ + "import upet", + "upet.save_upet(model='pet-mad', size='s', version='1.0.2')" + ] + }, + "url": "https://huggingface.co/metatensor/metatomic-regtests/resolve/8e38cf611b0da4f71391e884b4aa878c20c28d01/pet-mad-s-v1.0.2.pt" + }, + "pet-mad-s-v1.5.0": { + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "PyYAML==6.0.3", + "annotated-types==0.8.0", + "antlr4-python3-runtime==4.9.3", + "anyio==4.15.1", + "ase==3.29.0", + "attrs==26.1.0", + "certifi==2026.7.22", + "click==8.5.0", + "contourpy==1.4.0", + "ctypes-dlpack==0.1.0", + "cycler==0.12.1", + "filelock==4.0.1", + "fonttools==4.65.0", + "fsspec==2026.9.0", + "h11==0.16.0", + "hf-xet==1.6.0", + "httpcore==1.0.9", + "httpx==0.28.1", + "huggingface_hub==1.32.0", + "idna==3.20", + "jsonschema-specifications==2025.9.1", + "jsonschema==4.26.0", + "kiwisolver==1.5.1", + "matplotlib==3.11.2", + "metatensor-core==0.2.5", + "metatensor-learn==0.5.0", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "metatrain==2026.4", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "nvalchemi-toolkit-ops==0.4.1", + "omegaconf==2.3.1", + "packaging==26.3", + "pillow==12.3.0", + "platformdirs==4.11.12", + "pydantic==2.13.5", + "pydantic_core==2.46.5", + "pyparsing==3.3.3", + "python-dateutil==2.9.0.post0", + "python_hostlist==2.3.0", + "referencing==0.37.0", + "rpds-py==2026.6.3", + "scipy==1.18.1", + "setuptools==84.0.0", + "six==1.17.0", + "sympy==1.14.0", + "torch==2.13.0", + "tqdm==4.70.1", + "typing-inspection==0.4.4", + "typing_extensions==4.16.0", + "upet==0.3.0", + "vesin==0.6.1", + "warp-lang==1.17.0", + "wigners==0.4.1" + ], + "python": "3.13.15", + "sha256": "4045d466112e3a6750d2b430fa3409502286df1aadd75a5d2e10cac96798b306", + "size": 110815710, + "spec": { + "dependencies": [ + "upet" + ], + "source": [ + "import upet", + "upet.save_upet(model='pet-mad', size='s', version='1.5.0')" + ] + }, + "url": "https://huggingface.co/metatensor/metatomic-regtests/resolve/1943b1b663489a71d0a45488e71d1815f0bc744a/pet-mad-s-v1.5.0.pt" + }, + "shiftml4": { + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "ase==3.29.0", + "certifi==2026.7.22", + "charset-normalizer==3.5.1", + "contourpy==1.4.0", + "ctypes-dlpack==0.1.0", + "cycler==0.12.1", + "filelock==4.0.3", + "fonttools==4.66.0", + "fsspec==2026.9.0", + "idna==3.20", + "kiwisolver==1.5.1", + "matplotlib==3.11.2", + "metatensor-core==0.2.5", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "packaging==26.3", + "pillow==12.3.0", + "platformdirs==4.11.13", + "pyparsing==3.3.3", + "python-dateutil==2.9.0.post0", + "requests==2.34.2", + "scipy==1.18.1", + "setuptools==84.0.0", + "shiftml==0.2.0", + "six==1.17.0", + "sympy==1.14.0", + "torch==2.13.0", + "typing_extensions==4.16.0", + "urllib3==2.8.0", + "vesin==0.6.1", + "wigners==0.4.1" + ], + "python": "3.14.4", + "sha256": "1ad198bd75b0bb4b4098276be7e99abca9a43fbe2f4978493d5d83201205b42c", + "size": 10987853, + "spec": { + "dependencies": [ + "shiftml" + ], + "source": [ + "import shiftml.ase.calculator as calculator", + "from shiftml.utils.loading import load_model", + "url = calculator.url_resolve['ShiftML41']", + "calculator.download_with_retry(url, 'downloaded')", + "load_model('downloaded').save('shiftml4.pt')" + ] + }, + "url": "https://huggingface.co/metatensor/metatomic-regtests/resolve/c984ef0d5e431ae04dd4bd0557dff201ecbc7d5e/shiftml4.pt" + } +} diff --git a/regtests/pyproject.toml b/regtests/pyproject.toml new file mode 100644 index 000000000..2cff342cf --- /dev/null +++ b/regtests/pyproject.toml @@ -0,0 +1,7 @@ +# This is not an installable package: this file only carries the pytest configuration. + +[tool.pytest.ini_options] +python_files = ["*.py"] +testpaths = ["tests"] + +filterwarnings = ["ignore"] diff --git a/regtests/references/flashmd-pet-omatpes-v2-16fs/input.json b/regtests/references/flashmd-pet-omatpes-v2-16fs/input.json new file mode 100644 index 000000000..1df3aaa0f --- /dev/null +++ b/regtests/references/flashmd-pet-omatpes-v2-16fs/input.json @@ -0,0 +1,57 @@ +{ + "model": "flashmd-pet-omatpes-v2-16fs", + "length_unit": "angstrom", + "systems": [ + { + "types": [14, 7, 8, 28, 62, 10, 11, 47], + "positions": [ + [-0.045205, 0.029973, 0.093936], + [1.109695, 1.31844, 1.22018], + [0.110888, 2.624246, 2.74494], + [1.448454, 3.973221, 4.006242], + [2.81727, -0.08222, 2.721388], + [4.045614, 1.352814, 4.086579], + [2.752344, 2.967519, -0.14145], + [4.134036, 3.924733, 1.341168] + ], + "cell": [ + [5.55, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.07, 5.48] + ], + "pbc": [true, false, true], + "inputs": { + "momentum": [ + [0.038592, 0.577446, 0.693609], + [-0.136316, -0.728368, 0.429785], + [0.454582, -0.257445, -0.402374], + [-0.835319, -0.032925, -0.411871], + [0.146886, -0.276867, 0.131116], + [0.179059, -0.00977, 0.143577], + [0.513286, -0.216795, -0.273349], + [-0.36077, 0.944725, -0.310494] + ], + "mass": [28.085, 14.007, 15.999, 58.6934, 150.36, 20.1797, 22.989769, 107.8682] + } + } + ], + "selected_atoms": null, + "outputs": { + "position": { + "unit": "A", + "sample_kind": "atom", + "gradients": [], + "reference": "position.mts", + "rtol": 1e-6, + "atol": 1e-5 + }, + "momentum": { + "unit": "(eV*u)^(1/2)", + "sample_kind": "atom", + "gradients": [], + "reference": "momentum.mts", + "rtol": 1e-6, + "atol": 1e-5 + } + } +} diff --git a/regtests/references/flashmd-pet-omatpes-v2-16fs/momentum.mts b/regtests/references/flashmd-pet-omatpes-v2-16fs/momentum.mts new file mode 100644 index 000000000..6d02c899d Binary files /dev/null and b/regtests/references/flashmd-pet-omatpes-v2-16fs/momentum.mts differ diff --git a/regtests/references/flashmd-pet-omatpes-v2-16fs/position.mts b/regtests/references/flashmd-pet-omatpes-v2-16fs/position.mts new file mode 100644 index 000000000..97631effc Binary files /dev/null and b/regtests/references/flashmd-pet-omatpes-v2-16fs/position.mts differ diff --git a/regtests/references/flashmd-pet-omatpes-v2-16fs/references.lock b/regtests/references/flashmd-pet-omatpes-v2-16fs/references.lock new file mode 100644 index 000000000..80eded4b0 --- /dev/null +++ b/regtests/references/flashmd-pet-omatpes-v2-16fs/references.lock @@ -0,0 +1,29 @@ +{ + "input_sha256": "1d6fcebdc4821b1731bee9ba026b842fe57acdc8f18bc4a2b286decd31affe7d", + "references_sha256": { + "momentum.mts": "fdd8421c8c54b61768ede5b6c98852d525a69ecf61a6261af294c7e5dd01c29b", + "position.mts": "95836ecd4da4e8932880d9c174e44f5d983aca5190626c0bb991585884b60cff" + }, + "python": "3.14", + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "ctypes-dlpack==0.1.0", + "filelock==4.0.3", + "fsspec==2026.9.0", + "metatensor-core==0.2.5", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "setuptools==84.0.0", + "sympy==1.14.0", + "torch==2.13.0", + "typing_extensions==4.16.0", + "vesin==0.6.1", + "wigners==0.4.1" + ] +} diff --git a/regtests/references/pet-mad-dos-v1.0/dos.mts b/regtests/references/pet-mad-dos-v1.0/dos.mts new file mode 100644 index 000000000..0efa00f64 Binary files /dev/null and b/regtests/references/pet-mad-dos-v1.0/dos.mts differ diff --git a/regtests/references/pet-mad-dos-v1.0/input.json b/regtests/references/pet-mad-dos-v1.0/input.json new file mode 100644 index 000000000..7e20d4883 --- /dev/null +++ b/regtests/references/pet-mad-dos-v1.0/input.json @@ -0,0 +1,32 @@ +{ + "model": "pet-mad-dos-v1.0", + "length_unit": "angstrom", + "systems": [ + { + "types": [56, 22, 8, 8], + "positions": [ + [0.074946, 0.021462, -0.001003], + [2.036883, 2.014057, 1.882192], + [1.992929, 1.991172, 0.038224], + [1.920012, 0.002418, 2.012168] + ], + "cell": [ + [4.0, 0.0, 0.0], + [0.0, 4.0, 0.0], + [0.0, 0.0, 4.0] + ], + "pbc": [true, true, true] + } + ], + "selected_atoms": null, + "outputs": { + "mtt::dos": { + "unit": "", + "sample_kind": "atom", + "gradients": [], + "reference": "dos.mts", + "rtol": 1e-6, + "atol": 1e-5 + } + } +} diff --git a/regtests/references/pet-mad-dos-v1.0/references.lock b/regtests/references/pet-mad-dos-v1.0/references.lock new file mode 100644 index 000000000..7cf7ecce7 --- /dev/null +++ b/regtests/references/pet-mad-dos-v1.0/references.lock @@ -0,0 +1,28 @@ +{ + "input_sha256": "10725b2cee5d3465644aa3f48440c8ce384c6f4bd3715a95adeda50f0f8ee5f1", + "references_sha256": { + "dos.mts": "250f686e91aac8b8fe979464ce14480d324afdee1a556c32e8a4a9f5661fc4f2" + }, + "python": "3.14", + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "ctypes-dlpack==0.1.0", + "filelock==4.0.3", + "fsspec==2026.9.0", + "metatensor-core==0.2.5", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "setuptools==84.0.0", + "sympy==1.14.0", + "torch==2.13.0", + "typing_extensions==4.16.0", + "vesin==0.6.1", + "wigners==0.4.1" + ] +} diff --git a/regtests/references/pet-mad-featurizer-v1.0.0/feature.mts b/regtests/references/pet-mad-featurizer-v1.0.0/feature.mts new file mode 100644 index 000000000..26286ded4 Binary files /dev/null and b/regtests/references/pet-mad-featurizer-v1.0.0/feature.mts differ diff --git a/regtests/references/pet-mad-featurizer-v1.0.0/input.json b/regtests/references/pet-mad-featurizer-v1.0.0/input.json new file mode 100644 index 000000000..77da6c246 --- /dev/null +++ b/regtests/references/pet-mad-featurizer-v1.0.0/input.json @@ -0,0 +1,36 @@ +{ + "model": "pet-mad-featurizer-v1.0.0", + "length_unit": "angstrom", + "systems": [ + { + "types": [14, 7, 8, 28, 62, 10, 11, 47], + "positions": [ + [-0.045205, 0.029973, 0.093936], + [1.109695, 1.31844, 1.22018], + [0.110888, 2.624246, 2.74494], + [1.448454, 3.973221, 4.006242], + [2.81727, -0.08222, 2.721388], + [4.045614, 1.352814, 4.086579], + [2.752344, 2.967519, -0.14145], + [4.134036, 3.924733, 1.341168] + ], + "cell": [ + [5.55, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.07, 5.48] + ], + "pbc": [true, false, true] + } + ], + "selected_atoms": null, + "outputs": { + "feature": { + "unit": "", + "sample_kind": "system", + "gradients": [], + "reference": "feature.mts", + "rtol": 1e-6, + "atol": 1e-5 + } + } +} diff --git a/regtests/references/pet-mad-featurizer-v1.0.0/references.lock b/regtests/references/pet-mad-featurizer-v1.0.0/references.lock new file mode 100644 index 000000000..33fde290a --- /dev/null +++ b/regtests/references/pet-mad-featurizer-v1.0.0/references.lock @@ -0,0 +1,28 @@ +{ + "input_sha256": "e794ed023396d4026f077ae42d08fd1dbd76f2dc0900d5bc2a9f0bf9b629870f", + "references_sha256": { + "feature.mts": "dbd49d94dd02b0da689a9cc40324abe97baf3c8783644ce7c98cdd52678e0275" + }, + "python": "3.14", + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "ctypes-dlpack==0.1.0", + "filelock==4.0.3", + "fsspec==2026.9.0", + "metatensor-core==0.2.5", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "setuptools==84.0.0", + "sympy==1.14.0", + "torch==2.13.0", + "typing_extensions==4.16.0", + "vesin==0.6.1", + "wigners==0.4.1" + ] +} diff --git a/regtests/references/pet-mad-s-v1.0.2/energy.mts b/regtests/references/pet-mad-s-v1.0.2/energy.mts new file mode 100644 index 000000000..d4ad4c032 Binary files /dev/null and b/regtests/references/pet-mad-s-v1.0.2/energy.mts differ diff --git a/regtests/references/pet-mad-s-v1.0.2/feature.mts b/regtests/references/pet-mad-s-v1.0.2/feature.mts new file mode 100644 index 000000000..d04a311c7 Binary files /dev/null and b/regtests/references/pet-mad-s-v1.0.2/feature.mts differ diff --git a/regtests/references/pet-mad-s-v1.0.2/input.json b/regtests/references/pet-mad-s-v1.0.2/input.json new file mode 100644 index 000000000..d7cc2379e --- /dev/null +++ b/regtests/references/pet-mad-s-v1.0.2/input.json @@ -0,0 +1,47 @@ +{ + "model": "pet-mad-s-v1.0.2", + "length_unit": "angstrom", + "systems": [ + { + "types": [14, 7, 8, 28, 62, 10, 11, 47], + "positions": [ + [-0.045205, 0.029973, 0.093936], + [1.109695, 1.31844, 1.22018], + [0.110888, 2.624246, 2.74494], + [1.448454, 3.973221, 4.006242], + [2.81727, -0.08222, 2.721388], + [4.045614, 1.352814, 4.086579], + [2.752344, 2.967519, -0.14145], + [4.134036, 3.924733, 1.341168] + ], + "cell": [ + [5.55, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.07, 5.48] + ], + "pbc": [true, false, true] + } + ], + "selected_atoms": null, + "outputs": { + "energy": { + "unit": "eV", + "sample_kind": "system", + "gradients": [ + "positions", + "strain" + ], + "reference": "energy.mts", + "rtol": 1e-6, + "atol": 1e-5 + }, + "feature": { + "unit": "", + "sample_kind": "atom", + "gradients": [], + "reference": "feature.mts", + "rtol": 1e-4, + "atol": 1e-4 + } + } +} diff --git a/regtests/references/pet-mad-s-v1.0.2/references.lock b/regtests/references/pet-mad-s-v1.0.2/references.lock new file mode 100644 index 000000000..80141f79d --- /dev/null +++ b/regtests/references/pet-mad-s-v1.0.2/references.lock @@ -0,0 +1,29 @@ +{ + "input_sha256": "bdae7858b7cff00c73685a26b83ef2ebf8f6fc2babf75ebc69d8c8910f74d9ed", + "references_sha256": { + "energy.mts": "7ba198261c5f15a350372acc6b4344d7ce26595617937775996e1657547a01ff", + "feature.mts": "99eba0b0092f64caa4414f42510ca01af3a1ce1975e0f4726534b8eba485d84f" + }, + "python": "3.14", + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "ctypes-dlpack==0.1.0", + "filelock==4.0.3", + "fsspec==2026.9.0", + "metatensor-core==0.2.5", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "setuptools==84.0.0", + "sympy==1.14.0", + "torch==2.13.0", + "typing_extensions==4.16.0", + "vesin==0.6.1", + "wigners==0.4.1" + ] +} diff --git a/regtests/references/pet-mad-s-v1.5.0/energy.mts b/regtests/references/pet-mad-s-v1.5.0/energy.mts new file mode 100644 index 000000000..1c64a4ef1 Binary files /dev/null and b/regtests/references/pet-mad-s-v1.5.0/energy.mts differ diff --git a/regtests/references/pet-mad-s-v1.5.0/feature.mts b/regtests/references/pet-mad-s-v1.5.0/feature.mts new file mode 100644 index 000000000..5c3d25875 Binary files /dev/null and b/regtests/references/pet-mad-s-v1.5.0/feature.mts differ diff --git a/regtests/references/pet-mad-s-v1.5.0/input.json b/regtests/references/pet-mad-s-v1.5.0/input.json new file mode 100644 index 000000000..d4c046a35 --- /dev/null +++ b/regtests/references/pet-mad-s-v1.5.0/input.json @@ -0,0 +1,63 @@ +{ + "model": "pet-mad-s-v1.5.0", + "length_unit": "angstrom", + "systems": [ + { + "types": [14, 7, 8, 28, 62, 10, 11, 47], + "positions": [ + [-0.045205, 0.029973, 0.093936], + [1.109695, 1.31844, 1.22018], + [0.110888, 2.624246, 2.74494], + [1.448454, 3.973221, 4.006242], + [2.81727, -0.08222, 2.721388], + [4.045614, 1.352814, 4.086579], + [2.752344, 2.967519, -0.14145], + [4.134036, 3.924733, 1.341168] + ], + "cell": [ + [5.55, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.07, 5.48] + ], + "pbc": [true, false, true] + } + ], + "selected_atoms": null, + "outputs": { + "energy": { + "unit": "eV", + "sample_kind": "system", + "gradients": [ + "positions", + "strain" + ], + "reference": "energy.mts", + "rtol": 1e-6, + "atol": 1e-5 + }, + "non_conservative_force": { + "unit": "eV/A", + "sample_kind": "atom", + "gradients": [], + "reference": "nc_force.mts", + "rtol": 1e-6, + "atol": 1e-5 + }, + "non_conservative_stress": { + "unit": "eV/A^3", + "sample_kind": "system", + "gradients": [], + "reference": "nc_stress.mts", + "rtol": 1e-6, + "atol": 1e-5 + }, + "feature": { + "unit": "", + "sample_kind": "atom", + "gradients": [], + "reference": "feature.mts", + "rtol": 1e-4, + "atol": 1e-4 + } + } +} diff --git a/regtests/references/pet-mad-s-v1.5.0/nc_force.mts b/regtests/references/pet-mad-s-v1.5.0/nc_force.mts new file mode 100644 index 000000000..84719c958 Binary files /dev/null and b/regtests/references/pet-mad-s-v1.5.0/nc_force.mts differ diff --git a/regtests/references/pet-mad-s-v1.5.0/nc_stress.mts b/regtests/references/pet-mad-s-v1.5.0/nc_stress.mts new file mode 100644 index 000000000..e8febda82 Binary files /dev/null and b/regtests/references/pet-mad-s-v1.5.0/nc_stress.mts differ diff --git a/regtests/references/pet-mad-s-v1.5.0/references.lock b/regtests/references/pet-mad-s-v1.5.0/references.lock new file mode 100644 index 000000000..d439016ba --- /dev/null +++ b/regtests/references/pet-mad-s-v1.5.0/references.lock @@ -0,0 +1,31 @@ +{ + "input_sha256": "ef2e3da98a5a76cc3ce1e06a6cbe9f9f581ad3118c7a3d17b2c64f0d48ad5d89", + "references_sha256": { + "energy.mts": "be025a6d04aa23eec0e257156ae6e94e5a2697415a761d70b510a0a13e142d5b", + "feature.mts": "139f5413c896ceeead9941fd1389d19a7018babaa779627716dad89ad116c44f", + "nc_force.mts": "db3d42f429a04b1ef6dacdca890f0b496a84e90703217a48c139a0ec03f5ddd9", + "nc_stress.mts": "39eb71237c8bf587db5690267be5cbf76b47b91b487959488a34f75961caf51f" + }, + "python": "3.14", + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "ctypes-dlpack==0.1.0", + "filelock==4.0.3", + "fsspec==2026.9.0", + "metatensor-core==0.2.5", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "setuptools==84.0.0", + "sympy==1.14.0", + "torch==2.13.0", + "typing_extensions==4.16.0", + "vesin==0.6.1", + "wigners==0.4.1" + ] +} diff --git a/regtests/references/shiftml4/cs_iso.mts b/regtests/references/shiftml4/cs_iso.mts new file mode 100644 index 000000000..27c8eb2a2 Binary files /dev/null and b/regtests/references/shiftml4/cs_iso.mts differ diff --git a/regtests/references/shiftml4/input.json b/regtests/references/shiftml4/input.json new file mode 100644 index 000000000..9f4336610 --- /dev/null +++ b/regtests/references/shiftml4/input.json @@ -0,0 +1,36 @@ +{ + "model": "shiftml4", + "length_unit": "angstrom", + "systems": [ + { + "types": [6, 1, 1, 7, 8, 6, 1, 8], + "positions": [ + [0.041078, -0.052167, 0.05958], + [1.09656, -0.002623, -0.005521], + [0.026761, 1.108259, -0.04118], + [2.603335, 2.545063, 0.382558], + [2.873938, 0.328984, 2.915441], + [0.34278, 3.147015, 3.005015], + [1.371167, 3.286564, 3.409837], + [3.204182, 2.930543, 3.254507] + ], + "cell": [ + [6.2, 0.0, 0.0], + [0.0, 6.0, 0.0], + [0.15, 0.0, 6.1] + ], + "pbc": [true, true, true] + } + ], + "selected_atoms": null, + "outputs": { + "mtt::cs_iso": { + "unit": "", + "sample_kind": "atom", + "gradients": [], + "reference": "cs_iso.mts", + "rtol": 1e-4, + "atol": 1e-4 + } + } +} diff --git a/regtests/references/shiftml4/references.lock b/regtests/references/shiftml4/references.lock new file mode 100644 index 000000000..a9e22eabf --- /dev/null +++ b/regtests/references/shiftml4/references.lock @@ -0,0 +1,28 @@ +{ + "input_sha256": "325d4f3eb87f23b9e2e05407f2c0020a6f0c06359c29414770f79f33e4fc0b1e", + "references_sha256": { + "cs_iso.mts": "56b6c6eaad27e0eee0ff81f57abb13285eed89ef4949c35c3bba9a396fcabb5a" + }, + "python": "3.14", + "dependencies": [ + "Jinja2==3.1.6", + "MarkupSafe==3.0.3", + "ctypes-dlpack==0.1.0", + "filelock==4.0.3", + "fsspec==2026.9.0", + "metatensor-core==0.2.5", + "metatensor-operations==0.5.0", + "metatensor-torch==0.10.6", + "metatomic-ase==0.1.4", + "metatomic-torch==0.1.17", + "mpmath==1.3.0", + "networkx==3.7", + "numpy==2.5.3", + "setuptools==84.0.0", + "sympy==1.14.0", + "torch==2.13.0", + "typing_extensions==4.16.0", + "vesin==0.6.1", + "wigners==0.4.1" + ] +} diff --git a/regtests/tests/regressions.py b/regtests/tests/regressions.py new file mode 100644 index 000000000..44a239db0 --- /dev/null +++ b/regtests/tests/regressions.py @@ -0,0 +1,39 @@ +"""Regression tests for models exported by previous versions of metatomic""" + +import glob +import os +import sys + +import metatensor.torch as mts +import pytest + + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + +CASES = sorted(glob.glob(os.path.join(ROOT, "references", "*", "input.json"))) +CASE_IDS = [os.path.basename(os.path.dirname(case)) for case in CASES] + +sys.path.append(ROOT) + +import metatomic_regtest # noqa E401 + + +@pytest.mark.parametrize("case", CASES, ids=CASE_IDS) +def test_regression(case): + input = metatomic_regtest.load_input(case) + model = metatomic_regtest.load_model(input.model) + outputs = metatomic_regtest.run_model(model, input) + expected = metatomic_regtest.load_expected(input) + + for name in input.outputs: + try: + mts.allclose_raise( + outputs[name], + expected[name], + rtol=input.outputs[name]["rtol"], + atol=input.outputs[name]["atol"], + ) + except Exception as e: + raise AssertionError( + f"Output '{name}' of model '{input.model}' changed in case '{case}'" + ) from e diff --git a/regtests/update-references.py b/regtests/update-references.py new file mode 100644 index 000000000..f7d3067a4 --- /dev/null +++ b/regtests/update-references.py @@ -0,0 +1,510 @@ +#!/usr/bin/env python3 +""" +Create or update the reference outputs of the regression tests. + + python update-references.py --list + python update-references.py pet-mad-s-v1.5.0 + python update-references.py --all --force + python update-references.py --all --relock + python update-references.py --check +""" + +import argparse +import glob +import hashlib +import json +import os +import shutil +import subprocess +import sys + + +ROOT = os.path.dirname(os.path.abspath(__file__)) +VENVS_PATH = os.path.join(ROOT, "build", "venvs", "references") + +# torch publishes CPU-only wheels on its own index, and they are much smaller +DEFAULT_EXTRA_INDEX = "https://download.pytorch.org/whl/cpu" + +# Maximal size of a reference file, in bytes +MAX_REFERENCE_SIZE = 100 * 1024 + +# What the engine needs to run a model. +DEPENDENCIES = [ + "metatomic-torch", + "metatensor-operations", + "vesin", + "torch==2.13.*", + "numpy", +] + + +def all_cases(): + cases = glob.glob(os.path.join(ROOT, "references", "*", "input.json")) + return sorted(os.path.basename(os.path.dirname(case)) for case in cases) + + +def case_directory(case: str) -> str: + return os.path.join(ROOT, "references", case) + + +def lock_path(case: str) -> str: + return os.path.join(case_directory(case), "references.lock") + + +def read_lock(case: str) -> dict: + path = lock_path(case) + if not os.path.exists(path): + return {} + + with open(path) as fd: + return json.load(fd) + + +def file_hash(path: str) -> str: + with open(path, "rb") as fd: + return hashlib.sha256(fd.read()).hexdigest() + + +def input_hash(case: str) -> str: + """ + Checksum of the ``input.json`` of a case. + + The tolerances are left out of it: how closely an output has to match its reference + says nothing about whether the reference itself is still the right answer, so + changing one should not ask for a regeneration. The rest is hashed in a canonical + form, so that reformatting the file does not either. + """ + with open(os.path.join(case_directory(case), "input.json")) as fd: + data = json.load(fd) + + for output in data["outputs"].values(): + output.pop("rtol", None) + output.pop("atol", None) + + canonical = json.dumps(data, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(canonical.encode()).hexdigest() + + +def reference_hashes(case: str) -> dict: + """Checksum of every reference ``.mts`` file of a case""" + directory = case_directory(case) + with open(os.path.join(directory, "input.json")) as fd: + outputs = json.load(fd)["outputs"] + + return { + name: file_hash(os.path.join(directory, name)) + for name in sorted({output["reference"] for output in outputs.values()}) + } + + +def python_version(python: str) -> str: + version = subprocess.run( + [python, "-c", "import platform; print(platform.python_version())"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + return ".".join(version.split(".")[:2]) + + +def venv_python(directory: str) -> str: + if os.name == "nt": + return os.path.join(directory, "Scripts", "python.exe") + + return os.path.join(directory, "bin", "python") + + +def environment_marker(dependencies: list, python: str) -> str: + """What an environment holds, as the string stored inside it.""" + return json.dumps({"python": python_version(python), "dependencies": dependencies}) + + +def venv_directory(dependencies: list, python: str) -> str: + """Where the environment holding ``dependencies`` lives.""" + digest = hashlib.sha256(environment_marker(dependencies, python).encode()) + return os.path.join(VENVS_PATH, digest.hexdigest()[:12]) + + +def create_venv(dependencies: list, python: str) -> str: + """ + Get a virtual environment with exactly ``dependencies`` installed in it, creating it + if we don't have it already. + """ + marker = environment_marker(dependencies, python) + directory = venv_directory(dependencies, python) + stamp = os.path.join(directory, "regtests-environment.json") + + if os.path.exists(stamp) and open(stamp).read() == marker: + return directory + + if os.path.exists(directory): + shutil.rmtree(directory) + + os.makedirs(VENVS_PATH, exist_ok=True) + print(f"+++++ creating a virtualenv in {directory}") + subprocess.run([python, "-m", "venv", directory], check=True) + + environment = dict(os.environ) + environment.setdefault("PIP_EXTRA_INDEX_URL", DEFAULT_EXTRA_INDEX) + + pip = [venv_python(directory), "-m", "pip", "--disable-pip-version-check"] + subprocess.run(pip + ["install", "--upgrade", "pip"], check=True, env=environment) + + print("+++++ installing dependencies") + subprocess.run(pip + ["install"] + list(dependencies), check=True, env=environment) + + with open(stamp, "w") as fd: + fd.write(marker) + + return directory + + +def settle_venv(directory: str, dependencies: list, python: str) -> str: + """ + Move an environment to the name its *installed* dependencies give it. + + An environment is built from what the lock asks for, which the first time around is + the unpinned ``DEPENDENCIES``; the lock then records what pip actually installed. + Those are different names for the same environment, so without this the one we just + built is not the one the lock points at: it would be pruned right after being used, + and built again from scratch on the next run. + """ + target = venv_directory(dependencies, python) + if target == directory: + return directory + + if os.path.exists(target): + # another case already settled on this exact environment + shutil.rmtree(directory) + return target + + os.replace(directory, target) + with open(os.path.join(target, "regtests-environment.json"), "w") as fd: + fd.write(environment_marker(dependencies, python)) + + return target + + +def freeze(directory: str) -> list: + """The exact version of everything installed in the environment.""" + output = subprocess.run( + [venv_python(directory), "-m", "pip", "freeze", "--disable-pip-version-check"], + check=True, + capture_output=True, + text=True, + ) + + return sorted(line for line in output.stdout.splitlines() if line.strip()) + + +def prune_venvs(): + """ + Delete the environments that no case refers to any more. + + Regenerating a case for the first time moves it from the unpinned dependencies to + the ones its lock now holds, which is a different environment; without this the + first one would stay around forever. + """ + keep = { + venv_directory(read_lock(case)["dependencies"], sys.executable) + for case in all_cases() + if read_lock(case) + } + + if not os.path.exists(VENVS_PATH): + return + + for name in sorted(os.listdir(VENVS_PATH)): + directory = os.path.join(VENVS_PATH, name) + if directory not in keep: + print(f"+++++ removing the unused environment {name}") + shutil.rmtree(directory, ignore_errors=True) + + +def up_to_date(case: str, locked: dict) -> bool: + """ + Whether the references of ``case`` are the ones its current ``input.json`` produced. + + This is what the lock's checksums are for: without them we could only tell that the + reference files exist, not that they still answer the question being asked. + """ + if "input_sha256" not in locked or "references_sha256" not in locked: + return False + + if locked["input_sha256"] != input_hash(case): + return False + + directory = case_directory(case) + for name, digest in locked["references_sha256"].items(): + path = os.path.join(directory, name) + if not os.path.exists(path) or file_hash(path) != digest: + return False + + return True + + +def regenerate(case: str, force: bool, relock: bool, python: str) -> bool: + """Generate the references of ``case`` in the environment its lock describes.""" + print(f"+++++ {case}:") + + locked = read_lock(case) + if locked and not force and not relock and up_to_date(case, locked): + print("+++++ already generated from the current input.json") + return True + + if locked and not relock: + dependencies = locked["dependencies"] + print(f"+++++ using the {len(dependencies)} dependencies in references.lock") + else: + dependencies = DEPENDENCIES + if locked: + print("+++++ re-resolving the dependencies (--relock)") + + directory = create_venv(dependencies, python) + + command = [venv_python(directory), os.path.abspath(__file__), "--worker", case] + result = subprocess.run(command, cwd=ROOT) + if result.returncode != 0: + return False + + installed = freeze(directory) + directory = settle_venv(directory, installed, python) + + with open(lock_path(case), "w") as fd: + json.dump( + { + "input_sha256": input_hash(case), + "references_sha256": reference_hashes(case), + "python": python_version(venv_python(directory)), + "dependencies": installed, + }, + fd, + indent=2, + ) + fd.write("\n") + + return True + + +def deviation(actual, expected) -> str: + """Summarize how far ``actual`` is from ``expected``.""" + lines = [] + for key, expected_block in expected.items(): + try: + actual_block = actual.block(key) + except Exception: + return "+++++ metadata changed, can not compare values" + + if actual_block.values.shape != expected_block.values.shape: + return "+++++ shape changed, can not compare values" + + arrays = [("values", actual_block.values, expected_block.values)] + for name, gradient in expected_block.gradients(): + try: + arrays.append( + ( + f"{name} gradients", + actual_block.gradient(name).values, + gradient.values, + ) + ) + except Exception: + lines.append(f"+++++ {name} gradients: missing in the new output") + + for what, new, old in arrays: + if new.shape != old.shape: + lines.append(f"+++++ {what}: shape changed") + continue + + delta = (new.double() - old.double()).abs() + scale = old.double().abs().clamp(min=1e-100) + lines.append( + f"+++++ {what}: max absolute {delta.max():.3e}, " + f"max relative {(delta / scale).max():.3e}" + ) + + return "\n".join(lines) + + +def check(case: str, models: dict) -> list: + """Check that the reference files match what ``input.json`` asks for""" + import metatomic_regtest as regtest + + directory = case_directory(case) + try: + input = regtest.load_input(os.path.join(directory, "input.json")) + except Exception as e: + return [f"input.json can not be read: {e}"] + + problems = [] + if input.model not in models: + problems.append(f"uses the model '{input.model}', which is not in models.json") + + regenerate_it = f"run `python regtests/update-references.py {case} --force`" + locked = read_lock(case) + if not locked: + problems.append( + "references.lock is missing, so there is no record of what generated the " + f"references: {regenerate_it}" + ) + elif locked["input_sha256"] != input_hash(case): + problems.append( + f"input.json has changed since the references were generated: " + f"{regenerate_it}" + ) + + recorded = locked.get("references_sha256", {}) + referenced = set() + for name, output in input.outputs.items(): + reference = output["reference"] + referenced.add(reference) + + path = os.path.join(directory, reference) + if not os.path.exists(path): + problems.append( + f"the reference for '{name}' is missing ({reference}): run " + f"`python regtests/update-references.py {case}`" + ) + continue + + if reference in recorded and recorded[reference] != file_hash(path): + problems.append( + f"{reference} is not the file that was generated: {regenerate_it}" + ) + + size = os.path.getsize(path) + if size > MAX_REFERENCE_SIZE: + problems.append( + f"{reference} is {size / 1024:.0f} kiB, more than the " + f"{MAX_REFERENCE_SIZE // 1024} kiB a reference may take: use a smaller " + "system, or ask for fewer outputs" + ) + + for path in sorted(glob.glob(os.path.join(directory, "*.mts"))): + name = os.path.basename(path) + if name not in referenced: + problems.append(f"{name} is not used by any output of input.json") + + return problems + + +def generate(case: str) -> bool: + import metatensor.torch as mts + + import metatomic_regtest as regtest + + directory = case_directory(case) + input = regtest.load_input(os.path.join(directory, "input.json")) + + existing = {} + for name, output in input.outputs.items(): + path = os.path.join(directory, output["reference"]) + if os.path.exists(path): + existing[name] = mts.load(path) + + print(f"+++++ {case}: running {input.model}") + model = regtest.load_model(input.model) + outputs = regtest.run_model(model, input) + + for name, output in input.outputs.items(): + path = os.path.join(directory, output["reference"]) + if name in existing: + print(f"+++++ {name}: deviation from the previous reference") + print(deviation(outputs[name], existing[name])) + else: + print(f"+++++ {name}: new reference") + + mts.save(path, outputs[name]) + + return True + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("cases", nargs="*", help="cases to regenerate") + parser.add_argument("--all", action="store_true", help="regenerate every case") + parser.add_argument("--list", action="store_true", help="list the known cases") + parser.add_argument( + "--check", + action="store_true", + help="check that every case has the reference data input.json describes, " + "without running any model (every case, unless some are named)", + ) + parser.add_argument( + "--force", + action="store_true", + help="regenerate the references even if they are already up to date", + ) + parser.add_argument( + "--relock", + action="store_true", + help="re-resolve the dependencies instead of using the locked ones", + ) + parser.add_argument( + "--python", + default=sys.executable, + help="python interpreter used to create the virtualenv", + ) + parser.add_argument( + # how `regenerate` runs the generation inside the environment it just prepared + "--worker", + action="store_true", + help=argparse.SUPPRESS, + ) + args = parser.parse_args() + + # show the messages from this process and any subprocess in the correct order + sys.stdout.reconfigure(line_buffering=True) + + if args.list: + for case in all_cases(): + print(case) + return 0 + + if args.all or (args.check and not args.cases): + cases = all_cases() + elif args.cases: + cases = args.cases + else: + parser.error("give at least one case, or --all") + + for case in cases: + if case not in all_cases(): + known = "\n".join(f" - {name}" for name in all_cases()) + parser.error(f"unknown case '{case}', references/ contains:\n{known}") + + if args.check: + with open(os.path.join(ROOT, "models.json")) as fd: + models = json.load(fd) + + failed = False + for case in cases: + for message in check(case, models): + failed = True + print(f"ERROR: {case}: {message}", file=sys.stderr) + + if failed: + print("\nsome reference data is missing or out of date", file=sys.stderr) + return 1 + + print(f"+++++ all {len(cases)} cases have their reference data") + return 0 + + if args.worker: + return 0 if all(generate(case) for case in cases) else 1 + + ok = True + for case in cases: + ok &= regenerate(case, args.force, args.relock, args.python) + + prune_venvs() + + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tox.ini b/tox.ini index 919350ccd..72af79e64 100644 --- a/tox.ini +++ b/tox.ini @@ -23,7 +23,7 @@ setenv = COVERAGE_FILE={env_dir}/.coverage package = skip -lint_folders = "{toxinidir}/python" "{toxinidir}/setup.py" +lint_folders = "{toxinidir}/python" "{toxinidir}/regtests" "{toxinidir}/setup.py" build_single_wheel = --no-deps --no-build-isolation --check-build-dependencies install_lj_tests = pip install {[testenv]build_single_wheel} git+https://github.com/metatensor/lj-test@a1464af @@ -233,6 +233,31 @@ commands = pytest --cov={env_site_packages_dir}/metatomic_torchsim --cov-report= --import-mode=append {posargs} +################################################################################ +##### Ecosystem regression tests ##### +################################################################################ + +[testenv:regtests] +description = Run the regression tests (load and run real models) +deps = + {[testenv]testing_deps} + {[testenv]packaging_deps} + {[testenv]metatomic_deps} + + torch=={env:METATOMIC_TESTS_TORCH_VERSION:2.13}.* + numpy + + vesin >=0.6.0,<0.7 + +changedir = regtests +commands = + pip install {[testenv]build_single_wheel} ../python/metatomic_torch + + python export-models.py --check + python update-references.py --check + + pytest --import-mode=append {posargs} + ################################################################################ ##### Linter and formatter setup #####