Molecular Dynamics Simulations in TensorFlow

Platforms such as TensorFlow, that are customized to execute machine learning tasks, are increasingly being used to design machine learning based surrogates, integrators, and force fields for enhancing molecular dynamics simulations aimed at understanding material phenomena. The emerging centrality and integration of machine learning in molecular dynamics simulations leads to the question: Can we use existing machine learning frameworks such as TensorFlow for developing molecular dynamics simulations? We introduce TensorFlow Molecular Dynamics (TFMD), a system to develop molecular dynamics simulations and seamlessly integrate them with machine learning based enhancements. The high-level, python-based programming model provides the ability to rapidly prototype molecular dynamics simulations, tightly integrate them with machine learning techniques, and seamlessly use the next generation of cloud and HPC hardware. The capabilities of TFMD are illustrated via several examples including molecular dynamics simulations of electrolyte ions in nanoconfinement. Preliminary empirical results indicate that TensorFlow allows for rapid prototyping of molecular dynamics simulations of soft matter, and can automatically use GPUs to improve performance by more than 20x compared to CPUs.