diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c8b3b64f..a21e0652 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -22,13 +22,14 @@ jobs: fail-fast: false matrix: include: - - name: Linux (CUDA 11.8, Python 3.10, PyTorch 2.1) + - name: Linux (CUDA 13.0, Python 3.14, PyTorch 2.12) os: ubuntu-22.04 - cuda-version: "11.8.0" + cuda-version: "13.0" gcc-version: "10.3.*" - nvcc-version: "11.8" - python-version: "3.10" - pytorch-version: "2.1.*" + nvcc-version: "13.0" + python-version: "3.14" + pytorch-version: "2.12.*" + opencl: true - name: MacOS ARM (Python 3.11, PyTorch 2.4) os: macos-latest @@ -37,14 +38,16 @@ jobs: nvcc-version: "" python-version: "3.11" pytorch-version: "2.4.*" + opencl: false - - name: MacOS Intel (Python 3.11, PyTorch 2.4) - os: macos-13 + - name: MacOS Intel (Python 3.13, PyTorch 2.5) + os: macos-15-intel cuda-version: "" gcc-version: "" nvcc-version: "" - python-version: "3.11" - pytorch-version: "2.4.*" + python-version: "3.13" + pytorch-version: "2.5.*" + opencl: false steps: @@ -52,12 +55,17 @@ jobs: uses: actions/checkout@v2 - name: "Install CUDA Toolkit on Linux (if needed)" - uses: Jimver/cuda-toolkit@v0.2.15 + uses: Jimver/cuda-toolkit@v0.2.35 + id: cuda-toolkit with: - cuda: ${{ matrix.cuda-version }} + cuda: ${{ matrix.cuda-version }}.0 linux-local-args: '["--toolkit", "--override"]' if: startsWith(matrix.os, 'ubuntu') + - name: "Install OpenCL on Ubuntu (if needed)" + if: matrix.opencl == true + run: source devtools/scripts/install_intel_opencl.sh + - name: Manage disk space if: startsWith(matrix.os, 'ubuntu') uses: jlumbroso/free-disk-space@main @@ -78,7 +86,7 @@ jobs: NVCC_VERSION: ${{ matrix.nvcc-version }} PYTORCH_VERSION: ${{ matrix.pytorch-version }} - - uses: conda-incubator/setup-miniconda@v2 + - uses: conda-incubator/setup-miniconda@v4 name: "Install dependencies with Mamba" with: activate-environment: build @@ -111,8 +119,8 @@ jobs: -DOPENMM_DIR=${CONDA_PREFIX} \ -DTorch_DIR=${CONDA_PREFIX}/lib/python${{ matrix.python-version }}/site-packages/torch/share/cmake/Torch \ -DNN_BUILD_OPENCL_LIB=ON \ - -DOPENCL_INCLUDE_DIR=${CONDA_PREFIX}/include \ - -DOPENCL_LIBRARY=${CONDA_PREFIX}/lib/libOpenCL${SHLIB_EXT} + -DCUDAToolkit_ROOT="${{steps.cuda-toolkit.outputs.CUDA_PATH}}" \ + -DCMAKE_CUDA_COMPILER="${{steps.cuda-toolkit.outputs.CUDA_PATH}}/bin/nvcc" - name: "Build" shell: bash -l {0} @@ -132,6 +140,9 @@ jobs: run: | export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib/python${{ matrix.python-version }}/site-packages/torch/lib:${LD_LIBRARY_PATH}" cd build + if [[ ${{ matrix.opencl }} == true ]]; then + source /opt/intel/oneapi/setvars.sh + fi ctest --output-on-failure --exclude-regex "TestCuda|TestOpenCL" - name: "Run Python test" @@ -139,4 +150,7 @@ jobs: run: | export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib/python${{ matrix.python-version }}/site-packages/torch/lib:${LD_LIBRARY_PATH}" cd python/tests + if [[ ${{ matrix.opencl }} == true ]]; then + source /opt/intel/oneapi/setvars.sh + fi pytest --verbose Test* diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index 16c26025..b1b4e651 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -21,7 +21,7 @@ jobs: uses: actions/checkout@v2 - name: "Install dependencies with Mamba" - uses: conda-incubator/setup-miniconda@v2 + uses: conda-incubator/setup-miniconda@v4 with: activate-environment: test environment-file: devtools/conda-envs/test-tutorials.yml diff --git a/CMakeLists.txt b/CMakeLists.txt index 75095427..2e06a866 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -5,7 +5,7 @@ # The version number is set here and copied to anywhere it is needed. SET(OPENMM_TORCH_VERSION 1.5 CACHE STRING "The version of OpenMM-Torch that will be built." FORCE) -CMAKE_MINIMUM_REQUIRED(VERSION 3.5) +CMAKE_MINIMUM_REQUIRED(VERSION 3.12) # We need to know where OpenMM is installed so we can access the headers and libraries. SET(OPENMM_DIR "/usr/local/openmm" CACHE PATH "Where OpenMM is installed") @@ -139,12 +139,12 @@ IF(NN_BUILD_OPENCL_LIB) ADD_SUBDIRECTORY(platforms/opencl) ENDIF(NN_BUILD_OPENCL_LIB) -FIND_PACKAGE(CUDA QUIET) -IF(CUDA_FOUND) +FIND_PACKAGE(CUDAToolkit QUIET) +IF(CUDAToolkit_FOUND) SET(NN_BUILD_CUDA_LIB ON CACHE BOOL "Build implementation for CUDA") -ELSE(CUDA_FOUND) +ELSE(CUDAToolkit_FOUND) SET(NN_BUILD_CUDA_LIB OFF CACHE BOOL "Build implementation for CUDA") -ENDIF(CUDA_FOUND) +ENDIF(CUDAToolkit_FOUND) IF(NN_BUILD_CUDA_LIB) ADD_SUBDIRECTORY(platforms/cuda) ENDIF(NN_BUILD_CUDA_LIB) diff --git a/README.md b/README.md index 3456c23f..6abd071f 100644 --- a/README.md +++ b/README.md @@ -5,364 +5,154 @@ OpenMM PyTorch Plugin ===================== -This is a plugin for [OpenMM](http://openmm.org) that allows [PyTorch](https://pytorch.org/) static computation graphs -to be used for defining an OpenMM `TorchForce` object, an [OpenMM `Force` class](http://docs.openmm.org/latest/api-python/library.html#forces) that computes a contribution to the potential energy or used as a collective variable via [`CustomCVForce`](http://docs.openmm.org/latest/api-python/generated/simtk.openmm.openmm.CustomCVForce.html#simtk.openmm.openmm.CustomCVForce). +This is a plugin for [OpenMM](http://openmm.org) that allows using [PyTorch](https://pytorch.org/) models to compute +forces and energy in a simulation. It provides two force classes, `TorchForce` and `PythonTorchForce`. `TorchForce` +is deprecated, since it relies on TorchScript which is no longer maintained. `PythonTorchForce` is the recommended one +to use in all cases. -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. +`PythonTorchForce` is very similar to OpenMM's built in [`PythonForce`](https://docs.openmm.org/latest/api-python/generated/openmm.openmm.PythonForce.html) +class, but it is specialized for use with PyTorch. In particular, the particles positions and forces are represented +with tensors instead of NumPy arrays. The benefit is reducing overhead and improving performance. When the OpenMM +simulation and PyTorch model both run on the same GPU, coordinates and forces can be copied between them directly on the +GPU without ever needing to transfer them to the host. Installation ============ -Installing with conda ---------------------- +Installing with pip +------------------- + +We provide packages for Linux and macOS, which can be installed with the command: -We provide [conda](https://docs.conda.io/) packages for Linux and MacOS via [`conda-forge`](https://conda-forge.org/), which can be installed from the [conda-forge channel](https://anaconda.org/conda-forge/openmm-torch): ```bash -conda install -c conda-forge openmm-torch +pip install openmmtorch ``` -If you don't have `conda` available, we recommend installing [Miniconda for Python 3](https://docs.conda.io/en/latest/miniconda.html) to provide the `conda` package manager. Building from source -------------------- -This plugin uses [CMake](https://cmake.org/) as its build system. -Before compiling you must install [LibTorch](https://pytorch.org/cppdocs/installing.html), which is the PyTorch C++ API, by following the instructions at https://pytorch.org. -You can then follow these steps: +This plugin uses [CMake](https://cmake.org/) as its build system. Before compiling you must install PyTorch by +following the instructions at https://pytorch.org. You can then follow these steps: 1. Create a directory in which to build the plugin. 2. 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. -3. Press "Configure". (Do not worry if it produces an error message about not being able to find PyTorch.) +3. Press "Configure". 4. 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. ```python -from simtk import openmm +import openmm import os print(os.path.dirname(openmm.version.openmm_library_path)) ``` -5. Set `PYTORCH_DIR` to point to the directory where you installed the LibTorch. +5. Usually PyTorch will be found automatically. If it is not, set `Torch_DIR` to point to the directory containing its +CMake configuration files (e.g. `/share/cmake/Torch`). 6. Set `CMAKE_INSTALL_PREFIX` to the directory where the plugin should be installed. Usually, this will be the same as `OPENMM_DIR`, so the plugin will be added to your OpenMM installation. -7. If you plan to build the OpenCL platform, make sure that `OPENCL_INCLUDE_DIR` and -`OPENCL_LIBRARY` are set correctly, and that `NN_BUILD_OPENCL_LIB` is selected. - -8. If you plan to build the CUDA platform, make sure that `CUDA_TOOLKIT_ROOT_DIR` is set correctly -and that `NN_BUILD_CUDA_LIB` is selected. - -9. Press "Configure" again if necessary, then press "Generate". - -10. 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, and `make PythonInstall` to -install the Python wrapper. - -Using the OpenMM PyTorch plugin -=============================== - -Tutorials ---------- - -- [A simple simulation of alanine dipeptide with ANI-2x using OpenMM-Torch and NNPOps](tutorials/openmm-torch-nnpops.ipynb) +7. If you plan to build the OpenCL, CUDA, or HIP platform, make sure that `NN_BUILD_OPENCL_LIB`, `NN_BUILD_CUDA_LIB`, +or `NN_BUILD_HIP_LIB` respectively is selected. If the installed location of OpenCL, CUDA, or HIP was not found +automatically, set the appropriate CMake variables to locate them. -Exporting a PyTorch model for use in OpenMM -------------------------------------------- +8. Press "Configure" again if necessary, then press "Generate". -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. +9. Type `make install` to install the plugin, then `make PythonInstall` to install the Python wrapper. -The model must then be converted to a [TorchScript](https://pytorch.org/docs/stable/jit.html) module and saved to a file. -Converting to TorchScript can usually be done with a single call to [`torch.jit.script()`](https://pytorch.org/docs/stable/generated/torch.jit.script.html#torch.jit.script) or [`torch.jit.trace()`](https://pytorch.org/docs/stable/generated/torch.jit.trace.html#torch.jit.trace), -although more complicated models can sometimes require extra steps. -See the [PyTorch documentation](https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html) for details. +Using PythonTorchForce +====================== -Here is a simple Python example that does this for a very simple potential---a harmonic force attracting every particle to the origin: +To use PythonTorchForce, define a Python function that computes the interaction. It should take two arguments, a +State object and a Tensor of shape `(# particles, 3)`. For example, ```python 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. -Alternatively, the scripted module can be provided directly. -For example, -```python -# Create the TorchForce from the serialized compute graph -from openmmtorch import TorchForce -# Construct using a serialized module: -torch_force = TorchForce('model.pt') -# or using an instance of the module: -torch_force = TorchForce(module) - -# Add the TorchForce to your System -system.addForce(torch_force) +def compute(state, positions): + energy = torch.sum(positions**2) + forces = -0.5*positions + return energy, forces ``` -Defining a model that uses periodic boundary conditions -------------------------------------------------------- - -When defining the model to perform a calculation, you may want to apply periodic boundary conditions. +The State contains global parameters and periodic box vectors. The Tensor contains particle positions. The function +should compute the potential energy and forces, returning them as its two return values. The energy should be a +scalar Tensor containing the value in kJ/mol. The forces should be a Tensor of shape `(# particles, 3)` containing +the value in kJ/mol/nm. -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: +Now create a PythonTorchForce, passing the function to the constructor. ```python -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) +from openmmtorch import PythonTorchForce +force = PythonTorchForce(compute) ``` -Note that this code assumes a rectangular box. Applying periodic boundary -conditions with a triclinic box requires a slightly more complicated calculation. - -Defining global parameters that can be modified within the Context ------------------------------------------------------------------- - -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()`. +Do not make any assumptions about either the dtype or the device of the tensor containing positions. Both of them may +vary depending on the platform and precision mode used for the simulation. If you require a particular dtype or device, +call `to()` to ensure they are correct: ```python -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) +positions = positions.to(dtype=torch.float32, device='cuda:0') ``` -When you create the `TorchForce`, call `addGlobalParameter()` once for each extra argument. - -```python -torch_force.addGlobalParameter('scale', 2.0) -``` +### Global Parameters -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`. +The force can optionally depend on global parameters stored in the Context. To do this, pass a dict to the constructor +containing the names and default values of the parameters: ```python -context.setParameter('scale', 5.0) +force = PythonTorchForce(compute, {'k':2.5}) ``` -Computing forces in the model ------------------------------ - -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. +The computation function can then retrieve the parameter values from the State: ```python -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) +def compute(state, positions): + k = state.getParameters()['k'] + energy = k*torch.sum(positions**2) + forces = -0.5*k*positions + return energy, forces ``` -When you create the `TorchForce`, call `setOutputsForces()` to tell it to expect the model -to return forces. +You can change the parameter value at any time by calling `setParameter()` on the Context: ```python -torch_force.setOutputsForces(True) +context.setParameter('k', 5.0) ``` -Computing energy derivatives with respect to global parameters --------------------------------------------------------------- +### Periodic Boundary Conditions -TorchForce can compute derivatives of the energy with respect to global parameters.. In order to do so the global parameters must be registered as energy derivatives. This is done by calling `addEnergyParameterDerivative()` for each parameter. - -The parameter derivatives can be queried by calling `getEnergyParameterDerivatives()` on the `State` object returned by `Context.getState()`. The result is a dictionary with the parameter names as keys and the derivatives as values. +If you want your force to depend on periodic boundary conditions, call `setUsesPeriodicBoundaryConditions(True)` on the +PythonTorchForce. This has two effects. First, `usesPeriodicBoundaryConditions()` will return True, signaling to +other code that your system is periodic. Second, the State passed to the computation function will contain periodic +box vectors. You can use them however you want in computing the force. For example, ```python -import torch as pt -from openmmtorch import TorchForce -import openmm as mm - - -class ForceWithParameters(pt.nn.Module): - - def __init__(self): - super(ForceWithParameters, self).__init__() - - def forward(self, positions: pt.Tensor, k: pt.Tensor) -> pt.Tensor: - return k * pt.sum(positions**2) - - -numParticles = 10 -system = mm.System() -for _ in range(numParticles): - system.addParticle(1.0) - -model = pt.jit.script(ForceWithParameters()) -tforce = TorchForce(model) -tforce.setOutputsForces(False) -tforce.addGlobalParameter("k", 2.0) -tforce.addEnergyParameterDerivative("k") -system.addForce(tforce) -context = mm.Context(system, mm.VerletIntegrator(1.0)) -context.setPositions(pt.rand(numParticles, 3).numpy()) -state = context.getState(getParameterDerivatives=True) -dEdk = state.getEnergyParameterDerivatives()["k"] +def compute2(state, positions): + vectors = state.getPeriodicBoxVectors().value_in_unit(nanometer) + boxsize = torch.tensor(vectors, dtype=positions.dtype, device=positions.device).diag() + positions = positions - torch.floor(positions/boxsize)*boxsize + energy = torch.sum(positions**2) + forces = -0.5*positions + return energy, forces ``` +### Restricting to a Subset of Particles - -Recording the model into a CUDA graph -------------------------------------- - -You can ask `TorchForce` to run the model using [CUDA graphs](https://pytorch.org/docs/stable/notes/cuda.html#cuda-graphs). Not every model will be compatible with this feature, but it can be a significant performance boost for some models. To enable it the CUDA platform must be used and an special property must be provided to `TorchForce`: +A PythonTorchForce can optionally be applied to only a subset of the particles in a system. To do +this, call `setParticles()` on it, providing the indices of the particles to apply it to. ```python -torch_force.setProperty("useCUDAGraphs", "true") -# The property can also be set at construction -torch_force = TorchForce('model.pt', {'useCUDAGraphs': 'true'}) +force.setParticles(list(range(50))) # Apply to only the first 50 particles ``` -The first time the model is run, it will be compiled (also known as recording) into a CUDA graph. Subsequent runs will use the compiled graph, which can be significantly faster. It is possible that compilation fails, in which case an `OpenMMException` will be raised. If that happens, you can disable CUDA graphs and try again. - -It is required to run the model at least once before recording, in what is known as warmup. -By default ```TorchForce``` will run the model just a few times before recording, but controlling warmup steps might be desired. In these cases one can set the property ```CUDAGraphWarmupSteps```: -```python -torch_force.setProperty("CUDAGraphWarmupSteps", "12") -``` - -List of available properties ----------------------------- - -Some ```TorchForce``` functionalities can be customized by setting properties on an instance of it. Properties can be set at construction or by using ```setProperty```. A property is a pair of key/value strings. For instance: - -```python -torch_force = TorchForce('model.pt', {'useCUDAGraphs': 'true'}) -#Alternatively setProperty can be used to configure an already created instance. -#torch_force.setProperty("useCUDAGraphs", "true") -print("Current properties:") -for property in torch_force.getProperties(): - print(property.key, property.value) -``` - -Currently, the following properties are available: - -1. useCUDAGraphs: Turns on the CUDA graph functionality -2. CUDAGraphWarmupSteps: When CUDA graphs are being used, controls the number of warmup calls to the model before recording. - -License -======= - -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. +The computation function should then proceed as if those particles were the entire system. The positions +passed to it will be a smaller Tensor containing only the positions of those particles, and the returned +forces should similarly contain only those particles. That is, `forces[i]` should be the force on the i'th +particle passed to `setParticles()`. When applying forces to only a small fraction of the particles in a +system, this can greatly improve performance. diff --git a/devtools/conda-envs/build-macos-13.yml b/devtools/conda-envs/build-macos-15-intel.yml similarity index 100% rename from devtools/conda-envs/build-macos-13.yml rename to devtools/conda-envs/build-macos-15-intel.yml diff --git a/devtools/conda-envs/build-ubuntu-22.04.yml b/devtools/conda-envs/build-ubuntu-22.04.yml index 1cc25bbc..3582f3e3 100644 --- a/devtools/conda-envs/build-ubuntu-22.04.yml +++ b/devtools/conda-envs/build-ubuntu-22.04.yml @@ -3,15 +3,13 @@ channels: - conda-forge dependencies: - cmake - - cudatoolkit @CUDATOOLKIT_VERSION@ + - cuda-version @CUDATOOLKIT_VERSION@ - gxx_linux-64 @GCC_VERSION@ - make - nnpops - - nvcc_linux-64 @NVCC_VERSION@ - ocl-icd - openmm >=8.1 - pip - - pocl - pytest - python - pytorch-gpu @PYTORCH_VERSION@ diff --git a/devtools/scripts/install_intel_opencl.sh b/devtools/scripts/install_intel_opencl.sh new file mode 100644 index 00000000..49afa751 --- /dev/null +++ b/devtools/scripts/install_intel_opencl.sh @@ -0,0 +1,8 @@ +# This script installs Intel's OpenCL for CPUs. + +set -euxo pipefail + +wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | sudo tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null +echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list +sudo apt update +sudo apt install intel-basekit intel-hpckit intel-oneapi-runtime-opencl opencl-headers ocl-icd-opencl-dev -y diff --git a/openmmapi/include/PythonTorchForce.h b/openmmapi/include/PythonTorchForce.h new file mode 100644 index 00000000..824f3fd4 --- /dev/null +++ b/openmmapi/include/PythonTorchForce.h @@ -0,0 +1,193 @@ +#ifndef OPENMM_PYTHONTORCHFORCE_H_ +#define OPENMM_PYTHONTORCHFORCE_H_ + +/* -------------------------------------------------------------------------- * + * OpenMM * + * -------------------------------------------------------------------------- * + * This is part of the OpenMM molecular simulation toolkit. * + * See https://openmm.org/development. * + * * + * Portions copyright (c) 2025-2026 Stanford University and the Authors. * + * Authors: Peter Eastman * + * Contributors: * + * * + * 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. * + * -------------------------------------------------------------------------- */ + +#include "openmm/Force.h" +#include "openmm/State.h" +#include "internal/windowsExportTorch.h" +#include +#include +#include + +namespace TorchPlugin { + +/** + * This abstract class represents an interface for performing a computation. It is not intended to + * be used or subclassed directly by users. The Python wrapper contains a subclass that implements + * the interface using a Python function. + * @private + */ +class OPENMM_EXPORT_NN PythonTorchForceComputation { +public: + PythonTorchForceComputation() { + } + virtual ~PythonTorchForceComputation() { + } + /** + * Compute forces and energy. The State contains particle parameters and optionally + * periodic box vectors. Implementations should store the potential energy into the + * energy argument and return a tensor of shape (particles, 3) containing the forces. + */ + virtual torch::Tensor compute(const OpenMM::State& state, const torch::Tensor& positions, double& energy) const = 0; +}; + +/** + * This class provides a mechanism for computing forces and energy with Python code. It is similar to the + * PythonForce class included with OpenMM, but it is specialized to give better performance when working with + * PyTorch models. + * + * To use it, define a Python function that takes two arguments: a State object and a Tensor of shape (# particles, 3). + * The State contains global parameters and periodic box vectors. The Tensor contains particle positions. The function + * should compute the potential energy and forces, returning them as its two return values. The energy should be a + * scalar Tensor containing the value in kJ/mol. The forces should be a Tensor of shape (# particles, 3) containing + * the value in kJ/mol/nm. For example, + * + * \verbatim embed:rst:leading-asterisk + * .. code-block:: python + * + * def compute(state, pos): + * k = state.getParameters()['k'] + * energy = k*torch.sum(pos*pos) + * force = -0.5*k*pos + * return energy, force + * + * \endverbatim + * + * Now create a PythonTorchForce, passing the function to the constructor. If you want the force + * to depend on global parameters, pass a dict as the second parameter with the names and default + * values of the parameters. + * + * \verbatim embed:rst:leading-asterisk + * .. code-block:: python + * + * force = PythonTorchForce(compute, {'k':2.5}) + * + * \endverbatim + * + * The default value of a parameter is its value in newly created Contexts. After a Context is + * created, you can change the values of parameters by calling setParameter() on it. + * + * The PythonTorchForce cannot tell whether the function you provide makes use of periodic boundary + * conditions, so you must tell it. To make the force periodic, call + * setUsesPeriodicBoundaryConditions(True). This will cause usesPeriodicBoundaryConditions() + * to return True, and the State passed to the computation function will contain periodic + * box vectors. The positions may also be wrapped into a different periodic box to keep them + * closer to the origin and improve accuracy. + * + * A PythonTorchForce can optionally be applied to only a subset of the particles in a system. To do + * this, call setParticles() on it, providing the indices of the particles to apply it to. The + * computation function should then proceed as if those particles were the entire system. The positions + * passed to it will be a smaller Tensor containing only the positions of those particles, and the returned + * forces should similarly contain only those particles. That is, forces[i] should be the force on the i'th + * particle passed to setParticles(). When applying forces to only a small fraction of the particles in a + * system, this can greatly improve performance. + * + * When using XmlSerializer to save a PythonTorchForce, it uses the Python pickle module to save + * the computation function. If it cannot be pickled, you will not be able to serialize the + * PythonTorchForce. Functions defined at the top level of a module can usually be pickled, but local + * functions defined inside another function cannot. + */ +class OPENMM_EXPORT PythonTorchForce : public OpenMM::Force { +public: + /** + * Create a PythonTorchForce. This constructor is used internally, and is not intended for use + * by users. The Python wrapper defines an alternate constructor that takes a Python + * function instead of a PythonTorchForceComputation. + * + * @param computation an object defining how the forces and energy should be computed + * @param globalParameters any global parameters used by the force. Keys are the parameter + * names, and the corresponding values are their default values. + * @param particles the indices of the particles to use when computing the force. If + * this is empty (the default), all particles in the system will be used. + * @private + */ + explicit PythonTorchForce(PythonTorchForceComputation* computation, const std::map& globalParameters, + const std::vector& particles=std::vector()); + ~PythonTorchForce(); + /** + * Get the PythonTorchForceComputation that defines the computation. + * @private + */ + const PythonTorchForceComputation& getComputation() const; + /** + * Get all global parameters defined by this force. Keys are the parameter names, and the + * corresponding values are their default values. + */ + const std::map& getGlobalParameters() const; + /** + * Get the indices of the particles to use when computing the force. If this + * is empty, all particles in the system will be used. + */ + const std::vector& getParticles() const { + return particles; + } + /** + * Set the indices of the particles to use when computing the force. If this + * is empty, all particles in the system will be used. + */ + void setParticles(const std::vector& particles); + /** + * Get the pickled representation of the computation function. If it cannot be pickled, + * this will be an empty vector. + */ + const std::vector& getPickledFunction() const; + /** + * Set the pickled representation of the computation function. This is called automatically + * by the Python constructor. + * @private + */ + void setPickledFunction(char* function, int length); + /** + * Returns whether or not this force makes use of periodic boundary + * conditions. + * + * @returns true if force uses PBC and false otherwise + */ + bool usesPeriodicBoundaryConditions() const; + /** + * Set whether or not this force makes use of periodic boundary conditions. + * If this is set to true, periodic box vectors can be retrieved from the + * State passed to the computation function. + */ + void setUsesPeriodicBoundaryConditions(bool periodic); +protected: + OpenMM::ForceImpl* createImpl() const; +private: + PythonTorchForceComputation* computation; + std::map globalParameters; + bool usePeriodic; + std::vector particles; + std::vector pickled; +}; + +} // namespace TorchPlugin + +#endif /*OPENMM_PYTHONTORCHFORCE_H_*/ diff --git a/openmmapi/include/TorchKernels.h b/openmmapi/include/TorchKernels.h index d5f986e7..8682cefd 100644 --- a/openmmapi/include/TorchKernels.h +++ b/openmmapi/include/TorchKernels.h @@ -9,7 +9,7 @@ * Biological Structures at Stanford, funded under the NIH Roadmap for * * Medical Research, grant U54 GM072970. See https://simtk.org. * * * - * Portions copyright (c) 2018 Stanford University and the Authors. * + * Portions copyright (c) 2018-2026 Stanford University and the Authors. * * Authors: Peter Eastman * * Contributors: * * * @@ -33,6 +33,7 @@ * -------------------------------------------------------------------------- */ #include "TorchForce.h" +#include "PythonTorchForce.h" #include "openmm/KernelImpl.h" #include "openmm/Platform.h" #include "openmm/System.h" @@ -70,6 +71,34 @@ class CalcTorchForceKernel : public OpenMM::KernelImpl { virtual double execute(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy) = 0; }; +/** + * This kernel is invoked by PythonTorchForce to calculate the forces acting on the system and the energy of the system. + */ +class CalcPythonTorchForceKernel : public OpenMM::KernelImpl { +public: + static std::string Name() { + return "CalcPythonTorchForce"; + } + CalcPythonTorchForceKernel(std::string name, const OpenMM::Platform& platform) : OpenMM::KernelImpl(name, platform) { + } + /** + * Initialize the kernel. + * + * @param context the ContextImpl this kernel will be applied to + * @param force the PythonTorchForce this kernel will be used for + */ + virtual void initialize(const OpenMM::ContextImpl& context, const PythonTorchForce& force) = 0; + /** + * Execute the kernel to calculate the forces and/or energy. + * + * @param context the context in which to execute this kernel + * @param includeForces true if forces should be calculated + * @param includeEnergy true if the energy should be calculated + * @return the potential energy due to the force + */ + virtual double execute(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy) = 0; +}; + } // namespace TorchPlugin #endif /*TORCH_KERNELS_H_*/ diff --git a/openmmapi/include/internal/PythonTorchForceImpl.h b/openmmapi/include/internal/PythonTorchForceImpl.h new file mode 100644 index 00000000..dfb28bbf --- /dev/null +++ b/openmmapi/include/internal/PythonTorchForceImpl.h @@ -0,0 +1,73 @@ +#ifndef OPENMM_PYTHONTORCHFORCEIMPL_H_ +#define OPENMM_PYTHONTORCHFORCEIMPL_H_ + +/* -------------------------------------------------------------------------- * + * OpenMM * + * -------------------------------------------------------------------------- * + * This is part of the OpenMM molecular simulation toolkit. * + * See https://openmm.org/development. * + * * + * Portions copyright (c) 2026 Stanford University and the Authors. * + * Authors: Peter Eastman * + * Contributors: * + * * + * 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. * + * -------------------------------------------------------------------------- */ + +#include "PythonTorchForce.h" +#include "openmm/Kernel.h" +#include "openmm/internal/ForceImpl.h" +#include +#include +#include + +namespace TorchPlugin { + +/** + * This is the internal implementation of PythonTorchForce. + */ + +class PythonTorchForceImpl : public OpenMM::ForceImpl { +public: + PythonTorchForceImpl(const PythonTorchForce& owner); + ~PythonTorchForceImpl(); + void initialize(OpenMM::ContextImpl& context); + const PythonTorchForce& getOwner() const { + return owner; + } + void updateContextState(OpenMM::ContextImpl& context, bool& forcesInvalid) { + // This force field doesn't update the state directly. + } + double calcForcesAndEnergy(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy, int groups); + std::map getDefaultParameters(); + std::vector getKernelNames(); + std::vector > getBondedParticles() const { + return {}; + } +private: + const PythonTorchForce& owner; + const PythonTorchForceComputation& computation; + std::map defaultParameters; + bool usePeriodic; + OpenMM::Kernel kernel; +}; + +} // namespace TorchPlugin + +#endif /*OPENMM_PYTHONTORCHFORCEIMPL_H_*/ diff --git a/openmmapi/src/PythonTorchForce.cpp b/openmmapi/src/PythonTorchForce.cpp new file mode 100644 index 00000000..258244f5 --- /dev/null +++ b/openmmapi/src/PythonTorchForce.cpp @@ -0,0 +1,77 @@ +/* -------------------------------------------------------------------------- * + * OpenMM * + * -------------------------------------------------------------------------- * + * This is part of the OpenMM molecular simulation toolkit. * + * See https://openmm.org/development. * + * * + * Portions copyright (c) 2026 Stanford University and the Authors. * + * Authors: Peter Eastman * + * Contributors: * + * * + * 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. * + * -------------------------------------------------------------------------- */ + +#include "PythonTorchForce.h" +#include "internal/PythonTorchForceImpl.h" + +using namespace TorchPlugin; +using namespace OpenMM; +using namespace std; + +PythonTorchForce::PythonTorchForce(PythonTorchForceComputation* computation, const map& globalParameters, const vector& particles) : + computation(computation), globalParameters(globalParameters), usePeriodic(false), particles(particles) { +} + +PythonTorchForce::~PythonTorchForce() { + delete computation; +} + +const PythonTorchForceComputation& PythonTorchForce::getComputation() const { + return *computation; +} + +const map& PythonTorchForce::getGlobalParameters() const { + return globalParameters; +} + +void PythonTorchForce::setParticles(const std::vector& particles) { + this->particles = particles; +} + +bool PythonTorchForce::usesPeriodicBoundaryConditions() const { + return usePeriodic; +} + +void PythonTorchForce::setUsesPeriodicBoundaryConditions(bool periodic) { + usePeriodic = periodic; +} + +const vector& PythonTorchForce::getPickledFunction() const { + return pickled; +} + +void PythonTorchForce::setPickledFunction(char* function, int length) { + pickled.clear(); + for (int i = 0; i < length; i++) + pickled.push_back(function[i]); +} + +ForceImpl* PythonTorchForce::createImpl() const { + return new PythonTorchForceImpl(*this); +} diff --git a/openmmapi/src/PythonTorchForceImpl.cpp b/openmmapi/src/PythonTorchForceImpl.cpp new file mode 100644 index 00000000..a1fbfaa9 --- /dev/null +++ b/openmmapi/src/PythonTorchForceImpl.cpp @@ -0,0 +1,66 @@ +/* -------------------------------------------------------------------------- * + * OpenMM * + * -------------------------------------------------------------------------- * + * This is part of the OpenMM molecular simulation toolkit. * + * See https://openmm.org/development. * + * * + * Portions copyright (c) 2026 Stanford University and the Authors. * + * Authors: Peter Eastman * + * Contributors: * + * * + * 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. * + * -------------------------------------------------------------------------- */ + +#include "internal/PythonTorchForceImpl.h" +#include "TorchKernels.h" +#include "openmm/OpenMMException.h" +#include "openmm/internal/ContextImpl.h" +#include "openmm/kernels.h" +#include + +using namespace TorchPlugin; +using namespace OpenMM; +using namespace std; + +PythonTorchForceImpl::PythonTorchForceImpl(const PythonTorchForce& owner) : owner(owner), computation(owner.getComputation()), + defaultParameters(owner.getGlobalParameters()), usePeriodic(owner.usesPeriodicBoundaryConditions()) { + forceGroup = owner.getForceGroup(); +} + +PythonTorchForceImpl::~PythonTorchForceImpl() { +} + +void PythonTorchForceImpl::initialize(ContextImpl& context) { + kernel = context.getPlatform().createKernel(CalcPythonTorchForceKernel::Name(), context); + kernel.getAs().initialize(context, owner); +} + +double PythonTorchForceImpl::calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups) { + if ((groups&(1<().execute(context, includeForces, includeEnergy); + return 0.0; +} + +vector PythonTorchForceImpl::getKernelNames() { + return {CalcCustomCPPForceKernel::Name()}; +} + +map PythonTorchForceImpl::getDefaultParameters() { + return defaultParameters; +} diff --git a/platforms/common/src/CommonTorchKernels.cpp b/platforms/common/src/CommonTorchKernels.cpp index 1c9186ce..e04071c7 100644 --- a/platforms/common/src/CommonTorchKernels.cpp +++ b/platforms/common/src/CommonTorchKernels.cpp @@ -6,7 +6,7 @@ * Biological Structures at Stanford, funded under the NIH Roadmap for * * Medical Research, grant U54 GM072970. See https://simtk.org. * * * - * Portions copyright (c) 2018-2024 Stanford University and the Authors. * + * Portions copyright (c) 2018-2026 Stanford University and the Authors. * * Authors: Peter Eastman * * Contributors: * * * @@ -31,6 +31,7 @@ #include "CommonTorchKernels.h" #include "CommonTorchKernelSources.h" +#include "openmm/common/ContextSelector.h" #include "openmm/internal/ContextImpl.h" #include @@ -147,3 +148,112 @@ double CommonCalcTorchForceKernel::execute(ContextImpl& context, bool includeFor return energyTensor.item(); } +class CommonCalcPythonTorchForceKernel::ReorderListener : public ComputeContext::ReorderListener { +public: + ReorderListener(CommonCalcPythonTorchForceKernel& owner) : owner(owner) { + } + void execute() { + owner.sortParticles(); + } +private: + CommonCalcPythonTorchForceKernel& owner; +}; + +void CommonCalcPythonTorchForceKernel::initialize(const ContextImpl& context, const PythonTorchForce& force) { + ContextSelector selector(cc); + computation = &force.getComputation(); + usePeriodic = force.usesPeriodicBoundaryConditions(); + particles = force.getParticles(); + numParticles = particles.size(); + if (numParticles == 0) + numParticles = context.getSystem().getNumParticles(); + positionsVec.resize(3*numParticles); + int elementSize = (cc.getUseDoublePrecision() ? sizeof(double) : sizeof(float)); + positionsArray.initialize(cc, 3*numParticles, elementSize, "positions"); + forcesArray.initialize(cc, 3*numParticles, elementSize, "forces"); + map defines; + defines["NUM_ATOMS"] = cc.intToString(numParticles); + defines["PADDED_NUM_ATOMS"] = cc.intToString(cc.getPaddedNumAtoms()); + ComputeProgram program = cc.compileProgram(CommonTorchKernelSources::pythonTorchForce, defines); + if (particles.size() > 0) { + reorderedParticles.initialize(cc, numParticles, "reorderedParticles"); + reorderedParticles.upload(particles); + cc.addReorderListener(new ReorderListener(*this)); + addForcesKernel = program->createKernel("addForcesSubset"); + addForcesKernel->addArg(forcesArray); + addForcesKernel->addArg(cc.getLongForceBuffer()); + addForcesKernel->addArg(cc.getAtomIndexArray()); + addForcesKernel->addArg(reorderedParticles); + addForcesKernel->addArg(numParticles); + copyPositionsKernel = program->createKernel("copyPositionsSubset"); + } + else { + addForcesKernel = program->createKernel("addForcesAll"); + addForcesKernel->addArg(forcesArray); + addForcesKernel->addArg(cc.getLongForceBuffer()); + addForcesKernel->addArg(cc.getAtomIndexArray()); + copyPositionsKernel = program->createKernel("copyPositionsAll"); + } + copyPositionsKernel->addArg(cc.getPosq()); + copyPositionsKernel->addArg(positionsArray); + copyPositionsKernel->addArg(particles.size() > 0 ? reorderedParticles : cc.getAtomIndexArray()); + copyPositionsKernel->addArg(numParticles); +} + +double CommonCalcPythonTorchForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) { + if (cc.getContextIndex() != 0) + return 0.0; + torch::Tensor posTensor = getPositions(); + State::StateBuilder builder(contextImpl.getTime(), contextImpl.getStepCount()); + builder.setParameters(contextImpl.getParameters()); + if (usePeriodic) { + Vec3 a, b, c; + contextImpl.getPeriodicBoxVectors(a, b, c); + builder.setPeriodicBoxVectors(a, b, c); + } + State state = builder.getState(); + torch::Tensor forceTensor = computation->compute(state, posTensor, energy); + if (includeForces) + addForces(forceTensor); + return includeEnergy ? energy : 0.0; +} + +torch::Tensor CommonCalcPythonTorchForceKernel::getPositions() { + ContextSelector selector(cc); + copyPositionsKernel->execute(numParticles); + positionsArray.download(positionsVec.data()); + auto dtype = (cc.getUseDoublePrecision() ? torch::kFloat64 : torch::kFloat32); + return torch::from_blob(positionsVec.data(), {numParticles, 3}, torch::TensorOptions().dtype(dtype).requires_grad(true)); +} + +void CommonCalcPythonTorchForceKernel::sortParticles() { + // Update the list of particles to account for reordering. + + const vector& order = cc.getAtomIndex(); + vector inverseOrder(order.size()); + for (int i = 0; i < cc.getNumAtoms(); i++) + inverseOrder[order[i]] = i; + vector reordered(particles.size()); + for (int i = 0; i < particles.size(); i++) + reordered[i] = inverseOrder[particles[i]]; + reorderedParticles.upload(reordered); +} + +void CommonCalcPythonTorchForceKernel::addForces(torch::Tensor forceTensor) { + // Add in the forces. + + ContextSelector selector(cc); + if (cc.getUseDoublePrecision()) { + if (!(forceTensor.dtype() == torch::kFloat64)) + forceTensor = forceTensor.to(torch::kFloat64); + double* data = forceTensor.data_ptr(); + forcesArray.upload(data); + } + else { + if (!(forceTensor.dtype() == torch::kFloat32)) + forceTensor = forceTensor.to(torch::kFloat32); + float* data = forceTensor.data_ptr(); + forcesArray.upload(data); + } + addForcesKernel->execute(cc.getNumAtoms()); +} diff --git a/platforms/common/src/CommonTorchKernels.h b/platforms/common/src/CommonTorchKernels.h index 2c79ddda..227103a7 100644 --- a/platforms/common/src/CommonTorchKernels.h +++ b/platforms/common/src/CommonTorchKernels.h @@ -9,7 +9,7 @@ * Biological Structures at Stanford, funded under the NIH Roadmap for * * Medical Research, grant U54 GM072970. See https://simtk.org. * * * - * Portions copyright (c) 2018-2024 Stanford University and the Authors. * + * Portions copyright (c) 2018-2026 Stanford University and the Authors. * * Authors: Peter Eastman * * Contributors: * * * @@ -76,6 +76,54 @@ class CommonCalcTorchForceKernel : public CalcTorchForceKernel { OpenMM::ComputeKernel addForcesKernel; }; +/** + * This kernel is invoked by PythonTorchForce to calculate the forces acting on the system and the energy of the system. + */ +class CommonCalcPythonTorchForceKernel : public CalcPythonTorchForceKernel { +public: + CommonCalcPythonTorchForceKernel(std::string name, const OpenMM::Platform& platform, OpenMM::ContextImpl& contextImpl, OpenMM::ComputeContext& cc) : + CalcPythonTorchForceKernel(name, platform), contextImpl(contextImpl), cc(cc) { + } + /** + * Initialize the kernel. + * + * @param context the ContextImpl this kernel will be applied to + * @param force the PythonTorchForce this kernel will be used for + */ + void initialize(const OpenMM::ContextImpl& context, const PythonTorchForce& force); + /** + * Execute the kernel to calculate the forces and/or energy. + * + * @param context the context in which to execute this kernel + * @param includeForces true if forces should be calculated + * @param includeEnergy true if the energy should be calculated + * @return the potential energy due to the force + */ + double execute(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy); + /** + * Retrieve the current positions as a Tensor. Subclasses can override this to do it + * more efficiently. + */ + virtual torch::Tensor getPositions(); + /** + * Add in the forces. Subclasses can override this to do it more efficiently. + */ + virtual void addForces(torch::Tensor forceTensor); +protected: + class ReorderListener; + void sortParticles(); + OpenMM::ContextImpl& contextImpl; + OpenMM::ComputeContext& cc; + const PythonTorchForceComputation* computation; + OpenMM::ComputeArray positionsArray, forcesArray, particlesArray, reorderedParticles; + OpenMM::ComputeKernel copyPositionsKernel, addForcesKernel; + std::vector positionsVec; + std::vector particles; + int numParticles; + double energy; + bool usePeriodic, useWorkerThread; +}; + } // namespace TorchPlugin #endif /*COMMON_TORCH_KERNELS_H_*/ diff --git a/platforms/common/src/kernels/pythonTorchForce.cc b/platforms/common/src/kernels/pythonTorchForce.cc new file mode 100644 index 00000000..27f8589b --- /dev/null +++ b/platforms/common/src/kernels/pythonTorchForce.cc @@ -0,0 +1,36 @@ +KERNEL void copyPositionsAll(GLOBAL const real4* RESTRICT posq, GLOBAL real* RESTRICT positions, GLOBAL int* RESTRICT particles, int numParticles) { + for (int i = GLOBAL_ID; i < numParticles; i += GLOBAL_SIZE) { + int index = particles[i]; + real4 pos = posq[i]; + positions[3*index] = pos.x; + positions[3*index+1] = pos.y; + positions[3*index+2] = pos.z; + } +} + +KERNEL void copyPositionsSubset(GLOBAL const real4* RESTRICT posq, GLOBAL real* RESTRICT positions, GLOBAL int* RESTRICT particles, int numParticles) { + for (int i = GLOBAL_ID; i < numParticles; i += GLOBAL_SIZE) { + real4 pos = posq[particles[i]]; + positions[3*i] = pos.x; + positions[3*i+1] = pos.y; + positions[3*i+2] = pos.z; + } +} + +KERNEL void addForcesAll(GLOBAL const real* RESTRICT forces, GLOBAL mm_long* RESTRICT forceBuffers, GLOBAL int* RESTRICT atomIndex) { + for (int atom = GLOBAL_ID; atom < NUM_ATOMS; atom += GLOBAL_SIZE) { + int index = atomIndex[atom]; + forceBuffers[atom] += (mm_long) (forces[3*index]*0x100000000); + forceBuffers[atom+PADDED_NUM_ATOMS] += (mm_long) (forces[3*index+1]*0x100000000); + forceBuffers[atom+2*PADDED_NUM_ATOMS] += (mm_long) (forces[3*index+2]*0x100000000); + } +} + +KERNEL void addForcesSubset(GLOBAL const real* RESTRICT forces, GLOBAL mm_long* RESTRICT forceBuffers, GLOBAL int* RESTRICT atomIndex, GLOBAL int* RESTRICT particles, int numParticles) { + for (int i = GLOBAL_ID; i < numParticles; i += GLOBAL_SIZE) { + int index = particles[i]; + forceBuffers[index] += (mm_long) (forces[3*i]*0x100000000); + forceBuffers[index+PADDED_NUM_ATOMS] += (mm_long) (forces[3*i+1]*0x100000000); + forceBuffers[index+2*PADDED_NUM_ATOMS] += (mm_long) (forces[3*i+2]*0x100000000); + } +} diff --git a/platforms/cuda/CMakeLists.txt b/platforms/cuda/CMakeLists.txt index 0eb9e831..064e30c9 100644 --- a/platforms/cuda/CMakeLists.txt +++ b/platforms/cuda/CMakeLists.txt @@ -35,6 +35,8 @@ INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/src) INCLUDE_DIRECTORIES(BEFORE ${CMAKE_SOURCE_DIR}/platforms/cuda/include) INCLUDE_DIRECTORIES(BEFORE ${CMAKE_SOURCE_DIR}/platforms/cuda/src) INCLUDE_DIRECTORIES(BEFORE ${CMAKE_BINARY_DIR}/platforms/cuda/src) +INCLUDE_DIRECTORIES(BEFORE ${CMAKE_SOURCE_DIR}/platforms/common/src) +INCLUDE_DIRECTORIES(BEFORE ${CMAKE_BINARY_DIR}/platforms/common/src) # Set variables needed for encoding kernel sources into a C++ class @@ -42,7 +44,8 @@ SET(CUDA_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/src) SET(CUDA_SOURCE_CLASS CudaTorchKernelSources) SET(CUDA_KERNELS_CPP ${CMAKE_CURRENT_BINARY_DIR}/src/${CUDA_SOURCE_CLASS}.cpp) SET(CUDA_KERNELS_H ${CMAKE_CURRENT_BINARY_DIR}/src/${CUDA_SOURCE_CLASS}.h) -SET(SOURCE_FILES ${SOURCE_FILES} ${CUDA_KERNELS_CPP} ${CUDA_KERNELS_H}) +SET(COMMON_KERNELS_CPP ${CMAKE_CURRENT_SOURCE_DIR}/../common/src/CommonTorchKernels.cpp ${CMAKE_CURRENT_BINARY_DIR}/../common/src/CommonTorchKernelSources.cpp) +SET(SOURCE_FILES ${SOURCE_FILES} ${CUDA_KERNELS_CPP} ${CUDA_KERNELS_H} ${COMMON_KERNELS_CPP}) INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_BINARY_DIR}/src) # Create the library @@ -55,10 +58,11 @@ ADD_CUSTOM_COMMAND(OUTPUT ${CUDA_KERNELS_CPP} ${CUDA_KERNELS_H} ARGS -D CUDA_SOURCE_DIR=${CUDA_SOURCE_DIR} -D CUDA_KERNELS_CPP=${CUDA_KERNELS_CPP} -D CUDA_KERNELS_H=${CUDA_KERNELS_H} -D CUDA_SOURCE_CLASS=${CUDA_SOURCE_CLASS} -P ${CMAKE_SOURCE_DIR}/platforms/cuda/EncodeCUDAFiles.cmake DEPENDS ${CUDA_KERNELS} ) -SET_SOURCE_FILES_PROPERTIES(${CUDA_KERNELS_CPP} ${CUDA_KERNELS_H} PROPERTIES GENERATED TRUE) +SET_SOURCE_FILES_PROPERTIES(${CUDA_KERNELS_CPP} ${CUDA_KERNELS_H} ${COMMON_KERNELS_CPP} PROPERTIES GENERATED TRUE) ADD_LIBRARY(${SHARED_TARGET} SHARED ${SOURCE_FILES} ${SOURCE_INCLUDE_FILES} ${API_INCLUDE_FILES}) +ADD_DEPENDENCIES(${SHARED_TARGET} CommonKernels) -TARGET_LINK_LIBRARIES(${SHARED_TARGET} ${CUDA_LIBRARIES}) +TARGET_LINK_LIBRARIES(${SHARED_TARGET} CUDA::cuda_driver) IF(WIN32) FIND_LIBRARY(CUDA_DRIVER_LIBRARY cuda HINTS ${CUDA_DRIVER_LIBRARY_PATH}) IF(NOT CUDA_DRIVER_LIBRARY) diff --git a/platforms/cuda/src/CudaTorchKernelFactory.cpp b/platforms/cuda/src/CudaTorchKernelFactory.cpp index 316a480f..f8f83ebc 100644 --- a/platforms/cuda/src/CudaTorchKernelFactory.cpp +++ b/platforms/cuda/src/CudaTorchKernelFactory.cpp @@ -50,6 +50,7 @@ extern "C" OPENMM_EXPORT void registerKernelFactories() { Platform& platform = Platform::getPlatformByName("CUDA"); CudaTorchKernelFactory* factory = new CudaTorchKernelFactory(); platform.registerKernelFactory(CalcTorchForceKernel::Name(), factory); + platform.registerKernelFactory(CalcPythonTorchForceKernel::Name(), factory); } catch (std::exception ex) { // Ignore @@ -70,5 +71,7 @@ KernelImpl* CudaTorchKernelFactory::createKernelImpl(std::string name, const Pla CudaContext& cu = *static_cast(context.getPlatformData())->contexts[0]; if (name == CalcTorchForceKernel::Name()) return new CudaCalcTorchForceKernel(name, platform, cu); + if (name == CalcPythonTorchForceKernel::Name()) + return new CudaCalcPythonTorchForceKernel(name, platform, context, cu); throw OpenMMException((std::string("Tried to create kernel with illegal kernel name '")+name+"'").c_str()); } diff --git a/platforms/cuda/src/CudaTorchKernels.cpp b/platforms/cuda/src/CudaTorchKernels.cpp index 8fa167b0..a6c69e8c 100644 --- a/platforms/cuda/src/CudaTorchKernels.cpp +++ b/platforms/cuda/src/CudaTorchKernels.cpp @@ -37,6 +37,7 @@ #include #include #include + using namespace TorchPlugin; using namespace OpenMM; using namespace std; @@ -284,3 +285,32 @@ double CudaCalcTorchForceKernel::execute(ContextImpl& context, bool includeForce assert(primaryContext == ctx); // Check that the correct context was popped return energy; } + +CudaCalcPythonTorchForceKernel::CudaCalcPythonTorchForceKernel(std::string name, const OpenMM::Platform& platform, OpenMM::ContextImpl& contextImpl, OpenMM::ComputeContext& cc) : + CommonCalcPythonTorchForceKernel(name, platform, contextImpl, cc), cu(dynamic_cast(cc)) { +} + +void CudaCalcPythonTorchForceKernel::initialize(const ContextImpl& context, const PythonTorchForce& force) { + CommonCalcPythonTorchForceKernel::initialize(context, force); + torch::Device device(torch::kCUDA, cu.getDeviceIndex()); + torch::TensorOptions options = torch::TensorOptions().device(device).dtype(cu.getUseDoublePrecision() ? torch::kFloat64 : torch::kFloat32); + posTensor = torch::empty({numParticles, 3}, options.requires_grad(true)); +} + +torch::Tensor CudaCalcPythonTorchForceKernel::getPositions() { + ContextSelector selector(cc); + copyPositionsKernel->execute(numParticles); + CUdeviceptr source = cu.unwrap(positionsArray).getDevicePointer(); + CUdeviceptr dest = (CUdeviceptr) getTensorPointer(cu, posTensor); + CHECK_RESULT(cuMemcpyDtoDAsync(dest, source, positionsArray.getSize()*positionsArray.getElementSize(), cu.getCurrentStream()), "Error copying positions"); + return posTensor; +} + +void CudaCalcPythonTorchForceKernel::addForces(torch::Tensor forceTensor) { + ContextSelector selector(cc); + forceTensor = forceTensor.to(posTensor.device()).to(posTensor.dtype()); + CUdeviceptr source = (CUdeviceptr) getTensorPointer(cu, forceTensor); + CUdeviceptr dest = cu.unwrap(forcesArray).getDevicePointer(); + CHECK_RESULT(cuMemcpyDtoDAsync(dest, source, forcesArray.getSize()*forcesArray.getElementSize(), cu.getCurrentStream()), "Error copying forces"); + addForcesKernel->execute(cc.getNumAtoms()); +} diff --git a/platforms/cuda/src/CudaTorchKernels.h b/platforms/cuda/src/CudaTorchKernels.h index 1f7b4c97..3f0ac8c7 100644 --- a/platforms/cuda/src/CudaTorchKernels.h +++ b/platforms/cuda/src/CudaTorchKernels.h @@ -9,7 +9,7 @@ * Biological Structures at Stanford, funded under the NIH Roadmap for * * Medical Research, grant U54 GM072970. See https://simtk.org. * * * - * Portions copyright (c) 2018-2024 Stanford University and the Authors. * + * Portions copyright (c) 2018-2026 Stanford University and the Authors. * * Authors: Peter Eastman * * Contributors: Raimondas Galvelis, Raul P. Pelaez * * * @@ -33,6 +33,7 @@ * -------------------------------------------------------------------------- */ #include "TorchKernels.h" +#include "CommonTorchKernels.h" #include "openmm/cuda/CudaContext.h" #include "openmm/cuda/CudaArray.h" #include @@ -85,6 +86,35 @@ class CudaCalcTorchForceKernel : public CalcTorchForceKernel { int warmupSteps; }; + +/** + * This kernel is invoked by PythonTorchForce to calculate the forces acting on the system and the energy of the system. + */ +class CudaCalcPythonTorchForceKernel : public CommonCalcPythonTorchForceKernel { +public: + CudaCalcPythonTorchForceKernel(std::string name, const OpenMM::Platform& platform, OpenMM::ContextImpl& contextImpl, OpenMM::ComputeContext& cc); + /** + * Initialize the kernel. + * + * @param context the ContextImpl this kernel will be applied to + * @param force the PythonTorchForce this kernel will be used for + */ + void initialize(const OpenMM::ContextImpl& context, const PythonTorchForce& force); + /** + * Retrieve the current positions as a Tensor. This overrides the superclass to use + * a tensor stored on the GPU. + */ + torch::Tensor getPositions(); + /** + * Add in the forces. This overrides the superclass to work with a tensor stored on the GPU. + */ + void addForces(torch::Tensor forceTensor); +private: + OpenMM::CudaContext& cu; + torch::Tensor posTensor; +}; + + } // namespace TorchPlugin #endif /*CUDA_TORCH_KERNELS_H_*/ diff --git a/platforms/cuda/tests/TestCudaPythonTorchForce.cpp b/platforms/cuda/tests/TestCudaPythonTorchForce.cpp new file mode 100644 index 00000000..97390f17 --- /dev/null +++ b/platforms/cuda/tests/TestCudaPythonTorchForce.cpp @@ -0,0 +1,127 @@ +/* -------------------------------------------------------------------------- * + * OpenMM * + * -------------------------------------------------------------------------- * + * This is part of the OpenMM molecular simulation toolkit. * + * See https://openmm.org/development. * + * * + * Portions copyright (c) 2025-2026 Stanford University and the Authors. * + * Authors: Peter Eastman * + * Contributors: * + * * + * 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. * + * -------------------------------------------------------------------------- */ + +#include "PythonTorchForce.h" +#include "openmm/internal/AssertionUtilities.h" +#include "openmm/Context.h" +#include "openmm/NonbondedForce.h" +#include "openmm/Platform.h" +#include "openmm/VerletIntegrator.h" +#include "sfmt/SFMT.h" +#include + +using namespace TorchPlugin; +using namespace OpenMM; +using namespace std; + +extern "C" OPENMM_EXPORT void registerTorchCudaKernelFactories(); + +void testForce(bool subsetParticles) { + class Computation : public PythonTorchForceComputation { + torch::Tensor compute(const State& state, const torch::Tensor& positions, double& energy) const { + ASSERT_EQUAL(5.0, state.getParameters().at("a")); + ASSERT_EQUAL(10.0, state.getParameters().at("b")); + Vec3 a, b, c; + state.getPeriodicBoxVectors(a, b, c); + ASSERT_EQUAL(Vec3(2, 0, 0), a); + ASSERT_EQUAL(Vec3(0.1, 2, 0), b); + ASSERT_EQUAL(Vec3(0.1, 0.1, 2), c); + energy = 25.0; + return positions*2; + } + }; + int numParticles = 5; + int totalParticles = (subsetParticles ? numParticles+10 : numParticles); + System system; + Vec3 a(2, 0, 0); + Vec3 b(0.1, 2, 0); + Vec3 c(0.1, 0.1, 2); + system.setDefaultPeriodicBoxVectors(a, b, c); + NonbondedForce* nonbonded = new NonbondedForce(); // To trigger reordering + nonbonded->setNonbondedMethod(NonbondedForce::PME); + system.addForce(nonbonded); + vector positions; + OpenMM_SFMT::SFMT sfmt; + init_gen_rand(0, sfmt); + for (int i = 0; i < totalParticles; i++) { + system.addParticle(1.0); + positions.push_back(Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))); + nonbonded->addParticle(0.0, 1.0, 0.0); + } + map params; + params["a"] = 5.0; + params["b"] = 10.0; + vector particles; + if (subsetParticles) + for (int i = 0; i < numParticles; i++) + particles.push_back(i+5); + PythonTorchForce* force = new PythonTorchForce(new Computation(), params, particles); + ASSERT(!force->usesPeriodicBoundaryConditions()); + force->setUsesPeriodicBoundaryConditions(true); + ASSERT(force->usesPeriodicBoundaryConditions()); + system.addForce(force); + VerletIntegrator integrator(0.01); + Platform& platform = Platform::getPlatformByName("CUDA"); + Context context(system, integrator, platform); + context.setPositions(positions); + State state = context.getState(State::Energy | State::Forces); + ASSERT_EQUAL_TOL(25.0, state.getPotentialEnergy(), 1e-6); + if (subsetParticles) { + for (int i : particles) + ASSERT_EQUAL_VEC(2*positions[i], state.getForces()[i], 1e-6) + Vec3 zero; + for (int i = 0; i < 5; i++) + ASSERT_EQUAL_VEC(zero, state.getForces()[i], 1e-6); + } + else { + for (int i = 0; i < numParticles; i++) + ASSERT_EQUAL_VEC(2*positions[i], state.getForces()[i], 1e-6) + } + + // Check that force groups are handled correctly. + + ASSERT_EQUAL_TOL(25.0, context.getState(State::Energy, false, 1).getPotentialEnergy(), 1e-6); + ASSERT_EQUAL_TOL(0.0, context.getState(State::Energy, false, 2).getPotentialEnergy(), 1e-6); +} + +void runPlatformTests(); + +int main(int argc, char* argv[]) { + try { + registerTorchCudaKernelFactories(); + testForce(false); + testForce(true); + } + catch(const exception& e) { + cout << "exception: " << e.what() << endl; + return 1; + } + cout << "Done" << endl; + return 0; +} diff --git a/platforms/hip/src/HipTorchKernelFactory.cpp b/platforms/hip/src/HipTorchKernelFactory.cpp index 23b42c2a..24ade325 100644 --- a/platforms/hip/src/HipTorchKernelFactory.cpp +++ b/platforms/hip/src/HipTorchKernelFactory.cpp @@ -6,7 +6,7 @@ * Biological Structures at Stanford, funded under the NIH Roadmap for * * Medical Research, grant U54 GM072970. See https://simtk.org. * * * - * Portions copyright (c) 2024 Stanford University and the Authors. * + * Portions copyright (c) 2024-2025 Stanford University and the Authors. * * Authors: Peter Eastman * * Contributors: * * * @@ -51,6 +51,7 @@ extern "C" OPENMM_EXPORT void registerKernelFactories() { Platform& platform = Platform::getPlatformByName("HIP"); HipTorchKernelFactory* factory = new HipTorchKernelFactory(); platform.registerKernelFactory(CalcTorchForceKernel::Name(), factory); + platform.registerKernelFactory(CalcPythonTorchForceKernel::Name(), factory); } catch (std::exception ex) { // Ignore @@ -71,5 +72,7 @@ KernelImpl* HipTorchKernelFactory::createKernelImpl(std::string name, const Plat HipContext& cc = *static_cast(context.getPlatformData())->contexts[0]; if (name == CalcTorchForceKernel::Name()) return new CommonCalcTorchForceKernel(name, platform, cc); + if (name == CalcPythonTorchForceKernel::Name()) + return new CommonCalcPythonTorchForceKernel(name, platform, context, cc); throw OpenMMException((std::string("Tried to create kernel with illegal kernel name '")+name+"'").c_str()); } diff --git a/platforms/hip/tests/TestHipPythonTorchForce.cpp b/platforms/hip/tests/TestHipPythonTorchForce.cpp new file mode 100644 index 00000000..7bee37c6 --- /dev/null +++ b/platforms/hip/tests/TestHipPythonTorchForce.cpp @@ -0,0 +1,127 @@ +/* -------------------------------------------------------------------------- * + * OpenMM * + * -------------------------------------------------------------------------- * + * This is part of the OpenMM molecular simulation toolkit. * + * See https://openmm.org/development. * + * * + * Portions copyright (c) 2025-2026 Stanford University and the Authors. * + * Authors: Peter Eastman * + * Contributors: * + * * + * 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. * + * -------------------------------------------------------------------------- */ + +#include "PythonTorchForce.h" +#include "openmm/internal/AssertionUtilities.h" +#include "openmm/Context.h" +#include "openmm/NonbondedForce.h" +#include "openmm/Platform.h" +#include "openmm/VerletIntegrator.h" +#include "sfmt/SFMT.h" +#include + +using namespace TorchPlugin; +using namespace OpenMM; +using namespace std; + +extern "C" OPENMM_EXPORT void registerTorchHipKernelFactories(); + +void testForce(bool subsetParticles) { + class Computation : public PythonTorchForceComputation { + torch::Tensor compute(const State& state, const torch::Tensor& positions, double& energy) const { + ASSERT_EQUAL(5.0, state.getParameters().at("a")); + ASSERT_EQUAL(10.0, state.getParameters().at("b")); + Vec3 a, b, c; + state.getPeriodicBoxVectors(a, b, c); + ASSERT_EQUAL(Vec3(2, 0, 0), a); + ASSERT_EQUAL(Vec3(0.1, 2, 0), b); + ASSERT_EQUAL(Vec3(0.1, 0.1, 2), c); + energy = 25.0; + return positions*2; + } + }; + int numParticles = 5; + int totalParticles = (subsetParticles ? numParticles+10 : numParticles); + System system; + Vec3 a(2, 0, 0); + Vec3 b(0.1, 2, 0); + Vec3 c(0.1, 0.1, 2); + system.setDefaultPeriodicBoxVectors(a, b, c); + NonbondedForce* nonbonded = new NonbondedForce(); // To trigger reordering + nonbonded->setNonbondedMethod(NonbondedForce::PME); + system.addForce(nonbonded); + vector positions; + OpenMM_SFMT::SFMT sfmt; + init_gen_rand(0, sfmt); + for (int i = 0; i < totalParticles; i++) { + system.addParticle(1.0); + positions.push_back(Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))); + nonbonded->addParticle(0.0, 1.0, 0.0); + } + map params; + params["a"] = 5.0; + params["b"] = 10.0; + vector particles; + if (subsetParticles) + for (int i = 0; i < numParticles; i++) + particles.push_back(i+5); + PythonTorchForce* force = new PythonTorchForce(new Computation(), params, particles); + ASSERT(!force->usesPeriodicBoundaryConditions()); + force->setUsesPeriodicBoundaryConditions(true); + ASSERT(force->usesPeriodicBoundaryConditions()); + system.addForce(force); + VerletIntegrator integrator(0.01); + Platform& platform = Platform::getPlatformByName("HIP"); + Context context(system, integrator, platform); + context.setPositions(positions); + State state = context.getState(State::Energy | State::Forces); + ASSERT_EQUAL_TOL(25.0, state.getPotentialEnergy(), 1e-6); + if (subsetParticles) { + for (int i : particles) + ASSERT_EQUAL_VEC(2*positions[i], state.getForces()[i], 1e-6) + Vec3 zero; + for (int i = 0; i < 5; i++) + ASSERT_EQUAL_VEC(zero, state.getForces()[i], 1e-6); + } + else { + for (int i = 0; i < numParticles; i++) + ASSERT_EQUAL_VEC(2*positions[i], state.getForces()[i], 1e-6) + } + + // Check that force groups are handled correctly. + + ASSERT_EQUAL_TOL(25.0, context.getState(State::Energy, false, 1).getPotentialEnergy(), 1e-6); + ASSERT_EQUAL_TOL(0.0, context.getState(State::Energy, false, 2).getPotentialEnergy(), 1e-6); +} + +void runPlatformTests(); + +int main(int argc, char* argv[]) { + try { + registerTorchHipKernelFactories(); + testForce(false); + testForce(true); + } + catch(const exception& e) { + cout << "exception: " << e.what() << endl; + return 1; + } + cout << "Done" << endl; + return 0; +} diff --git a/platforms/opencl/src/OpenCLTorchKernelFactory.cpp b/platforms/opencl/src/OpenCLTorchKernelFactory.cpp index c847fd91..70383f2d 100644 --- a/platforms/opencl/src/OpenCLTorchKernelFactory.cpp +++ b/platforms/opencl/src/OpenCLTorchKernelFactory.cpp @@ -51,6 +51,7 @@ extern "C" OPENMM_EXPORT void registerKernelFactories() { Platform& platform = Platform::getPlatformByName("OpenCL"); OpenCLTorchKernelFactory* factory = new OpenCLTorchKernelFactory(); platform.registerKernelFactory(CalcTorchForceKernel::Name(), factory); + platform.registerKernelFactory(CalcPythonTorchForceKernel::Name(), factory); } catch (std::exception ex) { // Ignore @@ -71,5 +72,7 @@ KernelImpl* OpenCLTorchKernelFactory::createKernelImpl(std::string name, const P OpenCLContext& cl = *static_cast(context.getPlatformData())->contexts[0]; if (name == CalcTorchForceKernel::Name()) return new CommonCalcTorchForceKernel(name, platform, cl); + if (name == CalcPythonTorchForceKernel::Name()) + return new CommonCalcPythonTorchForceKernel(name, platform, context, cl); throw OpenMMException((std::string("Tried to create kernel with illegal kernel name '")+name+"'").c_str()); } diff --git a/platforms/opencl/tests/TestOpenCLPythonTorchForce.cpp b/platforms/opencl/tests/TestOpenCLPythonTorchForce.cpp new file mode 100644 index 00000000..84c65710 --- /dev/null +++ b/platforms/opencl/tests/TestOpenCLPythonTorchForce.cpp @@ -0,0 +1,127 @@ +/* -------------------------------------------------------------------------- * + * OpenMM * + * -------------------------------------------------------------------------- * + * This is part of the OpenMM molecular simulation toolkit. * + * See https://openmm.org/development. * + * * + * Portions copyright (c) 2025-2026 Stanford University and the Authors. * + * Authors: Peter Eastman * + * Contributors: * + * * + * 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. * + * -------------------------------------------------------------------------- */ + +#include "PythonTorchForce.h" +#include "openmm/internal/AssertionUtilities.h" +#include "openmm/Context.h" +#include "openmm/NonbondedForce.h" +#include "openmm/Platform.h" +#include "openmm/VerletIntegrator.h" +#include "sfmt/SFMT.h" +#include + +using namespace TorchPlugin; +using namespace OpenMM; +using namespace std; + +extern "C" OPENMM_EXPORT void registerTorchOpenCLKernelFactories(); + +void testForce(bool subsetParticles) { + class Computation : public PythonTorchForceComputation { + torch::Tensor compute(const State& state, const torch::Tensor& positions, double& energy) const { + ASSERT_EQUAL(5.0, state.getParameters().at("a")); + ASSERT_EQUAL(10.0, state.getParameters().at("b")); + Vec3 a, b, c; + state.getPeriodicBoxVectors(a, b, c); + ASSERT_EQUAL(Vec3(2, 0, 0), a); + ASSERT_EQUAL(Vec3(0.1, 2, 0), b); + ASSERT_EQUAL(Vec3(0.1, 0.1, 2), c); + energy = 25.0; + return positions*2; + } + }; + int numParticles = 5; + int totalParticles = (subsetParticles ? numParticles+10 : numParticles); + System system; + Vec3 a(2, 0, 0); + Vec3 b(0.1, 2, 0); + Vec3 c(0.1, 0.1, 2); + system.setDefaultPeriodicBoxVectors(a, b, c); + NonbondedForce* nonbonded = new NonbondedForce(); // To trigger reordering + nonbonded->setNonbondedMethod(NonbondedForce::PME); + system.addForce(nonbonded); + vector positions; + OpenMM_SFMT::SFMT sfmt; + init_gen_rand(0, sfmt); + for (int i = 0; i < totalParticles; i++) { + system.addParticle(1.0); + positions.push_back(Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))); + nonbonded->addParticle(0.0, 1.0, 0.0); + } + map params; + params["a"] = 5.0; + params["b"] = 10.0; + vector particles; + if (subsetParticles) + for (int i = 0; i < numParticles; i++) + particles.push_back(i+5); + PythonTorchForce* force = new PythonTorchForce(new Computation(), params, particles); + ASSERT(!force->usesPeriodicBoundaryConditions()); + force->setUsesPeriodicBoundaryConditions(true); + ASSERT(force->usesPeriodicBoundaryConditions()); + system.addForce(force); + VerletIntegrator integrator(0.01); + Platform& platform = Platform::getPlatformByName("OpenCL"); + Context context(system, integrator, platform); + context.setPositions(positions); + State state = context.getState(State::Energy | State::Forces); + ASSERT_EQUAL_TOL(25.0, state.getPotentialEnergy(), 1e-6); + if (subsetParticles) { + for (int i : particles) + ASSERT_EQUAL_VEC(2*positions[i], state.getForces()[i], 1e-6) + Vec3 zero; + for (int i = 0; i < 5; i++) + ASSERT_EQUAL_VEC(zero, state.getForces()[i], 1e-6); + } + else { + for (int i = 0; i < numParticles; i++) + ASSERT_EQUAL_VEC(2*positions[i], state.getForces()[i], 1e-6) + } + + // Check that force groups are handled correctly. + + ASSERT_EQUAL_TOL(25.0, context.getState(State::Energy, false, 1).getPotentialEnergy(), 1e-6); + ASSERT_EQUAL_TOL(0.0, context.getState(State::Energy, false, 2).getPotentialEnergy(), 1e-6); +} + +void runPlatformTests(); + +int main(int argc, char* argv[]) { + try { + registerTorchOpenCLKernelFactories(); + testForce(false); + testForce(true); + } + catch(const exception& e) { + cout << "exception: " << e.what() << endl; + return 1; + } + cout << "Done" << endl; + return 0; +} diff --git a/platforms/reference/src/ReferenceTorchKernelFactory.cpp b/platforms/reference/src/ReferenceTorchKernelFactory.cpp index 8571bc21..103b0030 100644 --- a/platforms/reference/src/ReferenceTorchKernelFactory.cpp +++ b/platforms/reference/src/ReferenceTorchKernelFactory.cpp @@ -49,6 +49,7 @@ extern "C" OPENMM_EXPORT void registerKernelFactories() { if (dynamic_cast(&platform) != NULL) { ReferenceTorchKernelFactory* factory = new ReferenceTorchKernelFactory(); platform.registerKernelFactory(CalcTorchForceKernel::Name(), factory); + platform.registerKernelFactory(CalcPythonTorchForceKernel::Name(), factory); } } } @@ -61,5 +62,7 @@ KernelImpl* ReferenceTorchKernelFactory::createKernelImpl(std::string name, cons ReferencePlatform::PlatformData& data = *static_cast(context.getPlatformData()); if (name == CalcTorchForceKernel::Name()) return new ReferenceCalcTorchForceKernel(name, platform); + if (name == CalcPythonTorchForceKernel::Name()) + return new ReferenceCalcPythonTorchForceKernel(name, platform); throw OpenMMException((std::string("Tried to create kernel with illegal kernel name '")+name+"'").c_str()); } diff --git a/platforms/reference/src/ReferenceTorchKernels.cpp b/platforms/reference/src/ReferenceTorchKernels.cpp index 346846e0..4f0836f4 100644 --- a/platforms/reference/src/ReferenceTorchKernels.cpp +++ b/platforms/reference/src/ReferenceTorchKernels.cpp @@ -126,3 +126,53 @@ double ReferenceCalcTorchForceKernel::execute(ContextImpl& context, bool include } return energyTensor.item(); } + +void ReferenceCalcPythonTorchForceKernel::initialize(const ContextImpl& context, const PythonTorchForce& force) { + computation = &force.getComputation(); + particles = force.getParticles(); + numParticles = particles.size(); + if (numParticles == 0) + numParticles = context.getSystem().getNumParticles(); + else + positions.resize(numParticles); + usePeriodic = force.usesPeriodicBoundaryConditions(); +} + +double ReferenceCalcPythonTorchForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) { + vector& posData = extractPositions(context); + vector& forceData = extractForces(context); + State::StateBuilder builder(context.getTime(), context.getStepCount()); + torch::Tensor posTensor; + if (particles.size() == 0) + posTensor = torch::from_blob(posData.data(), {numParticles, 3}, torch::TensorOptions().dtype(torch::kFloat64).requires_grad(true)); + else { + for (int i = 0; i < particles.size(); i++) + positions[i] = posData[particles[i]]; + posTensor = torch::from_blob(positions.data(), {numParticles, 3}, torch::TensorOptions().dtype(torch::kFloat64).requires_grad(true)); + } + builder.setParameters(context.getParameters()); + if (usePeriodic) { + Vec3 a, b, c; + context.getPeriodicBoxVectors(a, b, c); + builder.setPeriodicBoxVectors(a, b, c); + } + double energy; + State state = builder.getState(); + torch::Tensor forceTensor = computation->compute(state, posTensor, energy); + if (includeForces) { + if (!(forceTensor.dtype() == torch::kFloat64)) + forceTensor = forceTensor.to(torch::kFloat64); + double* outputForces = forceTensor.data_ptr(); + if (particles.size() == 0) { + for (int i = 0; i < numParticles; i++) + for (int j = 0; j < 3; j++) + forceData[i][j] += outputForces[3*i+j]; + } + else { + for (int i = 0; i < numParticles; i++) + for (int j = 0; j < 3; j++) + forceData[particles[i]][j] += outputForces[3*i+j]; + } + } + return energy; +} diff --git a/platforms/reference/src/ReferenceTorchKernels.h b/platforms/reference/src/ReferenceTorchKernels.h index 4d080ab2..f4c61e50 100644 --- a/platforms/reference/src/ReferenceTorchKernels.h +++ b/platforms/reference/src/ReferenceTorchKernels.h @@ -72,6 +72,37 @@ class ReferenceCalcTorchForceKernel : public CalcTorchForceKernel { bool usePeriodic, outputsForces; }; +/** + * This kernel is invoked by PythonTorchForceImpl to calculate the forces acting on the system and the energy of the system. + */ +class ReferenceCalcPythonTorchForceKernel : public CalcPythonTorchForceKernel { +public: + ReferenceCalcPythonTorchForceKernel(std::string name, const OpenMM::Platform& platform) : CalcPythonTorchForceKernel(name, platform) { + } + /** + * Initialize the kernel. + * + * @param context the ContextImpl this kernel will be applied to + * @param force the PythonTorchForce this kernel will be used for + */ + void initialize(const OpenMM::ContextImpl& context, const PythonTorchForce& force); + /** + * Execute the kernel to calculate the forces and/or energy. + * + * @param context the context in which to execute this kernel + * @param includeForces true if forces should be calculated + * @param includeEnergy true if the energy should be calculated + * @return the potential energy due to the force + */ + double execute(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy); +private: + const PythonTorchForceComputation* computation; + std::vector positions; + std::vector particles; + int numParticles; + bool usePeriodic; +}; + } // namespace TorchPlugin #endif /*REFERENCE_TORCH_KERNELS_H_*/ diff --git a/platforms/reference/tests/TestReferencePythonTorchForce.cpp b/platforms/reference/tests/TestReferencePythonTorchForce.cpp new file mode 100644 index 00000000..2d26e104 --- /dev/null +++ b/platforms/reference/tests/TestReferencePythonTorchForce.cpp @@ -0,0 +1,127 @@ +/* -------------------------------------------------------------------------- * + * OpenMM * + * -------------------------------------------------------------------------- * + * This is part of the OpenMM molecular simulation toolkit. * + * See https://openmm.org/development. * + * * + * Portions copyright (c) 2025-2026 Stanford University and the Authors. * + * Authors: Peter Eastman * + * Contributors: * + * * + * 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. * + * -------------------------------------------------------------------------- */ + +#include "PythonTorchForce.h" +#include "openmm/internal/AssertionUtilities.h" +#include "openmm/Context.h" +#include "openmm/NonbondedForce.h" +#include "openmm/Platform.h" +#include "openmm/VerletIntegrator.h" +#include "sfmt/SFMT.h" +#include + +using namespace TorchPlugin; +using namespace OpenMM; +using namespace std; + +extern "C" OPENMM_EXPORT void registerTorchReferenceKernelFactories(); + +void testForce(bool subsetParticles) { + class Computation : public PythonTorchForceComputation { + torch::Tensor compute(const State& state, const torch::Tensor& positions, double& energy) const { + ASSERT_EQUAL(5.0, state.getParameters().at("a")); + ASSERT_EQUAL(10.0, state.getParameters().at("b")); + Vec3 a, b, c; + state.getPeriodicBoxVectors(a, b, c); + ASSERT_EQUAL(Vec3(2, 0, 0), a); + ASSERT_EQUAL(Vec3(0.1, 2, 0), b); + ASSERT_EQUAL(Vec3(0.1, 0.1, 2), c); + energy = 25.0; + return positions*2; + } + }; + int numParticles = 5; + int totalParticles = (subsetParticles ? numParticles+10 : numParticles); + System system; + Vec3 a(2, 0, 0); + Vec3 b(0.1, 2, 0); + Vec3 c(0.1, 0.1, 2); + system.setDefaultPeriodicBoxVectors(a, b, c); + NonbondedForce* nonbonded = new NonbondedForce(); // To trigger reordering + nonbonded->setNonbondedMethod(NonbondedForce::PME); + system.addForce(nonbonded); + vector positions; + OpenMM_SFMT::SFMT sfmt; + init_gen_rand(0, sfmt); + for (int i = 0; i < totalParticles; i++) { + system.addParticle(1.0); + positions.push_back(Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))); + nonbonded->addParticle(0.0, 1.0, 0.0); + } + map params; + params["a"] = 5.0; + params["b"] = 10.0; + vector particles; + if (subsetParticles) + for (int i = 0; i < numParticles; i++) + particles.push_back(i+5); + PythonTorchForce* force = new PythonTorchForce(new Computation(), params, particles); + ASSERT(!force->usesPeriodicBoundaryConditions()); + force->setUsesPeriodicBoundaryConditions(true); + ASSERT(force->usesPeriodicBoundaryConditions()); + system.addForce(force); + VerletIntegrator integrator(0.01); + Platform& platform = Platform::getPlatformByName("Reference"); + Context context(system, integrator, platform); + context.setPositions(positions); + State state = context.getState(State::Energy | State::Forces); + ASSERT_EQUAL_TOL(25.0, state.getPotentialEnergy(), 1e-6); + if (subsetParticles) { + for (int i : particles) + ASSERT_EQUAL_VEC(2*positions[i], state.getForces()[i], 1e-6) + Vec3 zero; + for (int i = 0; i < 5; i++) + ASSERT_EQUAL_VEC(zero, state.getForces()[i], 1e-6); + } + else { + for (int i = 0; i < numParticles; i++) + ASSERT_EQUAL_VEC(2*positions[i], state.getForces()[i], 1e-6) + } + + // Check that force groups are handled correctly. + + ASSERT_EQUAL_TOL(25.0, context.getState(State::Energy, false, 1).getPotentialEnergy(), 1e-6); + ASSERT_EQUAL_TOL(0.0, context.getState(State::Energy, false, 2).getPotentialEnergy(), 1e-6); +} + +void runPlatformTests(); + +int main(int argc, char* argv[]) { + try { + registerTorchReferenceKernelFactories(); + testForce(false); + testForce(true); + } + catch(const exception& e) { + cout << "exception: " << e.what() << endl; + return 1; + } + cout << "Done" << endl; + return 0; +} diff --git a/python/openmmtorch.i b/python/openmmtorch.i index a1529377..44b73fe6 100644 --- a/python/openmmtorch.i +++ b/python/openmmtorch.i @@ -1,4 +1,7 @@ +%newobject TorchPlugin::PythonTorchForce::PythonTorchForce; + %pythonbegin %{ +import torch import sys if sys.platform == 'win32': import os @@ -10,26 +13,32 @@ if sys.platform == 'win32': os.environ['PATH'] = r'%(lib)s;%(lib)s\plugins;%(path)s' % {'lib': openmmtorch_library_path, 'path': _path} os.add_dll_directory(openmmtorch_library_path) - %} %module openmmtorch %include "factory.i" -%import(module="simtk.openmm") "swig/OpenMMSwigHeaders.i" +%import(module="openmm") "swig/OpenMMSwigHeaders.i" %include "swig/typemaps.i" -%include %include +%include +%include %{ #include "TorchForce.h" +#include "PythonTorchForce.h" #include "OpenMM.h" #include "OpenMMAmoeba.h" #include "OpenMMDrude.h" #include "openmm/RPMDIntegrator.h" #include "openmm/RPMDMonteCarloBarostat.h" +#include #include #include + +namespace TorchPlugin { + PythonTorchForce* _createPythonTorchForce(PyObject* computation, const std::map& globalParameters={}, const std::vector& particles={}); +} %} /* @@ -69,6 +78,8 @@ if sys.platform == 'win32': namespace std { %template(property_map) map; + %template(parameter_map) map; + %template(particle_list) vector; } namespace TorchPlugin { @@ -109,4 +120,228 @@ public: } }; +class PythonTorchForce : public OpenMM::Force { +public: + ~PythonTorchForce(); + const std::map& getGlobalParameters() const; + const std::vector& getParticles() const; + void setParticles(const std::vector &particles); + const std::vector& getPickledFunction() const; + virtual bool usesPeriodicBoundaryConditions() const; + void setUsesPeriodicBoundaryConditions(bool periodic); + + /* + * Add methods for casting a Force to a PythonTorchForce. + */ + %extend { + static TorchPlugin::PythonTorchForce& cast(OpenMM::Force& force) { + return dynamic_cast(force); + } + + static bool isinstance(OpenMM::Force& force) { + return (dynamic_cast(&force) != NULL); + } + } +}; + } + +%inline %{ + +namespace TorchPlugin { + /** + * This is the PythonTorchForceComputation that performs the computation for a PythonTorchForce. It invokes the function + * provided by the user, validates the outputs, and converts them to the required format. + */ + class ComputationWrapper : public PythonTorchForceComputation { + public: + ComputationWrapper(PyObject* computation) : computation(computation) { + Py_INCREF(computation); + } + ~ComputationWrapper() { + Py_XDECREF(computation); + } + torch::Tensor compute(const OpenMM::State& state, const torch::Tensor& positions, double& energy) const { + PyGILState_STATE gstate; + gstate = PyGILState_Ensure(); + + // Invoke the function. + + swig_type_info* info = SWIGTYPE_p_OpenMM__State; + PyObject* wrappedState = SWIG_NewPointerObj((void*) &state, info, 0); + PyObject* wrappedPositions = THPVariable_Wrap(positions); + PyObject* result = PyObject_CallFunctionObjArgs(computation, wrappedState, wrappedPositions, NULL); + Py_XDECREF(wrappedState); + Py_XDECREF(wrappedPositions); + if (result == NULL) { + // The function raised an exception. Convert it to an OpenMMException. + +#if PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION < 12 + PyObject *type; + PyObject *exception; + PyObject *traceback; + PyErr_Fetch(&type, &exception, &traceback); +#else + PyObject *exception = PyErr_GetRaisedException(); +#endif + PyObject *message = PyObject_Str(exception); + std::string *ptr; + SWIG_AsPtr_std_string(message, &ptr); + Py_XDECREF(message); + PyGILState_Release(gstate); + throw OpenMM::OpenMMException(*ptr); + } + + // Extract the return values. + + if (!PyTuple_Check(result) || PyTuple_Size(result) != 2) { + PyGILState_Release(gstate); + throw OpenMM::OpenMMException("PythonTorchForce: Expected two return values"); + } + PyObject* pyenergy = PyTuple_GetItem(result, 0); + PyObject* pyforces = PyTuple_GetItem(result, 1); + if (!THPVariable_Check(pyenergy)) { + PyGILState_Release(gstate); + throw OpenMM::OpenMMException("PythonTorchForce: Expected the energy to be a Tensor"); + } + if (!THPVariable_Check(pyforces)) { + PyGILState_Release(gstate); + throw OpenMM::OpenMMException("PythonTorchForce: Expected the forces to be a Tensor"); + } + torch::Tensor forces = THPVariable_Unpack(pyforces); + energy = THPVariable_Unpack(pyenergy).item(); + + // Clean up before returning. + + Py_XDECREF(result); + PyGILState_Release(gstate); + return forces; + } + private: + PyObject* computation; + }; + + /** + * Construct a new PythonTorchForce. + */ + PythonTorchForce* _createPythonTorchForce(PyObject* computation, const std::map& globalParameters, const std::vector& particles) { + PythonTorchForce* force = new PythonTorchForce(new ComputationWrapper(computation), globalParameters, particles); + PyObject* pickle = PyImport_ImportModule("pickle"); + PyObject* dumps = PyUnicode_FromString("dumps"); + PyObject* result = PyObject_CallMethodOneArg(pickle, dumps, computation); + if (result == NULL) { + // It couldn't be pickled. It will still work, but can't be serialized. Clear the error flag. + PyErr_Clear(); + } + else { + char* buffer; + Py_ssize_t len; + if (PyBytes_AsStringAndSize(result, &buffer, &len) == 0) + force->setPickledFunction(buffer, len); + } + return force; + } + + /** + * This is the serialization proxy used to serialize PythonTorchForce objects. + */ + class PythonTorchForceProxy : public OpenMM::SerializationProxy { + public: + PythonTorchForceProxy() : OpenMM::SerializationProxy("PythonTorchForce") { + } + + static std::string hexEncode(const std::vector& input) { + std::stringstream ss; + ss << std::hex << std::setfill('0'); + for (unsigned char i : input) + ss << std::setw(2) << static_cast(i); + return ss.str(); + } + + static std::vector hexDecode(const std::string& input) { + std::vector res; + res.reserve(input.size() / 2); + for (size_t i = 0; i < input.length(); i += 2) { + std::istringstream iss(input.substr(i, 2)); + uint64_t temp; + iss >> std::hex >> temp; + res.push_back(static_cast(temp)); + } + return res; + } + + void serialize(const void* object, OpenMM::SerializationNode& node) const { + node.setIntProperty("version", 1); + const PythonTorchForce& force = *reinterpret_cast(object); + if (force.getPickledFunction().size() == 0) + throw OpenMM::OpenMMException("PythonTorchForceProxy: Could not serialize PythonTorchForce because its function could not be pickled."); + node.setStringProperty("function", hexEncode(force.getPickledFunction())); + node.setIntProperty("forceGroup", force.getForceGroup()); + node.setStringProperty("name", force.getName()); + node.setBoolProperty("usesPeriodic", force.usesPeriodicBoundaryConditions()); + OpenMM::SerializationNode& globalParams = node.createChildNode("GlobalParameters"); + for (auto param : force.getGlobalParameters()) + globalParams.createChildNode("Parameter").setStringProperty("name", param.first).setDoubleProperty("default", param.second); + OpenMM::SerializationNode& particlesNode = node.createChildNode("Particles"); + for (int i : force.getParticles()) + particlesNode.createChildNode("Particle").setIntProperty("index", i); + } + + void* deserialize(const OpenMM::SerializationNode& node) const { + int version = node.getIntProperty("version"); + if (version != 1) + throw OpenMM::OpenMMException("Unsupported version number"); + std::vector pickledFunction = hexDecode(node.getStringProperty("function")); + PyObject* pickle = PyImport_ImportModule("pickle"); + PyObject* loads = PyUnicode_FromString("loads"); + PyObject *pythonBytes = PyBytes_FromStringAndSize(pickledFunction.data(), pickledFunction.size()); + PyObject *function = PyObject_CallMethodOneArg(pickle, loads, pythonBytes); + Py_XDECREF(pythonBytes); + const OpenMM::SerializationNode& paramsNode = node.getChildNode("GlobalParameters"); + std::map params; + for (auto& parameter : paramsNode.getChildren()) + params[parameter.getStringProperty("name")] = parameter.getDoubleProperty("default"); + std::vector particles; + for (auto& particle : node.getChildNode("Particles").getChildren()) + particles.push_back(particle.getIntProperty("index")); + PythonTorchForce* force = _createPythonTorchForce(function, params, particles); + if (node.hasProperty("forceGroup")) + force->setForceGroup(node.getIntProperty("forceGroup", 0)); + force->setName(node.getStringProperty("name", force->getName())); + if (node.hasProperty("usesPeriodic")) + force->setUsesPeriodicBoundaryConditions(node.getBoolProperty("usesPeriodic")); + return force; + } + }; + + /** + * Register the serialization proxy. This function is invoked automatically when the openmm module is imported. + */ + void registerPythonTorchForceProxy() { + OpenMM::SerializationProxy::registerProxy(typeid(PythonTorchForce), new PythonTorchForceProxy()); + } +} + +%} + +%extend TorchPlugin::PythonTorchForce { + %feature("docstring") PythonTorchForce """Create a PythonTorchForce. + +Parameters +---------- +computation : function + A function that performs the computation. It should take two arguments: a State and a + Tensor containing positions. It should return two values: the potential energy and the forces, + both represented as Tensors. +globalParameters : dict + Any global parameters the function depends on. Keys are the parameter names, and the + corresponding values are their default values. +""" + PythonTorchForce(PyObject* computation, const std::map& globalParameters={}, const std::vector& particles={}) { + return TorchPlugin::_createPythonTorchForce(computation, globalParameters, particles); + } +} + +%pythoncode %{ + registerPythonTorchForceProxy() +%} \ No newline at end of file diff --git a/python/setup.py b/python/setup.py index 1a20eb43..1b4b248d 100644 --- a/python/setup.py +++ b/python/setup.py @@ -11,7 +11,7 @@ extra_compile_args = ['-std=c++17'] extra_link_args = [] -libraries = ['OpenMM', 'OpenMMTorch'] +libraries = ['OpenMM', 'OpenMMTorch', 'torch_python'] runtime_library_dirs = [os.path.join(openmm_dir, 'lib'), torch_dir] # For Windows change the compiler flag to /std:c++17 diff --git a/python/tests/TestInteroperability.py b/python/tests/TestInteroperability.py deleted file mode 100644 index 1c67c8a1..00000000 --- a/python/tests/TestInteroperability.py +++ /dev/null @@ -1,71 +0,0 @@ -import openmm as mm -import openmm.unit as unit -import openmmtorch as ot -import platform -import pytest -from tempfile import NamedTemporaryFile -import torch as pt - - -@pytest.mark.skipif(platform.system() == 'Darwin', reason='There is no NNPOps package for MacOS') -@pytest.mark.parametrize('use_cv_force', [True, False]) -@pytest.mark.parametrize('platform', ['Reference', 'CPU', 'CUDA', 'OpenCL']) -def testTorchANI(use_cv_force, platform): - - if pt.cuda.device_count() < 1 and platform == 'CUDA': - pytest.skip('A CUDA device is not available') - - import NNPOps # There is no NNPOps package for MacOS - import torchani - - class Model(pt.nn.Module): - - def __init__(self): - super().__init__() - self.register_buffer('atomic_numbers', pt.tensor([[1, 1]])) - self.model = torchani.models.ANI2x(periodic_table_index=True) - self.model = NNPOps.OptimizedTorchANI(self.model, self.atomic_numbers) - - def forward(self, positions): - positions = positions.float().unsqueeze(0) * 10 # nm --> Ang - return self.model((self.atomic_numbers, positions)).energies[0] * 2625.5 # Hartree --> kJ/mol - - # Create a system - system = mm.System() - for _ in range(2): - system.addParticle(1.0) - positions = pt.tensor([[-5, 0.0, 0.0], [5, 0.0, 0.0]], requires_grad=True) - - with NamedTemporaryFile() as model_file: - - # Save the model - pt.jit.script(Model()).save(model_file.name) - - # Compute reference energy and forces - model = pt.jit.load(model_file) - ref_energy = model(positions) - ref_energy.backward() - ref_forces = positions.grad - - # Create a force - force = ot.TorchForce(model_file.name) - if use_cv_force: - # Wrap TorchForce into CustomCVForce - cv_force = mm.CustomCVForce('force') - cv_force.addCollectiveVariable('force', force) - system.addForce(cv_force) - else: - system.addForce(force) - - # Compute energy and forces - integ = mm.VerletIntegrator(1.0) - platform = mm.Platform.getPlatformByName(platform) - context = mm.Context(system, integ, platform) - context.setPositions(positions.detach().numpy()) - state = context.getState(getEnergy=True, getForces=True) - energy = state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole) - forces = state.getForces(asNumpy=True).value_in_unit(unit.kilojoules_per_mole/unit.nanometers) - - # Check energy and forces - assert pt.allclose(ref_energy, pt.tensor(energy, dtype=ref_energy.dtype)) - assert pt.allclose(ref_forces, pt.tensor(forces, dtype=ref_forces.dtype)) \ No newline at end of file diff --git a/python/tests/TestPythonTorchForce.py b/python/tests/TestPythonTorchForce.py new file mode 100644 index 00000000..4d75e576 --- /dev/null +++ b/python/tests/TestPythonTorchForce.py @@ -0,0 +1,248 @@ +import unittest +from openmm import * +from openmm.unit import * +from openmmtorch import PythonTorchForce +import numpy as np +import torch +import copy + +def compute(state, pos): + """This is a computation function used by the test cases.""" + k = state.getParameters()['k'] + energy = k*torch.sum(pos*pos) + force = -0.5*k*pos + return energy, force + +class TestPythonTorchForce(unittest.TestCase): + """Test the PythonTorchForce class""" + + def testComputeForce(self): + """Test using PythonTorchForce to compute forces.""" + system = System() + for i in range(5): + system.addParticle(1.0) + force = PythonTorchForce(compute, {'k':2.5}) + system.addForce(force) + positions = np.random.rand(5, 3) + for i in range(Platform.getNumPlatforms()): + integrator = VerletIntegrator(0.001) + try: + context = Context(system, integrator, Platform.getPlatform(i)) + except OpenMMException: + if i == 0: + raise + else: + # This happens on CI when no GPU is available. + continue + context.setPositions(positions) + state = context.getState(energy=True, forces=True) + self.assertAlmostEqual(2.5*np.sum(positions*positions), state.getPotentialEnergy().value_in_unit(kilojoules_per_mole), places=5) + self.assertTrue(np.allclose(-1.25*positions, state.getForces(asNumpy=True).value_in_unit(kilojoules_per_mole/nanometer))) + + def testParticleSubset(self): + """Test a PythonTorchForce appled to a subset of particles.""" + system = System() + for i in range(10): + system.addParticle(1.0) + force = PythonTorchForce(compute, {'k':2.5}) + particles = [1,3,5,7,9] + force.setParticles(particles) + system.addForce(force) + positions = np.random.rand(10, 3) + for i in range(Platform.getNumPlatforms()): + integrator = VerletIntegrator(0.001) + try: + context = Context(system, integrator, Platform.getPlatform(i)) + except OpenMMException: + if i == 0: + raise + else: + # This happens on CI when no GPU is available. + continue + context.setPositions(positions) + state = context.getState(energy=True, forces=True) + filtered = np.zeros(positions.shape) + filtered[particles] = positions[particles] + self.assertAlmostEqual(2.5*np.sum(filtered*filtered), state.getPotentialEnergy().value_in_unit(kilojoules_per_mole), places=5) + self.assertTrue(np.allclose(-1.25*filtered, state.getForces(asNumpy=True).value_in_unit(kilojoules_per_mole/nanometer))) + + def testPeriodic(self): + """Test using PythonTorchForce with periodic boundary conditions.""" + def compute2(state, positions): + vectors = state.getPeriodicBoxVectors().value_in_unit(nanometer) + boxsize = torch.tensor(vectors, dtype=positions.dtype, device=positions.device).diag() + positions = positions - torch.floor(positions/boxsize)*boxsize + energy = torch.sum(positions**2) + force = -0.5*positions + return energy, force + + system = System() + system.setDefaultPeriodicBoxVectors(Vec3(2, 0, 0), Vec3(0, 2, 0), Vec3(0, 0, 2)) + for i in range(10): + system.addParticle(1.0) + force = PythonTorchForce(compute2) + system.addForce(force) + self.assertFalse(system.usesPeriodicBoundaryConditions()) + force.setUsesPeriodicBoundaryConditions(True) + self.assertTrue(system.usesPeriodicBoundaryConditions()) + positions = 10*np.random.rand(10, 3)-3 + for i in range(Platform.getNumPlatforms()): + integrator = VerletIntegrator(0.001) + try: + context = Context(system, integrator, Platform.getPlatform(i)) + except OpenMMException: + if i == 0: + raise + else: + # This happens on CI when no GPU is available. + continue + context.setPositions(positions) + state = context.getState(energy=True, forces=True, positions=True, enforcePeriodicBox=True) + periodicPositions = state.getPositions(asNumpy=True).value_in_unit(nanometer) + self.assertAlmostEqual(np.sum(periodicPositions**2), state.getPotentialEnergy().value_in_unit(kilojoules_per_mole), places=5) + self.assertTrue(np.allclose(-0.5*periodicPositions, state.getForces(asNumpy=True).value_in_unit(kilojoules_per_mole/nanometer))) + + def testExceptions(self): + """Test that PythonTorchForce handles exceptions correctly.""" + def compute2(state, pos): + raise ValueError('This should fail') + + system = System() + system.addParticle(1.0) + force = PythonTorchForce(compute2) + system.addForce(force) + positions = np.random.rand(1, 3) + for i in range(Platform.getNumPlatforms()): + integrator = VerletIntegrator(0.001) + try: + context = Context(system, integrator, Platform.getPlatform(i)) + except OpenMMException: + if i == 0: + raise + else: + # This happens on CI when no GPU is available. + continue + context.setPositions(positions) + with self.assertRaises(OpenMMException) as cm: + context.getState(energy=True) + self.assertEqual('This should fail', str(cm.exception)) + + def testSerialize(self): + """Test that PythonTorchForce can be serialized.""" + force1 = PythonTorchForce(compute, {'k':2.5}) + force1.setUsesPeriodicBoundaryConditions(True) + force1.setParticles([1,3,5]) + force1.setName("custom name") + + # Make a copy by serializing and the deserializing it. + + copied = copy.copy(force1) + force2 = PythonTorchForce.cast(copied) + + # They should be identical. + + self.assertEqual(XmlSerializer.serialize(force1), XmlSerializer.serialize(force2)) + self.assertEqual(dict(force2.getGlobalParameters()), {'k':2.5}) + self.assertEqual(force1.getParticles(), force2.getParticles()) + self.assertTrue(force2.usesPeriodicBoundaryConditions()) + self.assertEqual(force1.getName(), force2.getName()) + + # A locally defined function cannot be pickled. We should not be able to serialize a force + # that uses it. + + def compute2(state): + return 1.0, np.zeros(len(state.getPositions()), 3) + + force3 = PythonTorchForce(compute2) + with self.assertRaises(OpenMMException): + XmlSerializer.serialize(force3) + + def testMinimization(self): + """Test that PythonTorchForce works correctly with the minimizer.""" + system = System() + for i in range(5): + system.addParticle(1.0) + force = PythonTorchForce(compute, {'k':2.5}) + system.addForce(force) + positions = np.random.rand(5, 3) + for i in range(Platform.getNumPlatforms()): + integrator = VerletIntegrator(0.001) + try: + context = Context(system, integrator, Platform.getPlatform(i)) + except OpenMMException: + if i == 0: + raise + else: + # This happens on CI when no GPU is available. + continue + context.setPositions(positions) + + # The PythonTorchForce and the MinimizationReporter both involve calling back into Python code, + # possibly from different threads. Make sure it doesn't cause any problems. + + class Reporter(MinimizationReporter): + count = 0 + def report(self, iteration, x, grad, args): + self.count += 1 + return False + + reporter = Reporter() + LocalEnergyMinimizer.minimize(context, tolerance=1e-3, reporter=reporter) + self.assertTrue(reporter.count > 0) + state = context.getState(energy=True, positions=True) + self.assertAlmostEqual(0.0, state.getPotentialEnergy().value_in_unit(kilojoules_per_mole)) + + def testMemory(self): + """Test for memory leaks in the Python/C++ interface.""" + try: + import resource + except: + # The resource module is not available on Windows. + return + system = System() + for i in range(1000): + system.addParticle(1.0) + force = PythonTorchForce(compute, {'k':2.5}) + system.addForce(force) + positions = np.random.rand(1000, 3) + integrator = VerletIntegrator(0.001) + context = Context(system, integrator, Platform.getPlatform('Reference')) + context.setPositions(positions) + integrator.step(5000) + memory1 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + integrator.step(5000) + memory2 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + self.assertTrue(memory2 < 1.05*memory1) + + def testDtypes(self): + """Test returning forces with different types.""" + for dtype in [torch.float32, torch.float64]: + def compute2(state, pos): + return torch.tensor(1.2, dtype=dtype), torch.tensor([[1,2,3],[4,5,6]], dtype=dtype) + + system = System() + system.addParticle(1.0) + system.addParticle(1.0) + force = PythonTorchForce(compute2) + system.addForce(force) + positions = np.random.rand(2, 3) + for i in range(Platform.getNumPlatforms()): + integrator = VerletIntegrator(0.001) + try: + context = Context(system, integrator, Platform.getPlatform(i)) + except OpenMMException: + if i == 0: + raise + else: + # This happens on CI when no GPU is available. + continue + context.setPositions(positions) + state = context.getState(forces=True, energy=True) + forces = state.getForces().value_in_unit(kilojoules_per_mole/nanometer) + self.assertEqual(Vec3(1,2,3), forces[0]) + self.assertEqual(Vec3(4,5,6), forces[1]) + energy = state.getPotentialEnergy().value_in_unit(kilojoules_per_mole) + self.assertAlmostEqual(1.2, energy) + +if __name__ == '__main__': + unittest.main() diff --git a/python/tests/TestTorchForce.py b/python/tests/TestTorchForce.py index 6eefed6e..56c973ef 100644 --- a/python/tests/TestTorchForce.py +++ b/python/tests/TestTorchForce.py @@ -6,7 +6,7 @@ import torch as pt from tempfile import NamedTemporaryFile -@pytest.mark.parametrize('model_file,', +@pytest.mark.parametrize('model_file', ['../../tests/central.pt', '../../tests/forces.pt']) def testConstructors(model_file):