Search arXiv⌕ Search

arXiv · 2510.05961

msmJAX: Fast and Differentiable Electrostatics on the GPU in Python

Abstract

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.

Explore related subjects

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Florian Buchner, Johannes Schörghuber, Nico Unglert, Jesús Carrete, Georg K. H. Madsen. 2025-10-07. msmJAX: Fast and Differentiable Electrostatics on the GPU in Python. https://arxiv.org/abs/2510.05961

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

A Physics-Driven Framework for Parametric Periodic-Flow Modeling and Finite-Amplitude Aeroelastic Response Analysis

Periodic unsteady flows are common in forced-motion and fluid-structure interaction problems. Their parametric analysis typically requires repeated high-fidelity simulations, whereas existing reduced-order and surrogate models generally rely on pre-generated flow-field or aerodynamic data. This study proposes a purely physics-driven framework for solving parametric periodic flows and finite-amplitude aeroelastic responses. First, a Periodic Physics-Informed Neural Network (P-PINN) is developed to directly solve periodic flows by imposing temporal periodicity over a single motion cycle, thereby avoiding the need to resolve the long transient evolution preceding the establishment of the periodic state. The flow conditions and motion parameters are further incorporated as network inputs to construct continuous parametric representations of the periodic flow field and aerodynamic forces. On this basis, the parametric aerodynamic model is coupled with the structural dynamic equation through first-order harmonic balance to solve the response amplitude and frequency of a single-degree-of-freedom aeroelastic system. This framework is validated using forced-motion cases of a circular cylinder and an airfoil, demonstrating accurate reproduction of periodic aerodynamic forces, surface load distributions, and instantaneous flow fields under different flow and motion parameters compared with time-marching results. Furthermore, aeroelastic analysis is conducted for an elastically mounted circular cylinder at subcritical Reynolds numbers, and the resulting aeroelastic response agrees well with fully coupled CFD/CSD results. Once trained offline, the parametric model can be repeatedly evaluated for different structural parameter states, enabling the complete aeroelastic response curve to be obtained online within seconds, without repeated long-time fluid-structure interaction time marching.

physics.comp-ph↗

Monte Carlo sampling of first-order QED processes in laser and pulsar plasmas

Monte Carlo sampling of strong-field quantum electrodynamics processes underpins simulations of high-intensity laser experiments and of astrophysical compact-object magnetospheres. Sampling an event requires the total rate of the process together with the cumulative probability that determines how energy is partitioned between the produced particles. Simulations typically tabulate both in advance and invert the tabulated probability numerically. Here we replace this procedure with elementary-function approximations for synchrotron radiation and the nonlinear Breit--Wheeler process. For each process, we approximate the auxiliary function that sets the total rate, as well as the cumulative probability, with Padé approximants chosen so that the inversion reduces to a quartic equation. This yields the sampled quantum parameter---electron $χ_e$ or photon $χ_γ$---in closed form. The approximations and the particle spectra sampled from them agree with the exact results to within $1\%$. The procedure requires no lookup tables, no interpolation, and no numerical root finding, and can be inserted directly into radiative particle-in-cell codes.

physics.comp-ph↗

Optimal Bias Potentials via Ergodic Optimal Control and Generator Learning

We investigate the computation of optimal bias potentials for accelerating transitions between metastable states and for computation of equilibrium properties in molecular dynamics simulations. We formulate optimal biasing as an ergodic optimal control problem (OCP), which can be recast as a linear eigenvalue problem for the infinitesimal generator of the unbiased dynamics. We demonstrate that data-driven learning methods for the generator enable reliable solution of the OCP, computation of biasing potentials, extraction of equilibrium properties, and acceleration of state transitions. We also explore the relation of the control problem to coarse grained representations and learning of coarse grained dynamics.

physics.comp-ph↗