Learning to Simulate Time-Averaged Coarse-Grained Molecular Dynamics with Geometric Machine Learning

Molecular dynamics (MD) simulation is the workhorse of various scientific domains. However, simulating a physical system with many particles is tremendously computationally expensive. Learning-based force fields have made major progress in accelerating ab-initio MD simulation but are still significantly slower than classical force fields. Complex systems such as battery and protein take weeks to months to simulate, even with classical force fields. We adopt a different ML approach by learning time-averaged acceleration at a coarse-grained level from trajectory data generated by traditional MD simulation. We coarse-grain a physical system using graph clustering and then use a deep graph neural network to model the time-averaged evolution. Our model can simulate complex systems at a lower spatial/temporal resolution and preserve key statistics of interest. Despite only trained to make single-step predictions, our model can rollout for 100,000 steps and recover properties related to 1-10ns level long-time dynamics. Our model applies to a range of estimation problems for complex systems, including predicting the radius of gyration of single-chain coarse-grained polymers of more than 1000 beads in implicit solvent and Li diffusivity of multi-component Li-ion polymer electrolyte systems.