Add PET models support

[abmazitov]

ALCHEMI Toolkit Pull Request

Description

This PR adds support for PET models from the lab-cosmo/metatrain package directly in nvalchemi-toolkit. It introduces the PETWrapper class, which wraps metatrain.pet.modules.backend.PETBackend (a pure-torch module) with nvalchemi's systems management and IO handling, exposing it through the standard BaseModelMixin interface. PETWrapper supports all PET-MAD model checkpoints, as well as other PET-based models from lab-cosmo/upet, and can fetch named checkpoints directly from HuggingFace or load local checkpoint files. The interface is torch-native and torch-compileable (fullgraph=True) for both forward energies and autograd forces.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

No related issues.

Changes Made

• Added nvalchemi/models/pet.py implementing PETWrapper, a BaseModelMixin wrapper around metatrain's pure-torch PETBackend.
• Added PETWrapper.from_checkpoint, which loads a checkpoint either from a local path or by fetching a named PET-MAD/upet checkpoint from the lab-cosmo/upet HuggingFace repository.
• Registered PET in nvalchemi/_optional.py as a new optional dependency, and exposed PETWrapper via nvalchemi/models/__init__.py.
• Added a new pet extra in pyproject.toml (pip install 'nvalchemi-toolkit[pet]'), pulling in metatrain and upet. Pinned metatrain to the upstream main branch via [tool.uv.sources] until the pure-torch PETBackend (Extract pure-torch PET backend into a separate module metatensor/metatrain#1180) lands in a tagged PyPI release.
• Added a filterwarnings entry in pyproject.toml to suppress torch.jit.script/torch.jit.script_method deprecation warnings emitted by metatensor.torch at import time.
• Added a hot-patch in pet.py that stubs a fake hostlist module in sys.modules when the real package isn't installed, since metatrain.pet.__init__ transitively imports metatrain.utils.distributed.slurm, which requires hostlist — a package explicitly blocked in this toolkit's dependency overrides.
• Added test/models/test_pet.py with unit tests covering PETWrapper construction, forward/backward passes, and checkpoint loading.
• Added examples/advanced/10_pet_nvt.py, an example running NVT dynamics with a PET-MAD model, and linked it from examples/advanced/README.rst.
• Updated CHANGELOG.md with an entry describing the new PETWrapper.
• Regenerated uv.lock.

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

Known caveats / follow-ups left for later PRs:

1. hostlist dependency conflict is not resolved. metatrain pulls python-hostlist as a transitive dependency and imports it at first import. This package is explicitly blocked in nvalchemi-toolkit's dependency overrides, and even though PETWrapper never uses it, we still have to hot-patch sys.modules to stub it out so the import chain resolves.
2. Energy variants and non-conservative forces/stresses are not implemented. Only the default energy head and conservative (autograd) forces are supported at the moment.
3. LLPR uncertainty quantification (UQ) is not implemented. On-the-fly uncertainty prediction via the LLPR module is not yet available for PETWrapper.
4. uv.lock was regenerated. This was required to pick up the new metatrain/upet dependencies from the pet extra (and the metatrain git source pin) — without regenerating it, the environment produced dependency conflicts.
5. OptionalDependency Enum apparently supports only one dependency per model. In the case of PET, at this moment, two dependencies are required: metatrain and upet. This detail requires adding a workaround for a failed import of the upet package, which cannot be done using the OptionalDependency.PET utility. It should be possible to get rid of the metatrain and only depend on upet soon, but we need a release of both packages before that will happen.

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds PETWrapper, a BaseModelMixin-compatible wrapper around metatrain's pure-torch PETBackend, supporting PET-MAD and other upet-based checkpoints with autograd forces/stress and optional torch.compile. It also adds the pet optional extra in pyproject.toml, a full test suite, and an NVT example.

• nvalchemi/models/pet.py (~1 130 lines): implements checkpoint loading (local or HuggingFace via upet), hyper-parameter normalisation, state-dict filtering, adapt_input/forward/compute_embeddings, and export_model.
• pyproject.toml: adds the pet extra with a temporary git-pinned metatrain (upstream main) until the PETBackend surface lands in a tagged PyPI release.
• test/models/test_pet.py and examples/advanced/10_pet_nvt.py: comprehensive tests and an end-to-end NVT example with an LJ fallback for CI.

Important Files Changed

Filename	Overview
nvalchemi/models/pet.py	New 1129-line PETWrapper implementation; contains two interface-contract bugs in from_checkpoint: version=None is rejected by the guard despite being documented as valid, and checkpoint_path docstring claims HuggingFace-fetch support that the code never delivers.
test/models/test_pet.py	Comprehensive 1012-line test suite covering construction, forward/backward, batching, and checkpoint loading; duplicates the hostlist stub from pet.py (harmless but repetitive).
examples/advanced/10_pet_nvt.py	NVT example with LJ fallback; unconditionally applies the hostlist sys.modules stub at import time even when PET is never loaded.
pyproject.toml	Adds pet extra depending on a git-pinned metatrain (upstream main) and upet; pin is expected to be temporary until a tagged PyPI release lands.
nvalchemi/_optional.py	Adds PET = (metatrain, nvalchemi-toolkit[pet]) entry; straightforward and consistent with existing pattern.
nvalchemi/models/init.py	Adds lazy PETWrapper import via getattr; follows the existing pattern used by MACEWrapper and others.

_{Reviews (1): Last reviewed commit: "Updated docs" | Re-trigger Greptile}

[greptile-apps]

version=None is documented as equivalent to "latest" (see the version parameter docstring and _fetch_pet_checkpoint which maps None → requested_version=None), but the guard version is not None silently rejects it. A user who follows the documented interface — PETWrapper.from_checkpoint(model="pet-mad-s", version=None) — hits the else branch and receives a confusing ValueError telling them to provide checkpoint_path or model, even though they did provide model.

Suggested change

	elif model is not None and version is not None:
	checkpoint_path = _fetch_pet_checkpoint(model, version)
	elif model is not None:
	checkpoint_path = _fetch_pet_checkpoint(model, version)

[greptile-apps]

The checkpoint_path parameter docstring states it can be "a named model such as pet-mad-xs-v1.6.0 to fetch from HuggingFace", but the code never acts on a non-file checkpoint_path string — if the path does not exist as a local file, the branch falls straight to elif model is not None, completely ignoring checkpoint_path. A user who passes checkpoint_path="pet-mad-xs-v1.6.0" with model=None hits the else and gets a misleading ValueError. The docstring should be corrected to remove the HuggingFace-fetch promise for checkpoint_path (the model parameter is the correct handle for named HF models).

Suggested change

	checkpoint_path : Path | str | None
	Path to a local PET checkpoint file (``.ckpt`` / ``.pt``), or a
	named model such as ``"pet-mad-xs-v1.6.0"`` to fetch from
	HuggingFace. If ``None``, *model* must be given instead.
	checkpoint_path : Path | str | None
	Path to a local PET checkpoint file (``.ckpt`` / ``.pt``).
	If ``None``, *model* must be given instead to fetch a checkpoint
	by name from HuggingFace.

[greptile-apps]

The error message has an unclosed opening parenthesis: "(a local file path, or \model`...". The (anever closes, so the human-readable text reads awkwardly asrequires either `checkpoint_path` (a local file path, or `model` (e.g. …).`

Suggested change

	raise ValueError(
	"PETWrapper.from_checkpoint requires either `checkpoint_path` (a "
	"local file path, or `model` (e.g. 'pet-mad-s', optionally with `version`). "
	)
	raise ValueError(
	"PETWrapper.from_checkpoint requires either `checkpoint_path` (a "
	"local file path) or `model` (e.g. 'pet-mad-s', optionally with `version`). "
	)

[greptile-apps]

Global sys.modules mutation even in LJ-fallback path

The hostlist stub is injected unconditionally at module import time, so every user who imports or runs this example — even those who never load PET and fall back to LJ — silently get a fake hostlist entry registered in sys.modules. The guard if "hostlist" not in sys.modules mitigates the most common case, but the whole block could be moved inside the if PET_MODEL_PATH: branch so it only runs when PET is actually needed.