JAX-ReaxFF: A Gradient-Based Framework for Fast Optimization of Reactive Force Fields.

Area of Science:

• Computational Chemistry
• Materials Science
• Chemical Physics

Background:

• Reactive force fields (ReaxFF) are crucial for simulating dynamic chemical bonding and polarizability.
• Accurate ReaxFF simulations require extensive parameter optimization against high-fidelity quantum mechanical data.
• Current optimization methods (genetic algorithms, Monte Carlo) are computationally intensive and slow.

Purpose of the Study:

• To introduce JAX-ReaxFF, a new software tool for rapid ReaxFF parameter optimization.
• To leverage machine learning infrastructure for efficient ReaxFF force field development.
• To enable faster refinement of ReaxFF parameters for improved simulation accuracy.

Main Methods:

• Utilizes the JAX library to compute gradients of the loss function.
• Employs local optimization methods initiated from multiple starting points in the parameter space.
• Designed for efficient execution on CPUs, GPUs, and TPUs.

Main Results:

• Significantly reduces ReaxFF parameter optimization time from days to minutes.
• Achieves high-quality optimization results through gradient-based local optimization.
• Demonstrates efficient performance across various hardware accelerators.

Conclusions:

• JAX-ReaxFF offers a substantial speedup for ReaxFF parameter optimization.
• The tool facilitates rapid development and refinement of accurate ReaxFF force fields.
• Provides a flexible platform for exploring modifications to the ReaxFF functional form.