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

2025-03-27

AI Summary

This work addresses the slow convergence and energy instability in machine learning interatomic potentials arising from dynamic charge modeling and long-range electrostatic coupling. We propose a non-iterative charge equilibration dynamics scheme embedded within the shadow molecular dynamics framework. Methodologically, we integrate Particle Mesh Ewald (PME) with conventional Ewald summation, leverage GPU acceleration via Triton/PyTorch, employ a custom high-efficiency neighbor list, and adopt an explicit charge-relaxation integrator. To our knowledge, this is the first Python-based implementation enabling large-scale, high-stability, long-timescale charge-responsive MD simulations for systems comprising hundreds of thousands of atoms. Compared to conventional iterative charge equilibration methods, our approach achieves significantly faster convergence and superior energy conservation, while preserving accuracy in structural and dynamical observables—including radial distribution functions and diffusion coefficients. The framework establishes a new paradigm for scalable, physically consistent, ML-driven molecular dynamics.

Technology Category

• Machine Learning: Hardware-aware ML
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems
• Multiagent Systems: Mechanism Design

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Economics, Online Markets and Human Computation: Architectures and workflows that use LLMs for crowd work
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• Integrating machine learning potentials with molecular dynamics in Python
• Addressing neglect of long-range electrostatics in machine learning potentials
• Achieving stable charge equilibration without high computational cost

Innovation

• GPU-accelerated charge equilibration for scalable MD
• Python-based shadow molecular dynamics with stability
• Efficient neighbor list and Ewald summation techniques