Merge pull request #1889 from peastman/groupderiv

Fixed error using interaction groups with parameter derivatives

platforms/cuda/include/CudaKernels.h

@@ -778,7 +778,7 @@ private:
     std::vector<CudaArray*> tabulatedFunctions;
     double longRangeCoefficient;
     std::vector<double> longRangeCoefficientDerivs;
-    bool hasInitializedLongRangeCorrection, hasInitializedKernel;
+    bool hasInitializedLongRangeCorrection, hasInitializedKernel, hasParamDerivs;
     int numGroupThreadBlocks;
     CustomNonbondedForce* forceCopy;
     const System& system;

platforms/cuda/src/CudaKernels.cpp

@@ -2664,6 +2664,7 @@ void CudaCalcCustomNonbondedForceKernel::initInteractionGroups(const CustomNonbo
     
     // Create the kernel.
     
+    hasParamDerivs = (force.getNumEnergyParameterDerivatives() > 0);
     map<string, string> replacements;
     replacements["COMPUTE_INTERACTION"] = interactionSource;
     const string suffixes[] = {"x", "y", "z", "w"};
@@ -2687,6 +2688,8 @@ void CudaCalcCustomNonbondedForceKernel::initInteractionGroups(const CustomNonbo
         args << ", const " << tableTypes[i]<< "* __restrict__ table" << i;
     if (globals != NULL)
         args<<", const float* __restrict__ globals";
+    if (hasParamDerivs)
+        args << ", mixed* __restrict__ energyParamDerivs";
     replacements["PARAMETER_ARGUMENTS"] = args.str();
     stringstream load1;
     for (int i = 0; i < (int) buffers.size(); i++)
@@ -2718,6 +2721,19 @@ void CudaCalcCustomNonbondedForceKernel::initInteractionGroups(const CustomNonbo
         }
     }
     replacements["LOAD_ATOM2_PARAMETERS"] = load2.str();
+    stringstream initDerivs, saveDerivs;
+    const vector<string>& allParamDerivNames = cu.getEnergyParamDerivNames();
+    int numDerivs = allParamDerivNames.size();
+    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
+        string paramName = force.getEnergyParameterDerivativeName(i);
+        string derivVariable = cu.getNonbondedUtilities().addEnergyParameterDerivative(paramName);
+        initDerivs<<"mixed "<<derivVariable<<" = 0;\n";
+        for (int index = 0; index < numDerivs; index++)
+            if (allParamDerivNames[index] == paramName)
+                saveDerivs<<"energyParamDerivs[(blockIdx.x*blockDim.x+threadIdx.x)*"<<numDerivs<<"+"<<index<<"] += "<<derivVariable<<";\n";
+    }
+    replacements["INIT_DERIVATIVES"] = initDerivs.str();
+    replacements["SAVE_DERIVATIVES"] = saveDerivs.str();
     map<string, string> defines;
     if (force.getNonbondedMethod() != CustomNonbondedForce::NoCutoff)
         defines["USE_CUTOFF"] = "1";
@@ -2779,6 +2795,8 @@ double CudaCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool in
                 interactionGroupArgs.push_back(&function->getDevicePointer());
             if (globals != NULL)
                 interactionGroupArgs.push_back(&globals->getDevicePointer());
+            if (hasParamDerivs)
+                interactionGroupArgs.push_back(&cu.getEnergyParamDerivBuffer().getDevicePointer());
         }
         int forceThreadBlockSize = cu.getNonbondedUtilities().getForceThreadBlockSize();
         cu.executeKernel(interactionGroupKernel, &interactionGroupArgs[0], numGroupThreadBlocks*forceThreadBlockSize, forceThreadBlockSize);

platforms/cuda/src/kernels/customNonbondedGroups.cu

@@ -17,6 +17,7 @@ extern "C" __global__ void computeInteractionGroups(
     const unsigned int tgx = threadIdx.x & (TILE_SIZE-1); // index within the warp
     const unsigned int tbx = threadIdx.x - tgx;           // block warpIndex
     mixed energy = 0;
+    INIT_DERIVATIVES
     __shared__ AtomData localData[LOCAL_MEMORY_SIZE];
 
     const unsigned int startTile = FIRST_TILE+warp*(LAST_TILE-FIRST_TILE)/totalWarps;
@@ -58,6 +59,7 @@ extern "C" __global__ void computeInteractionGroups(
                 LOAD_ATOM2_PARAMETERS
                 real dEdR = 0.0f;
                 real tempEnergy = 0.0f;
+                const real interactionScale = 1.0f;
                 COMPUTE_INTERACTION
                 energy += tempEnergy;
                 delta *= dEdR;
@@ -82,4 +84,5 @@ extern "C" __global__ void computeInteractionGroups(
         atomicAdd(&forceBuffers[atom2+2*PADDED_NUM_ATOMS], static_cast<unsigned long long>((long long) (localData[threadIdx.x].fz*0x100000000)));
     }
     energyBuffer[blockIdx.x*blockDim.x+threadIdx.x] += energy;
+    SAVE_DERIVATIVES
 }

platforms/opencl/include/OpenCLKernels.h

@@ -758,7 +758,7 @@ private:
     std::vector<OpenCLArray*> tabulatedFunctions;
     double longRangeCoefficient;
     std::vector<double> longRangeCoefficientDerivs;
-    bool hasInitializedLongRangeCorrection, hasInitializedKernel;
+    bool hasInitializedLongRangeCorrection, hasInitializedKernel, hasParamDerivs;
     int numGroupThreadBlocks;
     CustomNonbondedForce* forceCopy;
     const System& system;

platforms/opencl/src/OpenCLKernels.cpp

@@ -2789,6 +2789,7 @@ void OpenCLCalcCustomNonbondedForceKernel::initInteractionGroups(const CustomNon
     
     // Create the kernel.
     
+    hasParamDerivs = (force.getNumEnergyParameterDerivatives() > 0);
     map<string, string> replacements;
     replacements["COMPUTE_INTERACTION"] = interactionSource;
     const string suffixes[] = {"x", "y", "z", "w"};
@@ -2812,6 +2813,8 @@ void OpenCLCalcCustomNonbondedForceKernel::initInteractionGroups(const CustomNon
         args << ", __global const " << tableTypes[i]<< "* restrict table" << i;
     if (globals != NULL)
         args<<", __global const float* restrict globals";
+    if (hasParamDerivs)
+        args << ", __global mixed* restrict energyParamDerivs";
     replacements["PARAMETER_ARGUMENTS"] = args.str();
     stringstream load1;
     for (int i = 0; i < (int) buffers.size(); i++)
@@ -2843,6 +2846,19 @@ void OpenCLCalcCustomNonbondedForceKernel::initInteractionGroups(const CustomNon
         }
     }
     replacements["LOAD_ATOM2_PARAMETERS"] = load2.str();
+    stringstream initDerivs, saveDerivs;
+    const vector<string>& allParamDerivNames = cl.getEnergyParamDerivNames();
+    int numDerivs = allParamDerivNames.size();
+    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
+        string paramName = force.getEnergyParameterDerivativeName(i);
+        string derivVariable = cl.getNonbondedUtilities().addEnergyParameterDerivative(paramName);
+        initDerivs<<"mixed "<<derivVariable<<" = 0;\n";
+        for (int index = 0; index < numDerivs; index++)
+            if (allParamDerivNames[index] == paramName)
+                saveDerivs<<"energyParamDerivs[get_global_id(0)*"<<numDerivs<<"+"<<index<<"] += "<<derivVariable<<";\n";
+    }
+    replacements["INIT_DERIVATIVES"] = initDerivs.str();
+    replacements["SAVE_DERIVATIVES"] = saveDerivs.str();
     map<string, string> defines;
     if (force.getNonbondedMethod() != CustomNonbondedForce::NoCutoff)
         defines["USE_CUTOFF"] = "1";
@@ -2902,6 +2918,8 @@ double OpenCLCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool
                 interactionGroupKernel.setArg<cl::Memory>(index++, function->getDeviceBuffer());
             if (globals != NULL)
                 interactionGroupKernel.setArg<cl::Buffer>(index++, globals->getDeviceBuffer());
+            if (hasParamDerivs)
+                interactionGroupKernel.setArg<cl::Memory>(index++, cl.getEnergyParamDerivBuffer().getDeviceBuffer());
         }
         setPeriodicBoxArgs(cl, interactionGroupKernel, 4);
         int forceThreadBlockSize = max(32, cl.getNonbondedUtilities().getForceThreadBlockSize());

platforms/opencl/src/kernels/customNonbondedGroups.cl

@@ -50,6 +50,7 @@ __kernel void computeInteractionGroups(
     const unsigned int tgx = get_local_id(0) & (TILE_SIZE-1); // index within the warp
     const unsigned int tbx = get_local_id(0) - tgx;           // block warpIndex
     mixed energy = 0;
+    INIT_DERIVATIVES
     __local AtomData localData[LOCAL_MEMORY_SIZE];
 
     const unsigned int startTile = FIRST_TILE+warp*(LAST_TILE-FIRST_TILE)/totalWarps;
@@ -93,6 +94,7 @@ __kernel void computeInteractionGroups(
                     LOAD_ATOM2_PARAMETERS
                     real dEdR = 0.0f;
                     real tempEnergy = 0.0f;
+                    const real interactionScale = 1.0f;
                     COMPUTE_INTERACTION
                     energy += tempEnergy;
                     delta *= dEdR;
@@ -125,4 +127,5 @@ __kernel void computeInteractionGroups(
 #endif
     }
     energyBuffer[get_global_id(0)] += energy;
+    SAVE_DERIVATIVES
 }

tests/TestCustomNonbondedForce.h

@@ -1188,6 +1188,47 @@ void testEnergyParameterDerivatives2() {
     ASSERT_EQUAL_TOL((energy1-energy2)/(2*delta), derivs["a"], 1e-4);
 }
 
+void testEnergyParameterDerivativesWithGroups() {
+    System system;
+    system.addParticle(1.0);
+    system.addParticle(1.0);
+    system.addParticle(1.0);
+    VerletIntegrator integrator(0.01);
+    CustomNonbondedForce* nonbonded = new CustomNonbondedForce("k*(r-r0)^2");
+    nonbonded->addGlobalParameter("r0", 0.0);
+    nonbonded->addGlobalParameter("k", 0.0);
+    nonbonded->addEnergyParameterDerivative("k");
+    nonbonded->addEnergyParameterDerivative("r0");
+    vector<double> parameters;
+    nonbonded->addParticle(parameters);
+    nonbonded->addParticle(parameters);
+    nonbonded->addParticle(parameters);
+    set<int> set1, set2;
+    set1.insert(1);
+    set2.insert(0);
+    set2.insert(2);
+    nonbonded->addInteractionGroup(set1, set2);
+    system.addForce(nonbonded);
+    Context context(system, integrator, platform);
+    vector<Vec3> positions(3);
+    positions[0] = Vec3(0, 2, 0);
+    positions[1] = Vec3(0, 0, 0);
+    positions[2] = Vec3(1, 0, 0);
+    context.setPositions(positions);
+    for (int i = 0; i < 10; i++) {
+        double r0 = 0.1*i;
+        double k = 10-i;
+        context.setParameter("r0", r0);
+        context.setParameter("k", k);
+        State state = context.getState(State::ParameterDerivatives);
+        map<string, double> derivs = state.getEnergyParameterDerivatives();
+        double dEdr0 = -2*k*((2-r0)+(1-r0));
+        double dEdk = (2-r0)*(2-r0) + (1-r0)*(1-r0);
+        ASSERT_EQUAL_TOL(dEdr0, derivs["r0"], 1e-5);
+        ASSERT_EQUAL_TOL(dEdk, derivs["k"], 1e-5);
+    }
+}
+
 void runPlatformTests();
 
 int main(int argc, char* argv[]) {
@@ -1217,6 +1258,7 @@ int main(int argc, char* argv[]) {
         testIllegalVariable();
         testEnergyParameterDerivatives();
         testEnergyParameterDerivatives2();
+        testEnergyParameterDerivativesWithGroups();
         runPlatformTests();
     }
     catch(const exception& e) {