NVIDIA ALCHEMI Toolkit

UPET models can be driven through nvalchemi-toolkit, NVIDIA’s GPU-native toolkit for batched inference and molecular dynamics. UPETWrapper wraps any UPET / PET-MAD checkpoint as an nvalchemi-toolkit BaseModelMixin model, so it can be driven through nvalchemi’s batched Batch data pipeline and its FIRE / NVE / NVTLangevin / NPT integrators.

Requires the optional nvalchemi extra (see Installation):

pip install "upet[nvalchemi]"

Usage

Single-structure evaluation

Convert an ASE Atoms object to an AtomicData instance with from_atoms(), promote it to a single-graph Batch, compute its neighbor list, and evaluate:

import torch
from ase.build import bulk
from nvalchemi.data import AtomicData, Batch
from nvalchemi.neighbors import compute_neighbors
from upet.nvalchemi import UPETWrapper

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UPETWrapper.from_checkpoint(model="pet-mad-s", version="1.6.0", device=device)

atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
data = AtomicData.from_atoms(atoms, device=device)
batch = Batch.from_data_list([data], device=device)
compute_neighbors(batch, config=model.model_config.neighbor_config)

outputs = model(batch)
energy = outputs["energy"]
forces = outputs["forces"]
stress = outputs["stress"]

model.model_config.neighbor_config carries the cutoff and neighbor-list format the model expects, so compute_neighbors() never needs the cutoff repeated by hand.

Batched evaluation

Passing several structures at once only requires collecting one AtomicData per structure and collating them with from_data_list(); a single forward pass then evaluates all of them together:

structures = [
    bulk("Si", cubic=True, a=5.43, crystalstructure="diamond"),
    bulk("C", cubic=True, a=3.57, crystalstructure="diamond"),
    bulk("Ge", cubic=True, a=5.66, crystalstructure="diamond"),
]
data_list = [AtomicData.from_atoms(atoms, device=device) for atoms in structures]
batch = Batch.from_data_list(data_list, device=device)
compute_neighbors(batch, config=model.model_config.neighbor_config)

outputs = model(batch)
energies = outputs["energy"]  # one row per structure, shape [3, 1]

Batch.from_data_list handles the differing atom counts transparently; outputs["energy"] comes back with one row per input structure, and outputs["forces"] is stacked over all atoms in the batch in the same order as data_list.

Examples

Runnable end-to-end workflows built on top of UPETWrapper, driving nvalchemi’s FIRE, NVE, NVTLangevin, and NPT integrators: