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.

翻译：随着机器学习在原子间势能领域的近期进展，Python已成为探索新思想的首选编程语言。虽然机器学习势能通常在基于Python的框架中开发，但现有的分子动力学软件主要使用低级语言编写。这种差异使得将机器学习势能集成到这些分子动力学库中变得复杂。此外，机器学习势能通常关注局部特征，常因计算复杂性而忽略长程静电作用。这是一个关键限制，因为实际应用可能需要长程静电作用甚至柔性电荷以达到所需精度。近期的电荷平衡模型可以解决这些问题，但它们需要迭代求解器为原子分配弛豫后的柔性电荷。传统实现还需要非常严格的收敛条件才能实现长期稳定性，这进一步增加了计算成本。在本工作中，我们提出了一种基于电荷平衡模型的阴影分子动力学方案的可扩展Python实现，该方案在保持长期能量稳定性和可观测量精度的同时避免了收敛问题。为了提供功能完善且用户友好的基于Python的库，我们实现了高效的邻居列表算法、粒子网格Ewald方法和传统Ewald求和技术，充分利用了Triton和PyTorch的GPU加速能力。我们将这些方法与基于Python的阴影分子动力学方案相结合，实现了涉及数十万原子体系的可扩展机器学习势能的快速电荷平衡。