JAX-ReaxFF

JAX-ReaxFF optimizes ReaxFF reactive force field parameters using JAX automatic differentiation to fit parameters to quantum mechanical (QM) training data for molecular simulations.

Key Features:

• Automatic differentiation (JAX): Uses the JAX library to compute efficient gradients of the loss function via automatic differentiation.
• Gradient-based local optimization: Employs gradient-based local optimization methods initiated from multiple starting points in high-dimensional parameter spaces.
• Comparison to stochastic methods: Serves as an alternative to genetic algorithms and Monte Carlo approaches that require millions of error evaluations.
• Hardware portability: Executes across multicore CPUs, GPUs, and TPUs for accelerated computation.
• Speed-up in optimization: Reduces ReaxFF parameter optimization time from days to minutes as reported in the description.
• ReaxFF model support: Targets optimization for the ReaxFF reactive force field, which incorporates dynamic bonding and polarizability features.
• Functional-form sandbox: Provides an environment for exploring custom modifications to the ReaxFF functional form.
• QM training-data fitting: Fits ReaxFF parameters against high-fidelity quantum mechanical (QM) training data.

Scientific Applications:

• Force field parameterization: Development and refinement of ReaxFF parameter sets for molecular dynamics simulations.
• QM-guided fitting: Generating parameter sets that reproduce quantum mechanical reference energetics and properties.
• Functional-form development: Testing and validating custom modifications to the ReaxFF functional form.
• Accelerated iterative development: Shortening turnaround time for iterative force field refinement and testing.

Methodology:

Computes gradients with JAX automatic differentiation and applies gradient-based local optimization initiated from multiple starting points in the parameter space; execution is portable across multicore CPUs, GPUs, and TPUs and is presented as an alternative to stochastic genetic-algorithm and Monte Carlo approaches that perform millions of error evaluations.