CUDA implementation of new features in AmoebaVdwForce

platforms/cuda/src/kernels/nonbonded.cu

@@ -18,6 +18,10 @@ static __inline__ __device__ float real_shfl(float var, int srcLane) {
     return SHFL(var, srcLane);
 }
 
+static __inline__ __device__ float real_shfl(int var, int srcLane) {
+    return SHFL(var, srcLane);
+}
+
 static __inline__ __device__ double real_shfl(double var, int srcLane) {
     int hi, lo;
     asm volatile("mov.b64 { %0, %1 }, %2;" : "=r"(lo), "=r"(hi) : "d"(var));

plugins/amoeba/platforms/cuda/src/AmoebaCudaKernels.cpp

@@ -2366,23 +2366,36 @@ CudaCalcAmoebaVdwForceKernel::~CudaCalcAmoebaVdwForceKernel() {
 
 void CudaCalcAmoebaVdwForceKernel::initialize(const System& system, const AmoebaVdwForce& force) {
     cu.setAsCurrent();
-    sigmaEpsilon.initialize<float2>(cu, cu.getPaddedNumAtoms(), "sigmaEpsilon");
-    bondReductionAtoms.initialize<int>(cu, cu.getPaddedNumAtoms(), "bondReductionAtoms");
-    bondReductionFactors.initialize<float>(cu, cu.getPaddedNumAtoms(), "bondReductionFactors");
-    tempPosq.initialize(cu, cu.getPaddedNumAtoms(), cu.getUseDoublePrecision() ? sizeof(double4) : sizeof(float4), "tempPosq");
-    tempForces.initialize<long long>(cu, 3*cu.getPaddedNumAtoms(), "tempForces");
+    int paddedNumAtoms = cu.getPaddedNumAtoms();
+    bondReductionAtoms.initialize<int>(cu, paddedNumAtoms, "bondReductionAtoms");
+    bondReductionFactors.initialize<float>(cu, paddedNumAtoms, "bondReductionFactors");
+    tempPosq.initialize(cu, paddedNumAtoms, cu.getUseDoublePrecision() ? sizeof(double4) : sizeof(float4), "tempPosq");
+    tempForces.initialize<long long>(cu, 3*paddedNumAtoms, "tempForces");
     
     // Record atom parameters.
-    vector<float2> sigmaEpsilonVec(cu.getPaddedNumAtoms(), make_float2(0, 1));
-    vector<float> isAlchemicalVec(cu.getPaddedNumAtoms(), 0);
-    vector<int> bondReductionAtomsVec(cu.getPaddedNumAtoms(), 0);
-    vector<float> bondReductionFactorsVec(cu.getPaddedNumAtoms(), 0);
+    vector<int> atomTypeVec;
+    vector<vector<double> > sigmaMatrix, epsilonMatrix;
+    AmoebaVdwForceImpl::createParameterMatrix(force, atomTypeVec, sigmaMatrix, epsilonMatrix);
+    atomTypeVec.resize(paddedNumAtoms, 0);
+    int numTypes = sigmaMatrix.size();
+    atomType.initialize<int>(cu, paddedNumAtoms, "atomType");
+    sigmaEpsilon.initialize<float2>(cu, numTypes*numTypes, "sigmaEpsilon");
+    vector<float2> sigmaEpsilonVec(sigmaEpsilon.getSize());
+    for (int i = 0; i < numTypes; i++)
+        for (int j = 0; j < numTypes; j++)
+            sigmaEpsilonVec[i*numTypes+j] = make_float2((float) sigmaMatrix[i][j], (float) epsilonMatrix[i][j]);
+    atomType.upload(atomTypeVec);
+    sigmaEpsilon.upload(sigmaEpsilonVec);
+    
+    vector<float> isAlchemicalVec(paddedNumAtoms, 0);
+    vector<int> bondReductionAtomsVec(paddedNumAtoms, 0);
+    vector<float> bondReductionFactorsVec(paddedNumAtoms, 0);
     vector<vector<int> > exclusions(cu.getNumAtoms());
 
     // Handle Alchemical parameters.
     hasAlchemical = force.getAlchemicalMethod() != AmoebaVdwForce::None;
     if (hasAlchemical) {
-       isAlchemical.initialize<float>(cu, cu.getPaddedNumAtoms(), "isAlchemical");
+       isAlchemical.initialize<float>(cu, paddedNumAtoms, "isAlchemical");
        vdwLambda.initialize<float>(cu, 1, "vdwLambda");
        CHECK_RESULT(cuMemHostAlloc(&vdwLambdaPinnedBuffer, sizeof(float), 0), "Error allocating vdwLambda pinned buffer");
     }
@@ -2392,14 +2405,12 @@ void CudaCalcAmoebaVdwForceKernel::initialize(const System& system, const Amoeba
         double sigma, epsilon, reductionFactor;
         bool alchemical;
         force.getParticleParameters(i, ivIndex, sigma, epsilon, reductionFactor, alchemical, type);
-        sigmaEpsilonVec[i] = make_float2((float) sigma, (float) epsilon);
         isAlchemicalVec[i] = (alchemical) ? 1.0f : 0.0f;
         bondReductionAtomsVec[i] = ivIndex;
         bondReductionFactorsVec[i] = (float) reductionFactor;
         force.getParticleExclusions(i, exclusions[i]);
         exclusions[i].push_back(i);
     }
-    sigmaEpsilon.upload(sigmaEpsilonVec);
     bondReductionAtoms.upload(bondReductionAtomsVec);
     bondReductionFactors.upload(bondReductionFactorsVec);
     if (force.getUseDispersionCorrection())
@@ -2412,7 +2423,8 @@ void CudaCalcAmoebaVdwForceKernel::initialize(const System& system, const Amoeba
     // this force, so it will have its own neighbor list and interaction kernel.
     
     nonbonded = new CudaNonbondedUtilities(cu);
-    nonbonded->addParameter(CudaNonbondedUtilities::ParameterInfo("sigmaEpsilon", "float", 2, sizeof(float2), sigmaEpsilon.getDevicePointer()));
+    nonbonded->addParameter(CudaNonbondedUtilities::ParameterInfo("atomType", "int", 1, sizeof(int), atomType.getDevicePointer()));
+    nonbonded->addArgument(CudaNonbondedUtilities::ParameterInfo("sigmaEpsilon", "float", 2, sizeof(float2), sigmaEpsilon.getDevicePointer()));
 
     if (hasAlchemical) {
        isAlchemical.upload(isAlchemicalVec);
@@ -2425,33 +2437,11 @@ void CudaCalcAmoebaVdwForceKernel::initialize(const System& system, const Amoeba
     
     // Create the interaction kernel.
     
-    map<string, string> replacements;
-    string sigmaCombiningRule = force.getSigmaCombiningRule();
-    if (sigmaCombiningRule == "ARITHMETIC")
-        replacements["SIGMA_COMBINING_RULE"] = "1";
-    else if (sigmaCombiningRule == "GEOMETRIC")
-        replacements["SIGMA_COMBINING_RULE"] = "2";
-    else if (sigmaCombiningRule == "CUBIC-MEAN")
-        replacements["SIGMA_COMBINING_RULE"] = "3";
-    else
-        throw OpenMMException("Illegal combining rule for sigma: "+sigmaCombiningRule);
-    string epsilonCombiningRule = force.getEpsilonCombiningRule();
-    if (epsilonCombiningRule == "ARITHMETIC")
-        replacements["EPSILON_COMBINING_RULE"] = "1";
-    else if (epsilonCombiningRule == "GEOMETRIC")
-        replacements["EPSILON_COMBINING_RULE"] = "2";
-    else if (epsilonCombiningRule =="HARMONIC")
-        replacements["EPSILON_COMBINING_RULE"] = "3";
-    else if (epsilonCombiningRule == "W-H")
-        replacements["EPSILON_COMBINING_RULE"] = "4";
-    else if (epsilonCombiningRule == "HHG")
-        replacements["EPSILON_COMBINING_RULE"] = "5";
-    else
-        throw OpenMMException("Illegal combining rule for sigma: "+sigmaCombiningRule);
-
     replacements["VDW_ALCHEMICAL_METHOD"] = cu.intToString(force.getAlchemicalMethod()); 
     replacements["VDW_SOFTCORE_POWER"] = cu.intToString(force.getSoftcorePower());
     replacements["VDW_SOFTCORE_ALPHA"] = cu.doubleToString(force.getSoftcoreAlpha()); 
+    replacements["POTENTIAL_FUNCTION"] = cu.intToString(force.getPotentialFunction());
+    replacements["NUM_TYPES"] = cu.intToString(numTypes);
 
     double cutoff = force.getCutoffDistance();
     double taperCutoff = cutoff*0.9;
@@ -2467,7 +2457,7 @@ void CudaCalcAmoebaVdwForceKernel::initialize(const System& system, const Amoeba
     // Create the other kernels.
     
     map<string, string> defines;
-    defines["PADDED_NUM_ATOMS"] = cu.intToString(cu.getPaddedNumAtoms());
+    defines["PADDED_NUM_ATOMS"] = cu.intToString(paddedNumAtoms);
     CUmodule module = cu.createModule(CudaAmoebaKernelSources::amoebaVdwForce1, defines);
     prepareKernel = cu.getKernel(module, "prepareToComputeForce");
     spreadKernel = cu.getKernel(module, "spreadForces");
@@ -2512,8 +2502,21 @@ void CudaCalcAmoebaVdwForceKernel::copyParametersToContext(ContextImpl& context,
     if (force.getNumParticles() != cu.getNumAtoms())
         throw OpenMMException("updateParametersInContext: The number of particles has changed");
     
+    vector<int> atomTypeVec;
+    vector<vector<double> > sigmaMatrix, epsilonMatrix;
+    AmoebaVdwForceImpl::createParameterMatrix(force, atomTypeVec, sigmaMatrix, epsilonMatrix);
+    atomTypeVec.resize(cu.getPaddedNumAtoms(), 0);
+    int numTypes = sigmaMatrix.size();
+    if (sigmaEpsilon.getSize() != numTypes*numTypes)
+        throw OpenMMException("updateParametersInContext: The number of particle types has changed");
+    vector<float2> sigmaEpsilonVec(sigmaEpsilon.getSize());
+    for (int i = 0; i < numTypes; i++)
+        for (int j = 0; j < numTypes; j++)
+            sigmaEpsilonVec[i*numTypes+j] = make_float2((float) sigmaMatrix[i][j], (float) epsilonMatrix[i][j]);
+    atomType.upload(atomTypeVec);
+    sigmaEpsilon.upload(sigmaEpsilonVec);
+
     // Record the per-particle parameters.
-    vector<float2> sigmaEpsilonVec(cu.getPaddedNumAtoms(), make_float2(0, 1));
     vector<float> isAlchemicalVec(cu.getPaddedNumAtoms(), 0);
     vector<int> bondReductionAtomsVec(cu.getPaddedNumAtoms(), 0);
     vector<float> bondReductionFactorsVec(cu.getPaddedNumAtoms(), 0);
@@ -2522,12 +2525,10 @@ void CudaCalcAmoebaVdwForceKernel::copyParametersToContext(ContextImpl& context,
         double sigma, epsilon, reductionFactor;
         bool alchemical;
         force.getParticleParameters(i, ivIndex, sigma, epsilon, reductionFactor, alchemical, type);
-        sigmaEpsilonVec[i] = make_float2((float) sigma, (float) epsilon);
         isAlchemicalVec[i] = (alchemical) ? 1.0f : 0.0f;
         bondReductionAtomsVec[i] = ivIndex;
         bondReductionFactorsVec[i] = (float) reductionFactor;
     }
-    sigmaEpsilon.upload(sigmaEpsilonVec);
     if (hasAlchemical) isAlchemical.upload(isAlchemicalVec);
     bondReductionAtoms.upload(bondReductionAtomsVec);
     bondReductionFactors.upload(bondReductionFactorsVec);

plugins/amoeba/platforms/cuda/src/AmoebaCudaKernels.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-2018 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2020 Stanford University and the Authors.      *
  * Authors: Mark Friedrichs, Peter Eastman                                    *
  * Contributors:                                                              *
  *                                                                            *
@@ -584,7 +584,7 @@ private:
     CudaArray isAlchemical;
 
     double dispersionCoefficient;
-    CudaArray sigmaEpsilon;
+    CudaArray sigmaEpsilon, atomType;
     CudaArray bondReductionAtoms;
     CudaArray bondReductionFactors;
     CudaArray tempPosq;

plugins/amoeba/platforms/cuda/src/kernels/amoebaVdwForce2.cu

@@ -5,34 +5,9 @@
     unsigned int includeInteraction = (!isExcluded);
 #endif
     real tempForce = 0.0f;
-#if SIGMA_COMBINING_RULE == 1
-    real sigma = sigmaEpsilon1.x + sigmaEpsilon2.x;
-#elif SIGMA_COMBINING_RULE == 2
-    real sigma = 2*SQRT(sigmaEpsilon1.x*sigmaEpsilon2.x);
-#else
-    real sigma1_2 = sigmaEpsilon1.x*sigmaEpsilon1.x;
-    real sigma2_2 = sigmaEpsilon2.x*sigmaEpsilon2.x;
-    real sigmasum = sigma1_2+sigma2_2;
-    real sigma = (sigmasum == 0.0f ? (real) 0 : 2*(sigmaEpsilon1.x*sigma1_2 + sigmaEpsilon2.x*sigma2_2)/(sigma1_2+sigma2_2));
-#endif
-#if EPSILON_COMBINING_RULE == 1
-    real epsilon = 0.5f*(sigmaEpsilon1.y + sigmaEpsilon2.y);
-#elif EPSILON_COMBINING_RULE == 2
-    real epsilon = SQRT(sigmaEpsilon1.y*sigmaEpsilon2.y);
-#elif EPSILON_COMBINING_RULE == 3
-    real epssum = sigmaEpsilon1.y+sigmaEpsilon2.y;
-    real epsilon = (epssum == 0.0f ? (real) 0 : 2*(sigmaEpsilon1.y*sigmaEpsilon2.y)/(sigmaEpsilon1.y+sigmaEpsilon2.y));
-#elif EPSILON_COMBINING_RULE == 4
-    real sigma1_3 = sigmaEpsilon1.x*sigmaEpsilon1.x*sigmaEpsilon1.x;
-    real sigma2_3 = sigmaEpsilon2.x*sigmaEpsilon2.x*sigmaEpsilon2.x;
-    real sigma1_6 = sigma1_3*sigma1_3; 
-    real sigma2_6 = sigma2_3*sigma2_3;
-    real eps_s = SQRT(sigmaEpsilon1.y*sigmaEpsilon2.y);
-    real epsilon = (eps_s == 0.0f ? (real) 0 : 2*eps_s*sigma1_3*sigma2_3/(sigma1_6 + sigma2_6));
-#else
-    real epsilon_s = SQRT(sigmaEpsilon1.y) + SQRT(sigmaEpsilon2.y);
-    real epsilon = (epsilon_s == 0.0f ? (real) 0 : 4*sigmaEpsilon1.y*sigmaEpsilon2.y/(epsilon_s*epsilon_s));
-#endif
+    real2 pairSigmaEpsilon = sigmaEpsilon[atomType1+atomType2*NUM_TYPES];
+    real sigma = pairSigmaEpsilon.x;
+    real epsilon = pairSigmaEpsilon.y;
     real softcore = 0.0f;
 #if VDW_ALCHEMICAL_METHOD == 1
     if (isAlchemical1 != isAlchemical2) { 
@@ -45,6 +20,15 @@
        softcore = VDW_SOFTCORE_ALPHA * (1.0f - lambda) * (1.0f - lambda);
     }
 #endif
+#if POTENTIAL_FUNCTION == 1
+    real pp1 = sigma / r;
+    real pp2 = pp1 * pp1;
+    real pp3 = pp2 * pp1;
+    real pp6 = pp3 * pp3;
+    real pp12 = pp6 * pp6;
+    real termEnergy = 4 * epsilon * (pp12 - pp6);
+    real deltaE = -24 * epsilon * (2*pp12 - pp6) / r;
+#else
     real dhal = 0.07f;
     real ghal = 0.12f;
     real dhal1 = 1.07f;
@@ -65,6 +49,7 @@
     real dt2 = -7.0f * rho6 * t2 * s2;
     real termEnergy = epsilon * t1 * t2min;
     real deltaE = epsilon * (dt1 * t2min + t1 * dt2) / sigma;
+#endif
 #ifdef USE_CUTOFF
     if (r > TAPER_CUTOFF) {
         real x = r-TAPER_CUTOFF;

plugins/amoeba/platforms/cuda/tests/TestCudaAmoebaVdwForce.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-2012 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2020 Stanford University and the Authors.      *
  * Authors: Mark Friedrichs                                                   *
  * Contributors:                                                              *
  *                                                                            *
@@ -35,6 +35,7 @@
 
 #include "openmm/internal/AssertionUtilities.h"
 #include "openmm/Context.h"
+#include "openmm/CustomNonbondedForce.h"
 #include "OpenMMAmoeba.h"
 #include "openmm/System.h"
 #include "openmm/AmoebaVdwForce.h"
@@ -2106,6 +2107,94 @@ void testTriclinic() {
     }
 }
 
+void testCompareToCustom(AmoebaVdwForce::PotentialFunction potential) {
+    const int numParticles = 10;
+    const double cutoff = 1.0;
+    System system;
+    AmoebaVdwForce* vdw = new AmoebaVdwForce();
+    vdw->setNonbondedMethod(AmoebaVdwForce::CutoffPeriodic);
+    vdw->setCutoff(cutoff);
+    vdw->setPotentialFunction(potential);
+    vdw->setSigmaCombiningRule("ARITHMETIC");
+    vdw->setEpsilonCombiningRule("GEOMETRIC");
+    vdw->setUseDispersionCorrection(true);
+    system.addForce(vdw);
+    CustomNonbondedForce* custom;
+    if (potential == AmoebaVdwForce::LennardJones)
+        custom = new CustomNonbondedForce("4*eps*((sigma/r)^12-(sigma/r)^6); sigma=sigma1+sigma2; eps=sqrt(eps1*eps2)");
+    else
+        custom = new CustomNonbondedForce("eps*((1.07/(p+0.07))^7 * ((1.12/(p^7+0.12))-2)); p=r/sigma; sigma=sigma1+sigma2; eps=sqrt(eps1*eps2)");
+    custom->addPerParticleParameter("sigma");
+    custom->addPerParticleParameter("eps");
+    custom->setNonbondedMethod(CustomNonbondedForce::CutoffPeriodic);
+    custom->setCutoffDistance(cutoff);
+    custom->setUseLongRangeCorrection(true);
+    custom->setUseSwitchingFunction(true);
+    custom->setSwitchingDistance(0.9*cutoff);
+    custom->setForceGroup(1);
+    system.addForce(custom);
+    vector<Vec3> positions;
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(1.0);
+        if (i%2 == 0) {
+            vdw->addParticle(i, 0.2, 1.0, 0.0);
+            custom->addParticle({0.2, 1.0});
+        }
+        else {
+            vdw->addParticle(i, 0.25, 0.5, 0.0);
+            custom->addParticle({0.25, 0.5});
+        }
+        positions.push_back(2*Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt)));
+    }
+    LangevinIntegrator integrator(0.0, 0.1, 0.01);
+    Context context(system, integrator, Platform::getPlatformByName("CUDA"));
+    context.setPositions(positions);
+    State state1 = context.getState(State::Energy | State::Forces, true, 1);
+    State state2 = context.getState(State::Energy | State::Forces, true, 2);
+    ASSERT_EQUAL_TOL(state1.getPotentialEnergy(), state2.getPotentialEnergy(), 1e-3);
+    for (int i = 0; i < numParticles; i++)
+        ASSERT_EQUAL_VEC(state1.getForces()[i], state2.getForces()[i], 1e-5);
+}
+
+void testParticleTypes() {
+    System system;
+    for (int i = 0; i < 4; i++)
+        system.addParticle(1.0);
+    AmoebaVdwForce* vdw = new AmoebaVdwForce();
+    system.addForce(vdw);
+    vdw->setPotentialFunction(AmoebaVdwForce::LennardJones);
+    vdw->setSigmaCombiningRule("ARITHMETIC");
+    vdw->setEpsilonCombiningRule("GEOMETRIC");
+    vdw->addParticle(0, 0, 1.0);
+    vdw->addParticle(1, 2, 1.0);
+    vdw->addParticle(2, 0, 1.0);
+    vdw->addParticle(3, 1, 1.0);
+    vdw->addParticleType(0.3, 1.0);
+    vdw->addParticleType(0.4, 1.1);
+    vdw->addParticleType(0.5, 1.2);
+    vdw->addTypePair(2, 0, 0.6, 1.5);
+    vector<Vec3> positions;
+    positions.push_back(Vec3(0, 0, 0));
+    positions.push_back(Vec3(1, 0, 0));
+    positions.push_back(Vec3(0, 1, 0));
+    positions.push_back(Vec3(1, 1, 0));
+    LangevinIntegrator integrator(0.0, 0.1, 0.01);
+    Context context(system, integrator, Platform::getPlatformByName("CUDA"));
+    context.setPositions(positions);
+    State state = context.getState(State::Energy);
+    vector<double> r = {1.0, 1.0, sqrt(2.0), sqrt(2.0), 1.0, 1.0};
+    vector<double> sigma = {0.6, 0.3+0.3, 0.3+0.4, 0.6, 0.5+0.4, 0.3+0.4};
+    vector<double> epsilon = {1.5, sqrt(1.0*1.0), sqrt(1.0*1.1), 1.5, sqrt(1.1*1.2), sqrt(1.0*1.1)};
+    double expectedEnergy = 0;
+    for (int i = 0; i < 6; i++) {
+        double p = sigma[i]/r[i];
+        expectedEnergy += 4*epsilon[i]*(pow(p, 12) - pow(p, 6));
+    }
+    ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5);
+}
+
 int main(int argc, char* argv[]) {
     try {
         std::cout << "TestCudaAmoebaVdwForce running test..." << std::endl;
@@ -2161,6 +2250,15 @@ int main(int argc, char* argv[]) {
         // test triclinic boxes
 
         testTriclinic();
+        
+        // Compare to an equivalent CustomNonbondedForce.
+
+        testCompareToCustom(AmoebaVdwForce::Buffered147);
+        testCompareToCustom(AmoebaVdwForce::LennardJones);
+        
+        // Test specifying parameters by particle type.
+        
+        testParticleTypes();
 
         // Set lambda and the softcore power (n) to any values, while softcore alpha set to 0. 
         // The energy and forces are equal to scaling those from testVdwAmmoniaCubicMeanHhg by lambda^n;