[1]:
# /// script
# dependencies = [
#     "metatrain[pet]==2025.7",
#     "metatomic-torch>=0.1.1,<0.2",
#     "vesin-torch>=0.3.7",
# ]
# ///

Using the PET-MAD model with TorchSim

This tutorial explains how to use the PET-MAD model (https://arxiv.org/abs/2503.14118) via TorchSim’s metatomic interface.

Loading the model

Loading the model is simple: you simply need to specify the model name (in this case “pet-mad”), as shown below. All other arguments are optional: for example, you could specify the device. (If the device is not specified, like in this case, the optimal device is chosen automatically.)

[2]:
from torch_sim.models.metatomic import MetatomicModel

model = MetatomicModel("pet-mad")
/opt/hostedtoolcache/Python/3.12.11/x64/lib/python3.12/site-packages/pydantic/_internal/_generate_schema.py:2249: UnsupportedFieldAttributeWarning: The 'repr' attribute with value False was provided to the `Field()` function, which has no effect in the context it was used. 'repr' is field-specific metadata, and can only be attached to a model field using `Annotated` metadata or by assignment. This may have happened because an `Annotated` type alias using the `type` statement was used, or if the `Field()` function was attached to a single member of a union type.
  warnings.warn(
/opt/hostedtoolcache/Python/3.12.11/x64/lib/python3.12/site-packages/pydantic/_internal/_generate_schema.py:2249: UnsupportedFieldAttributeWarning: The 'frozen' attribute with value True was provided to the `Field()` function, which has no effect in the context it was used. 'frozen' is field-specific metadata, and can only be attached to a model field using `Annotated` metadata or by assignment. This may have happened because an `Annotated` type alias using the `type` statement was used, or if the `Field()` function was attached to a single member of a union type.
  warnings.warn(
/opt/hostedtoolcache/Python/3.12.11/x64/lib/python3.12/site-packages/metatrain/utils/io.py:193: UserWarning: trying to upgrade an old model checkpoint with unknown version, this might fail and require manual modifications
  warnings.warn(
/opt/hostedtoolcache/Python/3.12.11/x64/lib/python3.12/site-packages/metatrain/pet/model.py:102: UserWarning: PET assumes that Cartesian tensors of rank 2 are stress-like, meaning that they are symmetric and intensive. If this is not the case, please use a different model.
  self._add_output(target_name, target_info)

Using the model to run a molecular dynamics simulations

Once the model is loaded, you can use it just like any other TorchSim model to run simulations. Here, we show how to run a simple MD simulation consisting of an initial NVT equilibration run followed by an NVE run.

[3]:
from ase.build import bulk
import torch_sim as ts

atoms = bulk("Si", "diamond", a=5.43, cubic=True)

equilibrated_state = ts.integrate(
    system=atoms,
    model=model,
    integrator=ts.Integrator.nvt_langevin,
    n_steps=100,
    temperature=300,  # K
    timestep=0.001,  # ps
)

final_state = ts.integrate(
    system=equilibrated_state,
    model=model,
    integrator=ts.Integrator.nve,
    n_steps=100,
    temperature=300,  # K
    timestep=0.001,  # ps
)
/opt/hostedtoolcache/Python/3.12.11/x64/lib/python3.12/site-packages/torch_sim/integrators/nvt.py:57: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
  c1 = torch.exp(torch.tensor(-gamma * dt))

Further steps

Of course, in reality, you would want to run the simulation for much longer, probably save trajectories, and much more. However, this is all you need to get started with metatomic and PET-MAD. For more details on how to use TorchSim, you can refer to the other tutorials in this section.