Attempt at fixing errors with barostat (#4106)

* Attempt at fixing errors with barostat

* Missing ContextSelector

olla/include/openmm/kernels.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:                                                              *
  *                                                                            *
@@ -1458,6 +1458,13 @@ public:
      * @param rigidMolecules  whether molecules should be kept rigid while scaling coordinates
      */
     virtual void initialize(const System& system, const Force& barostat, bool rigidMolecules=true) = 0;
+    /**
+     * Save the coordinates before attempting a Monte Carlo step.  This allows us to restore them
+     * if the step is rejected.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    virtual void saveCoordinates(ContextImpl& context) = 0;
     /**
      * Attempt a Monte Carlo step, scaling particle positions (or cluster centers) by a specified value.
      * This version scales the x, y, and z positions independently.
@@ -1472,8 +1479,8 @@ public:
      */
     virtual void scaleCoordinates(ContextImpl& context, double scaleX, double scaleY, double scaleZ) = 0;
     /**
-     * Reject the most recent Monte Carlo step, restoring the particle positions to where they were before
-     * scaleCoordinates() was last called.
+     * Reject the most recent Monte Carlo step, restoring the particle positions to where they were when
+     * saveCoordinates() was last called.
      *
      * @param context    the context in which to execute this kernel
      */

openmmapi/src/MonteCarloAnisotropicBarostatImpl.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2010-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2010-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman, Lee-Ping Wang                                      *
  * Contributors:                                                              *
  *                                                                            *
@@ -107,10 +107,11 @@ void MonteCarloAnisotropicBarostatImpl::updateContextState(ContextImpl& context,
     double newVolume = volume+deltaVolume;
     Vec3 lengthScale(1.0, 1.0, 1.0);
     lengthScale[axis] = newVolume/volume;
-    kernel.getAs<ApplyMonteCarloBarostatKernel>().scaleCoordinates(context, lengthScale[0], lengthScale[1], lengthScale[2]);
+    kernel.getAs<ApplyMonteCarloBarostatKernel>().saveCoordinates(context);
     context.getOwner().setPeriodicBoxVectors(Vec3(box[0][0]*lengthScale[0], box[0][1]*lengthScale[1], box[0][2]*lengthScale[2]),
                                              Vec3(box[1][0]*lengthScale[0], box[1][1]*lengthScale[1], box[1][2]*lengthScale[2]),
                                              Vec3(box[2][0]*lengthScale[0], box[2][1]*lengthScale[1], box[2][2]*lengthScale[2]));
+    kernel.getAs<ApplyMonteCarloBarostatKernel>().scaleCoordinates(context, lengthScale[0], lengthScale[1], lengthScale[2]);
 
     // Compute the energy of the modified system.
     

openmmapi/src/MonteCarloBarostatImpl.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2010-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2010-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -78,8 +78,9 @@ void MonteCarloBarostatImpl::updateContextState(ContextImpl& context, bool& forc
     double deltaVolume = volumeScale*2*(SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber()-0.5);
     double newVolume = volume+deltaVolume;
     double lengthScale = pow(newVolume/volume, 1.0/3.0);
-    kernel.getAs<ApplyMonteCarloBarostatKernel>().scaleCoordinates(context, lengthScale, lengthScale, lengthScale);
+    kernel.getAs<ApplyMonteCarloBarostatKernel>().saveCoordinates(context);
     context.getOwner().setPeriodicBoxVectors(box[0]*lengthScale, box[1]*lengthScale, box[2]*lengthScale);
+    kernel.getAs<ApplyMonteCarloBarostatKernel>().scaleCoordinates(context, lengthScale, lengthScale, lengthScale);
 
     // Compute the energy of the modified system.
     

openmmapi/src/MonteCarloFlexibleBarostatImpl.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2010-2021 Stanford University and the Authors.      *
+ * Portions copyright (c) 2010-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman, Sander Vandenhaute                                 *
  * Contributors:                                                              *
  *                                                                            *
@@ -95,8 +95,9 @@ void MonteCarloFlexibleBarostatImpl::updateContextState(ContextImpl& context, bo
 
     // Scale particle coordinates and update box vectors in context.
 
-    kernel.getAs<ApplyMonteCarloBarostatKernel>().scaleCoordinates(context, trial[0][0]/box[0][0], trial[1][1]/box[1][1], trial[2][2]/box[2][2]);
+    kernel.getAs<ApplyMonteCarloBarostatKernel>().saveCoordinates(context);
     context.getOwner().setPeriodicBoxVectors(trial[0], trial[1], trial[2]);
+    kernel.getAs<ApplyMonteCarloBarostatKernel>().scaleCoordinates(context, trial[0][0]/box[0][0], trial[1][1]/box[1][1], trial[2][2]/box[2][2]);
 
     // Compute the energy of the modified system.
 

openmmapi/src/MonteCarloMembraneBarostatImpl.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2010-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2010-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman, Lee-Ping Wang                                      *
  * Contributors:                                                              *
  *                                                                            *
@@ -108,10 +108,11 @@ void MonteCarloMembraneBarostatImpl::updateContextState(ContextImpl& context, bo
         deltaVolume = 0;
     }
     double deltaArea = box[0][0]*lengthScale[0]*box[1][1]*lengthScale[1] - box[0][0]*box[1][1];
-    kernel.getAs<ApplyMonteCarloBarostatKernel>().scaleCoordinates(context, lengthScale[0], lengthScale[1], lengthScale[2]);
+    kernel.getAs<ApplyMonteCarloBarostatKernel>().saveCoordinates(context);
     context.getOwner().setPeriodicBoxVectors(Vec3(box[0][0]*lengthScale[0], box[0][1]*lengthScale[1], box[0][2]*lengthScale[2]),
                                              Vec3(box[1][0]*lengthScale[0], box[1][1]*lengthScale[1], box[1][2]*lengthScale[2]),
                                              Vec3(box[2][0]*lengthScale[0], box[2][1]*lengthScale[1], box[2][2]*lengthScale[2]));
+    kernel.getAs<ApplyMonteCarloBarostatKernel>().scaleCoordinates(context, lengthScale[0], lengthScale[1], lengthScale[2]);
 
     // Compute the energy of the modified system.
 

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

@@ -1547,6 +1547,13 @@ public:
      * @param rigidMolecules  whether molecules should be kept rigid while scaling coordinates
      */
     void initialize(const System& system, const Force& barostat, bool rigidMolecules=true);
+    /**
+     * Save the coordinates before attempting a Monte Carlo step.  This allows us to restore them
+     * if the step is rejected.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    void saveCoordinates(ContextImpl& context);
     /**
      * Attempt a Monte Carlo step, scaling particle positions (or cluster centers) by a specified value.
      * This version scales the x, y, and z positions independently.
@@ -1561,8 +1568,8 @@ public:
      */
     void scaleCoordinates(ContextImpl& context, double scaleX, double scaleY, double scaleZ);
     /**
-     * Reject the most recent Monte Carlo step, restoring the particle positions to where they were before
-     * scaleCoordinates() was last called.
+     * Reject the most recent Monte Carlo step, restoring the particle positions to where they were when
+     * saveCoordinates() was last called.
      *
      * @param context    the context in which to execute this kernel
      */

platforms/common/src/CommonKernels.cpp

@@ -7713,6 +7713,15 @@ void CommonApplyMonteCarloBarostatKernel::initialize(const System& system, const
     kernel = program->createKernel("scalePositions");
 }
 
+void CommonApplyMonteCarloBarostatKernel::saveCoordinates(ContextImpl& context) {
+    ContextSelector selector(cc);
+    cc.getPosq().copyTo(savedPositions);
+    cc.getLongForceBuffer().copyTo(savedLongForces);
+    if (savedFloatForces.isInitialized())
+        cc.getFloatForceBuffer().copyTo(savedFloatForces);
+    lastPosCellOffsets = cc.getPosCellOffsets();
+}
+
 void CommonApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context, double scaleX, double scaleY, double scaleZ) {
     ContextSelector selector(cc);
     if (!hasInitializedKernels) {
@@ -7755,11 +7764,6 @@ void CommonApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context,
         kernel->addArg(moleculeAtoms);
         kernel->addArg(moleculeStartIndex);
     }
-    cc.getPosq().copyTo(savedPositions);
-    cc.getLongForceBuffer().copyTo(savedLongForces);
-    if (savedFloatForces.isInitialized())
-        cc.getFloatForceBuffer().copyTo(savedFloatForces);
-    lastPosCellOffsets = cc.getPosCellOffsets();
     kernel->setArg(0, (float) scaleX);
     kernel->setArg(1, (float) scaleY);
     kernel->setArg(2, (float) scaleZ);

platforms/reference/include/ReferenceKernels.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:                                                              *
  *                                                                            *
@@ -1575,6 +1575,13 @@ public:
      * @param rigidMolecules  whether molecules should be kept rigid while scaling coordinates
      */
     void initialize(const System& system, const Force& barostat, bool rigidMolecules=true);
+    /**
+     * Save the coordinates before attempting a Monte Carlo step.  This allows us to restore them
+     * if the step is rejected.
+     *
+     * @param context    the context in which to execute this kernel
+     */
+    void saveCoordinates(ContextImpl& context);
     /**
      * Attempt a Monte Carlo step, scaling particle positions (or cluster centers) by a specified value.
      * This version scales the x, y, and z positions independently.
@@ -1589,8 +1596,8 @@ public:
      */
     void scaleCoordinates(ContextImpl& context, double scaleX, double scaleY, double scaleZ);
     /**
-     * Reject the most recent Monte Carlo step, restoring the particle positions to where they were before
-     * scaleCoordinates() was last called.
+     * Reject the most recent Monte Carlo step, restoring the particle positions to where they were when
+     * saveCoordinates() was last called.
      *
      * @param context    the context in which to execute this kernel
      */

platforms/reference/include/ReferenceMonteCarloBarostat.h

@@ -56,6 +56,16 @@ class ReferenceMonteCarloBarostat {
 
        ~ReferenceMonteCarloBarostat();
 
+      /**---------------------------------------------------------------------------------------
+
+         Save the positions before applying the barostat.
+
+         @param atomPositions      atom positions
+
+         --------------------------------------------------------------------------------------- */
+
+      void savePositions(std::vector<OpenMM::Vec3>& atomPositions);
+
       /**---------------------------------------------------------------------------------------
 
          Apply the barostat at the start of a time step, scaling x, y, and z coordinates independently.

platforms/reference/src/ReferenceKernels.cpp

@@ -6,7 +6,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:                                                              *
  *                                                                            *
@@ -35,7 +35,6 @@
 #include "ReferenceAngleBondIxn.h"
 #include "ReferenceBondForce.h"
 #include "ReferenceBrownianDynamics.h"
-#include "ReferenceCCMAAlgorithm.h"
 #include "ReferenceCMAPTorsionIxn.h"
 #include "ReferenceConstraints.h"
 #include "ReferenceCustomAngleIxn.h"
@@ -2897,7 +2896,7 @@ void ReferenceApplyMonteCarloBarostatKernel::initialize(const System& system, co
     this->rigidMolecules = rigidMolecules;
 }
 
-void ReferenceApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context, double scaleX, double scaleY, double scaleZ) {
+void ReferenceApplyMonteCarloBarostatKernel::saveCoordinates(ContextImpl& context) {
     if (barostat == NULL) {
         if (rigidMolecules)
             barostat = new ReferenceMonteCarloBarostat(context.getSystem().getNumParticles(), context.getMolecules());
@@ -2908,6 +2907,11 @@ void ReferenceApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& conte
             barostat = new ReferenceMonteCarloBarostat(context.getSystem().getNumParticles(), molecules);
         }
     }
+    vector<Vec3>& posData = extractPositions(context);
+    barostat->savePositions(posData);
+}
+
+void ReferenceApplyMonteCarloBarostatKernel::scaleCoordinates(ContextImpl& context, double scaleX, double scaleY, double scaleZ) {
     vector<Vec3>& posData = extractPositions(context);
     Vec3* boxVectors = extractBoxVectors(context);
     barostat->applyBarostat(posData, boxVectors, scaleX, scaleY, scaleZ);

platforms/reference/src/SimTKReference/ReferenceMonteCarloBarostat.cpp

@@ -52,6 +52,21 @@ ReferenceMonteCarloBarostat::ReferenceMonteCarloBarostat(int numAtoms, const vec
 ReferenceMonteCarloBarostat::~ReferenceMonteCarloBarostat() {
 }
 
+/**---------------------------------------------------------------------------------------
+
+  Save the positions before applying the barostat.
+
+  @param atomPositions      atom positions
+
+  --------------------------------------------------------------------------------------- */
+
+void ReferenceMonteCarloBarostat::savePositions(vector<Vec3>& atomPositions) {
+    int numAtoms = savedAtomPositions[0].size();
+    for (int i = 0; i < numAtoms; i++)
+        for (int j = 0; j < 3; j++)
+            savedAtomPositions[j][i] = atomPositions[i][j];
+}
+
 /**---------------------------------------------------------------------------------------
 
   Apply the barostat at the start of a time step.
@@ -65,11 +80,6 @@ ReferenceMonteCarloBarostat::~ReferenceMonteCarloBarostat() {
   --------------------------------------------------------------------------------------- */
 
 void ReferenceMonteCarloBarostat::applyBarostat(vector<Vec3>& atomPositions, const Vec3* boxVectors, double scaleX, double scaleY, double scaleZ) {
-    int numAtoms = savedAtomPositions[0].size();
-    for (int i = 0; i < numAtoms; i++)
-        for (int j = 0; j < 3; j++)
-            savedAtomPositions[j][i] = atomPositions[i][j];
-
     // Loop over molecules.
 
     for (auto& molecule : molecules) {
@@ -98,7 +108,7 @@ void ReferenceMonteCarloBarostat::applyBarostat(vector<Vec3>& atomPositions, con
 
 /**---------------------------------------------------------------------------------------
 
-  Restore atom positions to what they were before applyBarostat() was called.
+  Restore atom positions to what they were before savePositions() was called.
 
   @param atomPositions      atom positions
 

plugins/rpmd/openmmapi/src/RPMDMonteCarloBarostatImpl.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2010-2015 Stanford University and the Authors.      *
+ * Portions copyright (c) 2010-2023 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -107,8 +107,9 @@ void RPMDMonteCarloBarostatImpl::updateRPMDState(ContextImpl& context) {
     double newVolume = volume+deltaVolume;
     double lengthScale = std::pow(newVolume/volume, 1.0/3.0);
     context.setPositions(centroid);
-    kernel.getAs<ApplyMonteCarloBarostatKernel>().scaleCoordinates(context, lengthScale, lengthScale, lengthScale);
+    kernel.getAs<ApplyMonteCarloBarostatKernel>().saveCoordinates(context);
     context.getOwner().setPeriodicBoxVectors(box[0]*lengthScale, box[1]*lengthScale, box[2]*lengthScale);
+    kernel.getAs<ApplyMonteCarloBarostatKernel>().scaleCoordinates(context, lengthScale, lengthScale, lengthScale);
     State scaledState = context.getOwner().getState(State::Positions);
 
     // Now apply the same offset to all the copies.