Completed CPU implementation of LJPME

olla/include/openmm/kernels.h

@@ -1436,7 +1436,6 @@ public:
  */
 class CalcDispersionPmeReciprocalForceKernel : public KernelImpl {
 public:
-    class IO;
     static std::string Name() {
         return "CalcDispersionPmeReciprocalForce";
     }
@@ -1460,14 +1459,14 @@ public:
      * @param periodicBoxVectors  the vectors defining the periodic box (measured in nm)
      * @param includeEnergy       true if potential energy should be computed
      */
-    virtual void beginComputation(IO& io, const Vec3* periodicBoxVectors, bool includeEnergy) = 0;
+    virtual void beginComputation(CalcPmeReciprocalForceKernel::IO& io, const Vec3* periodicBoxVectors, bool includeEnergy) = 0;
     /**
      * Finish computing the force and energy.
      * 
      * @param io   an object that coordinates data transfer
      * @return the potential energy due to the PME reciprocal space interactions
      */
-    virtual double finishComputation(IO& io) = 0;
+    virtual double finishComputation(CalcPmeReciprocalForceKernel::IO& io) = 0;
     /**
      * Get the parameters being used for PME.
      * 

platforms/cpu/include/CpuKernels.h

@@ -268,7 +268,7 @@ private:
     class PmeIO;
     void computeParameters(ContextImpl& context, bool offsetsOnly);
     CpuPlatform::PlatformData& data;
-    int numParticles, num14, posqIndex;
+    int numParticles, num14, chargePosqIndex, ljPosqIndex;
     std::vector<std::vector<int> > bonded14IndexArray;
     std::vector<std::vector<double> > bonded14ParamArray;
     double nonbondedCutoff, switchingDistance, rfDielectric, ewaldAlpha, ewaldDispersionAlpha, ewaldSelfEnergy, dispersionCoefficient;

platforms/cpu/src/CpuKernels.cpp

@@ -490,7 +490,8 @@ CpuCalcNonbondedForceKernel::~CpuCalcNonbondedForceKernel() {
 }
 
 void CpuCalcNonbondedForceKernel::initialize(const System& system, const NonbondedForce& force) {
-    posqIndex = data.requestPosqIndex();
+    chargePosqIndex = data.requestPosqIndex();
+    ljPosqIndex = data.requestPosqIndex();
 
     // Identify which exceptions are 1-4 interactions.
 
@@ -637,6 +638,7 @@ double CpuCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeFo
 
             vector<string> kernelNames;
             kernelNames.push_back("CalcPmeReciprocalForce");
+            kernelNames.push_back("CalcDispersionPmeReciprocalForce");
             useOptimizedPme = getPlatform().supportsKernels(kernelNames);
             if (useOptimizedPme) {
                 optimizedPme = getPlatform().createKernel(CalcPmeReciprocalForceKernel::Name(), context);
@@ -648,7 +650,7 @@ double CpuCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeFo
         }
     }
     computeParameters(context, true);
-    copyChargesToPosq(context, charges, posqIndex);
+    copyChargesToPosq(context, charges, chargePosqIndex);
     AlignedArray<float>& posq = data.posq;
     vector<Vec3>& posData = extractPositions(context);
     vector<Vec3>& forceData = extractForces(context);
@@ -685,6 +687,11 @@ double CpuCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeFo
             Vec3 periodicBoxVectors[3] = {boxVectors[0], boxVectors[1], boxVectors[2]};
             optimizedPme.getAs<CalcPmeReciprocalForceKernel>().beginComputation(io, periodicBoxVectors, includeEnergy);
             nonbondedEnergy += optimizedPme.getAs<CalcPmeReciprocalForceKernel>().finishComputation(io);
+            if (nonbondedMethod == LJPME) {
+                copyChargesToPosq(context, C6params, ljPosqIndex);
+                optimizedDispersionPme.getAs<CalcDispersionPmeReciprocalForceKernel>().beginComputation(io, periodicBoxVectors, includeEnergy);
+                nonbondedEnergy += optimizedDispersionPme.getAs<CalcDispersionPmeReciprocalForceKernel>().finishComputation(io);
+            }
         }
         else
             nonbonded->calculateReciprocalIxn(numParticles, &posq[0], posData, particleParams, C6params, exclusions, forceData, includeEnergy ? &nonbondedEnergy : NULL);
@@ -715,7 +722,8 @@ void CpuCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& context,
 
     // Record the values.
 
-    posqIndex = data.requestPosqIndex();
+    chargePosqIndex = data.requestPosqIndex();
+    ljPosqIndex = data.requestPosqIndex();
     for (int i = 0; i < numParticles; ++i)
        force.getParticleParameters(i, baseParticleParams[i][0], baseParticleParams[i][1], baseParticleParams[i][2]);
     for (int i = 0; i < num14; ++i) {

plugins/cpupme/src/CpuPmeKernelFactory.cpp

@@ -35,8 +35,10 @@ using namespace OpenMM;
 extern "C" OPENMM_EXPORT_PME void registerKernelFactories() {
     if (CpuCalcPmeReciprocalForceKernel::isProcessorSupported()) {
         CpuPmeKernelFactory* factory = new CpuPmeKernelFactory();
-        for (int i = 0; i < Platform::getNumPlatforms(); i++)
+        for (int i = 0; i < Platform::getNumPlatforms(); i++) {
             Platform::getPlatform(i).registerKernelFactory(CalcPmeReciprocalForceKernel::Name(), factory);
+            Platform::getPlatform(i).registerKernelFactory(CalcDispersionPmeReciprocalForceKernel::Name(), factory);
+        }
     }
 }
 

plugins/cpupme/src/CpuPmeKernels.cpp

@@ -903,7 +903,7 @@ void CpuCalcDispersionPmeReciprocalForceKernel::runWorkerThread(ThreadPool& thre
     interpolateForces(posq, &force[0], realGrid, gridx, gridy, gridz, numParticles, periodicBoxVectors, recipBoxVectors, atomicCounter, epsilonFactor);
 }
 
-void CpuCalcDispersionPmeReciprocalForceKernel::beginComputation(IO& io, const Vec3* periodicBoxVectors, bool includeEnergy) {
+void CpuCalcDispersionPmeReciprocalForceKernel::beginComputation(CalcPmeReciprocalForceKernel::IO& io, const Vec3* periodicBoxVectors, bool includeEnergy) {
     this->io = &io;
     this->periodicBoxVectors[0] = periodicBoxVectors[0];
     this->periodicBoxVectors[1] = periodicBoxVectors[1];
@@ -927,7 +927,7 @@ void CpuCalcDispersionPmeReciprocalForceKernel::beginComputation(IO& io, const V
     pthread_mutex_unlock(&lock);
 }
 
-double CpuCalcDispersionPmeReciprocalForceKernel::finishComputation(IO& io) {
+double CpuCalcDispersionPmeReciprocalForceKernel::finishComputation(CalcPmeReciprocalForceKernel::IO& io) {
     pthread_mutex_lock(&lock);
     while (!isFinished) {
         pthread_cond_wait(&endCondition, &lock);

plugins/cpupme/src/CpuPmeKernels.h

@@ -142,9 +142,9 @@ private:
  * vectorized (requiring SSE 4.1) and multithreaded.  It uses FFTW to perform the FFTs.
  */
 
-class OPENMM_EXPORT_PME CpuCalcDispersionPmeReciprocalForceKernel : public CalcPmeReciprocalForceKernel {
+class OPENMM_EXPORT_PME CpuCalcDispersionPmeReciprocalForceKernel : public CalcDispersionPmeReciprocalForceKernel {
 public:
-    CpuCalcDispersionPmeReciprocalForceKernel(std::string name, const Platform& platform) : CalcPmeReciprocalForceKernel(name, platform),
+    CpuCalcDispersionPmeReciprocalForceKernel(std::string name, const Platform& platform) : CalcDispersionPmeReciprocalForceKernel(name, platform),
             hasCreatedPlan(false), isDeleted(false), realGrid(NULL), complexGrid(NULL)  {
     }
     /**
@@ -166,14 +166,14 @@ public:
      * @param periodicBoxVectors  the vectors defining the periodic box (measured in nm)
      * @param includeEnergy       true if potential energy should be computed
      */
-    void beginComputation(IO& io, const Vec3* periodicBoxVectors, bool includeEnergy);
+    void beginComputation(CalcPmeReciprocalForceKernel::IO& io, const Vec3* periodicBoxVectors, bool includeEnergy);
     /**
      * Finish computing the force and energy.
      * 
      * @param io   an object that coordinates data transfer
      * @return the potential energy due to the PME reciprocal space interactions
      */
-    double finishComputation(IO& io);
+    double finishComputation(CalcPmeReciprocalForceKernel::IO& io);
     /**
      * This routine contains the code executed by the main thread.
      */
@@ -221,7 +221,7 @@ private:
     pthread_mutex_t lock;
     pthread_t mainThread;
     // The following variables are used to store information about the calculation currently being performed.
-    IO* io;
+    CalcPmeReciprocalForceKernel::IO* io;
     float energy;
     float* posq;
     Vec3 periodicBoxVectors[3], recipBoxVectors[3];

plugins/cpupme/tests/TestCpuPme.cpp

@@ -637,6 +637,75 @@ void testPME(bool triclinic) {
         ASSERT_EQUAL_VEC(refState.getForces()[i], Vec3(io.force[4*i], io.force[4*i+1], io.force[4*i+2]), 1e-3);
 }
 
+void testLJPME(bool triclinic) {
+    // Create a cloud of random LJ particles.
+
+    const int numParticles = 51;
+    const double boxWidth = 5.0;
+    const double cutoff = 1.0;
+    const double alpha = 2.91842;
+    Vec3 boxVectors[3];
+    if (triclinic) {
+        boxVectors[0] = Vec3(boxWidth, 0, 0);
+        boxVectors[1] = Vec3(0.2*boxWidth, boxWidth, 0);
+        boxVectors[2] = Vec3(-0.3*boxWidth, -0.1*boxWidth, boxWidth);
+    }
+    else {
+        boxVectors[0] = Vec3(boxWidth, 0, 0);
+        boxVectors[1] = Vec3(0, boxWidth, 0);
+        boxVectors[2] = Vec3(0, 0, boxWidth);
+    }
+    System system;
+    system.setDefaultPeriodicBoxVectors(boxVectors[0], boxVectors[1], boxVectors[2]);
+    NonbondedForce* force = new NonbondedForce();
+    system.addForce(force);
+    vector<Vec3> positions(numParticles);
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(1.0);
+        force->addParticle(0, 0.5, 1.0);
+        positions[i] = Vec3(boxWidth*genrand_real2(sfmt), boxWidth*genrand_real2(sfmt), boxWidth*genrand_real2(sfmt));
+    }
+    force->setNonbondedMethod(NonbondedForce::LJPME);
+    force->setCutoffDistance(cutoff);
+    force->setReciprocalSpaceForceGroup(1);
+    force->setLJPMEParameters(alpha, 64, 64, 64);
+    
+    // Compute the reciprocal space forces with the reference platform.
+    
+    Platform& platform = Platform::getPlatformByName("Reference");
+    VerletIntegrator integrator(0.01);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    State refState = context.getState(State::Forces | State::Energy, false, 1<<1);
+    
+    // Now compute them with the optimized kernel.
+    
+    CpuCalcDispersionPmeReciprocalForceKernel pme(CalcDispersionPmeReciprocalForceKernel::Name(), platform);
+    IO io;
+    double ewaldSelfEnergy = 0;
+    for (int i = 0; i < numParticles; i++) {
+        io.posq.push_back(positions[i][0]);
+        io.posq.push_back(positions[i][1]);
+        io.posq.push_back(positions[i][2]);
+        double charge, sigma, epsilon;
+        force->getParticleParameters(i, charge, sigma, epsilon);
+        io.posq.push_back(pow(sigma, 3.0) * 2.0*sqrt(epsilon));
+        ewaldSelfEnergy += pow(alpha*sigma, 6.0) * epsilon / 3.0;
+    }
+    pme.initialize(64, 64, 64, numParticles, alpha, true);
+    pme.beginComputation(io, boxVectors, true);
+    double energy = pme.finishComputation(io);
+
+    // See if they match.
+    
+    ASSERT_EQUAL_TOL(refState.getPotentialEnergy(), energy+ewaldSelfEnergy, 1e-3);
+    for (int i = 0; i < numParticles; i++)
+        ASSERT_EQUAL_VEC(refState.getForces()[i], Vec3(io.force[4*i], io.force[4*i+1], io.force[4*i+2]), 1e-3);
+}
+
 int main(int argc, char* argv[]) {
     try {
         if (!CpuCalcPmeReciprocalForceKernel::isProcessorSupported()) {
@@ -645,6 +714,8 @@ int main(int argc, char* argv[]) {
         }
         testPME(false);
         testPME(true);
+        testLJPME(false);
+        testLJPME(true);
         test_water2_dpme_energies_forces_no_exclusions();
     }
     catch(const exception& e) {

wrappers/python/simtk/testInstallation.py

@@ -36,7 +36,7 @@ def run_tests():
     data_dir = os.path.join(os.path.abspath(os.path.split(__file__)[0]), 'openmm', 'app', 'data')
     pdb = PDBFile(os.path.join(data_dir, 'test.pdb'))
     forcefield = ForceField('amber99sb.xml', 'tip3p.xml')
-    system = forcefield.createSystem(pdb.topology, nonbondedMethod=PME, nonbondedCutoff=1*nanometer, constraints=HBonds)
+    system = forcefield.createSystem(pdb.topology, nonbondedMethod=LJPME, nonbondedCutoff=1*nanometer, constraints=HBonds, ewaldErrorTolerance=1e-4)
 
     # List all installed platforms and compute forces with each one.