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-torch
If you don't have conda
available, we recommend installing Miniconda for Python 3 to provide the conda
package manager.
This plugin uses CMake as its build system.
Before compiling you must install LibTorch, which is the PyTorch C++ API, by following the instructions at https://pytorch.org.
You can then follow these steps:
-
Create a directory in which to build the plugin.
-
Run the CMake GUI or
ccmake
, specifying your new directory as the build directory and the top level directory of this project as the source directory. -
Press "Configure". (Do not worry if it produces an error message about not being able to find PyTorch.)
-
Set
OPENMM_DIR
to point to the directory where OpenMM is installed. This is needed to locate the OpenMM header files and libraries. If you are unsure of what directory this is, the following script will print it out.
from simtk import openmm
import os
print(os.path.dirname(openmm.version.openmm_library_path))
-
Set
PYTORCH_DIR
to point to the directory where you installed the LibTorch. -
Set
CMAKE_INSTALL_PREFIX
to the directory where the plugin should be installed. Usually, this will be the same asOPENMM_DIR
, so the plugin will be added to your OpenMM installation. -
If you plan to build the OpenCL platform, make sure that
OPENCL_INCLUDE_DIR
andOPENCL_LIBRARY
are set correctly, and thatNN_BUILD_OPENCL_LIB
is selected. -
If you plan to build the CUDA platform, make sure that
CUDA_TOOLKIT_ROOT_DIR
is set correctly and thatNN_BUILD_CUDA_LIB
is selected. -
Press "Configure" again if necessary, then press "Generate".
-
Use the build system you selected to build and install the plugin. For example, if you selected Unix Makefiles, type
make install
to install the plugin, andmake PythonInstall
to install the Python wrapper.
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)
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.