msmJAX: Fast and Differentiable Electrostatics on the GPU in Python

Florian Buchner, Johannes Schörghuber, Nico Unglert, Jesús Carrete, Georg K. H. Madsen

arXiv:2510.05961·physics.comp-ph·Published 2025-10-07

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.

TopicsQuantum Chemistry & Force Fields

Tagsmaterials-science

arXiv categoriesphysics.comp-ph

arXiv abstract pagePDF