TorchSim: A Next-Generation PyTorch-Native Atomistic Simulation Engine for the MLIP Era
Radical AI has released TorchSim, a next-generation PyTorch-native atomistic simulation engine for the MLIP era. It accelerates materials simulation by orders of magnitude, transforming traditional scientific approaches. Current materials research requires large teams focused on single problems, resulting in slow progress and high costs. Radical AI aims to revolutionize this paradigm by enabling individual scientists to tackle multiple challenges simultaneously through AI and autonomous systems. TorchSim serves as the first public demonstration of this scientific approach, which allows real-time correlation between measured material properties and simulations at an unprecedented scale when integrated with self-driving laboratories.
TorchSim transforms atomistic simulation within PyTorch, delivering 100 times speedup compared to ASE and 100,000,000 times acceleration over DFT. TorchSim reimplements the most popular molecular dynamics and optimization algorithms, including NVE, NVT, NPT, gradient descent, and Frechet cell FIRE, while offering a user-friendly API with trajectory reporting, automatic memory management, and integration with established materials software and machine learning libraries. Radical AI released TorchSim as open-source software while maintaining it as one component of their Materials Flywheel™ ecosystem. The company aims to supply advanced materials to critical industries while accelerating materials development through this simulation revolution.
TorchSim simplifies atomistic simulation through a comprehensive high-level API featuring three primary “runner” functions: integrate for molecular dynamics, optimize for relaxation, and static for static evaluation. These functions share similar signatures while supporting auto batching, trajectory reporting, diverse models, and compatibility with popular libraries. The framework accommodates various simulation types, including NVT/NPT integration and gradient descent/FIRE optimization methods. The SimState is the core atomistic representation for the TorchSim package, containing atoms, atomic numbers, cell data, and all necessary simulation elements. SimState uses PyTorch tensors as attributes and employs a batched structure capable of representing single or multiple systems simultaneously.
TorchSim addresses the complex challenge of efficient GPU memory utilization during batched operations. Different models require varying memory allocations for identical systems, while memory footprint scaling depends on neighbor list computation methods. For instance, MACE models scale with the product of atom count and number density, whereas Fairchem models scale only with atom count. TorchSim dynamically determines model memory requirements and optimally arranges simulations to maximize available memory utilization. This intelligent management works across molecular dynamics simulations and optimization processes, ensuring computational resources are used efficiently throughout different simulation types.
TorchSim introduces a novel trajectory format designed for native integration with its batched state system, supporting binary encoding of diverse properties and real-time compression. Despite recognizing the current abundance of trajectory formats, developers determined that creating a new format was necessary to fulfill project requirements. The resulting TorchSimTrajectory is built on HDF5 and works as an efficient container for arbitrary arrays with utilities optimized for atomistic simulation. It utilizes consistent binary encoding and compression across all properties, including temperature, forces, per-atom energies, and electric fields, enabling comprehensive and efficient data management.
TorchSim welcomes community feedback as an experimental library. Contributors must first sign Radical AI’s Contributor License Agreement (CLA), a one-time requirement covering all Radical AI open source projects. This agreement allows contributors to retain ownership of their work while granting Radical AI necessary usage rights. The CLA-bot automatically verifies signatures on pull requests. All code submissions undergo mandatory review by project maintainers before merging. Contributors should submit changes through GitHub pull requests, with even maintainers’ submissions requiring review from other maintainers. Prompt responses to feedback and requested changes are expected throughout the review process.
Check out the Technical Details and GitHub Page. All credit for this research goes to the researchers of this project. Also, feel free to follow us on Twitter and don’t forget to join our 85k+ ML SubReddit.
Sajjad Ansari is a final year undergraduate from IIT Kharagpur. As a Tech enthusiast, he delves into the practical applications of AI with a focus on understanding the impact of AI technologies and their real-world implications. He aims to articulate complex AI concepts in a clear and accessible manner.
