TorchSim

TorchSim is a next-generation open-source atomistic simulation engine for the MLIP era. By rewriting the core primitives of atomistic simulation in Pytorch, it allows orders of magnitude acceleration of popular machine learning potentials.

• Automatic batching and GPU memory management allowing significant simulation speedup
• Support for MACE, Fairchem, SevenNet, ORB, MatterSim, graph-pes, and metatensor MLIP models
• Support for classical lennard jones, morse, and soft-sphere potentials
• Molecular dynamics integration schemes like NVE, NVT Langevin, and NPT Langevin
• Relaxation of atomic positions and cell with gradient descent and FIRE
• Swap monte carlo and hybrid swap monte carlo algorithm
• An extensible binary trajectory writing format with support for arbitrary properties
• A simple and intuitive high-level API for new users
• Integration with ASE, Pymatgen, and Phonopy
• and more: differentiable simulation, elastic properties, custom workflows...

Quick Start

Here is a quick demonstration of many of the core features of TorchSim: native support for GPUs, MLIP models, ASE integration, simple API, autobatching, and trajectory reporting, all in under 40 lines of code.

Running batched MD

import torch
import torch_sim as ts

# run natively on gpus
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# easily load the model from mace-mp
from mace.calculators.foundations_models import mace_mp
from torch_sim.models import MaceModel
mace = mace_mp(model="small", return_raw_model=True)
mace_model = MaceModel(model=mace, device=device)

from ase.build import bulk
cu_atoms = bulk("Cu", "fcc", a=3.58, cubic=True).repeat((2, 2, 2))
many_cu_atoms = [cu_atoms] * 50
trajectory_files = [f"Cu_traj_{i}.h5md" for i in range(len(many_cu_atoms))]

# run them all simultaneously with batching
final_state = ts.integrate(
    system=many_cu_atoms,
    model=mace_model,
    n_steps=50,
    timestep=0.002,
    temperature=1000,
    integrator=ts.nvt_langevin,
    trajectory_reporter=dict(filenames=trajectory_files, state_frequency=10),
)
final_atoms_list = final_state.to_atoms()

# extract the final energy from the trajectory file
final_energies = []
for filename in trajectory_files:
    with ts.TorchSimTrajectory(filename) as traj:
        final_energies.append(traj.get_array("potential_energy")[-1])

print(final_energies)

Running batched relaxation

To then relax those structures with FIRE is just a few more lines.

# relax all of the high temperature states
relaxed_state = ts.optimize(
    system=final_state,
    model=mace_model,
    optimizer=ts.frechet_cell_fire,
    autobatcher=True,
)

print(relaxed_state.energy)

Speedup

TorchSim achieves up to 100x speedup compared to ASE with popular MLIPs.

This figure compares the time per atom of ASE and torch_sim. Time per atom is defined as the number of atoms / total time. While ASE can only run a single system of n_atoms (on the $x$ axis), torch_sim can run as many systems as will fit in memory. On an H100 80 GB card, the max atoms that could fit in memory was ~8,000 for GemNet, ~10,000 for MACE, and ~2,500 for SevenNet. This metric describes model performance by capturing speed and memory usage simultaneously.

Installation

PyPI Installation

pip install torch-sim-atomistic

Installing from source

git clone https://github.com/radical-ai/torch-sim
cd torch-sim
pip install .

Examples

To understand how TorchSim works, start with the comprehensive tutorials in the documentation.

Core Modules

TorchSim's package structure is summarized in the API reference documentation and drawn as a treemap below.

License

TorchSim is released under an MIT license.

Citation

A manuscript is in preparation. Meanwhile, if you use TorchSim in your research, please cite the Zenodo archive.