(411f) HOOMD-TF: Online Machine Learning in Molecular Dynamics Simulations

We present a novel tool for low-latency online machine learning and collective variable calculation in the HOOMD-blue molecular dynamics (MD) engine. A common issue when pairing machine learning (ML) methods with MD simulations is a high barrier to entry for chemical physicists. This difficulty arises in both offline and online learning situations, in terms of pre- and post-processing of data in the former, or implementation and modification of simulation engine code in the latter. Here, we introduce a solution to this problem that bridges TensorFlow, Google's GPU-accelerated ML library, with HOOMD-blue, a GPU-accelerated MD engine. Our software uses a low-latency GPU-GPU communication scheme to expose the HOOMD neighbor lists of atoms to TensorFlow for reading, and optionally allows TensorFlow to output forces or collective variables calculated from these neighbor lists to HOOMD-blue. This also enables online learning using any value that can be specified as a tensor operation on the HOOMD-blue neighbor list. This gives rise to a wide variety of ML applications, such as coarse-grain force-field learning, calculation of arbitrary collective variables, biasing along arbitrary collective variables, and force-field calculation. In this work, we show four examples of these applications. We train a neural network model to reproduce the Lennard-Jones forces experienced by 10000 Lennard-Jones particles; we train a coarse-grained force field for one-bead methanol on-the-fly during an all-atom simulation; we calculate the neutron scattering profile of 5000 water molecules on-the-fly; and we bias a simulation of Lennard-Jones particles to have a target distance from the center of mass of the system.