msmJAX: Fast and differentiable electrostatics on the GPU in Python

We present msmJAX , a Python package implementing the multilevel summation method with B-spline interpolation, a linear-scaling algorithm for efficiently evaluating electrostatic and other long-range interactions in particle-based simulations. Built on the JAX framework, msmJAX integrates naturally with the machine-learning methods that are transforming chemistry and materials science, while also serving as a powerful tool in its own right. It combines high performance with Python’s accessibility, offers easy deployment on GPUs, and supports automatic differentiation. We outline the modular design of msmJAX , enabling users to adapt or extend the code, and present benchmarks and examples, including a verification of linear scaling, and demonstrations of its stability in molecular-dynamics simulations. PROGRAM SUMMARY Program Title: msmJAX CPC Library link to program files: https://doi.org/10.17632/8brhbk75vz.1 Developer’s repository link: https://github.com/Madsen-s-research-group/msmJAX Licensing provisions: Apache-2.0 Programming language: Python 3 Nature of problem: The problem of evaluating long-range pairwise interactions, epitomized by electrostatics, arises routinely in atomistic simulations, both in the context of classical force fields and machine-learned interaction potentials. To make it not just efficient but, in the case of periodic systems, possible at all, algorithms more sophisticated than direct summation are required. Existing implementations in established molecular-dynamics engines, developed with classical force fields in mind, do not interface well with machine-learning frameworks and workflows. Solution method: A Python and JAX implementation of the multilevel summation method (MSM) with B-spline interpolation [1], an efficient algorithm for long-range interactions that achieves linear scaling in the number of particles, is developed and released as the msmJAX package. The use of JAX brings the benefits of GPU support, automatic differentiation capabilities, compatibility with machine-learning frameworks, and a favorable balance between performance and comfort of use. Reference [1] D. J. Hardy, M. A. Wolff, J. Xia, K. Schulten, R. D. Skeel, Multilevel summation with B-spline interpolation for pairwise interactions in molecular dynamics simulations, The Journal of Chemical Physics 144 (11) (2016) 114112. doi:10.1063/1.4943868 .