Ensure valid atom order after loading a checkpoint (#3771)

platforms/common/include/openmm/common/ComputeContext.h

@@ -496,6 +496,11 @@ public:
      * definitions and order to be revalidated.
      */
     bool invalidateMolecules(ComputeForceInfo* force);
+    /**
+     * Make sure the current atom order is valid, based on the forces.  If not, perform reordering
+     * to generate a new valid order.  This method is only needed in very unusual situations.
+     */
+    void validateAtomOrder();
     /**
      * Wait until all work that has been queued (kernel executions, asynchronous data transfers, etc.)
      * has been submitted to the device.  This does not mean it has necessarily been completed.
@@ -508,6 +513,7 @@ protected:
     struct MoleculeGroup;
     class VirtualSiteInfo;
     void findMoleculeGroups();
+    void resetAtomOrder();
     /**
      * This is the internal implementation of reorderAtoms(), templatized by the numerical precision in use.
      */

platforms/common/src/ComputeContext.cpp

@@ -376,6 +376,13 @@ bool ComputeContext::invalidateMolecules(ComputeForceInfo* force) {
     // atoms to their original order, rebuild the list of identical molecules, and sort them
     // again.
 
+    resetAtomOrder();
+    findMoleculeGroups();
+    reorderAtoms();
+    return true;
+}
+
+void ComputeContext::resetAtomOrder() {
     ContextSelector selector(*this);
     vector<mm_int4> newCellOffsets(numAtoms);
     if (getUseDoublePrecision()) {
@@ -436,12 +443,25 @@ bool ComputeContext::invalidateMolecules(ComputeForceInfo* force) {
         posCellOffsets[i] = newCellOffsets[i];
     }
     getAtomIndexArray().upload(atomIndex);
-    findMoleculeGroups();
     for (auto listener : reorderListeners)
         listener->execute();
     forceNextReorder = true;
+}
+
+void ComputeContext::validateAtomOrder() {
+    for (auto& mol : moleculeGroups) {
+        for (int atom : mol.atoms) {
+            set<int> identical;
+            for (int offset : mol.offsets)
+                identical.insert(atom+offset);
+            for (int i : identical)
+                if (identical.find(atomIndex[i]) == identical.end()) {
+                    resetAtomOrder();
                     reorderAtoms();
-    return true;
+                    return;
+                }
+        }
+    }
 }
 
 void ComputeContext::forceReorder() {

platforms/cuda/src/CudaKernels.cpp

@@ -424,6 +424,7 @@ void CudaUpdateStateDataKernel::loadCheckpoint(ContextImpl& context, istream& st
     SimTKOpenMMUtilities::loadCheckpoint(stream);
     for (auto listener : cu.getReorderListeners())
         listener->execute();
+    cu.validateAtomOrder();
 }
 
 class CudaCalcNonbondedForceKernel::ForceInfo : public CudaForceInfo {

platforms/opencl/src/OpenCLKernels.cpp

@@ -441,6 +441,7 @@ void OpenCLUpdateStateDataKernel::loadCheckpoint(ContextImpl& context, istream&
     SimTKOpenMMUtilities::loadCheckpoint(stream);
     for (auto listener : cl.getReorderListeners())
         listener->execute();
+    cl.validateAtomOrder();
 }
 
 class OpenCLCalcNonbondedForceKernel::ForceInfo : public OpenCLForceInfo {