Fixed bug in PME when atoms were far outside the periodic box (see bug 1144)

platforms/cuda/src/kernels/kCalculatePME.cu

@@ -108,9 +108,12 @@ void kUpdateGridIndexAndFraction_kernel()
     for (int i = tid; i < cSim.atoms; i += tnb)
     {
         float4 posq = cSim.pPosq[i];
-        float3 t = make_float3((posq.x*cSim.invPeriodicBoxSizeX+1.0f)*cSim.pmeGridSize.x,
-                               (posq.y*cSim.invPeriodicBoxSizeY+1.0f)*cSim.pmeGridSize.y,
-                               (posq.z*cSim.invPeriodicBoxSizeZ+1.0f)*cSim.pmeGridSize.z);
+        posq.x -= floor(posq.x*cSim.invPeriodicBoxSizeX)*cSim.periodicBoxSizeX;
+        posq.y -= floor(posq.y*cSim.invPeriodicBoxSizeY)*cSim.periodicBoxSizeY;
+        posq.z -= floor(posq.z*cSim.invPeriodicBoxSizeZ)*cSim.periodicBoxSizeZ;
+        float3 t = make_float3((posq.x*cSim.invPeriodicBoxSizeX)*cSim.pmeGridSize.x,
+                               (posq.y*cSim.invPeriodicBoxSizeY)*cSim.pmeGridSize.y,
+                               (posq.z*cSim.invPeriodicBoxSizeZ)*cSim.pmeGridSize.z);
         int3 gridIndex = make_int3(((int) t.x) % cSim.pmeGridSize.x,
                               ((int) t.y) % cSim.pmeGridSize.y,
                               ((int) t.z) % cSim.pmeGridSize.z);
@@ -192,9 +195,12 @@ void kUpdateBsplines_kernel()
         }
 
         float4 posq = cSim.pPosq[i];
-        float3 t = make_float3((posq.x*cSim.invPeriodicBoxSizeX+1.0f)*cSim.pmeGridSize.x,
-                               (posq.y*cSim.invPeriodicBoxSizeY+1.0f)*cSim.pmeGridSize.y,
-                               (posq.z*cSim.invPeriodicBoxSizeZ+1.0f)*cSim.pmeGridSize.z);
+        posq.x -= floor(posq.x*cSim.invPeriodicBoxSizeX)*cSim.periodicBoxSizeX;
+        posq.y -= floor(posq.y*cSim.invPeriodicBoxSizeY)*cSim.periodicBoxSizeY;
+        posq.z -= floor(posq.z*cSim.invPeriodicBoxSizeZ)*cSim.periodicBoxSizeZ;
+        float3 t = make_float3((posq.x*cSim.invPeriodicBoxSizeX)*cSim.pmeGridSize.x,
+                               (posq.y*cSim.invPeriodicBoxSizeY)*cSim.pmeGridSize.y,
+                               (posq.z*cSim.invPeriodicBoxSizeZ)*cSim.pmeGridSize.z);
         float4 dr = make_float4(t.x-(int) t.x, t.y-(int) t.y, t.z-(int) t.z, 0.0f);
 
         data[PME_ORDER - 1] = make_float4(0.0f);
@@ -341,9 +347,12 @@ void kGridInterpolateForce_kernel()
     {
         float3 force = make_float3(0.0f, 0.0f, 0.0f);
         float4 posq = cSim.pPosq[atom];
-        float3 t = make_float3((posq.x*cSim.invPeriodicBoxSizeX+1.0f)*cSim.pmeGridSize.x,
-                               (posq.y*cSim.invPeriodicBoxSizeY+1.0f)*cSim.pmeGridSize.y,
-                               (posq.z*cSim.invPeriodicBoxSizeZ+1.0f)*cSim.pmeGridSize.z);
+        posq.x -= floor(posq.x*cSim.invPeriodicBoxSizeX)*cSim.periodicBoxSizeX;
+        posq.y -= floor(posq.y*cSim.invPeriodicBoxSizeY)*cSim.periodicBoxSizeY;
+        posq.z -= floor(posq.z*cSim.invPeriodicBoxSizeZ)*cSim.periodicBoxSizeZ;
+        float3 t = make_float3((posq.x*cSim.invPeriodicBoxSizeX)*cSim.pmeGridSize.x,
+                               (posq.y*cSim.invPeriodicBoxSizeY)*cSim.pmeGridSize.y,
+                               (posq.z*cSim.invPeriodicBoxSizeZ)*cSim.pmeGridSize.z);
         int3 gridIndex = make_int3(((int) t.x) % cSim.pmeGridSize.x,
                               ((int) t.y) % cSim.pmeGridSize.y,
                               ((int) t.z) % cSim.pmeGridSize.z);

platforms/opencl/src/OpenCLKernels.cpp

@@ -1296,9 +1296,8 @@ void OpenCLCalcNonbondedForceKernel::executeForces(ContextImpl& context) {
             pmeInterpolateForceKernel = cl::Kernel(program, "gridInterpolateForce");
             pmeGridIndexKernel.setArg<cl::Buffer>(0, cl.getPosq().getDeviceBuffer());
             pmeGridIndexKernel.setArg<cl::Buffer>(1, pmeAtomGridIndex->getDeviceBuffer());
-            pmeAtomRangeKernel.setArg<cl::Buffer>(0, cl.getPosq().getDeviceBuffer());
-            pmeAtomRangeKernel.setArg<cl::Buffer>(1, pmeAtomGridIndex->getDeviceBuffer());
-            pmeAtomRangeKernel.setArg<cl::Buffer>(2, pmeAtomRange->getDeviceBuffer());
+            pmeAtomRangeKernel.setArg<cl::Buffer>(0, pmeAtomGridIndex->getDeviceBuffer());
+            pmeAtomRangeKernel.setArg<cl::Buffer>(1, pmeAtomRange->getDeviceBuffer());
             pmeUpdateBsplinesKernel.setArg<cl::Buffer>(0, cl.getPosq().getDeviceBuffer());
             pmeUpdateBsplinesKernel.setArg<cl::Buffer>(1, pmeBsplineTheta->getDeviceBuffer());
             pmeUpdateBsplinesKernel.setArg<cl::Buffer>(2, pmeBsplineDtheta->getDeviceBuffer());
@@ -1336,11 +1335,13 @@ void OpenCLCalcNonbondedForceKernel::executeForces(ContextImpl& context) {
     if (pmeGrid != NULL) {
         mm_float4 boxSize = cl.getPeriodicBoxSize();
         mm_float4 invBoxSize = cl.getInvPeriodicBoxSize();
-        pmeGridIndexKernel.setArg<mm_float4>(2, invBoxSize);
+        pmeGridIndexKernel.setArg<mm_float4>(2, boxSize);
+        pmeGridIndexKernel.setArg<mm_float4>(3, invBoxSize);
         cl.executeKernel(pmeGridIndexKernel, cl.getNumAtoms());
         sort->sort(*pmeAtomGridIndex);
         cl.executeKernel(pmeAtomRangeKernel, cl.getNumAtoms());
-        pmeUpdateBsplinesKernel.setArg<mm_float4>(5, invBoxSize);
+        pmeUpdateBsplinesKernel.setArg<mm_float4>(5, boxSize);
+        pmeUpdateBsplinesKernel.setArg<mm_float4>(6, invBoxSize);
         cl.executeKernel(pmeUpdateBsplinesKernel, cl.getNumAtoms());
         cl.executeKernel(pmeSpreadChargeKernel, cl.getNumAtoms());
         fft->execFFT(*pmeGrid, true);
@@ -1348,7 +1349,8 @@ void OpenCLCalcNonbondedForceKernel::executeForces(ContextImpl& context) {
         pmeConvolutionKernel.setArg<cl_float>(6, (float) (1.0/(M_PI*boxSize.x*boxSize.y*boxSize.z)));
         cl.executeKernel(pmeConvolutionKernel, cl.getNumAtoms());
         fft->execFFT(*pmeGrid, false);
-        pmeInterpolateForceKernel.setArg<mm_float4>(5, invBoxSize);
+        pmeInterpolateForceKernel.setArg<mm_float4>(5, boxSize);
+        pmeInterpolateForceKernel.setArg<mm_float4>(6, invBoxSize);
         cl.executeKernel(pmeInterpolateForceKernel, cl.getNumAtoms());
     }
 }

platforms/opencl/src/kernels/pme.cl

-__kernel void updateGridIndexAndFraction(__global float4* posq, __global float2* pmeAtomGridIndex, float4 invPeriodicBoxSize) {
+__kernel void updateGridIndexAndFraction(__global float4* posq, __global float2* pmeAtomGridIndex, float4 periodicBoxSize, float4 invPeriodicBoxSize) {
     for (int i = get_global_id(0); i < NUM_ATOMS; i += get_global_size(0)) {
         float4 pos = posq[i];
-        float4 t = (float4) ((pos.x*invPeriodicBoxSize.x+1.0f)*GRID_SIZE_X,
-                             (pos.y*invPeriodicBoxSize.y+1.0f)*GRID_SIZE_Y,
-                             (pos.z*invPeriodicBoxSize.z+1.0f)*GRID_SIZE_Z, 0.0f);
+        pos.x -= floor(pos.x*invPeriodicBoxSize.x)*periodicBoxSize.x;
+        pos.y -= floor(pos.y*invPeriodicBoxSize.y)*periodicBoxSize.y;
+        pos.z -= floor(pos.z*invPeriodicBoxSize.z)*periodicBoxSize.z;
+        float4 t = (float4) ((pos.x*invPeriodicBoxSize.x)*GRID_SIZE_X,
+                             (pos.y*invPeriodicBoxSize.y)*GRID_SIZE_Y,
+                             (pos.z*invPeriodicBoxSize.z)*GRID_SIZE_Z, 0.0f);
         int4 gridIndex = (int4) (((int) t.x) % GRID_SIZE_X,
                                  ((int) t.y) % GRID_SIZE_Y,
                                  ((int) t.z) % GRID_SIZE_Z, 0);
@@ -15,7 +18,7 @@ __kernel void updateGridIndexAndFraction(__global float4* posq, __global float2*
  * For each grid point, find the range of sorted atoms associated with that point.
  */
 
-__kernel void findAtomRangeForGrid(__global float4* posq, __global float2* pmeAtomGridIndex, __global int* pmeAtomRange) {
+__kernel void findAtomRangeForGrid(__global float2* pmeAtomGridIndex, __global int* pmeAtomRange) {
     int start = (NUM_ATOMS*get_global_id(0))/get_global_size(0);
     int end = (NUM_ATOMS*(get_global_id(0)+1))/get_global_size(0);
     int last = (start == 0 ? -1 : pmeAtomGridIndex[start-1].y);
@@ -38,7 +41,7 @@ __kernel void findAtomRangeForGrid(__global float4* posq, __global float2* pmeAt
     }
 }
 
-__kernel void updateBsplines(__global float4* posq, __global float4* pmeBsplineTheta, __global float4* pmeBsplineDTheta, __local float4* bsplinesCache, __global float2* pmeAtomGridIndex, float4 invPeriodicBoxSize) {
+__kernel void updateBsplines(__global float4* posq, __global float4* pmeBsplineTheta, __global float4* pmeBsplineDTheta, __local float4* bsplinesCache, __global float2* pmeAtomGridIndex, float4 periodicBoxSize, float4 invPeriodicBoxSize) {
     const float4 scale = 1.0f/(PME_ORDER-1);
     for (int i = get_global_id(0); i < NUM_ATOMS; i += get_global_size(0)) {
         __local float4* data = &bsplinesCache[get_local_id(0)*PME_ORDER];
@@ -48,9 +51,12 @@ __kernel void updateBsplines(__global float4* posq, __global float4* pmeBsplineT
             ddata[j] = 0.0f;
         }
         float4 pos = posq[i];
-        float4 t = (float4) ((pos.x*invPeriodicBoxSize.x+1.0f)*GRID_SIZE_X,
-                             (pos.y*invPeriodicBoxSize.y+1.0f)*GRID_SIZE_Y,
-                             (pos.z*invPeriodicBoxSize.z+1.0f)*GRID_SIZE_Z, 0.0f);
+        pos.x -= floor(pos.x*invPeriodicBoxSize.x)*periodicBoxSize.x;
+        pos.y -= floor(pos.y*invPeriodicBoxSize.y)*periodicBoxSize.y;
+        pos.z -= floor(pos.z*invPeriodicBoxSize.z)*periodicBoxSize.z;
+        float4 t = (float4) ((pos.x*invPeriodicBoxSize.x)*GRID_SIZE_X,
+                             (pos.y*invPeriodicBoxSize.y)*GRID_SIZE_Y,
+                             (pos.z*invPeriodicBoxSize.z)*GRID_SIZE_Z, 0.0f);
         float4 dr = (float4) (t.x-(int) t.x, t.y-(int) t.y, t.z-(int) t.z, 0.0f);
         data[PME_ORDER-1] = 0.0f;
         data[1] = dr;
@@ -148,13 +154,16 @@ __kernel void reciprocalConvolution(__global float2* pmeGrid, __global float* en
     energyBuffer[get_global_id(0)] += 0.5f*energy;
 }
 
-__kernel void gridInterpolateForce(__global float4* posq, __global float4* forceBuffers, __global float4* pmeBsplineTheta, __global float4* pmeBsplineDTheta, __global float2* pmeGrid, float4 invPeriodicBoxSize) {
+__kernel void gridInterpolateForce(__global float4* posq, __global float4* forceBuffers, __global float4* pmeBsplineTheta, __global float4* pmeBsplineDTheta, __global float2* pmeGrid, float4 periodicBoxSize, float4 invPeriodicBoxSize) {
     for (int atom = get_global_id(0); atom < NUM_ATOMS; atom += get_global_size(0)) {
         float4 force = 0.0f;
         float4 pos = posq[atom];
-        float4 t = (float4) ((pos.x*invPeriodicBoxSize.x+1.0f)*GRID_SIZE_X,
-                             (pos.y*invPeriodicBoxSize.y+1.0f)*GRID_SIZE_Y,
-                             (pos.z*invPeriodicBoxSize.z+1.0f)*GRID_SIZE_Z, 0.0f);
+        pos.x -= floor(pos.x*invPeriodicBoxSize.x)*periodicBoxSize.x;
+        pos.y -= floor(pos.y*invPeriodicBoxSize.y)*periodicBoxSize.y;
+        pos.z -= floor(pos.z*invPeriodicBoxSize.z)*periodicBoxSize.z;
+        float4 t = (float4) ((pos.x*invPeriodicBoxSize.x)*GRID_SIZE_X,
+                             (pos.y*invPeriodicBoxSize.y)*GRID_SIZE_Y,
+                             (pos.z*invPeriodicBoxSize.z)*GRID_SIZE_Z, 0.0f);
         int4 gridIndex = (int4) (((int) t.x) % GRID_SIZE_X,
                                  ((int) t.y) % GRID_SIZE_Y,
                                  ((int) t.z) % GRID_SIZE_Z, 0);

platforms/reference/src/SimTKReference/PME.cpp

@@ -241,7 +241,9 @@ pme_update_grid_index_and_fraction(pme_t    pme,
              *    numerical problems, so this shouldnt cause any problems.
              *    (And, by adding 100.0 box lengths, we would lose a bit of numerical accuracy here!)
              */
-            t  = (atomCoordinates[i][d] / periodicBoxSize[d] + 1)*pme->ngrid[d];
+            RealOpenMM coord = atomCoordinates[i][d];
+            coord -= floor(coord/periodicBoxSize[d])*periodicBoxSize[d];
+            t  = (coord / periodicBoxSize[d])*pme->ngrid[d];
             ti = (int) t;
 
             pme->particlefraction[i][d] = t - ti;