Source-linked AI summary

OpenMM 8: Molecular Dynamics Simulation with Machine Learning Potentials

Peter Eastman, Raimondas Galvelis, Raúl P. Peláez, Charlles R. A. Abreu, Stephen E. Farr, Emilio Gallicchio, Anton Gorenko, Michael M. Henry, Frank Hu, Jing Huang, Andreas Krämer, Julien Michel, Joshua A. Mitchell, Vijay S. Pande, João PGLM Rodrigues, Jaime Rodriguez-Guerra, Andrew C. Simmonett, Sukrit Singh, Jason Swails, Philip Turner, Yuanqing Wang, Ivy Zhang, John D. Chodera, Gianni De Fabritiis, Thomas E. Markland

arXiv:2310.03121v2physics.chem-phcs.LG

TL;DR

Molecular simulation seeks higher accuracy at manageable cost, but machine learning potentials require practical software support and efficient execution. OpenMM 8 integrates arbitrary and pretrained PyTorch potentials with optimized operations, demonstrating agreement with AIMD and large speedups. These capabilities make ML-enhanced simulations practical at modest additional cost, within current applicability limits.

  • Problem

    Molecular simulation needs methods that improve accuracy over classical force fields without incurring the full cost of high-level quantum chemistry.

  • Method

    OpenMM 8 adds arbitrary PyTorch force and energy models, pretrained-potential interfaces, optimized CUDA kernels, and custom PyTorch operations for molecular dynamics.

  • Results

    MLP structural distributions agree with AIMD, while one GFP chromophore force evaluation is 1517 times faster than AIMD and optimized ANI-2x ligand modeling is 2.7 times slower than classical simulation.

  • Takeaways & Limitations

    OpenMM 8 makes it practical to use machine learning potentials for improved simulation accuracy at a modest computational cost.

  • Takeaways & Limitations

    Current pretrained potentials have limited applicability: ANI-2x supports seven elements, ANI-1ccx supports four, and neither supports charged molecules.

Abstract

from arXiv · show

Machine learning plays an important and growing role in molecular simulation. The newest version of the OpenMM molecular dynamics toolkit introduces new features to support the use of machine learning potentials. Arbitrary PyTorch models can be added to a simulation and used to compute forces and energy. A higher-level interface allows users to easily model their molecules of interest with general purpose, pretrained potential functions. A collection of optimized CUDA kernels and custom PyTorch operations greatly improves the speed of simulations. We demonstrate these features on simulations of cyclin-dependent kinase 8 (CDK8) and the green fluorescent protein (GFP) chromophore in water. Taken together, these features make it practical to use machine learning to improve the accuracy of simulations at only a modest increase in cost.

Introduction

Molecular simulation research seeks higher accuracy at lower cost, with machine learning potentials positioned between quantum chemistry and classical force fields. OpenMM 8 adds flexible machine-learning support for both methodological research and production simulations.

  • OpenMM 8: OpenMM 8 introduces features emphasizing machine learning while retaining the package’s performance, flexibility, and extensibility for molecular simulation.The release is intended for both methodological research and production simulations.
  • Motivation: Machine learning potentials aim to approach high-level quantum chemistry accuracy at a cost between quantum chemistry and classical force fields.Unlike designer-chosen force fields, they use flexible models, often neural networks, trained on molecular properties.
  • Machine Learning Potentials: New MLP architectures target improved accuracy, speed, transferability, and data efficiency through physical symmetries and embedded physics.Examples include equivariance under translations and rotations and explicit Coulomb, dispersion, and nuclear-repulsion terms.
  • Hybrid Simulations: ML/MM simulations combine machine learning for part of a system with classical force fields for the remainder, addressing the high cost of applying MLPs to large systems.This follows the same partitioning idea as QM/MM simulations.
  • OpenMM 8: OpenMM 8 supports arbitrary PyTorch models for computing forces and energies and provides validated pretrained MLPs for end users.Its interfaces are intended to make ML and ML/MM simulations no more difficult to set up than conventional force-field simulations.

PyTorch Model Support

OpenMM-Torch embeds PyTorch models in simulations, while OpenMM-ML provides higher-level access to pretrained and mixed ML/MM potentials. Current pretrained-potentials support remains limited in elemental and charge coverage.

  • OpenMM-Torch: OpenMM-Torch provides TorchForce, allowing arbitrary PyTorch models to be embedded in simulations to calculate forces and energies.PyTorch support also preserves the framework used by most recently published MLPs.
  • OpenMM-Torch: TorchForce accepts a model that maps atomic positions to potential energy, with forces obtained by backpropagation or returned explicitly.The model must be compilable to TorchScript, while other PyTorch features are supported.
  • OpenMM-ML: OpenMM-ML provides a convenient interface for running simulations with pretrained potentials such as ANI-2x.This makes pretrained MLPs simpler to use through a higher-level interface.
  • Mixed ML/MM Simulations: Mixed simulations assign selected atoms to an MLP and the remainder to a conventional force field.The interface handles setup details, including mechanical embedding between the ML and MM regions.
  • Scope: Currently supported pretrained potentials are ANI-1ccx and ANI-2x, whose applicability is limited by elemental coverage and lack of charged-molecule support.ANI-2x supports seven elements and ANI-1ccx supports four.
  • Training Data: The SPICE dataset provides high-level quantum-chemistry molecular energies and forces across 15 elements, charged and neutral molecules, conformations, and interaction types.It was created to facilitate development of new models relevant to drug-like molecules interacting with proteins.

Optimization Through NNPOps

NNPOps accelerates machine-learning potentials by supplying optimized CUDA kernels and PyTorch code for computational bottlenecks.

  • NNPOps: NNPOps accelerates important MLP operations with optimized CUDA kernels and PyTorch code.Its examples include ANI symmetry functions, neighbor-list construction, and Particle Mesh Ewald electrostatics.

Other Features

OpenMM 8 expands molecular simulation capabilities with flexible collective-variable, sampling, alchemical, force-field, integration, hardware, and machine-learning features. Demonstrations show optimized machine-learning simulations can approach classical-simulation practicality while retaining agreement with reference trajectories, although hybrid ML/MM tools remain a starting point for future development.

  • Enhanced sampling: OpenMM supports metadynamics through PLUMED or a native mechanism that uses arbitrary OpenMM energy terms as collective variables.Custom force classes can define collective variables from user-defined functional forms, including weighted distance and dihedral sums.
  • Enhanced sampling: Simulated tempering accelerates barrier crossing by temporarily increasing simulation temperature and complements aMD, replica exchange, and REST.The method is intended for situations in which specific collective variables are not known in advance.
  • Free-energy calculations: The Alchemical Transfer Method computes free energy differences by averaging interaction forces before and after coordinate transformations.It supports sampling of absolute or relative free-energy pathways without requiring soft-core interactions.
  • Force fields, integration, and hardware: OpenMM 8 adds newer Amber, CHARMM, AMOEBA, CHARMM Drude, GLYCAM, OpenFF, and GAFF force fields, plus LF-Middle integration and broader integrator support.LF-Middle can allow doubling the step size with no loss in accuracy in many cases, while new hardware support includes ARM, PowerPC, HIP, multi-GPU, and faster CPU execution.
  • Machine-learning performance: NNPOps makes ANI-2x 5.7 times faster than TorchANI, while using one model instead of the ensemble adds a further 1.9-fold speed improvement.Together, the optimizations make ANI-2x ligand modeling 2.7 times slower than a fully classical simulation at the same step size.
  • Machine-learning demonstrations: For the GFP chromophore, an equivariant-transformer MLP trained on AIMD and AI-PIMD data agrees with AIMD RDF and dihedral distributions within trajectory error bars.The MLP force test-set RMSE was 45.7 meV/A, and force evaluation was 1517 times faster than AIMD in the reported benchmark.
  • Scope and limitations: OpenMM 8 provides a foundation for machine-learning simulations, but pretrained general-purpose potentials remain scarce and limited in applicability.Hybrid ML/MM development still needs electrostatic embedding and machine-learning models that predict atomic partial charges as well as energy.

Disclosures

The authors disclose advisory, board, equity, and financial interests, along with extensive external funding received by the Chodera laboratory.

  • JDC reports advisory-board memberships and equity interests in Interline Therapeutics.
  • The Chodera laboratory receives or has received funding from numerous public, nonprofit, and industry sources.
  • JM reports membership on the Scientific Advisory Board of Cresset.
  • VSP reports board memberships at multiple companies and service on Schrodinger’s Scientific Advisory Board.
  • YW reports limited financial interests in Flagship Pioneering and its subsidiaries.
Loading 2310.03121v2…