Abstract:The reactive force field (ReaxFF) model bridges the gap between traditional classical models and quantum mechanical (QM) models by incorporating dynamic bonding and polarizability. To achieve realistic simulations using ReaxFF, model parameters must be optimized against high fidelity training data which typically come from QM calculations. Existing parameter optimization methods for ReaxFF consist of black box techniques using genetic algorithms or Monte Carlo methods. Due to the stochastic behavior of these methods, the optimization process oftentimes requires millions of error evaluations for complex parameter fitting tasks, thereby significantly hampering the rapid development of high quality parameter sets. Rapid optimization of the parameters is essential for developing and refining Reax force fields because producing a force field which exhibits empirical accuracy in terms of dynamics typically requires multiple refinements to the training data as well as to the parameters under optimization. In this work, we present JAX-ReaxFF, a novel software tool that leverages modern machine learning infrastructure to enable fast optimization of ReaxFF parameters. By calculating gradients of the loss function using the JAX library, JAX-ReaxFF utilizes highly effective local optimization methods that are initiated from multiple guesses in the high dimensional optimization space to obtain high quality results. Leveraging the architectural portability of the JAX framework, JAX-ReaxFF can execute efficiently on multicore CPUs, graphics processing units (GPUs), or even tensor processing units (TPUs). As a result of using the gradient information and modern hardware accelerators, we are able to decrease ReaxFF parameter optimization time from days to mere minutes. Furthermore, the JAX-ReaxFF framework can also serve as a sandbox environment for domain scientists to explore customizing the ReaxFF functional form for more accurate modeling.

JAX FDM: A differentiable solver for inverse form-finding

JAX-SSO: Differentiable Finite Element Analysis Solver for Structural Optimization and Seamless Integration with Neural Networks

Fabrication-Constrained Conjugate Adjoint Gradient Inverse Design for Large-Scale Aperiodic Frequency Selective Surface

JAX-FLUIDS: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows

JAX-based differentiable fluid dynamics on GPU and end-to-end optimization

Inverse Design Tool for Asymmetrical Self-Rising Surfaces with Color Texture.

A framework for structural shape optimization based on automatic differentiation, the adjoint method and accelerated linear algebra

JAX-SPH: A Differentiable Smoothed Particle Hydrodynamics Framework

JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows

Data-driven Design of Well-Behaved Nonlinear Structures: A Case Study on the Von Mises Truss

Differentiable hybrid neural modeling for fluid-structure interaction

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

jinns: a JAX Library for Physics-Informed Neural Networks

MPAX: Mathematical Programming in JAX

JAXFit: Trust Region Method for Nonlinear Least-Squares Curve Fitting on the GPU

Differentiable Conservative Radially Symmetric Fluid Simulations and Stellar Winds -- jf1uids

SynJax: Structured Probability Distributions for JAX

Rieoptax: Riemannian Optimization in JAX

EDACFEM: A linear truss and beam solver in MATLAB

A differentiable structural analysis framework for high-performance design optimization

Highly Efficient Feasible Direction Method (hefdim) for Structural Topology Optimization