Dependencies

/// script dependencies = [ “metatrain[pet]>=2025.4”, “metatensor-torch>=0.7,<0.8”] ///

Using the PET-MAD model with metatensor¶

This tutorial explains how to use the PET-MAD model (https://arxiv.org/abs/2503.14118) via TorchSim’s metatensor 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.)

[1]:

from torch_sim.models import MetatensorModel

model = MetatensorModel("pet-mad")

/opt/hostedtoolcache/Python/3.11.12/x64/lib/python3.11/site-packages/e3nn/o3/_wigner.py:10: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  _Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))

cuequivariance or cuequivariance_torch is not available. Cuequivariance acceleration will be disabled.

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.

[2]:

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.nvt_langevin,
    n_steps=100,
    temperature=300,  # K
    timestep=0.001,  # ps
)

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

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 metatensor and PET-MAD. For more details on how to use TorchSim, you can refer to the other tutorials in this section.