Simulations with TorchSim

TorchSim is a GPU-accelerated molecular simulation engine built on PyTorch. MatterSim provides a TorchSimWrapper that adapts MatterSim potentials as a TorchSim ModelInterface, so you can use TorchSim’s optimizers and integrators directly with MatterSim models.

Creating a TorchSim wrapper

The easiest way to create a wrapper is via the get_torchsim_wrapper helper, which accepts a model identifier, a checkpoint path, or an already-loaded Potential object.

import torch
from mattersim.torchsim import get_torchsim_wrapper

device = "cuda" if torch.cuda.is_available() else "cpu"

# From a model identifier (downloads the checkpoint automatically)
wrapper = get_torchsim_wrapper(potential="mattersim-v1.0.0-1M", device=device)

# Or from an already-loaded Potential
from mattersim.forcefield import Potential

potential = Potential.from_checkpoint(device=device)
wrapper = get_torchsim_wrapper(potential=potential, device=device)

You can also construct the wrapper directly if you need full control:

from mattersim.torchsim import TorchSimWrapper
from mattersim.forcefield import Potential

potential = Potential.from_checkpoint(device=device)
wrapper = TorchSimWrapper(model=potential, device=device)

Structure relaxation

TorchSim provides several optimizers for structure relaxation, including fire, lbfgs, bfgs, and gradient_descent.

import numpy as np
import torch
import torch_sim as ts
from ase.build import bulk
from mattersim.torchsim import get_torchsim_wrapper

device = "cuda" if torch.cuda.is_available() else "cpu"

# Create the wrapper
wrapper = get_torchsim_wrapper(potential="mattersim-v1.0.0-1M", device=device)

# Set up the structure
si = bulk("Si", "diamond", a=5.43)
si.positions += 0.1 * np.random.randn(len(si), 3)

# Initialize a TorchSim state from the ASE Atoms object
state = ts.initialize_state([si], device=device)

# Run relaxation with the FIRE optimizer
relaxed_state = ts.optimize(
    system=state,
    model=wrapper,
    optimizer=ts.Optimizer.fire,
    max_steps=500,
    pbar=True,
)

# Convert back to ASE Atoms
relaxed_atoms = relaxed_state.to_atoms()[0]
print(f"Relaxed energy: {relaxed_state.energy[0].item():.4f} eV")

You can also customize convergence criteria using ts.generate_force_convergence_fn:

convergence_fn = ts.generate_force_convergence_fn(force_tol=0.01)

relaxed_state = ts.optimize(
    system=state,
    model=wrapper,
    optimizer=ts.Optimizer.fire,
    max_steps=500,
    convergence_fn=convergence_fn,
    pbar=True,
)

Molecular dynamics

TorchSim supports various integrators for molecular dynamics simulations. Below we demonstrate NVT dynamics using a Langevin thermostat.

import torch
import torch_sim as ts
from ase.build import bulk
from mattersim.torchsim import get_torchsim_wrapper

device = "cuda" if torch.cuda.is_available() else "cpu"

# Create the wrapper
wrapper = get_torchsim_wrapper(potential="mattersim-v1.0.0-1M", device=device)

# Set up the structure
si = bulk("Si", "diamond", a=5.43)
state = ts.initialize_state([si], device=device)

# Run NVT MD at 300 K for 1000 steps
final_state = ts.integrate(
    system=state,
    model=wrapper,
    integrator=ts.Integrator.nvt_langevin,
    n_steps=1000,
    temperature=300.0,
    timestep=1e-3,
    pbar=True,
)

# Convert back to ASE Atoms
final_atoms = final_state.to_atoms()[0]

Available integrators include:

• NVE: ts.Integrator.nve — microcanonical ensemble

• NVT Langevin: ts.Integrator.nvt_langevin — Langevin thermostat

• NVT Nosé–Hoover: ts.Integrator.nvt_nose_hoover — Nosé–Hoover thermostat

• NPT Langevin: ts.Integrator.npt_langevin_isotropic — Langevin barostat (isotropic)

• NPT Nosé–Hoover: ts.Integrator.npt_nose_hoover_isotropic — Nosé–Hoover barostat (isotropic)

Saving trajectories

TorchSim can save trajectory frames to HDF5 files during the simulation using a TrajectoryReporter.

import torch
import torch_sim as ts
from ase.build import bulk
from mattersim.torchsim import get_torchsim_wrapper

device = "cuda" if torch.cuda.is_available() else "cpu"
wrapper = get_torchsim_wrapper(potential="mattersim-v1.0.0-1M", device=device)

si = bulk("Si", "diamond", a=5.43)
state = ts.initialize_state([si], device=device)

# Configure a trajectory reporter to save every 100 steps
reporter = ts.TrajectoryReporter(
    filenames=["md_trajectory.h5md"],
    state_frequency=100,
    state_kwargs=dict(save_velocities=True, save_forces=True),
)

final_state = ts.integrate(
    system=state,
    model=wrapper,
    integrator=ts.Integrator.nvt_langevin,
    n_steps=1000,
    temperature=300.0,
    timestep=1e-3,
    trajectory_reporter=reporter,
    pbar=True,
)

Note

For more details on TorchSim’s API, including autobatching, advanced trajectory handling, and custom integrators, please refer to the TorchSim documentation.