Fixed bugs in CustomGBForce

platforms/opencl/src/OpenCLKernels.cpp

@@ -1812,7 +1812,7 @@ void OpenCLCalcCustomGBForceKernel::initialize(const System& system, const Custo
     // Record derivatives of expressions needed for the chain rule terms.
 
     vector<vector<Lepton::ParsedExpression> > valueGradientExpressions(force.getNumComputedValues());
-    bool needParameterGradient = false;
+    needParameterGradient = false;
     for (int i = 1; i < force.getNumComputedValues(); i++) {
         Lepton::ParsedExpression ex = Lepton::Parser::parse(computedValueExpressions[i], functions).optimize();
         valueGradientExpressions[i].push_back(ex.differentiate("x").optimize());
@@ -2116,8 +2116,58 @@ void OpenCLCalcCustomGBForceKernel::initialize(const System& system, const Custo
         cl::Program program = cl.createProgram(cl.replaceStrings(OpenCLKernelSources::customGBEnergyPerParticle, replacements), defines);
         perParticleEnergyKernel = cl::Kernel(program, "computePerParticleEnergy");
     }
+    if (needParameterGradient) {
+        // Create the kernel to compute chain rule terms for computed values that depend explicitly on particle coordinates.
+
+        stringstream compute, extraArgs;
+        if (force.getNumGlobalParameters() > 0)
+            extraArgs << ", __constant float* globals";
+        for (int i = 0; i < (int) params->getBuffers().size(); i++) {
+            const OpenCLNonbondedUtilities::ParameterInfo& buffer = params->getBuffers()[i];
+            string paramName = "params"+intToString(i+1);
+            extraArgs << ", __global " << buffer.getType() << "* " << paramName;
+        }
+        for (int i = 0; i < (int) computedValues->getBuffers().size(); i++) {
+            const OpenCLNonbondedUtilities::ParameterInfo& buffer = computedValues->getBuffers()[i];
+            string valueName = "values"+intToString(i+1);
+            extraArgs << ", __global " << buffer.getType() << "* " << valueName;
+        }
+        for (int i = 0; i < (int) energyDerivs->getBuffers().size(); i++) {
+            const OpenCLNonbondedUtilities::ParameterInfo& buffer = energyDerivs->getBuffers()[i];
+            string index = intToString(i+1);
+            extraArgs << ", __global " << buffer.getType() << "* derivBuffers" << index;
+            compute << buffer.getType() << " deriv" << index << " = derivBuffers" << index << "[index];\n";
+        }
+        map<string, string> variables;
+        variables["x"] = "pos.x";
+        variables["y"] = "pos.y";
+        variables["z"] = "pos.z";
+        for (int i = 0; i < force.getNumPerParticleParameters(); i++)
+            variables[force.getPerParticleParameterName(i)] = "params"+params->getParameterSuffix(i, "[index]");
+        for (int i = 0; i < force.getNumGlobalParameters(); i++)
+            variables[force.getGlobalParameterName(i)] = "globals["+intToString(i)+"]";
+        for (int i = 0; i < force.getNumComputedValues(); i++)
+            variables[computedValueNames[i]] = "values"+computedValues->getParameterSuffix(i, "[index]");
+        map<string, Lepton::ParsedExpression> gradientExpressions;
+        for (int i = 1; i < force.getNumComputedValues(); i++) {
+            if (!isZeroExpression(valueGradientExpressions[i][0]))
+                gradientExpressions["force.x -= deriv"+energyDerivs->getParameterSuffix(i)+"*"] = valueGradientExpressions[i][0];
+            if (!isZeroExpression(valueGradientExpressions[i][1]))
+                gradientExpressions["force.y -= deriv"+energyDerivs->getParameterSuffix(i)+"*"] = valueGradientExpressions[i][1];
+            if (!isZeroExpression(valueGradientExpressions[i][2]))
+                gradientExpressions["force.z -= deriv"+energyDerivs->getParameterSuffix(i)+"*"] = valueGradientExpressions[i][2];
+        }
+        compute << OpenCLExpressionUtilities::createExpressions(gradientExpressions, variables, functionDefinitions, "temp", prefix+"functionParams");
+        map<string, string> replacements;
+        replacements["PARAMETER_ARGUMENTS"] = extraArgs.str()+tableArgs.str();
+        replacements["COMPUTE_FORCES"] = compute.str();
+        map<string, string> defines;
+        defines["NUM_ATOMS"] = intToString(cl.getNumAtoms());
+        cl::Program program = cl.createProgram(cl.replaceStrings(OpenCLKernelSources::customGBGradientChainRule, replacements), defines);
+        gradientChainRuleKernel = cl::Kernel(program, "computeGradientChainRuleTerms");
+    }
     {
-        // Create the code to calculate chain rules terms (possibly as part of the default nonbonded kernel).
+        // Create the code to calculate chain rules terms as part of the default nonbonded kernel.
 
         map<string, string> globalVariables;
         for (int i = 0; i < force.getNumGlobalParameters(); i++) {
@@ -2141,32 +2191,12 @@ void OpenCLCalcCustomGBForceKernel::initialize(const System& system, const Custo
         derivExpressions["float dV0dR1 = "] = dVdR;
         derivExpressions["float dV0dR2 = "] = dVdR.renameVariables(rename);
         chainSource << OpenCLExpressionUtilities::createExpressions(derivExpressions, variables, functionDefinitions, prefix+"temp0_", prefix+"functionParams");
-        if (needParameterGradient) {
-            chainSource << "float4 grad1_0_1 = (float4) 0;\n";
-            chainSource << "float4 grad1_0_2 = (float4) 0;\n";
-            chainSource << "float4 grad2_0_1 = (float4) 0;\n";
-            chainSource << "float4 grad2_0_2 = (float4) 0;\n";
-            if (useExclusionsForValue)
-                chainSource << "if (!isExcluded) {\n";
-            chainSource << "grad1_0_1 = dV0dR1*delta*invR;\n";
-            chainSource << "grad1_0_2 = dV0dR2*delta*invR;\n";
-            chainSource << "grad2_0_1 = grad1_0_1;\n";
-            chainSource << "grad2_0_2 = grad1_0_2;\n";
-            chainSource << "tempForce1 -= grad1_0_1*" << prefix << "dEdV" << energyDerivs->getParameterSuffix(0, "1") << ";\n";
-            chainSource << "tempForce1 -= grad1_0_2*" << prefix << "dEdV" << energyDerivs->getParameterSuffix(0, "2") << ";\n";
-            chainSource << "tempForce2 -= grad2_0_1*" << prefix << "dEdV" << energyDerivs->getParameterSuffix(0, "1") << ";\n";
-            chainSource << "tempForce2 -= grad2_0_2*" << prefix << "dEdV" << energyDerivs->getParameterSuffix(0, "2") << ";\n";
-            if (useExclusionsForValue)
-                chainSource << "}\n";
-        }
-        else {
         if (useExclusionsForValue)
             chainSource << "if (!isExcluded) {\n";
         chainSource << "tempForce -= dV0dR1*" << prefix << "dEdV" << energyDerivs->getParameterSuffix(0, "1") << ";\n";
         chainSource << "tempForce -= dV0dR2*" << prefix << "dEdV" << energyDerivs->getParameterSuffix(0, "2") << ";\n";
         if (useExclusionsForValue)
             chainSource << "}\n";
-        }
         variables = globalVariables;
         map<string, string> rename1;
         map<string, string> rename2;
@@ -2198,43 +2228,6 @@ void OpenCLCalcCustomGBForceKernel::initialize(const System& system, const Custo
             if (i == 0)
                 continue;
             string is = intToString(i);
-            if (needParameterGradient) {
-                chainSource << "float4 grad1_"+is+"_1 = (float4) 0;\n";
-                chainSource << "float4 grad1_"+is+"_2 = (float4) 0;\n";
-                chainSource << "float4 grad2_"+is+"_1 = (float4) 0;\n";
-                chainSource << "float4 grad2_"+is+"_2 = (float4) 0;\n";
-                for (int j = 0; j < i; j++) {
-                    string js = intToString(j);
-                    Lepton::ParsedExpression dVdV = Lepton::Parser::parse(computedValueExpressions[i], functions).differentiate(computedValueNames[j]).optimize();
-                    derivExpressions.clear();
-                    derivExpressions["float dV"+is+"dV"+js+"_1 = "] = dVdV.renameVariables(rename1);
-                    derivExpressions["float dV"+is+"dV"+js+"_2 = "] = dVdV.renameVariables(rename2);
-                    chainSource << OpenCLExpressionUtilities::createExpressions(derivExpressions, variables, functionDefinitions, prefix+"temp"+is+"_"+js+"_", prefix+"functionParams");
-                    chainSource << "grad1_"+is+"_1 += dV"+is+"dV"+js+"_1*grad1_"+js+"_1;\n";
-                    chainSource << "grad2_"+is+"_1 += dV"+is+"dV"+js+"_1*grad2_"+js+"_1;\n";
-                    chainSource << "grad1_"+is+"_2 += dV"+is+"dV"+js+"_2*grad1_"+js+"_2;\n";
-                    chainSource << "grad2_"+is+"_2 += dV"+is+"dV"+js+"_2*grad2_"+js+"_2;\n";
-                }
-                derivExpressions.clear();
-                if (!isZeroExpression(valueGradientExpressions[i][0])) {
-                    derivExpressions["grad1_"+is+"_1.x -= "] = valueGradientExpressions[i][0].renameVariables(rename1);
-                    derivExpressions["grad2_"+is+"_2.x -= "] = valueGradientExpressions[i][0].renameVariables(rename2);
-                }
-                if (!isZeroExpression(valueGradientExpressions[i][1])) {
-                    derivExpressions["grad1_"+is+"_1.y -= "] = valueGradientExpressions[i][1].renameVariables(rename1);
-                    derivExpressions["grad2_"+is+"_2.y -= "] = valueGradientExpressions[i][1].renameVariables(rename2);
-                }
-                if (!isZeroExpression(valueGradientExpressions[i][2])) {
-                    derivExpressions["grad1_"+is+"_1.z -= "] = valueGradientExpressions[i][2].renameVariables(rename1);
-                    derivExpressions["grad2_"+is+"_2.z -= "] = valueGradientExpressions[i][2].renameVariables(rename2);
-                }
-                chainSource << OpenCLExpressionUtilities::createExpressions(derivExpressions, variables, functionDefinitions, prefix+"temp"+is+"_", prefix+"functionParams");
-                chainSource << "tempForce1 -= grad1_"<<is<<"_1*"<<prefix<<"dEdV"<<energyDerivs->getParameterSuffix(i, "1")<<";\n";
-                chainSource << "tempForce2 -= grad2_"<<is<<"_1*"<<prefix<<"dEdV"<<energyDerivs->getParameterSuffix(i, "1")<<";\n";
-                chainSource << "tempForce1 -= grad1_"<<is<<"_2*"<<prefix<<"dEdV"<<energyDerivs->getParameterSuffix(i, "2")<<";\n";
-                chainSource << "tempForce2 -= grad2_"<<is<<"_2*"<<prefix<<"dEdV"<<energyDerivs->getParameterSuffix(i, "2")<<";\n";
-            }
-            else {
             chainSource << "float dV"+is+"dR1 = 0;\n";
             chainSource << "float dV"+is+"dR2 = 0;\n";
             for (int j = 0; j < i; j++) {
@@ -2248,7 +2241,6 @@ void OpenCLCalcCustomGBForceKernel::initialize(const System& system, const Custo
             chainSource << "tempForce -= dV"<< is << "dR1*" << prefix << "dEdV" << energyDerivs->getParameterSuffix(i, "1") << ";\n";
             chainSource << "tempForce -= dV"<< is << "dR2*" << prefix << "dEdV" << energyDerivs->getParameterSuffix(i, "2") << ";\n";
         }
-        }
         map<string, string> replacements;
         replacements["COMPUTE_FORCE"] = chainSource.str();
         string source = cl.replaceStrings(OpenCLKernelSources::customGBChainRule, replacements);
@@ -2273,21 +2265,11 @@ void OpenCLCalcCustomGBForceKernel::initialize(const System& system, const Custo
             globals->upload(globalParamValues);
             arguments.push_back(OpenCLNonbondedUtilities::ParameterInfo(prefix+"globals", "float", 1, sizeof(cl_float), globals->getDeviceBuffer()));
         }
-        if (needParameterGradient) {
-            chainRuleParameters = parameters;
-            chainRuleArguments = arguments;
-            chainRuleSource = source;
-            separateChainRuleKernel = true;
-            cl.getNonbondedUtilities().addInteraction(useCutoff, usePeriodic, force.getNumExclusions() > 0, force.getCutoffDistance(), exclusionList, "");
-        }
-        else {
         cl.getNonbondedUtilities().addInteraction(useCutoff, usePeriodic, force.getNumExclusions() > 0, force.getCutoffDistance(), exclusionList, source);
         for (int i = 0; i < (int) parameters.size(); i++)
             cl.getNonbondedUtilities().addParameter(parameters[i]);
         for (int i = 0; i < (int) arguments.size(); i++)
             cl.getNonbondedUtilities().addArgument(arguments[i]);
-            separateChainRuleKernel = false;
-        }
     }
     cl.addForce(new OpenCLCustomGBForceInfo(cl.getNonbondedUtilities().getNumForceBuffers(), force));
     for (int i = 0; i < (int) energyDerivs->getBuffers().size(); i++) {
@@ -2409,7 +2391,19 @@ void OpenCLCalcCustomGBForceKernel::executeForces(ContextImpl& context) {
                 perParticleEnergyKernel.setArg<cl::Buffer>(index++, tabulatedFunctions[i]->getDeviceBuffer());
             perParticleEnergyKernel.setArg<cl::Buffer>(index++, tabulatedFunctionParams->getDeviceBuffer());
         }
-        chainRuleKernel = nb.createInteractionKernel(chainRuleSource, chainRuleParameters, chainRuleArguments, true, false);
+        if (needParameterGradient) {
+            index = 0;
+            gradientChainRuleKernel.setArg<cl::Buffer>(index++, cl.getForceBuffers().getDeviceBuffer());
+            gradientChainRuleKernel.setArg<cl::Buffer>(index++, cl.getPosq().getDeviceBuffer());
+            if (globals != NULL)
+                gradientChainRuleKernel.setArg<cl::Buffer>(index++, globals->getDeviceBuffer());
+            for (int i = 0; i < (int) params->getBuffers().size(); i++)
+                gradientChainRuleKernel.setArg<cl::Memory>(index++, params->getBuffers()[i].getMemory());
+            for (int i = 0; i < (int) computedValues->getBuffers().size(); i++)
+                gradientChainRuleKernel.setArg<cl::Memory>(index++, computedValues->getBuffers()[i].getMemory());
+            for (int i = 0; i < (int) energyDerivs->getBuffers().size(); i++)
+                gradientChainRuleKernel.setArg<cl::Memory>(index++, energyDerivs->getBuffers()[i].getMemory());
+        }
     }
     if (globals != NULL) {
         bool changed = false;
@@ -2427,17 +2421,13 @@ void OpenCLCalcCustomGBForceKernel::executeForces(ContextImpl& context) {
         pairValueKernel.setArg<mm_float4>(11, cl.getInvPeriodicBoxSize());
         pairEnergyKernel.setArg<mm_float4>(11, cl.getPeriodicBoxSize());
         pairEnergyKernel.setArg<mm_float4>(12, cl.getInvPeriodicBoxSize());
-        if (separateChainRuleKernel) {
-            chainRuleKernel.setArg<mm_float4>(10, cl.getPeriodicBoxSize());
-            chainRuleKernel.setArg<mm_float4>(11, cl.getInvPeriodicBoxSize());
-        }
     }
     cl.executeKernel(pairValueKernel, nb.getTiles().getSize()*OpenCLContext::TileSize);
     cl.executeKernel(perParticleValueKernel, cl.getPaddedNumAtoms());
     cl.executeKernel(pairEnergyKernel, nb.getTiles().getSize()*OpenCLContext::TileSize);
     cl.executeKernel(perParticleEnergyKernel, cl.getPaddedNumAtoms());
-    if (separateChainRuleKernel)
-        cl.executeKernel(chainRuleKernel, nb.getTiles().getSize()*OpenCLContext::TileSize);
+    if (needParameterGradient)
+        cl.executeKernel(gradientChainRuleKernel, cl.getPaddedNumAtoms());
 }
 
 double OpenCLCalcCustomGBForceKernel::executeEnergy(ContextImpl& context) {

platforms/opencl/src/OpenCLKernels.h

@@ -624,7 +624,7 @@ public:
      */
     double executeEnergy(ContextImpl& context);
 private:
-    bool hasInitializedKernels, separateChainRuleKernel;
+    bool hasInitializedKernels, needParameterGradient;
     OpenCLContext& cl;
     OpenCLParameterSet* params;
     OpenCLParameterSet* computedValues;
@@ -635,11 +635,8 @@ private:
     std::vector<std::string> globalParamNames;
     std::vector<cl_float> globalParamValues;
     std::vector<OpenCLArray<mm_float4>*> tabulatedFunctions;
-    std::vector<OpenCLNonbondedUtilities::ParameterInfo> chainRuleParameters;
-    std::vector<OpenCLNonbondedUtilities::ParameterInfo> chainRuleArguments;
-    std::string chainRuleSource;
     System& system;
-    cl::Kernel pairValueKernel, perParticleValueKernel, pairEnergyKernel, perParticleEnergyKernel, chainRuleKernel;
+    cl::Kernel pairValueKernel, perParticleValueKernel, pairEnergyKernel, perParticleEnergyKernel, gradientChainRuleKernel;
 };
 
 /**

platforms/opencl/src/kernels/customGBGradientChainRule.cl

+/**
+ * Compute chain rule terms for computed values that depend explicitly on particle coordinates.
+ */
+
+__kernel void computeGradientChainRuleTerms(__global float4* forceBuffers, __global float4* posq
+        PARAMETER_ARGUMENTS) {
+    unsigned int index = get_global_id(0);
+    while (index < NUM_ATOMS) {
+        float4 pos = posq[index];
+        float4 force = forceBuffers[index];
+        COMPUTE_FORCES
+        forceBuffers[index] = force;
+        index += get_global_size(0);
+    }
+}

platforms/opencl/tests/TestOpenCLCustomGBForce.cpp

@@ -141,6 +141,95 @@ void testOBC(GBSAOBCForce::NonbondedMethod obcMethod, CustomGBForce::NonbondedMe
     }
 }
 
+void testMembrane() {
+    const int numMolecules = 70;
+    const int numParticles = numMolecules*2;
+    const double boxSize = 10.0;
+    OpenCLPlatform platform;
+
+    // Create a system with an implicit membrane.
+
+    System system;
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(1.0);
+    }
+    system.setDefaultPeriodicBoxVectors(Vec3(boxSize, 0.0, 0.0), Vec3(0.0, boxSize, 0.0), Vec3(0.0, 0.0, boxSize));
+    CustomGBForce* custom = new CustomGBForce();
+    custom->setCutoffDistance(2.0);
+    custom->addPerParticleParameter("q");
+    custom->addPerParticleParameter("radius");
+    custom->addPerParticleParameter("scale");
+    custom->addGlobalParameter("thickness", 3);
+    custom->addGlobalParameter("solventDielectric", 78.3);
+    custom->addGlobalParameter("soluteDielectric", 1);
+    custom->addComputedValue("Imol", "step(r+sr2-or1)*0.5*(1/L-1/U+0.25*(1/U^2-1/L^2)*(r-sr2*sr2/r)+0.5*log(L/U)/r+C);"
+                             "U=r+sr2;"
+                             "C=2*(1/or1-1/L)*step(sr2-r-or1);"
+                             "L=max(or1, D);"
+                             "D=abs(r-sr2);"
+                             "sr2 = scale2*or2;"
+                             "or1 = radius1-0.009; or2 = radius2-0.009", CustomGBForce::ParticlePairNoExclusions);
+    custom->addComputedValue("Imem", "(1/radius+2*log(2)/thickness)/(1+exp(7.2*(abs(z)+radius-0.5*thickness)))", CustomGBForce::SingleParticle);
+    custom->addComputedValue("B", "1/(1/or-tanh(1*psi-0.8*psi^2+4.85*psi^3)/radius);"
+                             "psi=max(Imol,Imem)*or; or=radius-0.009", CustomGBForce::SingleParticle);
+    custom->addEnergyTerm("28.3919551*(radius+0.14)^2*(radius/B)^6-0.5*138.935456*(1/soluteDielectric-1/solventDielectric)*q^2/B", CustomGBForce::SingleParticle);
+    custom->addEnergyTerm("-138.935456*(1/soluteDielectric-1/solventDielectric)*q1*q2/f;"
+                          "f=sqrt(r^2+B1*B2*exp(-r^2/(4*B1*B2)))", CustomGBForce::ParticlePairNoExclusions);
+    vector<Vec3> positions(numParticles);
+    vector<Vec3> velocities(numParticles);
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    vector<double> params(3);
+    for (int i = 0; i < numMolecules; i++) {
+        if (i < numMolecules/2) {
+            params[0] = 1.0;
+            params[1] = 0.2;
+            params[2] = 0.5;
+            custom->addParticle(params);
+            params[0] = -1.0;
+            params[1] = 0.1;
+            custom->addParticle(params);
+        }
+        else {
+            params[0] = 1.0;
+            params[1] = 0.2;
+            params[2] = 0.8;
+            custom->addParticle(params);
+            params[0] = -1.0;
+            params[1] = 0.1;
+            custom->addParticle(params);
+        }
+        positions[2*i] = Vec3(boxSize*genrand_real2(sfmt), boxSize*genrand_real2(sfmt), boxSize*genrand_real2(sfmt));
+        positions[2*i+1] = Vec3(positions[2*i][0]+1.0, positions[2*i][1], positions[2*i][2]);
+        velocities[2*i] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt));
+        velocities[2*i+1] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt));
+    }
+    system.addForce(custom);
+    VerletIntegrator integrator(0.01);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+    State state = context.getState(State::Forces | State::Energy);
+    const vector<Vec3>& forces = state.getForces();
+
+    // Take a small step in the direction of the energy gradient and see whether the potential energy changes by the expected amount.
+
+    double norm = 0.0;
+    for (int i = 0; i < (int) forces.size(); ++i)
+        norm += forces[i].dot(forces[i]);
+    norm = std::sqrt(norm);
+    const double stepSize = 1e-3;
+    double step = stepSize/norm;
+    for (int i = 0; i < (int) positions.size(); ++i) {
+        Vec3 p = positions[i];
+        Vec3 f = forces[i];
+        positions[i] = Vec3(p[0]-f[0]*step, p[1]-f[1]*step, p[2]-f[2]*step);
+    }
+    context.setPositions(positions);
+    State state2 = context.getState(State::Energy);
+    ASSERT_EQUAL_TOL(norm, (state2.getPotentialEnergy()-state.getPotentialEnergy())/stepSize, 1e-2);
+}
+
 void testTabulatedFunction(bool interpolating) {
     OpenCLPlatform platform;
     System system;
@@ -334,6 +423,7 @@ int main() {
         testOBC(GBSAOBCForce::NoCutoff, CustomGBForce::NoCutoff);
         testOBC(GBSAOBCForce::CutoffNonPeriodic, CustomGBForce::CutoffNonPeriodic);
         testOBC(GBSAOBCForce::CutoffPeriodic, CustomGBForce::CutoffPeriodic);
+        testMembrane();
         testTabulatedFunction(true);
         testTabulatedFunction(false);
         testMultipleChainRules();

platforms/reference/src/SimTKReference/ReferenceCustomGBIxn.cpp

@@ -349,6 +349,24 @@ void ReferenceCustomGBIxn::calculateChainRuleForces(int numAtoms, RealOpenMM** a
            }
         }
     }
+
+    // Compute chain rule terms for computed values that depend explicitly on particle coordinates.
+
+    map<string, double> variables = globalParameters;
+    for (int i = 0; i < numAtoms; i++) {
+        variables["x"] = atomCoordinates[i][0];
+        variables["y"] = atomCoordinates[i][1];
+        variables["z"] = atomCoordinates[i][2];
+        for (int j = 0; j < (int) paramNames.size(); j++)
+            variables[paramNames[j]] = atomParameters[i][j];
+        for (int j = 1; j < (int) valueNames.size(); j++) {
+            variables[valueNames[j-1]] = values[j-1][i];
+            for (int k = 0; k < 3; k++) {
+                RealOpenMM gradient = (RealOpenMM) valueGradientExpressions[j][k].evaluate(variables);
+                forces[i][k] -= dEdV[j][i]*gradient;
+            }
+        }
+    }
 }
 
 void ReferenceCustomGBIxn::calculateOnePairChainRule(int atom1, int atom2, RealOpenMM** atomCoordinates, RealOpenMM** atomParameters,
@@ -378,16 +396,18 @@ void ReferenceCustomGBIxn::calculateOnePairChainRule(int atom1, int atom2, RealO
 
     // Evaluate the derivative of each parameter with respect to position and apply forces.
 
-    vector<vector<RealOpenMM> > gradient1(valueDerivExpressions.size(), vector<RealOpenMM>(3, 0.0));
-    vector<vector<RealOpenMM> > gradient2(valueDerivExpressions.size(), vector<RealOpenMM>(3, 0.0));
-    if (!isExcluded || valueTypes[0] != OpenMM::CustomGBForce::ParticlePair) {
-        RealOpenMM dVdR = (RealOpenMM) valueDerivExpressions[0][0].evaluate(variables);
     RealOpenMM rinv = 1/r;
+    deltaR[0] *= rinv;
+    deltaR[1] *= rinv;
+    deltaR[2] *= rinv;
+    vector<RealOpenMM> dVdR1(valueDerivExpressions.size(), 0.0);
+    vector<RealOpenMM> dVdR2(valueDerivExpressions.size(), 0.0);
+    if (!isExcluded || valueTypes[0] != OpenMM::CustomGBForce::ParticlePair) {
+        dVdR1[0] = (RealOpenMM) valueDerivExpressions[0][0].evaluate(variables);;
+        dVdR2[0] = -dVdR1[0];
         for (int i = 0; i < 3; i++) {
-            gradient1[0][i] = dVdR*deltaR[i]*rinv;
-            gradient2[0][i] = -gradient1[0][i];
-            forces[atom1][i] -= dEdV[0][atom1]*gradient1[0][i];
-            forces[atom2][i] -= dEdV[0][atom1]*gradient2[0][i];
+            forces[atom1][i] -= dEdV[0][atom1]*dVdR1[0]*deltaR[i];
+            forces[atom2][i] -= dEdV[0][atom1]*dVdR2[0]*deltaR[i];
         }
     }
     variables = globalParameters;
@@ -401,15 +421,12 @@ void ReferenceCustomGBIxn::calculateOnePairChainRule(int atom1, int atom2, RealO
         variables["z"] = atomCoordinates[atom1][2];
         for (int j = 0; j < i; j++) {
             RealOpenMM dVdV = (RealOpenMM) valueDerivExpressions[i][j].evaluate(variables);
-            for (int k = 0; k < 3; k++) {
-                gradient1[i][k] += dVdV*gradient1[j][k];
-                gradient2[i][k] += dVdV*gradient2[j][k];
-            }
+            dVdR1[i] += dVdV*dVdR1[j];
+            dVdR2[i] += dVdV*dVdR2[j];
         }
         for (int k = 0; k < 3; k++) {
-            gradient1[i][k] += (RealOpenMM) valueGradientExpressions[i][k].evaluate(variables);
-            forces[atom1][k] -= dEdV[i][atom1]*gradient1[i][k];
-            forces[atom2][k] -= dEdV[i][atom1]*gradient2[i][k];
+            forces[atom1][k] -= dEdV[i][atom1]*dVdR1[i]*deltaR[k];
+            forces[atom2][k] -= dEdV[i][atom1]*dVdR2[i]*deltaR[k];
         }
     }
 }

platforms/reference/tests/TestReferenceCustomGBForce.cpp

@@ -143,6 +143,95 @@ void testOBC(GBSAOBCForce::NonbondedMethod obcMethod, CustomGBForce::NonbondedMe
     }
 }
 
+void testMembrane() {
+    const int numMolecules = 70;
+    const int numParticles = numMolecules*2;
+    const double boxSize = 10.0;
+    ReferencePlatform platform;
+
+    // Create a system with an implicit membrane.
+
+    System system;
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(1.0);
+    }
+    system.setDefaultPeriodicBoxVectors(Vec3(boxSize, 0.0, 0.0), Vec3(0.0, boxSize, 0.0), Vec3(0.0, 0.0, boxSize));
+    CustomGBForce* custom = new CustomGBForce();
+    custom->setCutoffDistance(2.0);
+    custom->addPerParticleParameter("q");
+    custom->addPerParticleParameter("radius");
+    custom->addPerParticleParameter("scale");
+    custom->addGlobalParameter("thickness", 3);
+    custom->addGlobalParameter("solventDielectric", 78.3);
+    custom->addGlobalParameter("soluteDielectric", 1);
+    custom->addComputedValue("Imol", "step(r+sr2-or1)*0.5*(1/L-1/U+0.25*(1/U^2-1/L^2)*(r-sr2*sr2/r)+0.5*log(L/U)/r+C);"
+                             "U=r+sr2;"
+                             "C=2*(1/or1-1/L)*step(sr2-r-or1);"
+                             "L=max(or1, D);"
+                             "D=abs(r-sr2);"
+                             "sr2 = scale2*or2;"
+                             "or1 = radius1-0.009; or2 = radius2-0.009", CustomGBForce::ParticlePairNoExclusions);
+    custom->addComputedValue("Imem", "(1/radius+2*log(2)/thickness)/(1+exp(7.2*(abs(z)+radius-0.5*thickness)))", CustomGBForce::SingleParticle);
+    custom->addComputedValue("B", "1/(1/or-tanh(1*psi-0.8*psi^2+4.85*psi^3)/radius);"
+                             "psi=max(Imol,Imem)*or; or=radius-0.009", CustomGBForce::SingleParticle);
+    custom->addEnergyTerm("28.3919551*(radius+0.14)^2*(radius/B)^6-0.5*138.935456*(1/soluteDielectric-1/solventDielectric)*q^2/B", CustomGBForce::SingleParticle);
+    custom->addEnergyTerm("-138.935456*(1/soluteDielectric-1/solventDielectric)*q1*q2/f;"
+                          "f=sqrt(r^2+B1*B2*exp(-r^2/(4*B1*B2)))", CustomGBForce::ParticlePairNoExclusions);
+    vector<Vec3> positions(numParticles);
+    vector<Vec3> velocities(numParticles);
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    vector<double> params(3);
+    for (int i = 0; i < numMolecules; i++) {
+        if (i < numMolecules/2) {
+            params[0] = 1.0;
+            params[1] = 0.2;
+            params[2] = 0.5;
+            custom->addParticle(params);
+            params[0] = -1.0;
+            params[1] = 0.1;
+            custom->addParticle(params);
+        }
+        else {
+            params[0] = 1.0;
+            params[1] = 0.2;
+            params[2] = 0.8;
+            custom->addParticle(params);
+            params[0] = -1.0;
+            params[1] = 0.1;
+            custom->addParticle(params);
+        }
+        positions[2*i] = Vec3(boxSize*genrand_real2(sfmt), boxSize*genrand_real2(sfmt), boxSize*genrand_real2(sfmt));
+        positions[2*i+1] = Vec3(positions[2*i][0]+1.0, positions[2*i][1], positions[2*i][2]);
+        velocities[2*i] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt));
+        velocities[2*i+1] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt));
+    }
+    system.addForce(custom);
+    VerletIntegrator integrator(0.01);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+    State state = context.getState(State::Forces | State::Energy);
+    const vector<Vec3>& forces = state.getForces();
+
+    // Take a small step in the direction of the energy gradient and see whether the potential energy changes by the expected amount.
+
+    double norm = 0.0;
+    for (int i = 0; i < (int) forces.size(); ++i)
+        norm += forces[i].dot(forces[i]);
+    norm = std::sqrt(norm);
+    const double stepSize = 1e-3;
+    double step = stepSize/norm;
+    for (int i = 0; i < (int) positions.size(); ++i) {
+        Vec3 p = positions[i];
+        Vec3 f = forces[i];
+        positions[i] = Vec3(p[0]-f[0]*step, p[1]-f[1]*step, p[2]-f[2]*step);
+    }
+    context.setPositions(positions);
+    State state2 = context.getState(State::Energy);
+    ASSERT_EQUAL_TOL(norm, (state2.getPotentialEnergy()-state.getPotentialEnergy())/stepSize, 1e-2);
+}
+
 void testTabulatedFunction(bool interpolating) {
     ReferencePlatform platform;
     System system;
@@ -775,6 +864,7 @@ int main() {
         testOBC(GBSAOBCForce::NoCutoff, CustomGBForce::NoCutoff);
         testOBC(GBSAOBCForce::CutoffNonPeriodic, CustomGBForce::CutoffNonPeriodic);
         testOBC(GBSAOBCForce::CutoffPeriodic, CustomGBForce::CutoffPeriodic);
+        testMembrane();
         testTabulatedFunction(true);
         testTabulatedFunction(false);
         testMultipleChainRules();