skip to main content


The NSF Public Access Repository (NSF-PAR) system and access will be unavailable from 11:00 PM ET on Thursday, June 13 until 2:00 AM ET on Friday, June 14 due to maintenance. We apologize for the inconvenience.

Title: End-to-End Differentiable Reactive Molecular Dynamics Simulations Using JAX
The reactive force field (ReaxFF) interatomic potential is a powerful tool for simulating the behavior of molecules in a wide range of chemical and physical systems at the atomic level. Unlike traditional classical force fields, ReaxFF employs dynamic bonding and polarizability to enable the study of reactive systems. Over the past couple decades, highly optimized parallel implementations have been developed for ReaxFF to efficiently utilize modern hardware such as multi-core processors and graphics processing units (GPUs). However, the complexity of the ReaxFF potential poses challenges in terms of portability to new architectures (AMD and Intel GPUs, RISC-V processors, etc.), and limits the ability of computational scientists to tailor its functional form to their target systems. In this regard, the convergence of cyber-infrastructure for high performance computing (HPC) and machine learning (ML) presents new opportunities for customization, programmer productivity and performance portability. In this paper, we explore the benefits and limitations of JAX, a modern ML library in Python representing a prime example of the convergence of HPC and ML software, for implementing ReaxFF. We demonstrate that by leveraging auto-differentiation, just-in-time compilation, and vectorization capabilities of JAX, one can attain a portable, performant, and easy to maintain ReaxFF software. Beyond enabling MD simulations, end-to-end differentiability of trajectories produced by ReaxFF implemented with JAX makes it possible to perform related tasks such as force field parameter optimization and meta-analysis without requiring any significant software developments. We also discuss scalability limitations using the current version of JAX for ReaxFF simulations.  more » « less
Award ID(s):
Author(s) / Creator(s):
; ; ; ; ;
Bhatele, A.; Hammond, J.; Baboulin, M.; Kruse, C.
Date Published:
Journal Name:
Lecture notes in computer science
Page Range / eLocation ID:
Medium: X
Sponsoring Org:
National Science Foundation
More Like this
  1. null (Ed.)
    Reactive force elds provide an a ordable model for simulating chemical reactions at a fraction of the cost of quantum mechanical approaches. However classically accounting for chemical reactivity often comes at the expense of accuracy and transferability, while computational cost is still large relative to non-reactive force elds. In this Perspective we summarize recent e orts for improving the performance of reactive force elds in these three areas with a focus on the ReaxFF theoretical model. To improve accuracy we describe recent reformulations of charge equilibration schemes to overcome unphysical long-range charge transfer, new ReaxFF models that account for explicit electrons, and corrections for energy conservation issues of the ReaxFF model. To enhance transferability we also highlight new advances to include explicit treatment of electrons in the ReaxFF and hybrid non-reactive/reactive simulations that make it possible to model charge transfer, redox chemistry, and large systems such as reverse micelles within the framework of a reactive force eld. To address the computational cost we review recent work in extended Lagrangian schemes and matrix preconditioners for accelerating the charge equilibration method component of ReaxFF and improvements in its software performance in LAMMPs. 
    more » « less
  2. null (Ed.)
    Abstract: Solid-state ion conduction (SSIC) is a mechanism of ionic current that has garnered increasing attention for applications in all-solid-state batteries and atomic switches. The Ag/S SSIC system in β-Ag S, possessing the highest ionic conductivity of any known material, provides a unique opportunity to better understand the fundamental nature of SSIC. β-Ag S is topographically similar to binary perovskites except that it is cubic, leading to isotropic SSIC exceeding 4 S/cm. The dynamic nature of SSIC makes it difficult to study by observational means, where inherent time-averaging obscures correlations among atomic transit routes.Molecular dynamics (MD) is a tool ideally suited for gaining insight into large atomic systems with subnanosecond time resolutions. However, traditional MD potentials lack a description of bond-breaking/forming reactions, which are an essential aspect of SSIC and related memristic properties. This limitation can be overcome by using a reactive force field (ReaxFF), which enables the simulation of bonding reactions with DFT-level accuracy. In this study, we present a ReaxFF force field for the Ag/S system, optimized for simulating SSIC in β-Ag S. Training data consisted of crystal structures, Bader partial charges, and energies of various Ag/S clusters calculated at the DFT-level. Energies were obtained with Gaussian 16, using the PBEh1PBE hybrid functional with a triple-zeta correlation-consistent basis set. Multiobjective parameter optimization was accomplished with an updated form of the Genetic Algorithm for Reactive Force Fields (GARFfield). The force field was validated with potential energy and ion conductivity calculations, along with relevant structural features. Results were compared with equivalent simulations from other established potentials. This new ReaxFF force field will enable modeling of realistic SSIC configurations for Ag/S-based materials and provides a viable approach for extending ReaxFF to other SSIC systems in the future. This work was supported by the National Science Foundation under grant #2025319. 
    more » « less

    Models of stellar population synthesis (SPS) are the fundamental tool that relates the physical properties of a galaxy to its spectral energy distribution (SED). In this paper, we present DSPS: a python package for SPS. All of the functionality in DSPS is implemented natively in the JAX library for automatic differentiation, and so our predictions for galaxy photometry are fully differentiable, and directly inherit the performance benefits of JAX, including portability onto GPUs. DSPS also implements several novel features, such as i) a flexible empirical model for stellar metallicity that incorporates correlations with stellar age, ii) support for the Diffstar model that provides a physically-motivated connection between the star formation history of a galaxy (SFH) and the mass assembly of its underlying dark matter halo. We detail a set of theoretical techniques for using autodiff to calculate gradients of predictions for galaxy SEDs with respect to SPS parameters that control a range of physical effects, including SFH, stellar metallicity, nebular emission, and dust attenuation. When forward modelling the colours of a synthetic galaxy population, we find that DSPS can provide a factor of 5 speed-up over standard SPS codes on a CPU, and a factor of 300-400 on a modern GPU. When coupled with gradient-based techniques for optimization and inference, DSPS makes it practical to conduct expansive likelihood analyses of simulation-based models of the galaxy–halo connection that fully forward model galaxy spectra and photometry.

    more » « less
  4. null (Ed.)
    The Interface force field (IFF) enables accurate simulations of bulk and interfacial properties of compounds and multiphase materials. However, the simulation of reactions and mechanical properties up to failure remains challenging and expensive. Here we introduce the Reactive Interface Force Field (IFF-R) to analyze bond breaking and failure of complex materials using molecular dynamics simulations. IFF-R uses a Morse potential instead of a harmonic potential as typically employed in molecular dynamics force fields to describe the bond energy, which can render any desired bond reactive by specification of the curve shape of the potential energy and the bond dissociation energy. This facile extension of IFF and other force fields that utilize a harmonic bond energy term allows the description of bond breaking without loss in functionality, accuracy, and speed. The method enables quantitative, on-the-fly computations of bond breaking and stress-strain curves up to failure in any material. We illustrate accurate predictions of mechanical behavior for a variety of material systems, including metals (iron), ceramics (carbon nanotubes), polymers (polyacrylonitrile and cellulose I\b{eta}), and include sample parameters for common bonds based on using experimental and high-level (MP2) quantum mechanical reference data. Computed structures, surface energies, elastic moduli, and tensile strengths are in excellent agreement with available experimental data. Non-reactive properties are shown to be essentially identical to IFF values. Computations are approximately 50 times faster than using ReaxFF and require only a single set of parameters. Compatibility of IFF and IFF-R with biomolecular force fields allows the quantitative analysis of the mechanics of proteins, DNA, and other biological molecules. 
    more » « less
  5. Abstract Background Bioinformatic workflows frequently make use of automated genome assembly and protein clustering tools. At the core of most of these tools, a significant portion of execution time is spent in determining optimal local alignment between two sequences. This task is performed with the Smith-Waterman algorithm, which is a dynamic programming based method. With the advent of modern sequencing technologies and increasing size of both genome and protein databases, a need for faster Smith-Waterman implementations has emerged. Multiple SIMD strategies for the Smith-Waterman algorithm are available for CPUs. However, with the move of HPC facilities towards accelerator based architectures, a need for an efficient GPU accelerated strategy has emerged. Existing GPU based strategies have either been optimized for a specific type of characters (Nucleotides or Amino Acids) or for only a handful of application use-cases. Results In this paper, we present ADEPT, a new sequence alignment strategy for GPU architectures that is domain independent, supporting alignment of sequences from both genomes and proteins. Our proposed strategy uses GPU specific optimizations that do not rely on the nature of sequence. We demonstrate the feasibility of this strategy by implementing the Smith-Waterman algorithm and comparing it to similar CPU strategies as well as the fastest known GPU methods for each domain. ADEPT’s driver enables it to scale across multiple GPUs and allows easy integration into software pipelines which utilize large scale computational systems. We have shown that the ADEPT based Smith-Waterman algorithm demonstrates a peak performance of 360 GCUPS and 497 GCUPs for protein based and DNA based datasets respectively on a single GPU node (8 GPUs) of the Cori Supercomputer. Overall ADEPT shows 10x faster performance in a node-to-node comparison against a corresponding SIMD CPU implementation. Conclusions ADEPT demonstrates a performance that is either comparable or better than existing GPU strategies. We demonstrated the efficacy of ADEPT in supporting existing bionformatics software pipelines by integrating ADEPT in MetaHipMer a high-performance denovo metagenome assembler and PASTIS a high-performance protein similarity graph construction pipeline. Our results show 10% and 30% boost of performance in MetaHipMer and PASTIS respectively. 
    more » « less