JAX, MD A framework for differentiable physics
JAX, MD A framework for differentiable physics
复制标题
DOI：
10.1088/1742-5468/ac3ae9
复制
发表时间：
2021-12-01
影响因子：
2.4
通讯作者：
Cubuk, Ekin D.
中科院分区：
文献类型：
--
作者：
Schoenholz, Samuel S.;Cubuk, Ekin D.
我们介绍了JAX MD，这是一个用于执行可微物理模拟的软件包，重点是分子动力学。JAX MD包括许多物理模拟环境，以及可以集成到这些环境中的交互潜力和神经网络，而无需编写任何额外的代码。由于仿真本身是可微函数，因此可以对整个轨迹进行区分以执行元优化。这些功能建立在空间划分等基本操作的基础上，这些操作允许在单个GPU上将模拟扩展到数十万个粒子。这些原语足够灵活，可以用来扩大分子动力学之外的工作量。我们提供了几个突出JAX MD特征的例子，包括：将图神经网络集成到传统模拟中，通过最小化粒子填充进行元优化，以及多智能体群集模拟。JAX MD可在https://www.github.com/google/jax-md.Y上获得
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