This is a plugin for OpenMM that allows PyTorch static computation graphs
to be used for defining an OpenMM TorchForce object, an OpenMM Force class that computes a contribution to the potential energy or used as a collective variable via CustomCVForce.
To use it, you create a PyTorch model that takes a (nparticles,3) tensor of particle positions (in nanometers) as input and produces energy (in kJ/mol) or the value of the collective variable as output.
The TorchForce provided by this plugin can then use the model to compute energy contributions or apply forces to particles during a simulation.
TorchForce also supports the use of global context parameters that can be fed to the model and changed dynamically during runtime.
We provide conda packages for Linux and MacOS via conda-forge, which can be installed from the conda-forge channel:
conda install -c conda-forge openmm-torchIf you don't have conda available, we recommend installing Miniconda for Python 3 to provide the conda package manager.
If you need to build from source look at INSTALL.md.
The first step is to create a PyTorch model defining the calculation to perform.
It should take particle positions in nanometers (in the form of a torch.Tensor of shape (nparticles,3) as input,
and return the potential energy in kJ/mol as a torch.Scalar as output.
The model must then be converted to a TorchScript module and saved to a file.
Converting to TorchScript can usually be done with a single call to torch.jit.script() or torch.jit.trace(),
although more complicated models can sometimes require extra steps.
See the PyTorch documentation for details.
Here is a simple Python example that does this for a very simple potential---a harmonic force attracting every particle to the origin:
import torch
class ForceModule(torch.nn.Module):
"""A central harmonic potential as a static compute graph"""
def forward(self, positions):
"""The forward method returns the energy computed from positions.
Parameters
----------
positions : torch.Tensor with shape (nparticles,3)
positions[i,k] is the position (in nanometers) of spatial dimension k of particle i
Returns
-------
potential : torch.Scalar
The potential energy (in kJ/mol)
"""
return torch.sum(positions**2)
# Render the compute graph to a TorchScript module
module = torch.jit.script(ForceModule())
# Serialize the compute graph to a file
module.save('model.pt')To use the exported model in a simulation, create a TorchForce object and add it to your System.
The constructor takes the path to the saved model as an argument.
For example,
# Create the TorchForce from the serialized compute graph
from openmmtorch import TorchForce
torch_force = TorchForce('model.pt')
# Add the TorchForce to your System
system.addForce(torch_force)When defining the model to perform a calculation, you may want to apply periodic boundary conditions.
To do this, call setUsesPeriodicBoundaryConditions(True) on the TorchForce.
The graph is then expected to take a second input, which contains the current periodic box vectors.
You can make use of them in whatever way you want for computing the force.
For example, the following code applies periodic boundary conditions to each
particle position to translate all of them into a single periodic cell:
class ForceModule(torch.nn.Module):
"""A central harmonic force with periodic boundary conditions"""
def forward(self, positions, boxvectors):
"""The forward method returns the energy computed from positions.
Parameters
----------
positions : torch.Tensor with shape (nparticles,3)
positions[i,k] is the position (in nanometers) of spatial dimension k of particle i
boxvectors : torch.tensor with shape (3,3)
boxvectors[i,k] is the box vector component k (in nanometers) of box vector i
Returns
-------
potential : torch.Scalar
The potential energy (in kJ/mol)
"""
# Image articles in rectilinear box
# NOTE: This will not work for non-orthogonal boxes
boxsize = boxvectors.diag()
periodicPositions = positions - torch.floor(positions/boxsize)*boxsize
# Compute central harmonic potential
return torch.sum(periodicPositions**2)Note that this code assumes a rectangular box. Applying periodic boundary conditions with a triclinic box requires a slightly more complicated calculation.
The graph can also take arbitrary scalar arguments that are passed in at
runtime. For example, this model multiplies the energy by scale, which is
passed as an argument to forward().
class ForceModule(torch.nn.Module):
"""A central harmonic force with a user-defined global scale parameter"""
def forward(self, positions, scale):
"""The forward method returns the energy computed from positions.
Parameters
----------
positions : torch.Tensor with shape (nparticles,3)
positions[i,k] is the position (in nanometers) of spatial dimension k of particle i
scale : torch.Scalar
A scalar tensor defined by 'TorchForce.addGlobalParameter'.
Here, it scales the contribution to the potential.
Note that parameters are passed in the order defined by `TorchForce.addGlobalParameter`, not by name.
Returns
-------
potential : torch.Scalar
The potential energy (in kJ/mol)
"""
return scale*torch.sum(positions**2)When you create the TorchForce, call addGlobalParameter() once for each extra argument.
torch_force.addGlobalParameter('scale', 2.0)This specifies the name of the parameter and its initial value. The name
does not need to match the argument to forward(). All global parameters
are simply passed to the model in the order you define them. The advantage
of using global parameters is that you can change their values at any time
by calling setParameter() on the Context.
context.setParameter('scale', 5.0)In the examples above, the PyTorch model computes the potential energy. Backpropagation can be used to compute the corresponding forces. That always works, but sometimes you may have a more efficient way to compute the forces than the generic backpropagation algorithm. In that case, you can have the model directly compute forces as well as energy, returning both of them in a tuple. Remember that the force is the negative gradient of the energy.
import torch
class ForceModule(torch.nn.Module):
"""A central harmonic potential that computes both energy and forces."""
def forward(self, positions):
"""The forward method returns the energy and forces computed from positions.
Parameters
----------
positions : torch.Tensor with shape (nparticles,3)
positions[i,k] is the position (in nanometers) of spatial dimension k of particle i
Returns
-------
potential : torch.Scalar
The potential energy (in kJ/mol)
forces : torch.Tensor with shape (nparticles,3)
The force (in kJ/mol/nm) on each particle
"""
return (torch.sum(positions**2), -2*positions)When you create the TorchForce, call setOutputsForces() to tell it to expect the model
to return forces.
torch_force.setOutputsForces(True)This is part of the OpenMM molecular simulation toolkit originating from Simbios, the NIH National Center for Physics-Based Simulation of Biological Structures at Stanford, funded under the NIH Roadmap for Medical Research, grant U54 GM072970. See https://simtk.org.
Portions copyright (c) 2018-2020 Stanford University and the Authors.
Authors: Peter Eastman
Contributors: Raimondas Galvelis, Jaime Rodriguez-Guerra, Yaoyi Chen, John D. Chodera
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.