Abstract
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.
S. S. Schoenholz—Work done while at Google.
This is a preview of subscription content, log in via an institution to check access.
Access this chapter
Tax calculation will be finalised at checkout
Purchases are for personal use only
References
Abadi, M., et al.: TensorFlow: a system for large-scale machine learning. In: 12th USENIX Symposium on Operating Systems Design and Implementation (OSDI 2016), pp. 265–283 (2016)
Aktulga, H.M., Fogarty, J.C., Pandit, S.A., Grama, A.Y.: Parallel reactive molecular dynamics: numerical methods and algorithmic techniques. Parallel Comput. 38(4–5), 245–259 (2012)
Aktulga, H.M., Pandit, S.A., van Duin, A.C., Grama, A.Y.: Reactive molecular dynamics: numerical methods and algorithmic techniques. SIAM J. Sci. Comput. 34(1), C1–C23 (2012)
Babuschkin, I., et al.: The DeepMind JAX ecosystem (2020). http://github.com/deepmind/jax
Battaglia, P.W., et al.: Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv:1806.01261 (2018)
Batzner, S., et al.: E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nat. Commun. 13(1), 1–11 (2022)
Behler, J., Parrinello, M.: Generalized neural-network representation of high-dimensional potential-energy surfaces. Phys. Rev. Lett. 98(14), 146401 (2007)
Bitzek, E., Koskinen, P., Gähler, F., Moseler, M., Gumbsch, P.: Structural relaxation made simple. Phys. Rev. Lett. 97(17), 170201 (2006)
Bradbury, J., Frostig, R., et al.: JAX: composable transformations of Python+ NumPy programs. Version 0.2 5, 14-24 (2018)
Brenner, D.W., Shenderova, O.A., Harrison, J.A., Stuart, S.J., Ni, B., Sinnott, S.B.: A second-generation reactive empirical bond order (REBO) potential energy expression for hydrocarbons. J. Phys.: Condens. Matter 14(4), 783 (2002)
Fogarty, J.C., Aktulga, H.M., Grama, A.Y., Van Duin, A.C., Pandit, S.A.: A reactive molecular dynamics simulation of the silica-water interface. J. Chem. Phys. 132(17), 174704 (2010)
Guo, F., et al.: Intelligent-ReaxFF: evaluating the reactive force field parameters with machine learning. Comput. Mater. Sci. 172, 109393 (2020)
Kaymak, M.C., Rahnamoun, A., O’Hearn, K.A., Van Duin, A.C., Merz Jr., K.M., Aktulga, H.M.: JAX-ReaxFF: a gradient-based framework for fast optimization of reactive force fields. J. Chem. Theory Comput. 18(9), 5181–5194 (2022)
Kylasa, S.B., Aktulga, H.M., Grama, A.Y.: PuReMD-GPU: a reactive molecular dynamics simulation package for GPUs. J. Comput. Phys. 272, 343–359 (2014)
Loshchilov, I., Hutter, F.: Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 (2017)
Mortier, W.J., Ghosh, S.K., Shankar, S.: Electronegativity-equalization method for the calculation of atomic charges in molecules. J. Am. Chem. Soc. 108(15), 4315–4320 (1986)
Paszke, A., et al.: PyTorch: an imperative style, high-performance deep learning library. Adv. Neural Inf. Process. Syst. 32 (2019)
Qiao, Z., Welborn, M., Anandkumar, A., Manby, F.R., Miller III, T.F.: OrbNet: deep learning for quantum chemistry using symmetry-adapted atomic-orbital features. J. Chem. Phys. 153(12), 124111 (2020)
Rahnamoun, A., et al.: ReaxFF/AMBER-a framework for hybrid reactive/nonreactive force field molecular dynamics simulations. J. Chem. Theory Comput. 16(12), 7645–7654 (2020)
Ramakrishnan, R., Dral, P.O., Rupp, M., Von Lilienfeld, O.A.: Quantum chemistry structures and properties of 134 kilo molecules. Sci. Data 1(1), 1–7 (2014)
Rappe, A.K., Goddard III, W.A.: Charge equilibration for molecular dynamics simulations. J. Phys. Chem. 95(8), 3358–3363 (1991)
ReaxFF, S.: Theoretical chemistry (2020)
Sabne, A.: XLA: compiling machine learning for peak performance (2020)
Schoenholz, S., Cubuk, E.D.: JAX MD: a framework for differentiable physics. Adv. Neural. Inf. Process. Syst. 33, 11428–11441 (2020)
Schütt, K., Kindermans, P.J., Sauceda Felix, H.E., Chmiela, S., Tkatchenko, A., Müller, K.R.: SchNet: a continuous-filter convolutional neural network for modeling quantum interactions. Adv. Neural. Inf. Process. Syst. 30 (2017)
Senftle, T.P., et al.: The ReaxFF reactive force-field: development, applications and future directions. NPJ Comput. Mater. 2(1), 1–14 (2016)
Tersoff, J.: Modeling solid-state chemistry: interatomic potentials for multicomponent systems. Phys. Rev. B 39(8), 5566 (1989)
Thompson, A.P., Swiler, L.P., Trott, C.R., Foiles, S.M., Tucker, G.J.: Spectral neighbor analysis method for automated generation of quantum-accurate interatomic potentials. J. Comput. Phys. 285, 316–330 (2015)
Trott, C.R., et al.: Kokkos 3: programming model extensions for the exascale era. IEEE Trans. Parallel Distrib. Syst. 33(4), 805–817 (2021)
Van Duin, A.C., Dasgupta, S., Lorant, F., Goddard, W.A.: ReaxFF: a reactive force field for hydrocarbons. J. Phys. Chem. A 105(41), 9396–9409 (2001)
Vazquez, F., Garzón, E.M., Martinez, J., Fernandez, J.: The sparse matrix vector product on GPUs. In: Proceedings of the 2009 International Conference on Computational and Mathematical Methods in Science and Engineering, vol. 2, pp. 1081–1092. Computational and Mathematical Methods in Science and Engineering Gijón, Spain (2009)
Vázquez, F., Fernández, J.J., Garzón, E.M.: A new approach for sparse matrix vector product on NVIDIA GPUs. Concurr. Comput.: Pract. Exp. 23(8), 815–826 (2011)
Verstraelen, T., Ayers, P., Van Speybroeck, V., Waroquier, M.: ACKS2: atom-condensed Kohn-Sham DFT approximated to second order. J. Chem. Phys. 138(7), 074108 (2013)
Author information
Authors and Affiliations
Corresponding author
Editor information
Editors and Affiliations
Rights and permissions
Copyright information
© 2023 The Author(s), under exclusive license to Springer Nature Switzerland AG
About this paper
Cite this paper
Kaymak, M.C., Schoenholz, S.S., Cubuk, E.D., O’Hearn, K.A., Merz Jr., K.M., Aktulga, H.M. (2023). End-to-End Differentiable Reactive Molecular Dynamics Simulations Using JAX. In: Bhatele, A., Hammond, J., Baboulin, M., Kruse, C. (eds) High Performance Computing. ISC High Performance 2023. Lecture Notes in Computer Science, vol 13948. Springer, Cham. https://doi.org/10.1007/978-3-031-32041-5_11
Download citation
DOI: https://doi.org/10.1007/978-3-031-32041-5_11
Published:
Publisher Name: Springer, Cham
Print ISBN: 978-3-031-32040-8
Online ISBN: 978-3-031-32041-5
eBook Packages: Computer ScienceComputer Science (R0)