JAX, M.D.: End-to-End Differentiable, Hardware Accelerated, Molecular Dynamics in Pure Python

12/09/2019
by   Samuel S. Schoenholz, et al.
0

A large fraction of computational science involves simulating the dynamics of particles that interact via pairwise or many-body interactions. These simulations, called Molecular Dynamics (MD), span a vast range of subjects from physics and materials science to biochemistry and drug discovery. Most MD software involves significant use of handwritten derivatives and code reuse across C++, FORTRAN, and CUDA. This is reminiscent of the state of machine learning before automatic differentiation became popular. In this work we bring the substantial advances in software that have taken place in machine learning to MD with JAX, M.D. (JAX MD). JAX MD is an end-to-end differentiable MD package written entirely in Python that can be just-in-time compiled to CPU, GPU, or TPU. JAX MD allows researchers to iterate extremely quickly and lets researchers easily incorporate machine learning models into their workflows. Finally, since all of the simulation code is written in Python, researchers can have unprecedented flexibility in setting up experiments without having to edit any low-level C++ or CUDA code. In addition to making existing workloads easier, JAX MD allows researchers to take derivatives through whole-simulations as well as seamlessly incorporate neural networks into simulations. This paper explores the architecture of JAX MD and its capabilities through several vignettes. Code is available at www.github.com/google/jax-md. We also provide an interactive Colab notebook that goes through all of the experiments discussed in the paper.

READ FULL TEXT

page 8

page 14

research
12/22/2020

TorchMD: A deep learning framework for molecular simulations

Molecular dynamics simulations provide a mechanistic description of mole...
research
12/20/2021

NetKet 3: Machine Learning Toolbox for Many-Body Quantum Systems

We introduce version 3 of NetKet, the machine learning toolbox for many-...
research
07/17/2023

SBMLtoODEjax: efficient simulation and optimization of ODE SBML models in JAX

Developing methods to explore, predict and control the dynamic behavior ...
research
10/22/2021

DQC: a Python program package for Differentiable Quantum Chemistry

Automatic differentiation represents a paradigm shift in scientific prog...
research
07/09/2020

ACORNS: An Easy-To-Use Code Generator for Gradients and Hessians

The computation of first and second-order derivatives is a staple in man...
research
02/05/2019

An open-source multi-person virtual reality framework for interactive molecular dynamics: from quantum chemistry to drug binding

Conceptually, the mechanics of nanoscale molecular objects arise through...
research
11/07/2022

pyGSL: A Graph Structure Learning Toolkit

We introduce pyGSL, a Python library that provides efficient implementat...

Please sign up or login with your details

Forgot password? Click here to reset