CPU implementation of parameter derivatives for CustomGBForce

platforms/cpu/include/CpuCustomGBForce.h

@@ -47,14 +47,14 @@ private:
     const CpuNeighborList* neighborList;
     float periodicBoxSize[3];
     float cutoffDistance, cutoffDistance2;
+    int numValues, numParams;
     const std::vector<std::set<int> > exclusions;
-    std::vector<std::string> valueNames;
     std::vector<CustomGBForce::ComputationType> valueTypes;
-    std::vector<std::string> paramNames;
     std::vector<CustomGBForce::ComputationType> energyTypes;
     ThreadPool& threads;
     std::vector<ThreadData*> threadData;
     std::vector<double> threadEnergy;
+    std::vector<std::vector<std::vector<float> > > dValuedParam;
     // Workspace vectors
     std::vector<std::vector<float> > values, dEdV;
     // The following variables are used to make information accessible to the individual threads.
@@ -189,11 +189,13 @@ public:
                         const std::vector<Lepton::CompiledExpression>& valueExpressions,
                         const std::vector<std::vector<Lepton::CompiledExpression> >& valueDerivExpressions,
                         const std::vector<std::vector<Lepton::CompiledExpression> >& valueGradientExpressions,
+                        const std::vector<std::vector<Lepton::CompiledExpression> >& valueParamDerivExpressions,
                         const std::vector<std::string>& valueNames,
                         const std::vector<CustomGBForce::ComputationType>& valueTypes,
                         const std::vector<Lepton::CompiledExpression>& energyExpressions,
                         const std::vector<std::vector<Lepton::CompiledExpression> >& energyDerivExpressions,
                         const std::vector<std::vector<Lepton::CompiledExpression> >& energyGradientExpressions,
+                        const std::vector<std::vector<Lepton::CompiledExpression> >& energyParamDerivExpressions,
                         const std::vector<CustomGBForce::ComputationType>& energyTypes,
                         const std::vector<std::string>& parameterNames, ThreadPool& threads);
 
@@ -221,16 +223,17 @@ public:
     /**
      * Calculate custom GB ixn
      * 
-     * @param numberOfAtoms    number of atoms
-     * @param posq             atom coordinates
-     * @param atomParameters   atomParameters[atomIndex][paramterIndex]
-     * @param globalParameters the values of global parameters
-     * @param forces           force array (forces added)
-     * @param totalEnergy      total energy
+     * @param numberOfAtoms      number of atoms
+     * @param posq               atom coordinates
+     * @param atomParameters     atomParameters[atomIndex][paramterIndex]
+     * @param globalParameters   the values of global parameters
+     * @param forces             force array (forces added)
+     * @param totalEnergy        total energy
+     * @param energyParamDerivs  derivatives of the energy with respect to global parameters
      */
 
-    void calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters,
-                     std::map<std::string, double>& globalParameters, std::vector<AlignedArray<float> >& threadForce, bool includeForce, bool includeEnergy, double& totalEnergy);
+    void calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters, std::map<std::string, double>& globalParameters,
+            std::vector<AlignedArray<float> >& threadForce, bool includeForce, bool includeEnergy, double& totalEnergy, double* energyParamDerivs);
 };
 
 class CpuCustomGBForce::ThreadData {
@@ -239,19 +242,23 @@ public:
                const std::vector<Lepton::CompiledExpression>& valueExpressions,
                const std::vector<std::vector<Lepton::CompiledExpression> >& valueDerivExpressions,
                const std::vector<std::vector<Lepton::CompiledExpression> >& valueGradientExpressions,
+               const std::vector<std::vector<Lepton::CompiledExpression> >& valueParamDerivExpressions,
                const std::vector<std::string>& valueNames,
                const std::vector<Lepton::CompiledExpression>& energyExpressions,
                const std::vector<std::vector<Lepton::CompiledExpression> >& energyDerivExpressions,
                const std::vector<std::vector<Lepton::CompiledExpression> >& energyGradientExpressions,
+               const std::vector<std::vector<Lepton::CompiledExpression> >& energyParamDerivExpressions,
                const std::vector<std::string>& parameterNames);
     CompiledExpressionSet expressionSet;
     std::vector<Lepton::CompiledExpression> valueExpressions;
     std::vector<std::vector<Lepton::CompiledExpression> > valueDerivExpressions;
     std::vector<std::vector<Lepton::CompiledExpression> > valueGradientExpressions;
+    std::vector<std::vector<Lepton::CompiledExpression> > valueParamDerivExpressions;
     std::vector<int> valueIndex;
     std::vector<Lepton::CompiledExpression> energyExpressions;
     std::vector<std::vector<Lepton::CompiledExpression> > energyDerivExpressions;
     std::vector<std::vector<Lepton::CompiledExpression> > energyGradientExpressions;
+    std::vector<std::vector<Lepton::CompiledExpression> > energyParamDerivExpressions;
     std::vector<int> paramIndex;
     std::vector<int> particleParamIndex;
     std::vector<int> particleValueIndex;
@@ -260,6 +267,8 @@ public:
     // Workspace vectors
     std::vector<float> value0, dVdR1, dVdR2, dVdX, dVdY, dVdZ;
     std::vector<std::vector<float> > dEdV;
+    std::vector<std::vector<float> > dValue0dParam;
+    std::vector<float> energyParamDerivs;
 };
 
 } // namespace OpenMM

platforms/cpu/include/CpuKernels.h

@@ -398,7 +398,7 @@ private:
     RealOpenMM nonbondedCutoff;
     CpuCustomGBForce* ixn;
     std::vector<std::set<int> > exclusions;
-    std::vector<std::string> particleParameterNames, globalParameterNames, valueNames;
+    std::vector<std::string> particleParameterNames, globalParameterNames, energyParamDerivNames, valueNames;
     std::vector<OpenMM::CustomGBForce::ComputationType> valueTypes;
     std::vector<OpenMM::CustomGBForce::ComputationType> energyTypes;
     NonbondedMethod nonbondedMethod;

platforms/cpu/src/CpuCustomGBForce.cpp

@@ -47,13 +47,16 @@ CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threa
                       const vector<Lepton::CompiledExpression>& valueExpressions,
                       const vector<vector<Lepton::CompiledExpression> >& valueDerivExpressions,
                       const vector<vector<Lepton::CompiledExpression> >& valueGradientExpressions,
+                      const vector<vector<Lepton::CompiledExpression> >& valueParamDerivExpressions,
                       const vector<string>& valueNames,
                       const vector<Lepton::CompiledExpression>& energyExpressions,
                       const vector<vector<Lepton::CompiledExpression> >& energyDerivExpressions,
                       const vector<vector<Lepton::CompiledExpression> >& energyGradientExpressions,
+                      const vector<vector<Lepton::CompiledExpression> >& energyParamDerivExpressions,
                       const vector<string>& parameterNames) :
             valueExpressions(valueExpressions), valueDerivExpressions(valueDerivExpressions), valueGradientExpressions(valueGradientExpressions),
-            energyExpressions(energyExpressions), energyDerivExpressions(energyDerivExpressions), energyGradientExpressions(energyGradientExpressions) {
+            valueParamDerivExpressions(valueParamDerivExpressions), energyExpressions(energyExpressions), energyDerivExpressions(energyDerivExpressions),
+            energyGradientExpressions(energyGradientExpressions), energyParamDerivExpressions(energyParamDerivExpressions) {
     firstAtom = (threadIndex*(long long) numAtoms)/numThreads;
     lastAtom = ((threadIndex+1)*(long long) numAtoms)/numThreads;
     for (int i = 0; i < (int) valueExpressions.size(); i++)
@@ -64,6 +67,9 @@ CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threa
     for (int i = 0; i < (int) valueGradientExpressions.size(); i++)
         for (int j = 0; j < (int) valueGradientExpressions[i].size(); j++)
             expressionSet.registerExpression(this->valueGradientExpressions[i][j]);
+    for (int i = 0; i < (int) valueParamDerivExpressions.size(); i++)
+        for (int j = 0; j < (int) valueParamDerivExpressions[i].size(); j++)
+            expressionSet.registerExpression(this->valueParamDerivExpressions[i][j]);
     for (int i = 0; i < (int) energyExpressions.size(); i++)
         expressionSet.registerExpression(this->energyExpressions[i]);
     for (int i = 0; i < (int) energyDerivExpressions.size(); i++)
@@ -72,6 +78,9 @@ CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threa
     for (int i = 0; i < (int) energyGradientExpressions.size(); i++)
         for (int j = 0; j < (int) energyGradientExpressions[i].size(); j++)
             expressionSet.registerExpression(this->energyGradientExpressions[i][j]);
+    for (int i = 0; i < (int) energyParamDerivExpressions.size(); i++)
+        for (int j = 0; j < (int) energyParamDerivExpressions[i].size(); j++)
+            expressionSet.registerExpression(this->energyParamDerivExpressions[i][j]);
     xindex = expressionSet.getVariableIndex("x");
     yindex = expressionSet.getVariableIndex("y");
     zindex = expressionSet.getVariableIndex("z");
@@ -101,30 +110,37 @@ CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threa
     dVdZ.resize(valueDerivExpressions.size());
     dVdR1.resize(valueDerivExpressions.size());
     dVdR2.resize(valueDerivExpressions.size());
+    dValue0dParam.resize(valueParamDerivExpressions[0].size(), vector<float>(numAtoms));
+    energyParamDerivs.resize(valueParamDerivExpressions[0].size());
 }
 
 CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const std::vector<std::set<int> >& exclusions,
                      const vector<Lepton::CompiledExpression>& valueExpressions,
                      const vector<vector<Lepton::CompiledExpression> >& valueDerivExpressions,
                      const vector<vector<Lepton::CompiledExpression> >& valueGradientExpressions,
+                     const vector<vector<Lepton::CompiledExpression> >& valueParamDerivExpressions,
                      const vector<string>& valueNames,
                      const vector<CustomGBForce::ComputationType>& valueTypes,
                      const vector<Lepton::CompiledExpression>& energyExpressions,
                      const vector<vector<Lepton::CompiledExpression> >& energyDerivExpressions,
                      const vector<vector<Lepton::CompiledExpression> >& energyGradientExpressions,
+                     const vector<vector<Lepton::CompiledExpression> >& energyParamDerivExpressions,
                      const vector<CustomGBForce::ComputationType>& energyTypes,
                      const vector<string>& parameterNames, ThreadPool& threads) :
-            exclusions(exclusions), cutoff(false), periodic(false), valueNames(valueNames), valueTypes(valueTypes),
-            energyTypes(energyTypes), paramNames(parameterNames), threads(threads) {
+            exclusions(exclusions), cutoff(false), periodic(false), valueTypes(valueTypes), energyTypes(energyTypes), numValues(valueNames.size()),
+            numParams(parameterNames.size()), threads(threads) {
     for (int i = 0; i < threads.getNumThreads(); i++)
-        threadData.push_back(new ThreadData(numAtoms, threads.getNumThreads(), i, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueNames,
-                      energyExpressions, energyDerivExpressions, energyGradientExpressions, parameterNames));
-    values.resize(valueNames.size());
-    dEdV.resize(valueNames.size());
+        threadData.push_back(new ThreadData(numAtoms, threads.getNumThreads(), i, valueExpressions, valueDerivExpressions, valueGradientExpressions,
+                valueParamDerivExpressions, valueNames, energyExpressions, energyDerivExpressions, energyGradientExpressions, energyParamDerivExpressions, parameterNames));
+    values.resize(numValues);
+    dEdV.resize(numValues);
     for (int i = 0; i < (int) values.size(); i++) {
         values[i].resize(numAtoms);
         dEdV[i].resize(numAtoms);
     }
+    dValuedParam.resize(numValues);
+    for (int i = 0; i < numValues; i++)
+        dValuedParam[i].resize(valueParamDerivExpressions[0].size(), vector<float>(numAtoms));
 }
 
 CpuCustomGBForce::~CpuCustomGBForce() {
@@ -153,7 +169,7 @@ void CpuCustomGBForce::setPeriodic(RealVec& boxSize) {
 
 void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters,
                                            map<string, double>& globalParameters, vector<AlignedArray<float> >& threadForce,
-                                           bool includeForce, bool includeEnergy, double& totalEnergy) {
+                                           bool includeForce, bool includeEnergy, double& totalEnergy, double* energyParamDerivs) {
     // Record the parameters for the threads.
     
     this->numberOfAtoms = numberOfAtoms;
@@ -173,12 +189,20 @@ void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM**
     gmx_atomic_set(&counter, 0);
     threads.execute(task);
     threads.waitForThreads();
+
+    // Sum derivatives of the first computed value with respect to global parameters.
+
+    bool hasParamDerivs = (threadData[0]->dValue0dParam.size() > 0);
+    if (hasParamDerivs) {
+        threads.resumeThreads();
+        threads.waitForThreads();
+    }
     
     // Calculate the remaining computed values.
     
     threads.resumeThreads();
     threads.waitForThreads();
-
+    
     // Calculate the energy terms.
 
     for (int i = 0; i < (int) threadData[0]->energyExpressions.size(); i++) {
@@ -205,6 +229,10 @@ void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM**
         for (int i = 0; i < numThreads; i++)
             totalEnergy += threadEnergy[i];
     }
+    if (hasParamDerivs)
+        for (int i = 0; i < threads.getNumThreads(); i++)
+            for (int j = 0; j < threadData[j]->energyParamDerivs.size(); j++)
+                energyParamDerivs[j] += threadData[i]->energyParamDerivs[j];
 }
 
 void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex) {
@@ -224,11 +252,28 @@ void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex)
 
     for (int i = 0; i < (int) data.value0.size(); i++)
         data.value0[i] = 0.0f;
+    for (int i = 0; i < (int) data.dValue0dParam.size(); i++)
+        for (int j = 0; j < (int) data.dValue0dParam[i].size(); j++)
+            data.dValue0dParam[i][j] = 0.0;
     if (valueTypes[0] == CustomGBForce::ParticlePair)
         calculateParticlePairValue(0, data, numberOfAtoms, posq, atomParameters, true, boxSize, invBoxSize);
     else
         calculateParticlePairValue(0, data, numberOfAtoms, posq, atomParameters, false, boxSize, invBoxSize);
     threads.syncThreads();
+    
+    // Sum derivatives of the first computed value with respect to global parameters.
+    
+    bool hasParamDerivs = (data.dValue0dParam.size() > 0);
+    if (hasParamDerivs) {
+        for (int j = 0; j < data.dValue0dParam.size(); j++)
+            for (int k = data.firstAtom; k < data.lastAtom; k++) {
+                float sum = 0.0f;
+                for (int m = 0; m < threadData.size(); m++)
+                    sum += threadData[m]->dValue0dParam[j][k];
+                dValuedParam[0][j][k] = sum;
+            }
+        threads.syncThreads();
+    }
 
     // Sum the first computed value and calculate the remaining ones.
 
@@ -241,11 +286,23 @@ void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex)
         data.expressionSet.setVariable(data.xindex, posq[4*atom]);
         data.expressionSet.setVariable(data.yindex, posq[4*atom+1]);
         data.expressionSet.setVariable(data.zindex, posq[4*atom+2]);
-        for (int j = 0; j < (int) paramNames.size(); j++)
+        for (int j = 0; j < numParams; j++)
             data.expressionSet.setVariable(data.paramIndex[j], atomParameters[atom][j]);
         for (int i = 1; i < numValues; i++) {
             data.expressionSet.setVariable(data.valueIndex[i-1], values[i-1][atom]);
             values[i][atom] = (float) data.valueExpressions[i].evaluate();
+
+            // Calculate derivatives with respect to parameters.
+
+            if (hasParamDerivs) {
+                for (int j = 0; j < data.valueParamDerivExpressions[i].size(); j++)
+                    dValuedParam[i][j][atom] = data.valueParamDerivExpressions[i][j].evaluate();
+                for (int j = 0; j < i; j++) {
+                    float dVdV = data.valueDerivExpressions[i][j].evaluate();
+                    for (int k = 0; k < data.valueParamDerivExpressions[i].size(); k++)
+                        dValuedParam[i][k][atom] += dVdV*dValuedParam[j][k][atom];
+                }
+            }
         }
     }
     threads.syncThreads();
@@ -254,7 +311,9 @@ void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex)
 
     for (int i = 0; i < (int) data.dEdV.size(); i++)
         for (int j = 0; j < (int) data.dEdV[i].size(); j++)
-            data.dEdV[i][j] = 0.0;
+            data.dEdV[i][j] = 0.0f;
+    for (int i = 0; i < (int) data.energyParamDerivs.size(); i++)
+        data.energyParamDerivs[i] = 0.0f;
     for (int termIndex = 0; termIndex < (int) data.energyExpressions.size(); termIndex++) {
         if (energyTypes[termIndex] == CustomGBForce::SingleParticle)
             calculateSingleParticleEnergyTerm(termIndex, data, numberOfAtoms, posq, atomParameters, forces, energy);
@@ -339,7 +398,7 @@ void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, Th
     if (cutoff && r2 >= cutoffDistance2)
         return;
     float r = sqrtf(r2);
-    for (int i = 0; i < (int) paramNames.size(); i++) {
+    for (int i = 0; i < numParams; i++) {
         data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
         data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
     }
@@ -349,6 +408,11 @@ void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, Th
         data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
     }
     valueArray[atom1] += (float) data.valueExpressions[index].evaluate();
+    
+    // Calculate derivatives with respect to parameters.
+    
+    for (int i = 0; i < data.valueParamDerivExpressions[index].size(); i++)
+        data.dValue0dParam[i][atom1] += data.valueParamDerivExpressions[index][i].evaluate();
 }
 
 void CpuCustomGBForce::calculateSingleParticleEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq,
@@ -357,17 +421,22 @@ void CpuCustomGBForce::calculateSingleParticleEnergyTerm(int index, ThreadData&
         data.expressionSet.setVariable(data.xindex, posq[4*i]);
         data.expressionSet.setVariable(data.yindex, posq[4*i+1]);
         data.expressionSet.setVariable(data.zindex, posq[4*i+2]);
-        for (int j = 0; j < (int) paramNames.size(); j++)
+        for (int j = 0; j < numParams; j++)
             data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
-        for (int j = 0; j < (int) valueNames.size(); j++)
+        for (int j = 0; j < (int) values.size(); j++)
             data.expressionSet.setVariable(data.valueIndex[j], values[j][i]);
         if (includeEnergy)
             totalEnergy += (float) data.energyExpressions[index].evaluate();
-        for (int j = 0; j < (int) valueNames.size(); j++)
+        for (int j = 0; j < (int) values.size(); j++)
             data.dEdV[j][i] += (float) data.energyDerivExpressions[index][j].evaluate();
         forces[4*i+0] -= (float) data.energyGradientExpressions[index][0].evaluate();
         forces[4*i+1] -= (float) data.energyGradientExpressions[index][1].evaluate();
         forces[4*i+2] -= (float) data.energyGradientExpressions[index][2].evaluate();
+        
+        // Compute derivatives with respect to parameters.
+        
+        for (int k = 0; k < data.energyParamDerivExpressions[index].size(); k++)
+            data.energyParamDerivs[k] += data.energyParamDerivExpressions[index][k].evaluate();
     }
 }
 
@@ -428,12 +497,12 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom
 
     // Record variables for evaluating expressions.
 
-    for (int i = 0; i < (int) paramNames.size(); i++) {
+    for (int i = 0; i < numParams; i++) {
         data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
         data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
     }
     data.expressionSet.setVariable(data.rindex, r);
-    for (int i = 0; i < (int) valueNames.size(); i++) {
+    for (int i = 0; i < (int) values.size(); i++) {
         data.expressionSet.setVariable(data.particleValueIndex[i*2], values[i][atom1]);
         data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
     }
@@ -447,10 +516,15 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom
     fvec4 result = deltaR*dEdR;
     (fvec4(forces+4*atom1)-result).store(forces+4*atom1);
     (fvec4(forces+4*atom2)+result).store(forces+4*atom2);
-    for (int i = 0; i < (int) valueNames.size(); i++) {
+    for (int i = 0; i < (int) values.size(); i++) {
         data.dEdV[i][atom1] += (float) data.energyDerivExpressions[index][2*i+1].evaluate();
         data.dEdV[i][atom2] += (float) data.energyDerivExpressions[index][2*i+2].evaluate();
     }
+        
+    // Compute derivatives with respect to parameters.
+
+    for (int i = 0; i < data.energyParamDerivExpressions[index].size(); i++)
+        data.energyParamDerivs[i] += data.energyParamDerivExpressions[index][i].evaluate();
 }
 
 void CpuCustomGBForce::calculateChainRuleForces(ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
@@ -500,9 +574,9 @@ void CpuCustomGBForce::calculateChainRuleForces(ThreadData& data, int numAtoms,
         data.expressionSet.setVariable(data.xindex, posq[4*i]);
         data.expressionSet.setVariable(data.yindex, posq[4*i+1]);
         data.expressionSet.setVariable(data.zindex, posq[4*i+2]);
-        for (int j = 0; j < (int) paramNames.size(); j++)
+        for (int j = 0; j < numParams; j++)
             data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
-        for (int j = 1; j < (int) valueNames.size(); j++) {
+        for (int j = 1; j < (int) values.size(); j++) {
             data.expressionSet.setVariable(data.valueIndex[j-1], values[j-1][i]);
             data.dVdX[j] = 0.0;
             data.dVdY[j] = 0.0;
@@ -521,6 +595,13 @@ void CpuCustomGBForce::calculateChainRuleForces(ThreadData& data, int numAtoms,
             forces[4*i+2] -= dEdV[j][i]*data.dVdZ[j];
         }
     }
+        
+    // Compute chain rule terms for derivatives with respect to parameters.
+
+    for (int i = data.firstAtom; i < data.lastAtom; i++)
+        for (int j = 0; j < data.valueIndex.size(); j++)
+            for (int k = 0; k < dValuedParam[j].size(); k++)
+                data.energyParamDerivs[k] += dEdV[j][i]*dValuedParam[j][k][i];
 }
 
 void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
@@ -538,7 +619,7 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadDat
 
     // Record variables for evaluating expressions.
 
-    for (int i = 0; i < (int) paramNames.size(); i++) {
+    for (int i = 0; i < numParams; i++) {
         data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
         data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
         data.expressionSet.setVariable(data.paramIndex[i], atomParameters[atom1][i]);
@@ -562,7 +643,7 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadDat
         f1 -= deltaR*(dEdV[0][atom1]*data.dVdR1[0]);
         f2 -= deltaR*(dEdV[0][atom1]*data.dVdR2[0]);
     }
-    for (int i = 1; i < (int) valueNames.size(); i++) {
+    for (int i = 1; i < (int) values.size(); i++) {
         data.expressionSet.setVariable(data.valueIndex[i], values[i][atom1]);
         data.dVdR1[i] = 0.0;
         data.dVdR2[i] = 0.0;

platforms/cpu/src/CpuKernels.cpp

@@ -1045,6 +1045,7 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB
 
     vector<vector<Lepton::CompiledExpression> > valueDerivExpressions(force.getNumComputedValues());
     vector<vector<Lepton::CompiledExpression> > valueGradientExpressions(force.getNumComputedValues());
+    vector<vector<Lepton::CompiledExpression> > valueParamDerivExpressions(force.getNumComputedValues());
     vector<Lepton::CompiledExpression> valueExpressions;
     vector<Lepton::CompiledExpression> energyExpressions;
     set<string> particleVariables, pairVariables;
@@ -1079,6 +1080,11 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB
                 valueDerivExpressions[i].push_back(ex.differentiate(valueNames[j]).createCompiledExpression());
             validateVariables(ex.getRootNode(), particleVariables);
         }
+        for (int j = 0; j < force.getNumEnergyParameterDerivatives(); j++) {
+            string param = force.getEnergyParameterDerivativeName(j);
+            energyParamDerivNames.push_back(param);
+            valueParamDerivExpressions[i].push_back(ex.differentiate(param).createCompiledExpression());
+        }
         particleVariables.insert(name);
         pairVariables.insert(name+"1");
         pairVariables.insert(name+"2");
@@ -1088,6 +1094,7 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB
 
     vector<vector<Lepton::CompiledExpression> > energyDerivExpressions(force.getNumEnergyTerms());
     vector<vector<Lepton::CompiledExpression> > energyGradientExpressions(force.getNumEnergyTerms());
+    vector<vector<Lepton::CompiledExpression> > energyParamDerivExpressions(force.getNumEnergyTerms());
     for (int i = 0; i < force.getNumEnergyTerms(); i++) {
         string expression;
         CustomGBForce::ComputationType type;
@@ -1111,14 +1118,17 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB
                 validateVariables(ex.getRootNode(), pairVariables);
             }
         }
+        for (int j = 0; j < force.getNumEnergyParameterDerivatives(); j++)
+            energyParamDerivExpressions[i].push_back(ex.differentiate(force.getEnergyParameterDerivativeName(j)).createCompiledExpression());
     }
 
     // Delete the custom functions.
 
     for (map<string, Lepton::CustomFunction*>::iterator iter = functions.begin(); iter != functions.end(); iter++)
         delete iter->second;
-    ixn = new CpuCustomGBForce(numParticles, exclusions, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueNames, valueTypes, energyExpressions,
-        energyDerivExpressions, energyGradientExpressions, energyTypes, particleParameterNames, data.threads);
+    ixn = new CpuCustomGBForce(numParticles, exclusions, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueParamDerivExpressions,
+        valueNames, valueTypes, energyExpressions, energyDerivExpressions, energyGradientExpressions, energyParamDerivExpressions, energyTypes,
+        particleParameterNames, data.threads);
     data.isPeriodic = (force.getNonbondedMethod() == CustomGBForce::CutoffPeriodic);
 }
 
@@ -1135,7 +1145,11 @@ double CpuCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeFor
     map<string, double> globalParameters;
     for (int i = 0; i < (int) globalParameterNames.size(); i++)
         globalParameters[globalParameterNames[i]] = context.getParameter(globalParameterNames[i]);
-    ixn->calculateIxn(numParticles, &data.posq[0], particleParamArray, globalParameters, data.threadForce, includeForces, includeEnergy, energy);
+    vector<double> energyParamDerivValues(energyParamDerivNames.size()+1, 0.0);
+    ixn->calculateIxn(numParticles, &data.posq[0], particleParamArray, globalParameters, data.threadForce, includeForces, includeEnergy, energy, &energyParamDerivValues[0]);
+    map<string, double>& energyParamDerivs = extractEnergyParameterDerivatives(context);
+    for (int i = 0; i < energyParamDerivNames.size(); i++)
+        energyParamDerivs[energyParamDerivNames[i]] += energyParamDerivValues[i];
     return energy;
 }
 

platforms/reference/include/ReferenceCustomGBIxn.h

@@ -170,7 +170,7 @@ class ReferenceCustomGBIxn {
          --------------------------------------------------------------------------------------- */
 
       void calculateChainRuleForces(int numAtoms, std::vector<OpenMM::RealVec>& atomCoordinates, RealOpenMM** atomParameters,
-                                      const std::vector<std::set<int> >& exclusions, std::vector<OpenMM::RealVec>& forces);
+                                      const std::vector<std::set<int> >& exclusions, std::vector<OpenMM::RealVec>& forces, double* energyParamDerivs);
 
       /**---------------------------------------------------------------------------------------
 

platforms/reference/src/SimTKReference/ReferenceCustomGBIxn.cpp

@@ -190,7 +190,7 @@ void ReferenceCustomGBIxn::calculateIxn(int numberOfAtoms, vector<RealVec>& atom
 
     // Apply the chain rule to evaluate forces.
 
-    calculateChainRuleForces(numberOfAtoms, atomCoordinates, atomParameters, exclusions, forces);
+    calculateChainRuleForces(numberOfAtoms, atomCoordinates, atomParameters, exclusions, forces, energyParamDerivs);
 }
 
 void ReferenceCustomGBIxn::calculateSingleParticleValue(int index, int numAtoms, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters) {
@@ -204,11 +204,11 @@ void ReferenceCustomGBIxn::calculateSingleParticleValue(int index, int numAtoms,
         for (int j = 0; j < index; j++)
             expressionSet.setVariable(valueIndex[j], values[j][i]);
         values[index][i] = (RealOpenMM) valueExpressions[index].evaluate();
-        for (int j = 0; j < valueParamDerivExpressions[index].size(); j++)
-            dValuedParam[index][j][i] += valueParamDerivExpressions[index][j].evaluate();
 
         // Calculate derivatives with respect to parameters.
 
+        for (int j = 0; j < valueParamDerivExpressions[index].size(); j++)
+            dValuedParam[index][j][i] += valueParamDerivExpressions[index][j].evaluate();
         for (int j = 0; j < index; j++) {
             RealOpenMM dVdV = valueDerivExpressions[index][j].evaluate();
             for (int k = 0; k < valueParamDerivExpressions[index].size(); k++)
@@ -296,11 +296,8 @@ void ReferenceCustomGBIxn::calculateSingleParticleEnergyTerm(int index, int numA
         
         // Compute derivatives with respect to parameters.
         
-        for (int k = 0; k < energyParamDerivExpressions[index].size(); k++) {
+        for (int k = 0; k < energyParamDerivExpressions[index].size(); k++)
             energyParamDerivs[k] += energyParamDerivExpressions[index][k].evaluate();
-            for (int j = 0; j < (int) valueIndex.size(); j++)
-                energyParamDerivs[k] += dEdV[j][i]*dValuedParam[j][k][i];
-        }
     }
 }
 
@@ -376,7 +373,7 @@ void ReferenceCustomGBIxn::calculateOnePairEnergyTerm(int index, int atom1, int
 }
 
 void ReferenceCustomGBIxn::calculateChainRuleForces(int numAtoms, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters,
-        const vector<set<int> >& exclusions, vector<RealVec>& forces) {
+        const vector<set<int> >& exclusions, vector<RealVec>& forces, double* energyParamDerivs) {
     if (cutoff) {
         // Loop over all pairs in the neighbor list.
 
@@ -426,6 +423,13 @@ void ReferenceCustomGBIxn::calculateChainRuleForces(int numAtoms, vector<RealVec
             forces[i][2] -= dEdV[j][i]*dVdZ[j];
         }
     }
+        
+    // Compute chain rule terms for derivatives with respect to parameters.
+
+    for (int i = 0; i < numAtoms; i++)
+        for (int j = 0; j < (int) valueIndex.size(); j++)
+            for (int k = 0; k < dValuedParam[j].size(); k++)
+                energyParamDerivs[k] += dEdV[j][i]*dValuedParam[j][k][i];
 }
 
 void ReferenceCustomGBIxn::calculateOnePairChainRule(int atom1, int atom2, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters,