Title: END-TO-END DIFFERENTIABLE, HARDWARE ACCELERATED, MOLECULAR DYNAMICS IN PURE PYTHON

Abstract: 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/jaxmd/jax-md along with an interactive Colab notebook that goes through all of the experiments discussed in the paper.

1. INTRODUCTION:

Understanding complex many-body systems is a challenge that underlies many of the hard problems in the physical sciences. A ubiquitous tool at our disposal in trying to understand such systems is to posit interactions between the constituents and then simulate the resulting dynamics. If interactions can be identified such that the simulation captures macroscopic behaviors observed in experiments, then the simulation can be studied to gain insight into the physical system. Since one has access to the full microscopic state at each step, it is possible to test hypotheses and make measurements that would otherwise be impossible. Such techniques, generally called molecular dynamics (MD), have been used to understand a wide range of systems including molecules, crystals, glasses, proteins, polymers, and whole biological cells.

Significant effort has gone into a number of high quality MD packages such as LAMMPS (Plimpton, 1995) , HOOMD-Blue (Anderson et al., 2008; Glaser et al., 2015) , and OpenMM (Eastman et al., 2017) . Traditional simulation environments are large and specialized codebases written in C++ or FORTRAN, along with custom CUDA kernels for GPU acceleration. These packages include significant amounts of code duplication and hand written gradients. The state of affairs is reminiscent of Machine Learning (ML) before the popularization of Automatic Differentiation (AD). Researchers trying a new idea often have to spend significant effort computing derivatives and integrating them into these large and specialized codebases. Simultaneously, the amount of data produced from MD simulations has been rapidly increasing, in part due to ever increasing computational resources along with more efficient MD software. Furthermore, deep learning is becoming a popular tool both for making MD simulations more accurate and for analyzing data produced in the simulations. Unfortunately, the issues facing MD libraries are even more pronounced when combining MD with deep learning, which typically involves complicated derivatives that can take weeks to derive and implement.

Here we introduce JAX, M.D. (JAX MD) which is a new MD package that leverages the substantial progress made in ML software to improve this state of affairs. JAX MD is end-to-end differentiable, written in pure python, and is fast since simulations are just-in-time compiled to CPU, GPU, or TPU using XLA. Moreover, JAX MD is based on JAX (Bradbury et al., 2018; Frostig et al., 2018) which has a strong neural network ecosystem that can be used seamlessly with simulations. In addition to a strong neural network ecosystem and just-in-time compilation, JAX can automatically vectorize calculations over one-or multiple-devices. This makes it easy to simulate ensembles of systems in JAX MD. We will begin with a short introduction to simulations in JAX MD followed by a brief description of JAX and some architectural choices underlying JAX MD. We then explore the features of JAX MD through several experiments:

• Efficient generation of ensembles of systems.

• Using neural networks to do machine learning of a potential.

• Meta-optimization through a simulation to optimize physical parameters.

• Cooling a glass.

While these examples are designed to be illustrative, they are similar to problems faced in actual research. Moreover, all but the last example would be significantly more difficult using existing tools.

JAX MD has so far implemented simple pairwise potentials (Lennard-Jones, soft-sphere, Morse) and the embedded atom method (EAM) (Daw & Baskes, 1984) . It can also work with the Atomic Simulation Environment (Larsen et al., 2017) and other first-principles calculations that can be accessed from Python (e.g. Quantum Espresso (Giannozzi et al., 2009) ). Due to its efficient spatial partitioning strategy, it can simulate millions of particles on a single GPU. On the ML side, JAX MD has access to all of the ML developments in JAX, including state-of-the-art convolutional networks and graph networks.

2. RELATED WORK:

Automatic differentiation has enjoyed a rich history in machine learning as well as the physical sciences (Baydin et al., 2018) . Backprop has been the core algorithm in the recent explosion of ML research, enabled by packages like TensorFlow (Abadi et al., 2016) , Torch (Collobert et al., 2002) , and Theano (Bastien et al., 2012) . While still generally focused on ML, more recent packages have made automatic differentiation more generally available as a computational tool (e.g. Autograd (Maclaurin et al., 2015) , JAX (Bradbury et al., 2018; Frostig et al., 2018) , PyTorch (Paszke et al., 2017) , Chainer (Tokui et al., 2015) , and Zygote (Innes et al., 2019) ).

In the physical sciences, automatic differentiation has been applied to a large variety of problems in structural optimization (Hoyer et al., 2019) , quantum chemistry (Tamayo-Mendoza et al., 2018) , fluid dynamics (Müller & Cusdin, 2005; Thomas et al., 2010; Bischof et al., 2007) , computational finance (Capriotti, 2010) , atmospheric modelling (Charpentier & Ghemires, 2000; Carmichael et al., 1997 ), optimal control (Walther, 2007) , physics engines (de Avila Belbute-Peres et al., 2018) , protein modelling (Ingraham et al., 2018; AlQuraishi, 2019) , and quantum circuits (Schuld et al., 2019) . For further related work at the intersection of ML and MD, please see Appendix E. Despite significant work on the topic, the number of general purpose simulation environments that are integrated with AD is scarce.

3. WARM-UP: SIMULATING A BUBBLE RAFT:

We begin with a lightning introduction to MD simulations. As an example, we're going to imagine some bubbles floating on water so that they live on a two-dimensional interface between water and air. We describe N bubbles by positions, { r i } 1≤i≤N . Since the bubbles are confined to the water's surface, the positions will be 2-dimensional vectors, r i ∈ R 2 . For simplicity, we can assume that the bubbles all are the same size and let their diameter be 1 without a loss of generality. We now have to posit interactions between the simulated bubbles that qualitatively model how real bubbles behave. More accurate interactions will produce more accurate simulations, which will in turn capture more realistic phenomena. For the purposes of this example, we assume that bubbles interact with each other in pairs. We model pairs of bubbles by defining an energy function for the pair that depends only on the distance between them. We will choose an energy that is zero if the bubbles aren't touching and then increases gradually as they get pushed together. Specifically, if r ij is the distance between bubble i and j, we use a pairwise energy function,

U (r ij ) = (1 -r ij ) 2 if r ij < 1 0 if r ij > 1 (1)

Once an energy has been defined we can compute the forces on a bubble, F i , as the negative gradient of the energy. From their definition, we see that forces move bubbles to minimize their energy. From Eq. ( 1) low energy configurations will be those where bubbles are touching as little as possible.

However, bubbles are situated on water which is full of water molecules that are moving around. These water molecules bump into the bubbles and push them slightly. To model the interaction between the bubbles and the water we will assume that there are very small random forces from the water that push the bubbles. This is a model called Brownian motion and it is described by a first-order differential equation relating the velocity of bubbles to the forces on them along with random kicks coming from the water,

EQUATION

Here F i (t) are forces, ξ i ∼ N (0, 1) is i.i.d. Gaussian distributed noise, and k B T specifies the temperature of the water. Incidentally, this model of objects in water dates back to Einstein (1905) .

In Appendix A we show an example where we define a function, final_positions = simulation(rng_key), that takes a random number generator state and returns the final positions of the particles after simulating for some time. In this example, although we only simulated a small number of bubbles we were able to emulate a much large bubble raft by using what are known as "periodic boundary conditions" (which are used in the popular game, "Asteroids"). With periodic boundary conditions bubbles can wrap around the side of the simulation to the other side, this is a ubiquitous technique for simulating the "bulk" properties of a system. In Appendix A we also show figures from a real experiment compared with the results of the simulation which shows striking similarities despite the significant simplifying assumptions we made in defining our simulation.

4. ARCHITECTURE:

We begin by briefly describing JAX before discussing the architectural choices we made in designing JAX MD. JAX is the successor to Autograd and shares key design features. As with Autograd, the main user-facing API of JAX is in one-to-one correspondence with Numpy (Van Der Walt et al., 2011) , the ubiquitous numerical computing library for Python. On top of this, JAX implements sophisticated "tracing" machinery that takes arbitrary python functions and builds an Abstract Syntax Tree (AST) for the function called a "Jaxpr". JAX includes a number of transformations that can be applied to Jaxprs to produce new Jaxprs. Examples of such transformations are: automatic differentiation (grad), vectorization on a single device (vmap), parallelization across multiple devices (pmap), and just-in-time compilation (jit). To see an example of this see Appendix D, which shows that the grad function ingests a function and returns a new, transformed function that computes its gradient. This is emblematic of JAX's functional design; all transformations take functions and return transformed functions. These function transformations are arbitrarily composable so one can write e.g. jit(vmap(grad(f))) to just-in-time compile a function that computes per example gradients of a function f . As discussed above, JAX makes heavy use of the accelerated linear algebra library, XLA, which allows compiled functions to be run in a single call on CPU, GPU, or TPU. This effectively removes execution speed issues that generally face Python programs.

JAX MD adopts a similarly functional style with immutable data and first-class functions. In a further departure from most other MD software, JAX MD features no classes or object hierarchies and instead uses a data driven approach that involves transforming arrays of data. JAX MD often uses named tuples to organize collections of arrays. This functional and data-oriented approach complements JAX's style and makes it easy to apply the range of function transformations that JAX provides. JAX MD makes extensive use of automatic differentiation and automatic vectorization to concisely express ideas (e.g. force as the negative gradient of energy) that are challenging in more conventional libraries. Since JAX MD leverages XLA to compile whole simulations to single GPU calls, it can be entirely written in Python while still being extremely fast. Together this means that implementing simulations in JAX MD looks almost verbatim like textbook descriptions of the subject.

While our use of JAX and XLA provides a significantly more productive research environment, it does have several drawbacks. Most significantly, the primitives exposed by XLA are often at odds with computations that are commonplace in molecular dynamics. Notably, XLA requires that shapes be statically known at compile-time and it is often challenging to use complex data-structures. This is makes several operations, for example spatial partitioning using cell lists, suboptimal. While JAX MD is fast enough for many research applications, since XLA has been optimized for machine learning workloads JAX MD is still slower than production quality MD packages using custom CUDA kernels. Indeed, we benchmarked JAX MD against HOOMD Blue Anderson et al. (2008) on a 4096 particle, Lennard-Jones system and observe 112µs/step for HOOMD Blue and 3200µs/step for JAX MD which represents a slowdown of around 25×. We expect this gap to shrink as XLA (and other ML language frameworks such as MLIR) continue to improve along with JAX MD. We now go through the major systems underlying JAX MD.

4.1. SPACES:

In MD we simulate a collection of N particles in either two-or three-dimensions. In the simplest case, these particles are defined by a collection of position vectors, { r i } 1≤i≤N . Some simulations are performed with r i ∈ R D where D = 2, 3 is the spatial dimension of the system; this is implemented in JAX MD using the space.free() function. However, as discussed in Section 3, it is more common to use periodic boundary conditions in which case r i ∈ R D with the association that r i = r i +Lê k for basis vectors e k and some "box size" L. In this case the simulation space is homeomorphic to a D-torus; this is implemented in JAX MD using the space.periodic(box_size) function.

These boundary conditions can be implemented by defining two functions. First, a function d : R D × R D → R D that computes the displacement between two particles. This function can in turn be used to define a metric on the space d :

R D × R D → R by d( a, b) = | d( a, b)| 2 .

Note, that in systems with periodic boundary conditions d computes the displacement between a particle and the nearest "image" of the second particle. Second, a shift function µ : R D × R D → R D must be defined that moves a particle by a shift. Motivated by this, in JAX MD we implement the spaces outlined above by functions that return a "displacement" and "shift" functions. We show an example below.

import jax.numpy as np from jax_md import space r_1 = np.array([0.5, 0, 0]) r_2 = np.array([0, 0.5, 0]) displacement, shift = space.periodic(1.) dR = displacement(r_1, r_2, t=0.1)

4.2. ENERGY AND FORCES:

As discussed above, MD simulations often proceed by defining a potential energy function, U ({ r i }), between particles. The degree to which U ({ r i }) approximates reality has a significant influence of the fidelity of the simulation. For this reason, approximating potential energy functions has received significant attention from the ML community. JAX MD allows potential energy functions to be arbitrary "JAX traceable" functions, E : R N ×D → R, including arbitrary neural networks.

JAX MD provides a number of predefined, common, choices of energy functions including several pair potentials -Lennard-Jones, soft-sphere, and Morse -as well as the Embedded Atom (Daw & Baskes, 1984) many-body potential and soft-spring bonds. Functions to compute the energy of a system are constructed by providing a displacement function for example, energy_fun = energy.soft_sphere_pair(displacement). Forces can easily be computed using a helper function quantity.force(energy_fun) which is a thin wrapper around grad. In addition to the pre-defined energy functions, it is easy to add new potential energy functions to JAX MD. In Section 5.2 we show how to add a neural network many-body potential called the Behler-Perrinello (Behler, 2011) . In sec. B we describe some additional tools provided by JAX MD to easily define custom energy functions.

In many applications, the scalar function u(r) has compact support such that u(r) = 0 if r > r c for some cutoff, r c ∈ R. We say that particles are not interacting if r ij > r c . The pairwise function defined in Eq. ( 4) is wasteful in this case since the number of computations scales as O(N 2 ) even though the total number of interactions scales as O(N ). To improve calculations in this case we provide the function cell_list_fun = smap.cell_list(fun, box_size, r_c, example_positions) that takes a function and returns a new function that uses spatial partitioning to provide a speed up.

4.3. DYNAMICS AND SIMULATIONS:

Once an energy function has been defined, there are a number of simulations that can be run. JAX MD supports simple constant energy (NVE) simulation as well as Nose-Hoover (Martyna et al., 1992) , Langevin, and Brownian simulations at constant temperature (NVT). JAX MD also supports Fast Inertial Relaxation Engine (Fire) Descent (Bitzek et al., 2006) and Gradient Descent to minimize the energy of systems. All simulations in JAX MD follow a pattern that is inspired by JAX's optimizers; for simplicity, we will use Brownian motion as an example in this section. Simulations are pairs of two functions: an initialization function, state = init_fun(key, positions), that takes particle positions and returns an initial simulation state and an update function, state = update_fun(state), that takes a state and applies a single update step to the state. To see an example of this in the case of Brownian motion, see code from the warm-up in sec. A. Simulation functions can also feature time-dependent temperatures or spaces in which case a time parameter can be passed to the update function, state = update_fun(state, t=t).

5. FOUR VIGNETTES:



5.1. VECTORIZED GENERATION OF ENSEMBLES:

Increases in computing power are increasingly due to device parallelism rather compute speed. Indeed GPUs are designed to process significant amounts of data in parallel and TPUs move futher in this direction by offering high speed interconnects between chips. This parallelism is often used to simulate ever larger systems. However, there are other interesting uses of parallelism that have received less attention. Many of these methods (e.g. replica exchange MCMC sampling (Swendsen & Wang, 1986) or nudged elastic band (Henkelman et al., 2000) ) involve simulating an ensemble of states simultaneously.

Thanks to JAX, ensembling can be done automatically in JAX MD. For small systems, the amount of necessary compute can be sub-linear in the number of replicas since it can otherwise be difficult to saturate the parallelism of accelerators. Here we go through an example where we use automatic ensembling to quickly compute statistics of a simulation. Suppose we have a function simulate(key) that simulates a single system given a random key and returns its final positions using code similar to Section 3. As discussed in Section 4 JAX includes the function vmap that automatically vectorizes computations. Here to run an ensemble of simulations we simply define, vectorized_simulation = vmap(simulation)

Fig. 1 (a) shows some example simulations of small, 32-particle systems that were performed in parallel on a single GPU. These simulations are too small to saturate the compute on a single GPU and Fig 1 (b) shows that the time-per-simulation decreases with the number of simulations being performed in parallel. This scaling continues until a batch size of about 100 when the GPU compute becomes saturated.

a b:

Figure 1 : On the left are 6 of the configurations produced by the vectorized simulation function. On the right is the time-per-simulation using the vmap functionality of JAX.

5.2. EASY MACHINE LEARNED POTENTIALS:

Historically energy functions were often derived by hand based on coarse heuristics and scarce experimental results were used to fit parameters. More recently, energy functions with a larger number of fitting parameters (e.g. ReaxFF (Van Duin et al., 2001) ) have become popular due to their ability to accurately describe certain systems. However, these methods traditionally involve significant expert knowledge and fail for systems that deviate too much from those that they were designed for. A natural progression of this trend is to use neural networks and large datasets to learn energy functions. There were a number of early efforts that received mixed success; however, it was not until 2007 when Behler and Parrinello (Behler & Parrinello, 2007) published their general purpose neural network architecture that learned energy functions emerged as a viable alternative to traditional approaches.

Since then large amounts of work has been done on this topic, however the substantial progress in machine learned potentials has not seen as much use as might be expected. At the root of this discrepancy are two points of friction at the intersection of ML and MD that prevent rapid prototyping and deployment of learned energies. First, simulation code and machine learning code are written in different languages. Second, due to the lack of automatic differentiation in molecular dynamics packages, including neural network potentials in physics simulations can require substantial work which often prohibits easy experimentation (see Eq. 3 below).

To address these issues, several projects developed adapters (Artrith & Urban, 2016; Artrith et al., 2017; Lot et al., 2019; Onat et al., 2018) between common ML languages, like Torch and Tensorflow, and common MD languages like LAMMPS. However, these solutions require researchers to be working in exactly the regime serviced by the adapter. One of the consequences of this is that the atomistic features which get fed into the neural network need to be differentiated by hand within the MD package to compute forces. Trying out a new set of features can easily take weeks or months of work to compute and implement these derivatives.

As an example, we will fit a neural network to the bubble potential defined in Eq. ( 1) and see how JAX MD gets around these issues easily. The Behler-Parrinello architecture describes the total energy of the system as a sum over individual contributions from per-atom neural networks,

U ({ r i }) = j E(θ; G j ({ r i }))

where f (θ; x) is a fully-connected neural network with parameters θ and G j ({ r i }) are hand-designed, many-body, features for a particle j. While many choices of features exist, one simple set are given by the local pair correlation function, g i (ρ) = j δ(r ij -ρ), which measures the the density of particles a distance ρ from a central particle. The Behler-Parrinello architecture can be described and initialized in two lines of python.

init_fun, E = stax.serial( stax.Dense(no_hidden_units), stax.Relu, # hidden layer 1 stax.Dense(no_hidden_units), stax.Relu, # hidden layer 2 stax.Dense(1)) # readout _, params = init_fun(key, (-1, number_of_features)) stax is JAX's native neural network library.

It is also easy to define the Behler-Parrinello loss using vmap and the JAX MD function pair_corr_fun = quantity.pair_correlation(displacement) as shown below. per_example_loss defines the MSE loss on a single state (atomic configuration) and loss is the total loss over a minibatch of states. We see a comparison between the learned energies and ground truth energies after training the above architecture for 20 seconds on 800 example states in Fig 5 .2 (a).

Figure 2 : Left panel shows the agreement between predicted energies and the correct energies of the bubble rafts in the test set. Right panel shows the distribution of the inner product between the correct force and the predicted force.

We now compute forces with JAX MD and consider how this would be implemented in a standard MD package. Recall that

F i = -∂ ri U ({ r i })

, where U is the potential energy of the system defined in eq (5.2). Thus,

EQUATION

using the chain rule. Since dE/dG j is the gradient of a neural network it is easy get in most neural network packages and feed into MD. However, traditionally dG j /dr i is a pain point and has to be coded up by hand. In JAX, MD, we get dG j /dr i for free without any extra work using JAX's grad function as grad(lambda params, r: -E(params, g(r)), argnums=1)

This energy function and force can now be used in any JAX MD simulation. In Fig. 2 (b) we see a comparison between the predictions of this network after 20 seconds of training on states generated in JAX MD. Despite the small amounts of compute involved, we see reasonable agreement for energies and forces between the machine-learned potential and the real potential.

5.3. OPTIMIZATION THROUGH DYNAMICS:

So far we have demonstrated how JAX MD can make common workloads easier. However, combining molecular dynamics with automatic differentiation opens the door for qualitatively new research. One such avenue involves differentiating through the simulation trajectory to optimize physical parameters. There have been several excellent applications so far in e.g. protein folding (AlQuraishi, 2019; Ingraham et al., 2018) , but until now this has involved significant amounts of specialized code. This vein of research is also similar to recent work in machine learning on meta-optimization (Andrychowicz et al., 2016; Metz et al., 2018) .

We revisit the bubble raft example above. In this case, we will control the structure of the bubble raft by differentiating through the simulation. As we saw in Section 3, bubble rafts form a hexagonal structure when all of the bubbles have the same size. However, when the bubbles have different sizes the situation can change considerably. To experiment with these changes, we're going to set up a simulation of a bubble raft with bubbles of two distinct sizes. To keep things simple, we'll let half of the bubbles have diameter 1 and half have diameter D. To control the conditions of the experiment, we will keep the total volume of the bubbles constant (see appendix C. Unlike the previous simulations, we will minimize the energy of the system using a function simulate(diameter, key) that returns the energy of a system given a diameter and a random key. Using vmap we can vectorize the simulation to compute the statistics for an ensemble of states at different diameters in parallel on the GPU.

Some example states at different diameters along with the energy as a function of diameter can be found in Fig 3 . We see that the hexagonal structure breaks down in the two-species case. Moreover, we see that there is a "most disordered" point when D = 0.8, which can be seen as the highest energy point in Fig. 3 (a). The study of such disordered systems is often referred to as the study of "Jammed (O'hern et al., 2003) " solids. However, this was a somewhat brute-force way to investigate the role of size-disparity in the structure of bubble rafts. Could we have seen the same result more directly? Since each energy calculation is a result of a differentiable simulation, we can differentiate through the minimization with respect to D. This would allow us to find extrema of the minimized-energy as a function of diameter using first-order optimization methods. This could be implemented in JAX MD as, dE_dD_fun = grad(simulate). Of course the dE_dD_fun function can be vectorized to aggregate statistics from an ensemble.

The gradient is plotted in Fig. 3 (b ). We see that the gradient is positive and constant for D < 0.8 corresponding to the linear increase in the average energy. Moreover, we see that the derivative crosses zero exactly at the maximum average energy. Finally, we observe that the gradient goes back to zero at D = 1. This suggests that D = 0.8 is the point of maximum disorder, as we found by brute force above. It also shows that D = 1 is the minimum energy configuration of the diameter. Although we hadn't hypothesized it, we realize this must be true since D < 1 states are symmetric with D > 1 as we keep the total packing fraction constant. On the right is the final energy as a function of the rate at which the glass was cooled.

5.4. COOLING A GLASS:

The above examples have mostly focused on small scale problems that were pedagogical; however, JAX MD is fast enough to do research at scale. To demonstrate this we consider a problem of ongoing study: the cooling of a liquid to form a glass. When liquids are cooled quickly enough, instead of crystallizing to form a solid they fall out-of-equilibrium into a glassy state. The transition of a liquid into a glass is one of the longest standing problems in condensed matter physics (Angell, 1995b) . One particularly interesting phenomenon within glass transition research is that the final properties of the glass depend on how slowly the glass itself is cooled (Angell, 1995a; Debenedetti & Stillinger, 2001; Schoenholz et al., 2017) . We consider a glass known as the Kob-Andersen Lennard-Jones mixture (Kob & Andersen, 1995) which continues to be widely studied and run several experiments cooling the liquid into a glass at different rates. We plot the results of this experiment in fig. 4 and we see that the final energy of a cooled glass gets lower as it gets cooled for longer. The difference between these two functions amounts to the choice of whether to use d pair or d bond defined above. The two-line examples above and in Section 5.2 should be contrasted with the significant undertaking that would be required to implement these features in traditional MD packages.

C OPTIMIZATION THROUGH DYNAMICS:

If there are N bubbles then the total volume of water filled by bubbles is,

V bubbles = N 8 π(D 2 + 1)

where the factor of 8 comes from the fact that our system is split into two halves and we are using diameters not radii. Since the volume of our simulation is V = L 2 if we want to keep the "packing fraction", φ = V bubbles /V constant then we will have to scale the size of the box to be, L = Vbubbles φ .

D EXAMPLE JAXPR AND ITS GRADIENT:

f (x) = x 3

Listing 2: Python 

:

CONCLUSIONWe have described recent work on a general purpose differentiable molecular dynamics physics package, JAX MD and demonstrated several instances where it simplifies the research process and enables new avenues of work.

Figures and Tables:

g = quantity.pair_correlation(displacement) U = lambda params, positions: np.sum(E(params, g(positions))) # Eq. 4. def per_example_loss(params, positions): return (U(params, positions) -energy_fun(positions)) ** 2 def loss(params, batch_positions): return np.mean(vmap(per_example_loss, in_axis=(None, 0))(params, batch_positions))

Figure 3: Panel a shows the average energy and the standard deviation of the energy at D. Panel b shows the derivative we calculate by differentiating through energy optimization by gradient descent as a function of D.

Figure4: On the left is the average energy plotted against time for a glasses cooled at different rates. On the right is the final energy as a function of the rate at which the glass was cooled.

Figure 5: An experiment of a bubble raft compared with the results of a simulation.

from jax.api import grad df_dx = grad(f) print(df_dx(2.0)) # 12.0 Listing 5: Jaxpr { lambda ; ; a. let b = pow a 2.0 c = mul 3.0 b d = safe_mul 1.0 c in [d] }

Peer review (iclr2020):
This paper announces a new software package for simulating molecular dynamics which includes close integration with a neural network / machine learning library--the first to do so. Straightforward access to hardware acceleration (e.g. GPU) is provided for both the simulation and machine learning. I lean toward accepting this submission. If it were only about simulation molecular dynamics using hardware accelerators, I would question the appropriateness of the venue, but because it is explicitly intended to support training and usage of learned potential functions, it seems suitable. Still might be better placed in a physics/chemistry venue, as where most of the references come from and likely where users would, too. The application area is no doubt an important research technique. The paper is clearly written, with enough specific examples to contrast previous pain points in this line of work against its smoother interface. All of these points are fine for a package-release/tutorial paper, but for a conference paper, might hope to see these addressed: Description of the elements of the design of JAX which are useful here are presented, and appear distinct from other AD libraries like Tensorflow or PyTorch, although the authors stop short of explicitly stating which functionality would be more difficult/impossible to support with the possible alternatives (automatic vectorization of the simulations seems like one?). Limitations of the library and drawbacks of any design decisions (something must be traded at some point?) are not explicitly mentioned. Despite mentioning numerous existing MD libraries, no performance comparison is drawn against any other. Could also show some demonstration of running an experiment which has complexity on par with state of the art research? The bubble raft example is great for illustrative purposes, but it could be better to save some of that for a tutorial and use space to exercise this library on a relevant problem and show performance there. I took a quick glance at the code on github; it is substantial but not huge, cleanly organized, and is well-documented.