# /// script
# dependencies = [
#     "graph-pes>=0.0.30",
#     "torch==2.5",
#     "vesin-torch>=0.3.7",
# ]
# ///

Integrating TorchSim with graph-pes

This brief tutorial demonstrates how to use models trained with the graph-pes package to drive MD simulations and geometry optimizations in TorchSim.

Step 1: loading a model

As an output of the graph-pes-train command, you receive a path to a .pt file containing your trained model. To use this model with TorchSim, pass the path to this .pt file, or the model itself, to the GraphPESWrapper constructor.

Below, we create a dummy TensorNet model with random weights as a demonstration:

from graph_pes.models import TensorNet, load_model

# if you had a model saved to disk, you could load it like this:
# model = load_model("path/to/model.pt")

# here, we just create a TensorNet model with random weights
model = TensorNet(cutoff=5.0)

print(f"Number of parameters: {sum(p.numel() for p in model.parameters()):,}")

Number of parameters: 56,280

Step 2: wrapping the model for use with TorchSim

We provide the GraphPESWrapper class to wrap a graph-pes model for use with TorchSim. If you intend to drive simulations that require stresses, you will need to specify the compute_stress argument to True.

from torch_sim.models.graphpes import GraphPESWrapper

# wrap the model for use with TorchSim
ts_model = GraphPESWrapper(model, compute_stress=False)

# or, alternatively, pass a model path directly:
# ts_model = GraphPESWrapper("path/to/model.pt", compute_stress=False)

Step 3: driving MD with the model

Now that we have a model, we can drive MD simulations with it. For this, we will use the integrate function.

from ase.build import molecule
import torch_sim as ts
from load_atoms import view

# NVT at 300K
atoms = molecule("H2O")

final_state = ts.integrate(
    system=atoms,
    model=ts_model,
    integrator=ts.Integrator.nvt_langevin,
    n_steps=50,
    temperature=300,
    timestep=0.001,
)

final_atoms = final_state.to_atoms()[0]
view(final_atoms, show_bonds=True)

/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))

Of course, this is a very simple example. However, you are now equipped to use any graph-pes model that you have trained to drive any of the functionality exposed by TorchSim!