OpenCL implementation of RMSDForce

platforms/opencl/include/OpenCLKernels.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) 2008-2017 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2018 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -1256,6 +1256,59 @@ private:
     cl::Kernel copyStateKernel, copyForcesKernel, addForcesKernel;
 };
 
+/**
+ * This kernel is invoked by RMSDForce to calculate the forces acting on the system and the energy of the system.
+ */
+class OpenCLCalcRMSDForceKernel : public CalcRMSDForceKernel {
+public:
+    OpenCLCalcRMSDForceKernel(std::string name, const Platform& platform, OpenCLContext& cl) : CalcRMSDForceKernel(name, platform),
+            cl(cl), referencePos(NULL), particles(NULL), buffer(NULL) {
+    }
+    ~OpenCLCalcRMSDForceKernel();
+    /**
+     * Initialize the kernel.
+     *
+     * @param system     the System this kernel will be applied to
+     * @param force      the RMSDForce this kernel will be used for
+     */
+    void initialize(const System& system, const RMSDForce& force);
+    /**
+     * Record the reference positions and particle indices.
+     */
+    void recordParameters(const RMSDForce& 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(ContextImpl& context, bool includeForces, bool includeEnergy);
+    /**
+     * This is the internal implementation of execute(), templatized on whether we're
+     * using single or double precision.
+     */
+    template <class REAL>
+    double executeImpl(ContextImpl& context);
+    /**
+     * Copy changed parameters over to a context.
+     *
+     * @param context    the context to copy parameters to
+     * @param force      the RMSDForce to copy the parameters from
+     */
+    void copyParametersToContext(ContextImpl& context, const RMSDForce& force);
+private:
+    class ForceInfo;
+    OpenCLContext& cl;
+    ForceInfo* info;
+    double sumNormRef;
+    OpenCLArray* referencePos;
+    OpenCLArray* particles;
+    OpenCLArray* buffer;
+    cl::Kernel kernel1, kernel2;
+};
+
 /**
  * This kernel is invoked by VerletIntegrator to take one time step.
  */

platforms/opencl/src/OpenCLKernelFactory.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) 2008-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2018 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -108,6 +108,8 @@ KernelImpl* OpenCLKernelFactory::createKernelImpl(std::string name, const Platfo
         return new OpenCLCalcCustomCompoundBondForceKernel(name, platform, cl, context.getSystem());
     if (name == CalcCustomCVForceKernel::Name())
         return new OpenCLCalcCustomCVForceKernel(name, platform, cl);
+    if (name == CalcRMSDForceKernel::Name())
+        return new OpenCLCalcRMSDForceKernel(name, platform, cl);
     if (name == CalcCustomManyParticleForceKernel::Name())
         return new OpenCLCalcCustomManyParticleForceKernel(name, platform, cl, context.getSystem());
     if (name == CalcGayBerneForceKernel::Name())

platforms/opencl/src/OpenCLKernels.cpp

@@ -51,6 +51,7 @@
 #include "ReferenceTabulatedFunction.h"
 #include "SimTKOpenMMRealType.h"
 #include "SimTKOpenMMUtilities.h"
+#include "jama_eig.h"
 #include <algorithm>
 #include <cmath>
 #include <set>
@@ -7064,6 +7065,203 @@ void OpenCLCalcCustomCVForceKernel::copyState(ContextImpl& context, ContextImpl&
         innerContext.setParameter(param.first, context.getParameter(param.first));
 }
 
+class OpenCLCalcRMSDForceKernel::ForceInfo : public OpenCLForceInfo {
+public:
+    ForceInfo(const RMSDForce& force) : OpenCLForceInfo(0), force(force) {
+        updateParticles();
+    }
+    void updateParticles() {
+        particles.clear();
+        for (int i : force.getParticles())
+            particles.insert(i);
+    }
+    bool areParticlesIdentical(int particle1, int particle2) {
+        bool include1 = (particles.find(particle1) != particles.end());
+        bool include2 = (particles.find(particle2) != particles.end());
+        return (include1 == include2);
+    }
+private:
+    const RMSDForce& force;
+    set<int> particles;
+};
+
+OpenCLCalcRMSDForceKernel::~OpenCLCalcRMSDForceKernel() {
+    if (referencePos != NULL)
+        delete referencePos;
+    if (particles != NULL)
+        delete particles;
+    if (buffer != NULL)
+        delete buffer;
+}
+
+void OpenCLCalcRMSDForceKernel::initialize(const System& system, const RMSDForce& force) {
+    // Create data structures.
+    
+    bool useDouble = cl.getUseDoublePrecision();
+    int elementSize = (useDouble ? sizeof(cl_double) : sizeof(cl_float));
+    int numParticles = force.getParticles().size();
+    if (numParticles == 0)
+        numParticles = system.getNumParticles();
+    referencePos = new OpenCLArray(cl, system.getNumParticles(), 4*elementSize, "referencePos");
+    particles = OpenCLArray::create<cl_int>(cl, numParticles, "particles");
+    buffer = new OpenCLArray(cl, 13, elementSize, "buffer");
+    recordParameters(force);
+    info = new ForceInfo(force);
+    cl.addForce(info);
+    
+    // Create the kernels.
+
+    cl::Program program = cl.createProgram(OpenCLKernelSources::rmsd);
+    kernel1 = cl::Kernel(program, "computeRMSDPart1");
+    kernel2 = cl::Kernel(program, "computeRMSDForces");
+}
+
+void OpenCLCalcRMSDForceKernel::recordParameters(const RMSDForce& force) {
+    // Record the parameters and center the reference positions.
+    
+    vector<int> particleVec = force.getParticles();
+    if (particleVec.size() == 0)
+        for (int i = 0; i < cl.getNumAtoms(); i++)
+            particleVec.push_back(i);
+    vector<Vec3> centeredPositions = force.getReferencePositions();
+    Vec3 center;
+    for (int i : particleVec)
+        center += centeredPositions[i];
+    center /= particleVec.size();
+    for (Vec3& p : centeredPositions)
+        p -= center;
+
+    // Upload them to the device.
+
+    particles->upload(particleVec);
+    if (cl.getUseDoublePrecision()) {
+        vector<mm_double4> pos;
+        for (Vec3 p : centeredPositions)
+            pos.push_back(mm_double4(p[0], p[1], p[2], 0));
+        referencePos->upload(pos);
+    }
+    else {
+        vector<mm_float4> pos;
+        for (Vec3 p : centeredPositions)
+            pos.push_back(mm_float4(p[0], p[1], p[2], 0));
+        referencePos->upload(pos);
+    }
+
+    // Record the sum of the norms of the reference positions.
+
+    sumNormRef = 0.0;
+    for (int i : particleVec) {
+        Vec3 p = centeredPositions[i];
+        sumNormRef += p.dot(p);
+    }
+}
+
+double OpenCLCalcRMSDForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
+    if (cl.getUseDoublePrecision())
+        return executeImpl<double>(context);
+    return executeImpl<float>(context);
+}
+
+template <class REAL>
+double OpenCLCalcRMSDForceKernel::executeImpl(ContextImpl& context) {
+    // Execute the first kernel.
+
+    int numParticles = particles->getSize();
+    int blockSize = 128;
+    kernel1.setArg<cl_int>(0, numParticles);
+    kernel1.setArg<cl::Buffer>(1, cl.getPosq().getDeviceBuffer());
+    kernel1.setArg<cl::Buffer>(2, referencePos->getDeviceBuffer());
+    kernel1.setArg<cl::Buffer>(3, particles->getDeviceBuffer());
+    kernel1.setArg<cl::Buffer>(4, buffer->getDeviceBuffer());
+    kernel1.setArg(5, blockSize*sizeof(REAL), NULL);
+    cl.executeKernel(kernel1, blockSize, blockSize);
+    
+    // Download the results, build the F matrix, and find the maximum eigenvalue
+    // and eigenvector.
+
+    vector<REAL> b;
+    buffer->download(b);
+    Array2D<double> F(4, 4);
+    F[0][0] =  b[0*3+0] + b[1*3+1] + b[2*3+2];
+    F[1][0] =  b[1*3+2] - b[2*3+1];
+    F[2][0] =  b[2*3+0] - b[0*3+2];
+    F[3][0] =  b[0*3+1] - b[1*3+0];
+    F[0][1] =  b[1*3+2] - b[2*3+1];
+    F[1][1] =  b[0*3+0] - b[1*3+1] - b[2*3+2];
+    F[2][1] =  b[0*3+1] + b[1*3+0];
+    F[3][1] =  b[0*3+2] + b[2*3+0];
+    F[0][2] =  b[2*3+0] - b[0*3+2];
+    F[1][2] =  b[0*3+1] + b[1*3+0];
+    F[2][2] = -b[0*3+0] + b[1*3+1] - b[2*3+2];
+    F[3][2] =  b[1*3+2] + b[2*3+1];
+    F[0][3] =  b[0*3+1] - b[1*3+0];
+    F[1][3] =  b[0*3+2] + b[2*3+0];
+    F[2][3] =  b[1*3+2] + b[2*3+1];
+    F[3][3] = -b[0*3+0] - b[1*3+1] + b[2*3+2];
+    JAMA::Eigenvalue<double> eigen(F);
+    Array1D<double> values;
+    eigen.getRealEigenvalues(values);
+    Array2D<double> vectors;
+    eigen.getV(vectors);
+
+    // Compute the RMSD.
+
+    double msd = (sumNormRef+b[9]-2*values[3])/numParticles;
+    msd = max(0.0, msd); // Protect against numerical error when the particles are perfectly aligned.
+    double rmsd = sqrt(msd);
+    b[9] = rmsd;
+
+    // Compute the rotation matrix.
+
+    double q[] = {vectors[0][3], vectors[1][3], vectors[2][3], vectors[3][3]};
+    double q00 = q[0]*q[0], q01 = q[0]*q[1], q02 = q[0]*q[2], q03 = q[0]*q[3];
+    double q11 = q[1]*q[1], q12 = q[1]*q[2], q13 = q[1]*q[3];
+    double q22 = q[2]*q[2], q23 = q[2]*q[3];
+    double q33 = q[3]*q[3];
+    b[0] = q00+q11-q22-q33;
+    b[1] = 2*(q12-q03);
+    b[2] = 2*(q13+q02);
+    b[3] = 2*(q12+q03);
+    b[4] = q00-q11+q22-q33;
+    b[5] = 2*(q23-q01);
+    b[6] = 2*(q13-q02);
+    b[7] = 2*(q23+q01);
+    b[8] = q00-q11-q22+q33;
+
+    // Upload it to the device and invoke the kernel to apply forces.
+    
+    buffer->upload(b);
+    kernel2.setArg<cl_int>(0, numParticles);
+    kernel2.setArg<cl::Buffer>(1, cl.getPosq().getDeviceBuffer());
+    kernel2.setArg<cl::Buffer>(2, referencePos->getDeviceBuffer());
+    kernel2.setArg<cl::Buffer>(3, particles->getDeviceBuffer());
+    kernel2.setArg<cl::Buffer>(4, buffer->getDeviceBuffer());
+    kernel2.setArg<cl::Buffer>(5, cl.getForceBuffers().getDeviceBuffer());
+    cl.executeKernel(kernel2, numParticles);
+    return rmsd;
+}
+
+void OpenCLCalcRMSDForceKernel::copyParametersToContext(ContextImpl& context, const RMSDForce& force) {
+    if (referencePos->getSize() != force.getReferencePositions().size())
+        throw OpenMMException("updateParametersInContext: The number of reference positions has changed");
+    int numParticles = force.getParticles().size();
+    if (numParticles == 0)
+        numParticles = context.getSystem().getNumParticles();
+    if (numParticles != particles->getSize()) {
+        // Recreate the particles array.
+        
+        delete particles;
+        particles = NULL;
+        particles = OpenCLArray::create<cl_int>(cl, numParticles, "particles");
+    }
+    recordParameters(force);
+    
+    // Mark that the current reordering may be invalid.
+    
+    info->updateParticles();
+    cl.invalidateMolecules(info);
+}
+
 OpenCLIntegrateVerletStepKernel::~OpenCLIntegrateVerletStepKernel() {
 }
 

platforms/opencl/src/OpenCLPlatform.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) 2008-2017 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2018 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -83,6 +83,7 @@ OpenCLPlatform::OpenCLPlatform() {
     registerKernelFactory(CalcCustomCentroidBondForceKernel::Name(), factory);
     registerKernelFactory(CalcCustomCompoundBondForceKernel::Name(), factory);
     registerKernelFactory(CalcCustomCVForceKernel::Name(), factory);
+    registerKernelFactory(CalcRMSDForceKernel::Name(), factory);
     registerKernelFactory(CalcCustomManyParticleForceKernel::Name(), factory);
     registerKernelFactory(CalcGayBerneForceKernel::Name(), factory);
     registerKernelFactory(IntegrateVerletStepKernel::Name(), factory);

platforms/opencl/src/kernels/rmsd.cl

+// This file contains kernels to compute the RMSD and its gradient using the algorithm described
+// in Coutsias et al, "Using quaternions to calculate RMSD" (doi: 10.1002/jcc.20110).
+
+/**
+ * Sum a value over all threads.
+ */
+real reduceValue(real value, __local real* temp) {
+    const int thread = get_local_id(0);
+    temp[thread] = value;
+    barrier(CLK_LOCAL_MEM_FENCE);
+    for (uint step = 1; step < get_local_size(0); step *= 2) {
+        if (thread+step < get_local_size(0) && thread%(2*step) == 0)
+            temp[thread] = temp[thread] + temp[thread+step];
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+    return temp[0];
+}
+
+/**
+ * Perform the first step of computing the RMSD.  This is executed as a single work group.
+ */
+__kernel void computeRMSDPart1(int numParticles, __global const real4* restrict posq, __global const real4* restrict referencePos,
+        __global const int* restrict particles, __global real* buffer, __local real* restrict temp) {
+    // Compute the center of the particle positions.
+    
+    real3 center = (real3) 0;
+    for (int i = get_local_id(0); i < numParticles; i += get_local_size(0))
+        center += posq[particles[i]].xyz;
+    center.x = reduceValue(center.x, temp)/numParticles;
+    center.y = reduceValue(center.y, temp)/numParticles;
+    center.z = reduceValue(center.z, temp)/numParticles;
+    
+    // Compute the correlation matrix.
+    
+    real R[3][3] = {{0, 0, 0}, {0, 0, 0}, {0, 0, 0}};
+    real sum = 0;
+    for (int i = get_local_id(0); i < numParticles; i += get_local_size(0)) {
+        int index = particles[i];
+        real3 pos = posq[index].xyz - center;
+        real3 refPos = referencePos[index].xyz;
+        R[0][0] += pos.x*refPos.x;
+        R[0][1] += pos.x*refPos.y;
+        R[0][2] += pos.x*refPos.z;
+        R[1][0] += pos.y*refPos.x;
+        R[1][1] += pos.y*refPos.y;
+        R[1][2] += pos.y*refPos.z;
+        R[2][0] += pos.z*refPos.x;
+        R[2][1] += pos.z*refPos.y;
+        R[2][2] += pos.z*refPos.z;
+        sum += dot(pos, pos);
+    }
+    for (int i = 0; i < 3; i++)
+        for (int j = 0; j < 3; j++)
+            R[i][j] = reduceValue(R[i][j], temp);
+    sum = reduceValue(sum, temp);
+
+    // Copy everything into the output buffer to send back to the host.
+    
+    if (get_local_id(0) == 0) {
+        for (int i = 0; i < 3; i++)
+            for (int j = 0; j < 3; j++)
+                buffer[3*i+j] = R[i][j];
+        buffer[9] = sum;
+        buffer[10] = center.x;
+        buffer[11] = center.y;
+        buffer[12] = center.z;
+    }
+}
+
+/**
+ * Apply forces based on the RMSD.
+ */
+__kernel void computeRMSDForces(int numParticles, __global const real4* restrict posq, __global const real4* restrict referencePos,
+        __global const int* restrict particles, __global const real* buffer, __global real4* restrict forceBuffers) {
+    real3 center = (real3) (buffer[10], buffer[11], buffer[12]);
+    real scale = 1 / (real) (buffer[9]*numParticles);
+    for (int i = get_global_id(0); i < numParticles; i += get_global_size(0)) {
+        int index = particles[i];
+        real3 pos = posq[index].xyz - center;
+        real3 refPos = referencePos[index].xyz;
+        real3 rotatedRef = (real3) (buffer[0]*refPos[0] + buffer[3]*refPos[1] + buffer[6]*refPos[2],
+                                    buffer[1]*refPos[0] + buffer[4]*refPos[1] + buffer[7]*refPos[2],
+                                    buffer[2]*refPos[0] + buffer[5]*refPos[1] + buffer[8]*refPos[2]);
+        forceBuffers[index].xyz -= (pos-rotatedRef)*scale;
+    }
+}

platforms/opencl/tests/TestOpenCLRMSDForce.cpp

+/* -------------------------------------------------------------------------- *
+ *                                   OpenMM                                   *
+ * -------------------------------------------------------------------------- *
+ * 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 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 "OpenCLTests.h"
+#include "TestRMSDForce.h"
+
+void runPlatformTests() {
+}

platforms/reference/src/ReferenceKernels.cpp

@@ -2058,6 +2058,9 @@ void ReferenceCalcCustomCVForceKernel::copyState(ContextImpl& context, ContextIm
 
 void ReferenceCalcRMSDForceKernel::initialize(const System& system, const RMSDForce& force) {
     particles = force.getParticles();
+    if (particles.size() == 0)
+        for (int i = 0; i < system.getNumParticles(); i++)
+            particles.push_back(i);
     referencePos = force.getReferencePositions();
     Vec3 center;
     for (int i : particles)
@@ -2078,6 +2081,9 @@ void ReferenceCalcRMSDForceKernel::copyParametersToContext(ContextImpl& context,
     if (referencePos.size() != force.getReferencePositions().size())
         throw OpenMMException("updateParametersInContext: The number of reference positions has changed");
     particles = force.getParticles();
+    if (particles.size() == 0)
+        for (int i = 0; i < referencePos.size(); i++)
+            particles.push_back(i);
     referencePos = force.getReferencePositions();
     Vec3 center;
     for (int i : particles)

platforms/reference/src/SimTKReference/ReferenceRMSDForce.cpp

@@ -96,7 +96,9 @@ double ReferenceRMSDForce::calculateIxn(vector<Vec3>& atomCoordinates, vector<Ve
         int index = particles[i];
         sum += positions[i].dot(positions[i]) + referencePos[index].dot(referencePos[index]);
     }
-    double rmsd = sqrt((sum-2*values[3])/numParticles);
+    double msd = (sum-2*values[3])/numParticles;
+    msd = max(0.0, msd); // Protect against numerical error when the particles are perfectly aligned.
+    double rmsd = sqrt(msd);
 
     // Compute the rotation matrix.
 

tests/TestRMSDForce.h

@@ -99,7 +99,7 @@ void testRMSD() {
     }
     context.setPositions(positions);
     state1 = context.getState(State::Energy | State::Forces);
-    ASSERT_EQUAL_TOL(rmsd, state1.getPotentialEnergy(), 1e-5);
+    ASSERT_EQUAL_TOL(rmsd, state1.getPotentialEnergy(), 1e-4);
 
     // Take a small step in the direction of the energy gradient and see whether the potential energy changes by the expected amount.
 
@@ -108,7 +108,7 @@ void testRMSD() {
     for (int i = 0; i < (int) forces.size(); ++i)
         norm += forces[i].dot(forces[i]);
     norm = std::sqrt(norm);
-    const double stepSize = 1e-3;
+    const double stepSize = 0.1;
     double step = 0.5*stepSize/norm;
     vector<Vec3> positions2(numParticles), positions3(numParticles);
     for (int i = 0; i < (int) positions.size(); ++i) {
@@ -121,7 +121,16 @@ void testRMSD() {
     State state2 = context.getState(State::Energy);
     context.setPositions(positions3);
     State state3 = context.getState(State::Energy);
-    ASSERT_EQUAL_TOL(norm, (state2.getPotentialEnergy()-state3.getPotentialEnergy())/stepSize, 1e-4);
+    ASSERT_EQUAL_TOL(norm, (state2.getPotentialEnergy()-state3.getPotentialEnergy())/stepSize, 1e-3);
+    
+    // Check that updateParametersInContext() works correctly.
+    
+    context.setPositions(positions);
+    force->setReferencePositions(positions);
+    force->updateParametersInContext(context);
+    ASSERT_EQUAL_TOL(0.0, context.getState(State::Energy).getPotentialEnergy(), 1e-2);
+    context.setPositions(referencePos);
+    ASSERT_EQUAL_TOL(rmsd, context.getState(State::Energy).getPotentialEnergy(), 1e-4);
 }
 
 void runPlatformTests();