Fixed bug where MonteCarloBarostat would cause molecules to be reordered (see bug 1753)

platforms/cuda/src/CudaKernels.cpp

@@ -5217,16 +5217,45 @@ void CudaApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context, d
     cu.executeKernel(kernel, args, cu.getNumAtoms());
     for (int i = 0; i < (int) cu.getPosCellOffsets().size(); i++)
         cu.getPosCellOffsets()[i] = make_int4(0, 0, 0, 0);
+    lastAtomOrder = cu.getAtomIndex();
 }
 
 void CudaApplyMonteCarloBarostatKernel::restoreCoordinates(ContextImpl& context) {
     cu.setAsCurrent();
-    int bytesToCopy = cu.getPosq().getSize()*(cu.getUseDoublePrecision() ? sizeof(double4) : sizeof(float4));
-    CUresult result = cuMemcpyDtoD(cu.getPosq().getDevicePointer(), savedPositions->getDevicePointer(), bytesToCopy);
-    if (result != CUDA_SUCCESS) {
-        std::stringstream m;
-        m<<"Error restoring positions for MC barostat: "<<cu.getErrorString(result)<<" ("<<result<<")";
-        throw OpenMMException(m.str());
+    if (cu.getAtomsWereReordered()) {
+        // The atoms were reordered since we saved the positions, so we need to fix them.
+        
+        const vector<int> atomOrder = cu.getAtomIndex();
+        int numAtoms = cu.getNumAtoms();
+        if (cu.getUseDoublePrecision()) {
+            double4* pos = (double4*) cu.getPinnedBuffer();
+            savedPositions->download(pos);
+            vector<double4> fixedPos(cu.getPaddedNumAtoms());
+            for (int i = 0; i < numAtoms; i++)
+                fixedPos[lastAtomOrder[i]] = pos[i];
+            for (int i = 0; i < numAtoms; i++)
+                pos[i] = fixedPos[atomOrder[i]];
+            cu.getPosq().upload(pos);
+        }
+        else {
+            float4* pos = (float4*) cu.getPinnedBuffer();
+            savedPositions->download(pos);
+            vector<float4> fixedPos(cu.getPaddedNumAtoms());
+            for (int i = 0; i < numAtoms; i++)
+                fixedPos[lastAtomOrder[i]] = pos[i];
+            for (int i = 0; i < numAtoms; i++)
+                pos[i] = fixedPos[atomOrder[i]];
+            cu.getPosq().upload(pos);
+        }
+    }
+    else {
+        int bytesToCopy = cu.getPosq().getSize()*(cu.getUseDoublePrecision() ? sizeof(double4) : sizeof(float4));
+        CUresult result = cuMemcpyDtoD(cu.getPosq().getDevicePointer(), savedPositions->getDevicePointer(), bytesToCopy);
+        if (result != CUDA_SUCCESS) {
+            std::stringstream m;
+            m<<"Error restoring positions for MC barostat: "<<cu.getErrorString(result)<<" ("<<result<<")";
+            throw OpenMMException(m.str());
+        }
     }
 }
 

platforms/cuda/src/CudaKernels.h

@@ -1273,6 +1273,7 @@ private:
     CudaArray* moleculeAtoms;
     CudaArray* moleculeStartIndex;
     CUfunction kernel;
+    std::vector<int> lastAtomOrder;
 };
 
 /**

platforms/opencl/src/OpenCLKernels.cpp

@@ -5469,10 +5469,38 @@ void OpenCLApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context,
     cl.executeKernel(kernel, cl.getNumAtoms());
     for (int i = 0; i < (int) cl.getPosCellOffsets().size(); i++)
         cl.getPosCellOffsets()[i] = mm_int4(0, 0, 0, 0);
+    lastAtomOrder = cl.getAtomIndex();
 }
 
 void OpenCLApplyMonteCarloBarostatKernel::restoreCoordinates(ContextImpl& context) {
-    cl.getQueue().enqueueCopyBuffer(savedPositions->getDeviceBuffer(), cl.getPosq().getDeviceBuffer(), 0, 0, cl.getPosq().getSize()*sizeof(mm_float4));
+    if (cl.getAtomsWereReordered()) {
+        // The atoms were reordered since we saved the positions, so we need to fix them.
+        
+        const vector<int> atomOrder = cl.getAtomIndex();
+        int numAtoms = cl.getNumAtoms();
+        if (cl.getUseDoublePrecision()) {
+            mm_double4* pos = (mm_double4*) cl.getPinnedBuffer();
+            savedPositions->download(pos);
+            vector<mm_double4> fixedPos(cl.getPaddedNumAtoms());
+            for (int i = 0; i < numAtoms; i++)
+                fixedPos[lastAtomOrder[i]] = pos[i];
+            for (int i = 0; i < numAtoms; i++)
+                pos[i] = fixedPos[atomOrder[i]];
+            cl.getPosq().upload(pos);
+        }
+        else {
+            mm_float4* pos = (mm_float4*) cl.getPinnedBuffer();
+            savedPositions->download(pos);
+            vector<mm_float4> fixedPos(cl.getPaddedNumAtoms());
+            for (int i = 0; i < numAtoms; i++)
+                fixedPos[lastAtomOrder[i]] = pos[i];
+            for (int i = 0; i < numAtoms; i++)
+                pos[i] = fixedPos[atomOrder[i]];
+            cl.getPosq().upload(pos);
+        }
+    }
+    else
+        cl.getQueue().enqueueCopyBuffer(savedPositions->getDeviceBuffer(), cl.getPosq().getDeviceBuffer(), 0, 0, cl.getPosq().getSize()*sizeof(mm_float4));
 }
 
 OpenCLRemoveCMMotionKernel::~OpenCLRemoveCMMotionKernel() {

platforms/opencl/src/OpenCLKernels.h

@@ -1282,6 +1282,7 @@ private:
     OpenCLArray* moleculeAtoms;
     OpenCLArray* moleculeStartIndex;
     cl::Kernel kernel;
+    std::vector<int> lastAtomOrder;
 };
 
 /**