Parts | Days | Selection | Search | Updates | Downloads | Help
AKPIK: Arbeitskreis Physik, moderne Informationstechnologie und Künstliche Intelligenz
AKPIK 2: Machine Learning & Physics
AKPIK 2.1: Talk
Wednesday, March 20, 2024, 15:00–15:15, MAR 0.002
Bringing long-ranged interactions to the JAX ecosystem with the multilevel summation method — •Florian Buchner1, Johannes Schörghuber1, Jesús Carrete2,1, and Georg K. H. Madsen1 — 1Institute of Materials Chemistry, TU Wien, 1060 Vienna, Austria — 2Instituto de Nanociencia y Materiales de Aragón (INMA), CSIC-Universidad de Zaragoza, 50009 Zaragoza, Spain
Despite the tremendous success of machine-learned force fields (MLFFs), their extension beyond the locality approximation remains a field of active research. In constructing such MLFFs, the efficient (ideally with linear scaling) evaluation of pairwise long-ranged interactions is a ubiquitous requirement. It is routinely solved by well-established algorithms such as Ewald summation.
While implementations of such algorithms are readily available, they tend not to interface well with modern machine-learning environments and workflows. This includes Google’s JAX framework, which is proving transformative to machine-learning research by providing high performance and general-purpose automatic differentiation.
Here, we present a JAX-based implementation of the multilevel summation method (MSM) [D. J. Hardy et al., J. Chem. Phys. 144, 114112 (2016)], a powerful linearly scaling algorithm for pairwise long-ranged interactions. Its notable features include support for mixed boundary conditions and freedom from artefactual force discontinuities encountered in competing methods. We introduce the basics of the MSM, discuss our design and implementation strategy, and highlight example applications.
Keywords: Electrostatics; Long-ranged interactions; Machine-Learned Force Fields; Automatic Differentiation; Atomistic Simulation