another attempted fix for MonteCarloBarostat (#4119)

* save atomIndex order in MCBarostat kernel

* Update device indices

* Only setAtomIndex if there was a reorder

platforms/common/include/openmm/common/CommonKernels.h

@@ -1576,7 +1576,7 @@ public:
     void restoreCoordinates(ContextImpl& context);
 private:
     ComputeContext& cc;
-    bool hasInitializedKernels, rigidMolecules;
+    bool hasInitializedKernels, rigidMolecules, atomsWereReordered;
     int numMolecules;
     ComputeArray savedPositions, savedFloatForces, savedLongForces;
     ComputeArray moleculeAtoms;

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

@@ -387,6 +387,10 @@ public:
     const std::vector<int>& getAtomIndex() const {
         return atomIndex;
     }
+    /**
+     * Set the vector which contains the index of each atom.
+     */
+    void setAtomIndex(std::vector<int>& index);
     /**
      * Get the array which contains the index of each atom.
      */

platforms/common/src/CommonKernels.cpp

@@ -7720,10 +7720,15 @@ void CommonApplyMonteCarloBarostatKernel::saveCoordinates(ContextImpl& context)
     if (savedFloatForces.isInitialized())
         cc.getFloatForceBuffer().copyTo(savedFloatForces);
     lastPosCellOffsets = cc.getPosCellOffsets();
+    lastAtomOrder = cc.getAtomIndex();
 }
 
 void CommonApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context, double scaleX, double scaleY, double scaleZ) {
     ContextSelector selector(cc);
+
+    // check if atoms were reordered from energy evaluation before scaling
+    atomsWereReordered = cc.getAtomsWereReordered();
+
     if (!hasInitializedKernels) {
         hasInitializedKernels = true;
 
@@ -7769,7 +7774,6 @@ void CommonApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context,
     kernel->setArg(2, (float) scaleZ);
     setPeriodicBoxArgs(cc, kernel, 4);
     kernel->execute(cc.getNumAtoms());
-    lastAtomOrder = cc.getAtomIndex();
 }
 
 void CommonApplyMonteCarloBarostatKernel::restoreCoordinates(ContextImpl& context) {
@@ -7779,4 +7783,8 @@ void CommonApplyMonteCarloBarostatKernel::restoreCoordinates(ContextImpl& contex
     cc.setPosCellOffsets(lastPosCellOffsets);
     if (savedFloatForces.isInitialized())
         savedFloatForces.copyTo(cc.getFloatForceBuffer());
+
+    // check if atoms were reordered from energy evaluation before or after scaling
+    if (atomsWereReordered || cc.getAtomsWereReordered())
+        cc.setAtomIndex(lastAtomOrder);
 }
\ No newline at end of file

platforms/common/src/ComputeContext.cpp

@@ -57,6 +57,13 @@ void ComputeContext::addForce(ComputeForceInfo* force) {
     forces.push_back(force);
 }
 
+void ComputeContext::setAtomIndex(std::vector<int>& index){
+    atomIndex = index;
+    getAtomIndexArray().upload(atomIndex);
+    for (auto listener : reorderListeners)
+        listener->execute();
+}
+
 string ComputeContext::replaceStrings(const string& input, const std::map<std::string, std::string>& replacements) const {
     static set<char> symbolChars;
     if (symbolChars.size() == 0) {