Apax: A Flexible and Performant Framework for the Development of Machine-Learned Interatomic Potentials.

Area of Science:

• Computational materials science
• Machine learning for scientific computing

Background:

• Machine-learned interatomic potentials (MLIPs) are crucial for simulating materials at the atomic scale.
• Existing MLIP frameworks can be computationally intensive and lack flexibility for advanced applications like active learning.

Purpose of the Study:

• Introduce Atomistic learned potentials in JAX (apax), an open-source software package for MLIPs.
• Demonstrate apax's capabilities in accelerating materials simulations and active learning workflows.

Main Methods:

• Developed apax leveraging the JAX framework for GPU acceleration and flexible model design.
• Implemented features for active learning: kernel-based data selection, uncertainty estimation, and enhanced sampling.
• Trained models on ionic liquid and Li3PO4 datasets, comparing performance against existing methods.

Main Results:

• Continuously learning models in apax reduced training times by up to 85% for active learning.
• A Gaussian Moment Neural Network in apax achieved higher accuracy and 10x faster inference than Allegro.
• Demonstrated good scalability in data-parallel training and modularity for complex model ensembles.

Conclusions:

• apax provides a flexible, efficient, and user-friendly platform for developing and deploying MLIPs.
• The software facilitates advanced applications like active learning, leading to faster and more accurate materials simulations.
• apax enables the development of novel MLIP models and accelerates the pace of materials discovery.