PME works correctly in CPU platform

platforms/cpu/include/CpuKernels.h

@@ -44,7 +44,7 @@ namespace OpenMM {
  */
 class CpuCalcNonbondedForceKernel : public CalcNonbondedForceKernel {
 public:
-    CpuCalcNonbondedForceKernel(std::string name, const Platform& platform) : CalcNonbondedForceKernel(name, platform) {
+    CpuCalcNonbondedForceKernel(std::string name, const Platform& platform) : CalcNonbondedForceKernel(name, platform), bonded14IndexArray(NULL), bonded14ParamArray(NULL) {
     }
     ~CpuCalcNonbondedForceKernel();
     /**
@@ -75,12 +75,12 @@ public:
 private:
     int numParticles, num14;
     int **bonded14IndexArray;
-    float **particleParamArray;
     double **bonded14ParamArray;
-    double nonbondedCutoff, switchingDistance, rfDielectric, ewaldAlpha, dispersionCoefficient;
+    double nonbondedCutoff, switchingDistance, rfDielectric, ewaldAlpha, ewaldSelfEnergy, dispersionCoefficient;
     int kmax[3], gridSize[3];
     bool useSwitchingFunction;
     std::vector<std::set<int> > exclusions;
+    std::vector<std::pair<float, float> > particleParams;
     std::vector<float> posq;
     std::vector<float> forces;
     NonbondedMethod nonbondedMethod;

platforms/cpu/include/CpuNonbondedForce.h

@@ -49,30 +49,7 @@ class CpuNonbondedForce {
       int numRx, numRy, numRz;
       int meshDim[3];
       __m128 boxSize, invBoxSize, half;
-
-      // parameter indices
-
-      static const int SigIndex = 0;
-      static const int EpsIndex = 1;
-      static const int   QIndex = 2;
-            
-      /**---------------------------------------------------------------------------------------
-      
-         Calculate LJ Coulomb pair ixn between two atoms
-      
-         @param atom1            the index of the first atom
-         @param atom2            the index of the second atom
-         @param atomCoordinates  atom coordinates
-         @param atomParameters   atom parameters (charges, c6, c12, ...)     atomParameters[atomIndex][paramterIndex]
-         @param forces           force array (forces added)
-         @param totalEnergy      total energy
-            
-         --------------------------------------------------------------------------------------- */
-          
-      void calculateOneIxn( int atom1, int atom2, float* atomCoordinates,
-                            float** atomParameters, float* forces,
-                            double* totalEnergy ) const;
-
+      static float TWO_OVER_SQRT_PI;
 
    public:
 
@@ -153,51 +130,83 @@ class CpuNonbondedForce {
 
       /**---------------------------------------------------------------------------------------
       
-         Calculate LJ Coulomb pair ixn
+         Calculate Ewald ixn
       
          @param numberOfAtoms    number of atoms
-         @param atomCoordinates  atom coordinates
-         @param atomParameters   atom parameters (charges, c6, c12, ...)     atomParameters[atomIndex][paramterIndex]
+         @param posq             atom coordinates and charges
+         @param atomCoordinates  atom coordinates (in format needed by PME)
+         @param atomParameters   atom parameters (sigma/2, 2*sqrt(epsilon))
          @param exclusions       atom exclusion indices
                                  exclusions[atomIndex] contains the list of exclusions for that atom
          @param fixedParameters  non atom parameters (not currently used)
          @param forces           force array (forces added)
          @param totalEnergy      total energy
-         @param includeDirect      true if direct space interactions should be included
-         @param includeReciprocal  true if reciprocal space interactions should be included
             
          --------------------------------------------------------------------------------------- */
           
-      void calculatePairIxn(int numberOfAtoms, float* atomCoordinates,
-                            float** atomParameters, std::vector<std::set<int> >& exclusions,
-                            float* fixedParameters, float* forces,
-                            float* totalEnergy, bool includeDirect, bool includeReciprocal) const;
+      void calculateReciprocalIxn(int numberOfAtoms, float* posq, std::vector<OpenMM::RealVec>& atomCoordinates,
+                            const std::vector<std::pair<float, float> >& atomParameters, const std::vector<std::set<int> >& exclusions,
+                            float* fixedParameters, std::vector<OpenMM::RealVec>& forces, float* totalEnergy) const;
       
-private:
       /**---------------------------------------------------------------------------------------
       
-         Calculate Ewald ixn
+         Calculate LJ Coulomb pair ixn
       
          @param numberOfAtoms    number of atoms
-         @param atomCoordinates  atom coordinates
-         @param atomParameters   atom parameters (charges, c6, c12, ...)     atomParameters[atomIndex][paramterIndex]
+         @param posq             atom coordinates and charges
+         @param atomParameters   atom parameters (sigma/2, 2*sqrt(epsilon))
          @param exclusions       atom exclusion indices
                                  exclusions[atomIndex] contains the list of exclusions for that atom
          @param fixedParameters  non atom parameters (not currently used)
          @param forces           force array (forces added)
          @param totalEnergy      total energy
-         @param includeDirect      true if direct space interactions should be included
-         @param includeReciprocal  true if reciprocal space interactions should be included
       
          --------------------------------------------------------------------------------------- */
           
-      void calculateEwaldIxn(int numberOfAtoms, float* atomCoordinates,
-                            float** atomParameters, std::vector<std::set<int> >& exclusions,
-                            float* fixedParameters, float* forces,
-                            float* totalEnergy, bool includeDirect, bool includeReciprocal) const;
+      void calculateDirectIxn(int numberOfAtoms, float* posq,
+                            const std::vector<std::pair<float, float> >& atomParameters, const std::vector<std::set<int> >& exclusions,
+                            float* fixedParameters, float* forces, float* totalEnergy) const;
+
+private:
+            
+      /**---------------------------------------------------------------------------------------
+      
+         Calculate LJ Coulomb pair ixn between two atoms
+      
+         @param atom1            the index of the first atom
+         @param atom2            the index of the second atom
+         @param posq             atom coordinates and charges
+         @param atomParameters   atom parameters (sigma/2, 2*sqrt(epsilon))
+         @param forces           force array (forces added)
+         @param totalEnergy      total energy
+            
+         --------------------------------------------------------------------------------------- */
+          
+      void calculateOneIxn( int atom1, int atom2, float* posq,
+                            const std::vector<std::pair<float, float> >& atomParameters, float* forces,
+                            double* totalEnergy ) const;
+            
+      /**---------------------------------------------------------------------------------------
+      
+         Calculate LJ Coulomb pair ixn between two atoms
+      
+         @param atom1            the index of the first atom
+         @param atom2            the index of the second atom
+         @param posq             atom coordinates and charges
+         @param atomParameters   atom parameters (sigma/2, 2*sqrt(epsilon))
+         @param forces           force array (forces added)
+         @param totalEnergy      total energy
+            
+         --------------------------------------------------------------------------------------- */
+          
+      void calculateOneEwaldIxn( int atom1, int atom2, float* posq,
+                            const std::vector<std::pair<float, float> >& atomParameters, float* forces,
+                            double* totalEnergy ) const;
+
       
       void getDeltaR(const __m128& posI, const __m128& posJ, __m128& deltaR, float& r2, bool periodic) const;
 
+      static float erfcApprox(float x);
 };
 
 // ---------------------------------------------------------------------------------------

platforms/cpu/src/CpuKernels.cpp

@@ -63,9 +63,14 @@ static RealVec& extractBoxSize(ContextImpl& context) {
 }
 
 CpuCalcNonbondedForceKernel::~CpuCalcNonbondedForceKernel() {
-    delete bonded14IndexArray; // Do this properly
-    delete bonded14ParamArray; // Do this properly
-    delete particleParamArray; // Do this properly
+    if (bonded14ParamArray != NULL) {
+        for (int i = 0; i < num14; i++) {
+            delete[] bonded14IndexArray[i];
+            delete[] bonded14ParamArray[i];
+        }
+        delete bonded14IndexArray;
+        delete bonded14ParamArray;
+    }
 }
 
 void CpuCalcNonbondedForceKernel::initialize(const System& system, const NonbondedForce& force) {
@@ -96,16 +101,14 @@ void CpuCalcNonbondedForceKernel::initialize(const System& system, const Nonbond
     bonded14ParamArray = new double*[num14];
     for (int i = 0; i < num14; i++)
         bonded14ParamArray[i] = new double[3];
-    particleParamArray = new float*[numParticles];
-    for (int i = 0; i < numParticles; i++)
-        particleParamArray[i] = new float[3];
+    particleParams.resize(numParticles);
+    double sumSquaredCharges = 0.0;
     for (int i = 0; i < numParticles; ++i) {
         double charge, radius, depth;
         force.getParticleParameters(i, charge, radius, depth);
         posq[4*i+3] = (float) charge;
-        particleParamArray[i][0] = (float) (0.5*radius);
-        particleParamArray[i][1] = (float) (2.0*sqrt(depth));
-        particleParamArray[i][2] = (float) (charge);
+        particleParams[i] = make_pair((float) (0.5*radius), (float) (2.0*sqrt(depth)));
+        sumSquaredCharges += charge*charge;
     }
     for (int i = 0; i < num14; ++i) {
         int particle1, particle2;
@@ -135,6 +138,10 @@ void CpuCalcNonbondedForceKernel::initialize(const System& system, const Nonbond
         NonbondedForceImpl::calcPMEParameters(system, force, alpha, gridSize[0], gridSize[1], gridSize[2]);
         ewaldAlpha = alpha;
     }
+    if (nonbondedMethod == Ewald || nonbondedMethod == PME)
+        ewaldSelfEnergy = -ONE_4PI_EPS0*ewaldAlpha*sumSquaredCharges/sqrt(M_PI);
+    else
+        ewaldSelfEnergy = 0.0;
     rfDielectric = force.getReactionFieldDielectric();
     if (force.getUseDispersionCorrection())
         dispersionCoefficient = NonbondedForceImpl::calcDispersionCorrection(system, force);
@@ -147,7 +154,7 @@ double CpuCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeFo
     vector<RealVec>& forceData = extractForces(context);
     RealVec boxSize = extractBoxSize(context);
     float floatBoxSize[3] = {(float) boxSize[0], (float) boxSize[1], (float) boxSize[2]};
-    double energy = 0;
+    double energy = ewaldSelfEnergy;
     CpuNonbondedForce clj;
     bool periodic = (nonbondedMethod == CutoffPeriodic);
     bool ewald  = (nonbondedMethod == Ewald);
@@ -186,9 +193,12 @@ double CpuCalcNonbondedForceKernel::execute(ContextImpl& context, bool includeFo
         clj.setUsePME(ewaldAlpha, gridSize);
     if (useSwitchingFunction)
         clj.setUseSwitchingFunction(switchingDistance);
-    float directEnergy = 0;
-    clj.calculatePairIxn(numParticles, &posq[0], particleParamArray, exclusions, 0, &forces[0], includeEnergy ? &directEnergy : NULL, includeDirect, includeReciprocal);
-    energy += directEnergy;
+    float nonbondedEnergy = 0;
+    if (includeDirect)
+        clj.calculateDirectIxn(numParticles, &posq[0], particleParams, exclusions, 0, &forces[0], includeEnergy ? &nonbondedEnergy : NULL);
+    if (includeReciprocal)
+        clj.calculateReciprocalIxn(numParticles, &posq[0], posData, particleParams, exclusions, 0, forceData, includeEnergy ? &nonbondedEnergy : NULL);
+    energy += nonbondedEnergy;
     for (int i = 0; i < numParticles; i++) {
         forceData[i][0] += forces[4*i];
         forceData[i][1] += forces[4*i+1];
@@ -220,13 +230,18 @@ void CpuCalcNonbondedForceKernel::copyParametersToContext(ContextImpl& context,
 
     // Record the values.
 
+    double sumSquaredCharges = 0.0;
     for (int i = 0; i < numParticles; ++i) {
         double charge, radius, depth;
         force.getParticleParameters(i, charge, radius, depth);
-        particleParamArray[i][0] = (float) (0.5*radius);
-        particleParamArray[i][1] = (float) (2.0*sqrt(depth));
-        particleParamArray[i][2] = (float) (charge);
+        posq[4*i+3] = (float) charge;
+        particleParams[i] = make_pair((float) (0.5*radius), (float) (2.0*sqrt(depth)));
+        sumSquaredCharges += charge*charge;
     }
+    if (nonbondedMethod == Ewald || nonbondedMethod == PME)
+        ewaldSelfEnergy = -ONE_4PI_EPS0*ewaldAlpha*sumSquaredCharges/sqrt(M_PI);
+    else
+        ewaldSelfEnergy = 0.0;
     for (int i = 0; i < num14; ++i) {
         int particle1, particle2;
         double charge, radius, depth;

platforms/reference/include/ReferencePME.h

@@ -77,7 +77,7 @@ int
 pme_exec(pme_t       pme,
          std::vector<OpenMM::RealVec>& atomCoordinates,
          std::vector<OpenMM::RealVec>& forces,
-         RealOpenMM **  atomParameters,
+         std::vector<RealOpenMM>& charges,
          const RealOpenMM  periodicBoxSize[3],
          RealOpenMM *    energy,
          RealOpenMM      pme_virial[3][3]);

platforms/reference/src/SimTKReference/ReferenceLJCoulombIxn.cpp

@@ -233,7 +233,10 @@ void ReferenceLJCoulombIxn::calculateEwaldIxn(int numberOfAtoms, vector<RealVec>
 
     pme_init(&pmedata,alphaEwald,numberOfAtoms,meshDim,5,1);
 
-    pme_exec(pmedata,atomCoordinates,forces,atomParameters,periodicBoxSize,&recipEnergy,virial);
+    vector<RealOpenMM> charges(numberOfAtoms);
+    for (int i = 0; i < numberOfAtoms; i++)
+        charges[i] = atomParameters[i][QIndex];
+    pme_exec(pmedata,atomCoordinates,forces,charges,periodicBoxSize,&recipEnergy,virial);
 
     if( totalEnergy )
        *totalEnergy += recipEnergy;

platforms/reference/src/SimTKReference/ReferencePME.cpp

@@ -317,10 +317,8 @@ pme_update_bsplines(pme_t    pme)
 
 
 static void
-pme_grid_spread_charge(pme_t      pme,
-                       RealOpenMM **   atomParameters)
+pme_grid_spread_charge(pme_t      pme, vector<RealOpenMM>& charges)
 {
-        static const int   QIndex = 2; // atom charges are stored in atomParameters[atomID][2]
     int       order;
     int       i;
     int       ix,iy,iz;
@@ -342,7 +340,7 @@ pme_grid_spread_charge(pme_t      pme,
 
     for(i=0;i<pme->natoms;i++)
     {
-        q = atomParameters[i][QIndex];
+        q = charges[i];
 
         /* Grid index for the actual atom position */
         x0index = pme->particleindex[i][0];
@@ -523,10 +521,9 @@ pme_reciprocal_convolution(pme_t     pme,
 static void
 pme_grid_interpolate_force(pme_t      pme,
                            const RealOpenMM     periodicBoxSize[3],
-                           RealOpenMM **   atomParameters,
+                           vector<RealOpenMM>& charges,
                            vector<RealVec>&   forces)
 {
-        static const int   QIndex = 2; // atom charges are stored in atomParameters[atomID][2]
     int       i;
     int       ix,iy,iz;
     int       x0index,y0index,z0index;
@@ -558,7 +555,7 @@ pme_grid_interpolate_force(pme_t      pme,
     {
         fx = fy = fz = 0;
 
-        q = atomParameters[i][QIndex];
+        q = charges[i];
 
         /* Grid index for the actual atom position */
         x0index = pme->particleindex[i][0];
@@ -671,7 +668,7 @@ pme_init(pme_t *       ppme,
 int pme_exec(pme_t       pme,
              vector<RealVec>&   atomCoordinates,
              vector<RealVec>&   forces,
-             RealOpenMM **   atomParameters,
+             vector<RealOpenMM>& charges,
              const RealOpenMM      periodicBoxSize[3],
              RealOpenMM *    energy,
              RealOpenMM      pme_virial[3][3])
@@ -692,7 +689,7 @@ int pme_exec(pme_t       pme,
     pme_update_bsplines(pme);
 
     /* Spread the charges on grid (using newly calculated bsplines in the pme structure) */
-    pme_grid_spread_charge(pme,atomParameters);
+    pme_grid_spread_charge(pme, charges);
 
     /* do 3d-fft */
     fftpack_exec_3d(pme->fftplan,FFTPACK_FORWARD,pme->grid,pme->grid);
@@ -704,7 +701,7 @@ int pme_exec(pme_t       pme,
     fftpack_exec_3d(pme->fftplan,FFTPACK_BACKWARD,pme->grid,pme->grid);
 
     /* Get the particle forces from the grid and bsplines in the pme structure */
-    pme_grid_interpolate_force(pme,periodicBoxSize,atomParameters,forces);
+    pme_grid_interpolate_force(pme,periodicBoxSize,charges,forces);
 
     return 0;
 }