期刊
JOURNAL OF STATISTICAL MECHANICS-THEORY AND EXPERIMENT
卷 2021, 期 12, 页码 -出版社
IOP Publishing Ltd
DOI: 10.1088/1742-5468/ac3ae9
关键词
machine learning; molecular dynamics; numerical simulations
JAX MD is a software package for differentiable physics simulations, focusing on molecular dynamics. It includes physics simulation environments, interaction potentials, and neural networks. The simulations are differentiable, allowing for meta-optimization.
We introduce JAX MD, a software package for performing differentiable physics simulations with a focus on molecular dynamics. JAX MD includes a number of physics simulation environments, as well as interaction potentials and neural networks that can be integrated into these environments without writing any additional code. Since the simulations themselves are differentiable functions, entire trajectories can be differentiated to perform meta-optimization. These features are built on primitive operations, such as spatial partitioning, that allow simulations to scale to hundreds-of-thousands of particles on a single GPU. These primitives are flexible enough that they can be used to scale up workloads outside of molecular dynamics. We present several examples that highlight the features of JAX MD including: integration of graph neural networks into traditional simulations, meta-optimization through minimization of particle packings, and a multi-agent flocking simulation. JAX MD is available at https://www.github.com/google/jax-md.Y
作者
我是这篇论文的作者
点击您的名字以认领此论文并将其添加到您的个人资料中。
推荐
暂无数据