GPU-Accelerated Charge-Equilibration for Shadow Molecular Dynamics in Python

TL;DR

This work addresses the challenge of integrating machine-learning interatomic potentials with classical MD by introducing a GPU-accelerated Python framework for shadow Born-Oppenheimer MD with a charge equilibration (QEQ) model to capture long-range electrostatics. It combines fast, scalable Python implementations of neighbor lists, regular Ewald, and PME electrostatics using PyTorch and Triton, with an extended Lagrangian (shadow) formulation that relaxes iterative convergence requirements yet preserves conservative forces. The authors demonstrate stability and accuracy across large systems (hundreds of thousands of atoms) and report substantial speedups (up to 3x) for charge equilibration in shadow MD, while maintaining energy conservation and matching observables such as radial distribution functions and charge distributions. This open-source framework lowers the barrier to deploying physically faithful electrostatics in ML-driven potentials and paves the way for multi-GPU scaling, multipole extensions, and broader MD integrations in Python-based workflows.

Abstract

With recent advancements in machine learning for interatomic potentials, Python has become the go-to programming language for exploring new ideas. While machine-learning potentials are often developed in Python-based frameworks, existing molecular dynamics software is predominantly written in lower-level languages. This disparity complicates the integration of machine learning potentials into these molecular dynamics libraries. Additionally, machine learning potentials typically focus on local features, often neglecting long-range electrostatics due to computational complexities. This is a key limitation as applications can require long-range electrostatics and even flexible charges to achieve the desired accuracy. Recent charge equilibration models can address these issues, but they require iterative solvers to assign relaxed flexible charges to the atoms. Conventional implementations also demand very tight convergence to achieve long-term stability, further increasing computational cost. In this work, we present a scalable Python implementation of a recently proposed shadow molecular dynamics scheme based on a charge equilibration model, which avoids the convergence problem while maintaining long-term energy stability and accuracy of observable properties. To deliver a functional and user-friendly Python-based library, we implemented an efficient neighbor list algorithm, Particle Mesh Ewald, and traditional Ewald summation techniques, leveraging the GPU-accelerated power of Triton and PyTorch. We integrated these approaches with the Python-based shadow molecular dynamics scheme, enabling fast charge equilibration for scalable machine learning potentials involving systems with hundreds of thousands of atoms.

Figures (9)

• Figure 1: Illustration of the cell based neighbor search method.
• Figure 2: Example NVE simulations from \ref{['table:solver_valid']}. Y-axis shows the deviation of total energy from the average energy ($E_{avg}$).
• Figure 3: Radial distribution functions for O-O (left) and O-H (right) pairs.
• Figure 4: Charge distribution of hydrogen (left) and oxygen (right) atoms.
• Figure 5: Performance of neighbor list generation for the COO format (left panel) and the ELLPACK format (right panel) on an A100 GPU with 40 GB VRAM using randomly generated water-like systems. Missing data points indicate GPU memory limitations.