Prevent molecules from accidentally getting wrapped by barostat (#4034)

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

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -1576,6 +1576,7 @@ private:
     ComputeArray moleculeStartIndex;
     ComputeKernel kernel;
     std::vector<int> lastAtomOrder;
+    std::vector<mm_int4> lastPosCellOffsets;
 };
 
 } // namespace OpenMM

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

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2019-2022 Stanford University and the Authors.      *
+ * Portions copyright (c) 2019-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -397,6 +397,12 @@ public:
     std::vector<mm_int4>& getPosCellOffsets() {
         return posCellOffsets;
     }
+    /**
+     * Set the number of cells by which the positions are offset.
+     */
+    void setPosCellOffsets(std::vector<mm_int4>& offsets) {
+        posCellOffsets = offsets;
+    }
     /**
      * Replace all occurrences of a list of substrings.
      *

platforms/common/src/CommonKernels.cpp

@@ -7787,12 +7787,14 @@ void CommonApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context,
     setPeriodicBoxArgs(cc, kernel, 4);
     kernel->execute(cc.getNumAtoms());
     lastAtomOrder = cc.getAtomIndex();
+    lastPosCellOffsets = cc.getPosCellOffsets();
 }
 
 void CommonApplyMonteCarloBarostatKernel::restoreCoordinates(ContextImpl& context) {
     ContextSelector selector(cc);
     savedPositions.copyTo(cc.getPosq());
     savedLongForces.copyTo(cc.getLongForceBuffer());
+    cc.setPosCellOffsets(lastPosCellOffsets);
     if (savedFloatForces.isInitialized())
         savedFloatForces.copyTo(cc.getFloatForceBuffer());
 }