From 7054e1b679aadef3aca426c671d66b542342d1ed Mon Sep 17 00:00:00 2001 From: peastman Date: Thu, 23 Apr 2026 11:51:24 -0700 Subject: [PATCH 01/29] Beginning of PythonTorchForce --- .github/workflows/CI.yml | 4 +- openmmapi/include/PythonTorchForce.h | 200 ++++++++++++++++++ openmmapi/include/TorchKernels.h | 31 ++- .../include/internal/PythonTorchForceImpl.h | 73 +++++++ openmmapi/src/PythonTorchForce.cpp | 77 +++++++ openmmapi/src/PythonTorchForceImpl.cpp | 66 ++++++ .../src/ReferenceTorchKernelFactory.cpp | 3 + .../reference/src/ReferenceTorchKernels.cpp | 50 +++++ .../reference/src/ReferenceTorchKernels.h | 31 +++ .../tests/TestReferencePythonTorchForce.cpp | 127 +++++++++++ python/openmmtorch.i | 2 +- 11 files changed, 660 insertions(+), 4 deletions(-) create mode 100644 openmmapi/include/PythonTorchForce.h create mode 100644 openmmapi/include/internal/PythonTorchForceImpl.h create mode 100644 openmmapi/src/PythonTorchForce.cpp create mode 100644 openmmapi/src/PythonTorchForceImpl.cpp create mode 100644 platforms/reference/tests/TestReferencePythonTorchForce.cpp diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c8b3b64f..fcc6afb7 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -38,8 +38,8 @@ jobs: python-version: "3.11" pytorch-version: "2.4.*" - - 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: "" diff --git a/openmmapi/include/PythonTorchForce.h b/openmmapi/include/PythonTorchForce.h new file mode 100644 index 00000000..dfa3fe1c --- /dev/null +++ b/openmmapi/include/PythonTorchForce.h @@ -0,0 +1,200 @@ +#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. To use it, + * define a Python function that takes a State object as its only argument. The State contains + * particle positions and global parameters. Based on it, the function should compute the + * potential energy and forces, returning them as its two return values. The forces should be + * represented as a NumPy array of shape (# particles, 3). For example, + * + * \verbatim embed:rst:leading-asterisk + * .. code-block:: python + * + * def compute(state): + * pos = state.getPositions(asNumpy=True).value_in_unit(nanometer) + * k = state.getParameters()['k'] + * energy = k*np.sum(pos*pos) + * force = -0.5*k*pos + * return energy*kilojoules_per_mole, force*kilojoules_per_mole/nanometer + * + * \endverbatim + * + * Attaching units to the return values is optional. If units are omitted, the values are assumed + * to be in the default units (energy in kJ/mol, forces in kJ/mol/nm). + * + * Now create a Python force, 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. + * state.getPositions() will return a smaller array containing only the positions of those + * particles, and the array of 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. + * + * Compared to other types of forces, computing a force with Python code is slow and has high + * overhead. When possible, using a different force class is usually preferred. For example, + * the Python force shown in the example code above (a harmonic force attracting every particle + * to the origin) could be implemented just as easily with a CustomExternalForce, and would + * execute much faster if done that way. + */ +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/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..9293020e 100644 --- a/python/openmmtorch.i +++ b/python/openmmtorch.i @@ -16,7 +16,7 @@ if sys.platform == 'win32': %module openmmtorch %include "factory.i" -%import(module="simtk.openmm") "swig/OpenMMSwigHeaders.i" +%import(module="openmm") "swig/OpenMMSwigHeaders.i" %include "swig/typemaps.i" %include %include From 781e643cd730f0964a8d3f0d9e55212cb8f966f3 Mon Sep 17 00:00:00 2001 From: peastman Date: Sat, 25 Apr 2026 16:04:29 -0700 Subject: [PATCH 02/29] Python API --- python/openmmtorch.i | 241 ++++++++++++++++++++++++++- python/tests/TestPythonTorchForce.py | 202 ++++++++++++++++++++++ 2 files changed, 441 insertions(+), 2 deletions(-) create mode 100644 python/tests/TestPythonTorchForce.py diff --git a/python/openmmtorch.i b/python/openmmtorch.i index 9293020e..482d461f 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,7 +13,6 @@ 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 @@ -18,18 +20,25 @@ if sys.platform == 'win32': %include "factory.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,230 @@ 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); + 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(wrappedState); + Py_XDECREF(wrappedPositions); + Py_XDECREF(result); + Py_XDECREF(pyenergy); + Py_XDECREF(pyforces); + 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.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)); + 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/tests/TestPythonTorchForce.py b/python/tests/TestPythonTorchForce.py new file mode 100644 index 00000000..15e896ba --- /dev/null +++ b/python/tests/TestPythonTorchForce.py @@ -0,0 +1,202 @@ +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 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]) + + # 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()) + + # 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) + integrator = VerletIntegrator(0.001) + context = Context(system, integrator, Platform.getPlatform('Reference')) + 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() From ccef319e1ea05fd1697238cefa188ac06d3eeefb Mon Sep 17 00:00:00 2001 From: peastman Date: Mon, 27 Apr 2026 13:42:06 -0700 Subject: [PATCH 03/29] OpenCL platform --- platforms/common/src/CommonTorchKernels.cpp | 146 +++++++++++++++++- platforms/common/src/CommonTorchKernels.h | 50 +++++- .../opencl/src/OpenCLTorchKernelFactory.cpp | 3 + .../tests/TestOpenCLPythonTorchForce.cpp | 127 +++++++++++++++ python/openmmtorch.i | 6 +- 5 files changed, 326 insertions(+), 6 deletions(-) create mode 100644 platforms/opencl/tests/TestOpenCLPythonTorchForce.cpp diff --git a/platforms/common/src/CommonTorchKernels.cpp b/platforms/common/src/CommonTorchKernels.cpp index 1c9186ce..034d6e21 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,8 @@ #include "CommonTorchKernels.h" #include "CommonTorchKernelSources.h" +#include "openmm/common/CommonKernelSources.h" +#include "openmm/common/ContextSelector.h" #include "openmm/internal/ContextImpl.h" #include @@ -147,3 +149,145 @@ 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(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(CommonKernelSources::pythonForce, defines); + if (particles.size() > 0) { + particlesArray.initialize(cc, numParticles, "particles"); + reorderedParticles.initialize(cc, numParticles, "reorderedParticles"); + particlesArray.upload(particles); + reorderedParticles.upload(particles); + cc.addReorderListener(new ReorderListener(*this)); + copyPositionsKernel = program->createKernel("copyPositions"); + copyPositionsKernel->addArg(cc.getPosq()); + copyPositionsKernel->addArg(positionsArray); + copyPositionsKernel->addArg(reorderedParticles); + copyPositionsKernel->addArg(numParticles); + addForcesKernel = program->createKernel("addForcesSubset"); + addForcesKernel->addArg(forcesArray); + addForcesKernel->addArg(cc.getLongForceBuffer()); + addForcesKernel->addArg(cc.getAtomIndexArray()); + addForcesKernel->addArg(reorderedParticles); + addForcesKernel->addArg(numParticles); + } + else { + addForcesKernel = program->createKernel("addForcesAll"); + addForcesKernel->addArg(forcesArray); + addForcesKernel->addArg(cc.getLongForceBuffer()); + addForcesKernel->addArg(cc.getAtomIndexArray()); + } +} + +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() { + // If the NonbondedUtilities uses periodic boundary conditions, the positions might have been + // wrapped to the periodic box. If this force also applies periodic boundary conditions, that's + // alright. Otherwise, we need to move them back. + + bool fixPeriodic = usePeriodic || !cc.getNonbondedUtilities().getUsePeriodic(); + if (particles.size() == 0) { + // The force applies to the whole system, so we can just use the standard getPositions(). + + contextImpl.getPositions(positionsVec, fixPeriodic); + } + else { + // Retrieve positions for the subset of particles the force is applied to. + + ContextSelector selector(cc); + copyPositionsKernel->execute(numParticles); + if (cc.getUseDoublePrecision()) { + vector pos(3*numParticles); + positionsArray.download(pos); + for (int i = 0; i < numParticles; i++) + positionsVec[i] = Vec3(pos[3*i], pos[3*i+1], pos[3*i+2]); + } + else { + vector pos(3*numParticles); + positionsArray.download(pos); + for (int i = 0; i < numParticles; i++) + positionsVec[i] = Vec3((double) pos[3*i], (double) pos[3*i+1], (double) pos[3*i+2]); + } + if (fixPeriodic) { + Vec3 boxVectors[3]; + cc.getPeriodicBoxVectors(boxVectors[0], boxVectors[1], boxVectors[2]); + for (int i = 0; i < numParticles; ++i) { + mm_int4 offset = cc.getPosCellOffsets()[particles[i]]; + positionsVec[i] -= boxVectors[0]*offset.x-boxVectors[1]*offset.y-boxVectors[2]*offset.z; + } + } + } + return torch::from_blob(positionsVec.data(), {numParticles, 3}, torch::TensorOptions().dtype(torch::kFloat64).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..1fd12d1c 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); +private: + 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/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/python/openmmtorch.i b/python/openmmtorch.i index 482d461f..d17de96b 100644 --- a/python/openmmtorch.i +++ b/python/openmmtorch.i @@ -146,10 +146,8 @@ public: } - %inline %{ - namespace TorchPlugin { /** * This is the PythonTorchForceComputation that performs the computation for a PythonTorchForce. It invokes the function @@ -173,6 +171,8 @@ namespace TorchPlugin { 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. @@ -213,8 +213,6 @@ namespace TorchPlugin { // Clean up before returning. - Py_XDECREF(wrappedState); - Py_XDECREF(wrappedPositions); Py_XDECREF(result); Py_XDECREF(pyenergy); Py_XDECREF(pyforces); From 69d8e7e8d399f4b40b32fcf8215160c5451cb8c7 Mon Sep 17 00:00:00 2001 From: peastman Date: Mon, 27 Apr 2026 16:46:43 -0700 Subject: [PATCH 04/29] Code simplification --- platforms/common/src/CommonTorchKernels.cpp | 56 ++++----------------- platforms/common/src/CommonTorchKernels.h | 2 +- 2 files changed, 12 insertions(+), 46 deletions(-) diff --git a/platforms/common/src/CommonTorchKernels.cpp b/platforms/common/src/CommonTorchKernels.cpp index 034d6e21..e2b7db43 100644 --- a/platforms/common/src/CommonTorchKernels.cpp +++ b/platforms/common/src/CommonTorchKernels.cpp @@ -168,7 +168,7 @@ void CommonCalcPythonTorchForceKernel::initialize(const ContextImpl& context, co numParticles = particles.size(); if (numParticles == 0) numParticles = context.getSystem().getNumParticles(); - positionsVec.resize(numParticles); + 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"); @@ -177,16 +177,9 @@ void CommonCalcPythonTorchForceKernel::initialize(const ContextImpl& context, co defines["PADDED_NUM_ATOMS"] = cc.intToString(cc.getPaddedNumAtoms()); ComputeProgram program = cc.compileProgram(CommonKernelSources::pythonForce, defines); if (particles.size() > 0) { - particlesArray.initialize(cc, numParticles, "particles"); reorderedParticles.initialize(cc, numParticles, "reorderedParticles"); - particlesArray.upload(particles); reorderedParticles.upload(particles); cc.addReorderListener(new ReorderListener(*this)); - copyPositionsKernel = program->createKernel("copyPositions"); - copyPositionsKernel->addArg(cc.getPosq()); - copyPositionsKernel->addArg(positionsArray); - copyPositionsKernel->addArg(reorderedParticles); - copyPositionsKernel->addArg(numParticles); addForcesKernel = program->createKernel("addForcesSubset"); addForcesKernel->addArg(forcesArray); addForcesKernel->addArg(cc.getLongForceBuffer()); @@ -200,6 +193,11 @@ void CommonCalcPythonTorchForceKernel::initialize(const ContextImpl& context, co addForcesKernel->addArg(cc.getLongForceBuffer()); addForcesKernel->addArg(cc.getAtomIndexArray()); } + copyPositionsKernel = program->createKernel("copyPositions"); + 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) { @@ -221,43 +219,11 @@ double CommonCalcPythonTorchForceKernel::execute(ContextImpl& context, bool incl } torch::Tensor CommonCalcPythonTorchForceKernel::getPositions() { - // If the NonbondedUtilities uses periodic boundary conditions, the positions might have been - // wrapped to the periodic box. If this force also applies periodic boundary conditions, that's - // alright. Otherwise, we need to move them back. - - bool fixPeriodic = usePeriodic || !cc.getNonbondedUtilities().getUsePeriodic(); - if (particles.size() == 0) { - // The force applies to the whole system, so we can just use the standard getPositions(). - - contextImpl.getPositions(positionsVec, fixPeriodic); - } - else { - // Retrieve positions for the subset of particles the force is applied to. - - ContextSelector selector(cc); - copyPositionsKernel->execute(numParticles); - if (cc.getUseDoublePrecision()) { - vector pos(3*numParticles); - positionsArray.download(pos); - for (int i = 0; i < numParticles; i++) - positionsVec[i] = Vec3(pos[3*i], pos[3*i+1], pos[3*i+2]); - } - else { - vector pos(3*numParticles); - positionsArray.download(pos); - for (int i = 0; i < numParticles; i++) - positionsVec[i] = Vec3((double) pos[3*i], (double) pos[3*i+1], (double) pos[3*i+2]); - } - if (fixPeriodic) { - Vec3 boxVectors[3]; - cc.getPeriodicBoxVectors(boxVectors[0], boxVectors[1], boxVectors[2]); - for (int i = 0; i < numParticles; ++i) { - mm_int4 offset = cc.getPosCellOffsets()[particles[i]]; - positionsVec[i] -= boxVectors[0]*offset.x-boxVectors[1]*offset.y-boxVectors[2]*offset.z; - } - } - } - return torch::from_blob(positionsVec.data(), {numParticles, 3}, torch::TensorOptions().dtype(torch::kFloat64).requires_grad(true)); + 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() { diff --git a/platforms/common/src/CommonTorchKernels.h b/platforms/common/src/CommonTorchKernels.h index 1fd12d1c..d4354d74 100644 --- a/platforms/common/src/CommonTorchKernels.h +++ b/platforms/common/src/CommonTorchKernels.h @@ -117,7 +117,7 @@ class CommonCalcPythonTorchForceKernel : public CalcPythonTorchForceKernel { const PythonTorchForceComputation* computation; OpenMM::ComputeArray positionsArray, forcesArray, particlesArray, reorderedParticles; OpenMM::ComputeKernel copyPositionsKernel, addForcesKernel; - std::vector positionsVec; + std::vector positionsVec; std::vector particles; int numParticles; double energy; From d505b2f1faa06702ead5d4ded5549792d4ad48ba Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Tue, 28 Apr 2026 18:04:24 -0700 Subject: [PATCH 05/29] Start of CUDA platform --- CMakeLists.txt | 8 +- platforms/common/src/CommonTorchKernels.cpp | 6 +- platforms/common/src/CommonTorchKernels.h | 2 +- .../common/src/kernels/pythonTorchForce.cc | 36 +++++ platforms/cuda/CMakeLists.txt | 10 +- platforms/cuda/src/CudaTorchKernelFactory.cpp | 3 + platforms/cuda/src/CudaTorchKernels.h | 40 +++++- .../cuda/tests/TestCudaPythonTorchForce.cpp | 127 ++++++++++++++++++ python/setup.py | 2 +- 9 files changed, 221 insertions(+), 13 deletions(-) create mode 100644 platforms/common/src/kernels/pythonTorchForce.cc create mode 100644 platforms/cuda/tests/TestCudaPythonTorchForce.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 75095427..11666d20 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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/platforms/common/src/CommonTorchKernels.cpp b/platforms/common/src/CommonTorchKernels.cpp index e2b7db43..e04071c7 100644 --- a/platforms/common/src/CommonTorchKernels.cpp +++ b/platforms/common/src/CommonTorchKernels.cpp @@ -31,7 +31,6 @@ #include "CommonTorchKernels.h" #include "CommonTorchKernelSources.h" -#include "openmm/common/CommonKernelSources.h" #include "openmm/common/ContextSelector.h" #include "openmm/internal/ContextImpl.h" #include @@ -175,7 +174,7 @@ void CommonCalcPythonTorchForceKernel::initialize(const ContextImpl& context, co map defines; defines["NUM_ATOMS"] = cc.intToString(numParticles); defines["PADDED_NUM_ATOMS"] = cc.intToString(cc.getPaddedNumAtoms()); - ComputeProgram program = cc.compileProgram(CommonKernelSources::pythonForce, defines); + ComputeProgram program = cc.compileProgram(CommonTorchKernelSources::pythonTorchForce, defines); if (particles.size() > 0) { reorderedParticles.initialize(cc, numParticles, "reorderedParticles"); reorderedParticles.upload(particles); @@ -186,14 +185,15 @@ void CommonCalcPythonTorchForceKernel::initialize(const ContextImpl& context, co 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 = program->createKernel("copyPositions"); copyPositionsKernel->addArg(cc.getPosq()); copyPositionsKernel->addArg(positionsArray); copyPositionsKernel->addArg(particles.size() > 0 ? reorderedParticles : cc.getAtomIndexArray()); diff --git a/platforms/common/src/CommonTorchKernels.h b/platforms/common/src/CommonTorchKernels.h index d4354d74..227103a7 100644 --- a/platforms/common/src/CommonTorchKernels.h +++ b/platforms/common/src/CommonTorchKernels.h @@ -109,7 +109,7 @@ class CommonCalcPythonTorchForceKernel : public CalcPythonTorchForceKernel { * Add in the forces. Subclasses can override this to do it more efficiently. */ virtual void addForces(torch::Tensor forceTensor); -private: +protected: class ReorderListener; void sortParticles(); OpenMM::ContextImpl& contextImpl; 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.h b/platforms/cuda/src/CudaTorchKernels.h index 1f7b4c97..793f3bdb 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,43 @@ 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) : + CommonCalcPythonTorchForceKernel(name, platform, contextImpl, 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); +}; + + } // 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/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 From 4cd399784bf96290286ce55ee9a51b3b41330043 Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Wed, 29 Apr 2026 13:32:26 -0700 Subject: [PATCH 06/29] Avoid copying data between host and device --- platforms/cuda/src/CudaTorchKernels.cpp | 30 ++++++++++++++++ platforms/cuda/src/CudaTorchKernels.h | 48 +++++++++++-------------- 2 files changed, 50 insertions(+), 28 deletions(-) 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 793f3bdb..3f0ac8c7 100644 --- a/platforms/cuda/src/CudaTorchKernels.h +++ b/platforms/cuda/src/CudaTorchKernels.h @@ -92,34 +92,26 @@ class CudaCalcTorchForceKernel : public CalcTorchForceKernel { */ class CudaCalcPythonTorchForceKernel : public CommonCalcPythonTorchForceKernel { public: - CudaCalcPythonTorchForceKernel(std::string name, const OpenMM::Platform& platform, OpenMM::ContextImpl& contextImpl, OpenMM::ComputeContext& cc) : - CommonCalcPythonTorchForceKernel(name, platform, contextImpl, 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); + 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; }; From 67e46be0ee9d8cf9327b09f2109555989fb49bf8 Mon Sep 17 00:00:00 2001 From: peastman Date: Thu, 30 Apr 2026 13:57:35 -0700 Subject: [PATCH 07/29] Documentation --- README.md | 364 ++++++--------------------- openmmapi/include/PythonTorchForce.h | 45 ++-- python/tests/TestPythonTorchForce.py | 36 +++ 3 files changed, 132 insertions(+), 313 deletions(-) diff --git a/README.md b/README.md index 3456c23f..1fe127a1 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 Tenor 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/openmmapi/include/PythonTorchForce.h b/openmmapi/include/PythonTorchForce.h index dfa3fe1c..9307ca16 100644 --- a/openmmapi/include/PythonTorchForce.h +++ b/openmmapi/include/PythonTorchForce.h @@ -60,28 +60,28 @@ class OPENMM_EXPORT_NN PythonTorchForceComputation { }; /** - * This class provides a mechanism for computing forces and energy with Python code. To use it, - * define a Python function that takes a State object as its only argument. The State contains - * particle positions and global parameters. Based on it, the function should compute the - * potential energy and forces, returning them as its two return values. The forces should be - * represented as a NumPy array of shape (# particles, 3). For example, + * 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 Tenor 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 = state.getPositions(asNumpy=True).value_in_unit(nanometer) + * def compute(state, pos): * k = state.getParameters()['k'] - * energy = k*np.sum(pos*pos) + * energy = k*torch.sum(pos*pos) * force = -0.5*k*pos - * return energy*kilojoules_per_mole, force*kilojoules_per_mole/nanometer - * + * return energy, force + * * \endverbatim * - * Attaching units to the return values is optional. If units are omitted, the values are assumed - * to be in the default units (energy in kJ/mol, forces in kJ/mol/nm). - * - * Now create a Python force, passing the function to the constructor. If you want the force + * 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. * @@ -104,23 +104,16 @@ class OPENMM_EXPORT_NN PythonTorchForceComputation { * * 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. - * state.getPositions() will return a smaller array containing only the positions of those - * particles, and the array of 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. + * 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. - * - * Compared to other types of forces, computing a force with Python code is slow and has high - * overhead. When possible, using a different force class is usually preferred. For example, - * the Python force shown in the example code above (a harmonic force attracting every particle - * to the origin) could be implemented just as easily with a CustomExternalForce, and would - * execute much faster if done that way. */ class OPENMM_EXPORT PythonTorchForce : public OpenMM::Force { public: diff --git a/python/tests/TestPythonTorchForce.py b/python/tests/TestPythonTorchForce.py index 15e896ba..94ffd1f3 100644 --- a/python/tests/TestPythonTorchForce.py +++ b/python/tests/TestPythonTorchForce.py @@ -66,6 +66,42 @@ def testParticleSubset(self): 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): From 984e3af9e8d28d294a93a5a6214c1377bbaa6aab Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Thu, 18 Jun 2026 13:41:17 -0700 Subject: [PATCH 08/29] Use common implementation for HIP --- platforms/hip/src/HipTorchKernelFactory.cpp | 5 +- .../hip/tests/TestHipPythonTorchForce.cpp | 127 ++++++++++++++++++ 2 files changed, 131 insertions(+), 1 deletion(-) create mode 100644 platforms/hip/tests/TestHipPythonTorchForce.cpp 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; +} From f5766215aa019c41e7ac7ffa6305181885246048 Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Thu, 18 Jun 2026 14:14:11 -0700 Subject: [PATCH 09/29] Fixed a memory error --- python/openmmtorch.i | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/openmmtorch.i b/python/openmmtorch.i index d17de96b..c8365cab 100644 --- a/python/openmmtorch.i +++ b/python/openmmtorch.i @@ -214,8 +214,6 @@ namespace TorchPlugin { // Clean up before returning. Py_XDECREF(result); - Py_XDECREF(pyenergy); - Py_XDECREF(pyforces); PyGILState_Release(gstate); return forces; } From 3c2f3cfb1150bb81c732442ff43d5624abf0d7c9 Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Thu, 18 Jun 2026 16:53:49 -0700 Subject: [PATCH 10/29] Use newer version of setup-miniconda --- .github/workflows/CI.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index fcc6afb7..89ec7c0e 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -78,7 +78,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 From a5abced6f2e57da3b780e799e4bc55c27a442e90 Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Thu, 18 Jun 2026 17:23:51 -0700 Subject: [PATCH 11/29] Work with newer version of pytest --- python/tests/TestTorchForce.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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): From 6312f21ab50df33d1fbd7ff26ef4af55eeac828a Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Thu, 18 Jun 2026 17:27:13 -0700 Subject: [PATCH 12/29] Fix name of environment file --- .../conda-envs/{build-macos-13.yml => build-macos-15-intel.yml} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename devtools/conda-envs/{build-macos-13.yml => build-macos-15-intel.yml} (100%) 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 From 933cf21025d444f51b449173d0933eb27a20410b Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Thu, 18 Jun 2026 17:31:38 -0700 Subject: [PATCH 13/29] Update setup-miniconda version in another workflow --- .github/workflows/tutorials.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From 127aece09a337f2275ad4fc8a33b6f4c3cbe050d Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Thu, 18 Jun 2026 19:30:25 -0700 Subject: [PATCH 14/29] Update CI to newer package versions --- .github/workflows/CI.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 89ec7c0e..c5b339ca 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -22,13 +22,13 @@ jobs: fail-fast: false matrix: include: - - name: Linux (CUDA 11.8, Python 3.10, PyTorch 2.1) + - name: Linux (CUDA 13.2, Python 3.14, PyTorch 2.12) os: ubuntu-22.04 - cuda-version: "11.8.0" + cuda-version: "13.2.0" gcc-version: "10.3.*" - nvcc-version: "11.8" - python-version: "3.10" - pytorch-version: "2.1.*" + nvcc-version: "13.2" + python-version: "3.14" + pytorch-version: "2.12.*" - name: MacOS ARM (Python 3.11, PyTorch 2.4) os: macos-latest From 49369db0415ed228d95e27f32885420d5b825dfd Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Thu, 18 Jun 2026 19:33:09 -0700 Subject: [PATCH 15/29] Newer version of cuda-toolkit action --- .github/workflows/CI.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c5b339ca..bdad0198 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -52,7 +52,7 @@ 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 with: cuda: ${{ matrix.cuda-version }} linux-local-args: '["--toolkit", "--override"]' From 603e4a89fd666a7b9b8a82419c6af811fe91f8a6 Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Thu, 18 Jun 2026 19:46:39 -0700 Subject: [PATCH 16/29] Try a different CUDA version --- .github/workflows/CI.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index bdad0198..6da70873 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -22,11 +22,11 @@ jobs: fail-fast: false matrix: include: - - name: Linux (CUDA 13.2, Python 3.14, PyTorch 2.12) + - name: Linux (CUDA 12.8, Python 3.14, PyTorch 2.12) os: ubuntu-22.04 - cuda-version: "13.2.0" + cuda-version: "12.8.0" gcc-version: "10.3.*" - nvcc-version: "13.2" + nvcc-version: "12.8" python-version: "3.14" pytorch-version: "2.12.*" From b5ddd8673c7f2e6661ff474fb71b8bdda123adb9 Mon Sep 17 00:00:00 2001 From: peastman Date: Fri, 19 Jun 2026 08:00:45 -0700 Subject: [PATCH 17/29] Use cuda-version instead of cudatoolkit --- .github/workflows/CI.yml | 4 ++-- devtools/conda-envs/build-ubuntu-22.04.yml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 6da70873..2a4d0b10 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -24,7 +24,7 @@ jobs: include: - name: Linux (CUDA 12.8, Python 3.14, PyTorch 2.12) os: ubuntu-22.04 - cuda-version: "12.8.0" + cuda-version: "12.8" gcc-version: "10.3.*" nvcc-version: "12.8" python-version: "3.14" @@ -54,7 +54,7 @@ jobs: - name: "Install CUDA Toolkit on Linux (if needed)" uses: Jimver/cuda-toolkit@v0.2.35 with: - cuda: ${{ matrix.cuda-version }} + cuda: ${{ matrix.cuda-version }}.0 linux-local-args: '["--toolkit", "--override"]' if: startsWith(matrix.os, 'ubuntu') diff --git a/devtools/conda-envs/build-ubuntu-22.04.yml b/devtools/conda-envs/build-ubuntu-22.04.yml index 1cc25bbc..6ff0897c 100644 --- a/devtools/conda-envs/build-ubuntu-22.04.yml +++ b/devtools/conda-envs/build-ubuntu-22.04.yml @@ -3,7 +3,7 @@ channels: - conda-forge dependencies: - cmake - - cudatoolkit @CUDATOOLKIT_VERSION@ + - cuda-version @CUDATOOLKIT_VERSION@ - gxx_linux-64 @GCC_VERSION@ - make - nnpops From a0cb6d030003e966557e312cbde50fc77ba55c27 Mon Sep 17 00:00:00 2001 From: peastman Date: Fri, 19 Jun 2026 08:18:09 -0700 Subject: [PATCH 18/29] Pytorch 2.12 requires newer CUDA --- .github/workflows/CI.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 2a4d0b10..c3d7d158 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -22,9 +22,9 @@ jobs: fail-fast: false matrix: include: - - name: Linux (CUDA 12.8, Python 3.14, PyTorch 2.12) + - name: Linux (CUDA 13.0, Python 3.14, PyTorch 2.12) os: ubuntu-22.04 - cuda-version: "12.8" + cuda-version: "13.0" gcc-version: "10.3.*" nvcc-version: "12.8" python-version: "3.14" From 3df60fe592aff214b872743988d68f86dbe88849 Mon Sep 17 00:00:00 2001 From: peastman Date: Fri, 19 Jun 2026 08:30:01 -0700 Subject: [PATCH 19/29] Fix nvcc version --- .github/workflows/CI.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c3d7d158..4e0a24d7 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -26,7 +26,7 @@ jobs: os: ubuntu-22.04 cuda-version: "13.0" gcc-version: "10.3.*" - nvcc-version: "12.8" + nvcc-version: "13.0" python-version: "3.14" pytorch-version: "2.12.*" From f66bf8754c5865a47970b43452a678de4f0c7347 Mon Sep 17 00:00:00 2001 From: peastman Date: Fri, 19 Jun 2026 08:39:33 -0700 Subject: [PATCH 20/29] Remove nvcc conda package --- devtools/conda-envs/build-ubuntu-22.04.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/devtools/conda-envs/build-ubuntu-22.04.yml b/devtools/conda-envs/build-ubuntu-22.04.yml index 6ff0897c..6c5c7d75 100644 --- a/devtools/conda-envs/build-ubuntu-22.04.yml +++ b/devtools/conda-envs/build-ubuntu-22.04.yml @@ -7,7 +7,6 @@ dependencies: - gxx_linux-64 @GCC_VERSION@ - make - nnpops - - nvcc_linux-64 @NVCC_VERSION@ - ocl-icd - openmm >=8.1 - pip From 231a0e074472183f4ad35b44d3bfb1518638553f Mon Sep 17 00:00:00 2001 From: peastman Date: Fri, 19 Jun 2026 10:50:08 -0700 Subject: [PATCH 21/29] Debugging --- .github/workflows/CI.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 4e0a24d7..696aaf8b 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -112,7 +112,8 @@ jobs: -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} + -DOPENCL_LIBRARY=${CONDA_PREFIX}/lib/libOpenCL${SHLIB_EXT} \ + -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc - name: "Build" shell: bash -l {0} From 4cdc87d4e1d8bad378f73dc18546dab266abdd18 Mon Sep 17 00:00:00 2001 From: peastman Date: Fri, 19 Jun 2026 11:01:11 -0700 Subject: [PATCH 22/29] Debugging --- .github/workflows/CI.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 696aaf8b..72846400 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -53,6 +53,7 @@ jobs: - name: "Install CUDA Toolkit on Linux (if needed)" uses: Jimver/cuda-toolkit@v0.2.35 + id: cuda-toolkit with: cuda: ${{ matrix.cuda-version }}.0 linux-local-args: '["--toolkit", "--override"]' @@ -113,7 +114,7 @@ jobs: -DNN_BUILD_OPENCL_LIB=ON \ -DOPENCL_INCLUDE_DIR=${CONDA_PREFIX}/include \ -DOPENCL_LIBRARY=${CONDA_PREFIX}/lib/libOpenCL${SHLIB_EXT} \ - -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc + -DCMAKE_CUDA_COMPILER="${{steps.cuda-toolkit.outputs.CUDA_PATH}}/bin/nvcc" - name: "Build" shell: bash -l {0} From 80ae7121f3b56bc0856fe5d75297a457416e7f70 Mon Sep 17 00:00:00 2001 From: peastman Date: Fri, 19 Jun 2026 11:18:15 -0700 Subject: [PATCH 23/29] Debugging --- .github/workflows/CI.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 72846400..dfda2fae 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -114,7 +114,7 @@ jobs: -DNN_BUILD_OPENCL_LIB=ON \ -DOPENCL_INCLUDE_DIR=${CONDA_PREFIX}/include \ -DOPENCL_LIBRARY=${CONDA_PREFIX}/lib/libOpenCL${SHLIB_EXT} \ - -DCMAKE_CUDA_COMPILER="${{steps.cuda-toolkit.outputs.CUDA_PATH}}/bin/nvcc" + -DCUDAToolkit_ROOT="${{steps.cuda-toolkit.outputs.CUDA_PATH}}" - name: "Build" shell: bash -l {0} From 6dd5a2d8318ccd616f54fa837ab958560bb685ab Mon Sep 17 00:00:00 2001 From: peastman Date: Fri, 19 Jun 2026 11:32:30 -0700 Subject: [PATCH 24/29] Debugging --- CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 11666d20..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") From add1ec7053b8b3e7d98123af0a8a613b4a5825c0 Mon Sep 17 00:00:00 2001 From: peastman Date: Fri, 19 Jun 2026 11:52:34 -0700 Subject: [PATCH 25/29] Debugging --- .github/workflows/CI.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index dfda2fae..25db31ac 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -114,7 +114,8 @@ jobs: -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}}" + -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} From 039c71eac787ed79c6ed8be76ccac3265d3edb8f Mon Sep 17 00:00:00 2001 From: peastman Date: Fri, 19 Jun 2026 12:20:21 -0700 Subject: [PATCH 26/29] Removed an obsolete test --- python/tests/TestInteroperability.py | 71 ---------------------------- 1 file changed, 71 deletions(-) delete mode 100644 python/tests/TestInteroperability.py 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 From 44498e9a8fa9f58f2baef633b842378623d660ec Mon Sep 17 00:00:00 2001 From: peastman Date: Tue, 23 Jun 2026 20:52:31 -0700 Subject: [PATCH 27/29] Changes based on suggestions --- .github/workflows/CI.yml | 4 +-- README.md | 2 +- openmmapi/include/PythonTorchForce.h | 2 +- python/openmmtorch.i | 2 ++ python/tests/TestPythonTorchForce.py | 40 +++++++++++++++++----------- 5 files changed, 31 insertions(+), 19 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 25db31ac..00c293c9 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -43,8 +43,8 @@ jobs: cuda-version: "" gcc-version: "" nvcc-version: "" - python-version: "3.11" - pytorch-version: "2.4.*" + python-version: "3.13" + pytorch-version: "2.5.*" steps: diff --git a/README.md b/README.md index 1fe127a1..6abd071f 100644 --- a/README.md +++ b/README.md @@ -82,7 +82,7 @@ def compute(state, positions): 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 Tenor containing the value in kJ/mol. The forces should be a Tensor of shape `(# particles, 3)` containing +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. Now create a PythonTorchForce, passing the function to the constructor. diff --git a/openmmapi/include/PythonTorchForce.h b/openmmapi/include/PythonTorchForce.h index 9307ca16..824f3fd4 100644 --- a/openmmapi/include/PythonTorchForce.h +++ b/openmmapi/include/PythonTorchForce.h @@ -67,7 +67,7 @@ class OPENMM_EXPORT_NN PythonTorchForceComputation { * 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 Tenor containing the value in kJ/mol. The forces should be a Tensor of shape (# particles, 3) containing + * 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 diff --git a/python/openmmtorch.i b/python/openmmtorch.i index c8365cab..44b73fe6 100644 --- a/python/openmmtorch.i +++ b/python/openmmtorch.i @@ -277,6 +277,7 @@ namespace TorchPlugin { 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()) @@ -306,6 +307,7 @@ namespace TorchPlugin { 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; diff --git a/python/tests/TestPythonTorchForce.py b/python/tests/TestPythonTorchForce.py index 94ffd1f3..4d75e576 100644 --- a/python/tests/TestPythonTorchForce.py +++ b/python/tests/TestPythonTorchForce.py @@ -132,6 +132,7 @@ def testSerialize(self): 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. @@ -144,6 +145,7 @@ def testSerialize(self): 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. @@ -163,24 +165,32 @@ def testMinimization(self): force = PythonTorchForce(compute, {'k':2.5}) system.addForce(force) positions = np.random.rand(5, 3) - integrator = VerletIntegrator(0.001) - context = Context(system, integrator, Platform.getPlatform('Reference')) - context.setPositions(positions) + 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. + # 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 + 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)) + 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.""" From 452dac06f43539432d9b3b0c28ce2bbb32b7a79a Mon Sep 17 00:00:00 2001 From: peastman Date: Wed, 24 Jun 2026 13:33:24 -0700 Subject: [PATCH 28/29] Use Intel OpenCL instead of pocl --- .github/workflows/CI.yml | 15 +++++++++++++-- devtools/conda-envs/build-ubuntu-22.04.yml | 1 - devtools/scripts/install_intel_opencl.sh | 8 ++++++++ 3 files changed, 21 insertions(+), 3 deletions(-) create mode 100644 devtools/scripts/install_intel_opencl.sh diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 00c293c9..f9202500 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -29,6 +29,7 @@ jobs: 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,6 +38,7 @@ jobs: nvcc-version: "" python-version: "3.11" pytorch-version: "2.4.*" + opencl: false - name: MacOS Intel (Python 3.13, PyTorch 2.5) os: macos-15-intel @@ -45,6 +47,7 @@ jobs: nvcc-version: "" python-version: "3.13" pytorch-version: "2.5.*" + opencl: false steps: @@ -59,6 +62,10 @@ jobs: linux-local-args: '["--toolkit", "--override"]' if: startsWith(matrix.os, 'ubuntu') + - name: "Install OpenCL on Ubuntu (if needed)" + if: matrix.opencl == true + run: source devtools/ci/gh-actions/scripts/install_intel_opencl.sh + - name: Manage disk space if: startsWith(matrix.os, 'ubuntu') uses: jlumbroso/free-disk-space@main @@ -112,8 +119,6 @@ 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" @@ -135,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" @@ -142,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/devtools/conda-envs/build-ubuntu-22.04.yml b/devtools/conda-envs/build-ubuntu-22.04.yml index 6c5c7d75..3582f3e3 100644 --- a/devtools/conda-envs/build-ubuntu-22.04.yml +++ b/devtools/conda-envs/build-ubuntu-22.04.yml @@ -10,7 +10,6 @@ dependencies: - 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 From b1b52965d93ce67f74fb54058b4b1d9d92044077 Mon Sep 17 00:00:00 2001 From: peastman Date: Wed, 24 Jun 2026 13:41:22 -0700 Subject: [PATCH 29/29] Fixed path --- .github/workflows/CI.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index f9202500..a21e0652 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -64,7 +64,7 @@ jobs: - name: "Install OpenCL on Ubuntu (if needed)" if: matrix.opencl == true - run: source devtools/ci/gh-actions/scripts/install_intel_opencl.sh + run: source devtools/scripts/install_intel_opencl.sh - name: Manage disk space if: startsWith(matrix.os, 'ubuntu')