TorchSim: An efficient atomistic simulation engine in PyTorch

TL;DR

TorchSim is a GPU-accelerated, batched atomistic simulation engine built in PyTorch that significantly speeds up molecular dynamics and structural relaxation for multiple systems simultaneously, supporting various potentials and differentiable simulation.

Contribution

TorchSim introduces a PyTorch-based framework for efficient batched atomistic simulations, enabling faster MLIP applications and seamless integration with materials informatics tools.

Findings

Achieves orders of magnitude acceleration over traditional methods.

Supports both machine-learned and classical interatomic potentials.

Enables differentiable simulation for advanced analysis.

Abstract

We introduce TorchSim, an open-source atomistic simulation engine tailored for the Machine Learned Interatomic Potential (MLIP) era. By rewriting core atomistic simulation primitives in PyTorch, TorchSim can achieve orders of magnitude acceleration for popular MLIPs. Unlike existing molecular dynamics packages, which simulate one system at a time, TorchSim performs batched simulations that efficiently utilize modern GPUs by evolving multiple systems concurrently. TorchSim supports molecular dynamics integrators, structural relaxation optimizers, both machine-learned and classical interatomic potentials (such as Lennard-Jones, Morse, soft-sphere), batching with automatic memory management, differentiable simulation, and integration with popular materials informatics tools.