Implemented velocity constraints for CustomIntegrator

openmmapi/src/CustomIntegrator.cpp

@@ -72,7 +72,6 @@ vector<string> CustomIntegrator::getKernelNames() {
 void CustomIntegrator::step(int steps) {
     globalsAreCurrent = false;
     for (int i = 0; i < steps; ++i) {
-        context->updateContextState();
         dynamic_cast<IntegrateCustomStepKernel&>(kernel.getImpl()).execute(*context, *this, forcesAreValid);
     }
 }

platforms/reference/src/ReferenceKernels.cpp

@@ -145,18 +145,22 @@ void ReferenceCalcForcesAndEnergyKernel::initialize(const System& system) {
 }
 
 void ReferenceCalcForcesAndEnergyKernel::beginComputation(ContextImpl& context, bool includeForces, bool includeEnergy) {
+    vector<RealVec>& forceData = extractForces(context);
     if (includeForces) {
         int numParticles = context.getSystem().getNumParticles();
-        vector<RealVec>& forceData = extractForces(context);
         for (int i = 0; i < numParticles; ++i) {
             forceData[i][0] = (RealOpenMM) 0.0;
             forceData[i][1] = (RealOpenMM) 0.0;
             forceData[i][2] = (RealOpenMM) 0.0;
         }
     }
+    else
+        savedForces = forceData;
 }
 
 double ReferenceCalcForcesAndEnergyKernel::finishComputation(ContextImpl& context, bool includeForces, bool includeEnergy) {
+    if (!includeForces)
+        extractForces(context) = savedForces; // Restore the forces so computing the energy doesn't overwrite the forces with incorrect values.
     return 0.0;
 }
 

platforms/reference/src/ReferenceKernels.h

@@ -89,6 +89,8 @@ public:
      * energy directly, <i>or</i> add it to an internal buffer so that it will be included here.
      */
     double finishComputation(ContextImpl& context, bool includeForce, bool includeEnergy);
+private:
+    std::vector<RealVec> savedForces;
 };
 
 /**

platforms/reference/src/SimTKReference/ReferenceCCMAAlgorithm.cpp

@@ -88,9 +88,7 @@ ReferenceCCMAAlgorithm::ReferenceCCMAAlgorithm( int numberOfAtoms,
    // work arrays
 
    if (_numberOfConstraints > 0) {
-       _r_ij                       = SimTKOpenMMUtilities::allocateTwoDRealOpenMMArray( numberOfConstraints, threeI, NULL,
-                                                                                        1, zero, "r_ij" );
-
+       _r_ij.resize(numberOfConstraints);
        _d_ij2                      = SimTKOpenMMUtilities::allocateOneDRealOpenMMArray( numberOfConstraints, NULL, 1, zero, "dij_2" );
        _distanceTolerance          = SimTKOpenMMUtilities::allocateOneDRealOpenMMArray( numberOfConstraints, NULL, 1, zero, "distanceTolerance" );
        _reducedMasses              = SimTKOpenMMUtilities::allocateOneDRealOpenMMArray( numberOfConstraints, NULL, 1, zero, "reducedMasses" );
@@ -225,8 +223,6 @@ ReferenceCCMAAlgorithm::~ReferenceCCMAAlgorithm( ){
    // ---------------------------------------------------------------------------------------
 
     if (_numberOfConstraints > 0) {
-       SimTKOpenMMUtilities::freeTwoDRealOpenMMArray( _r_ij,  "r_ij" );
-
        SimTKOpenMMUtilities::freeOneDRealOpenMMArray( _d_ij2, "d_ij2" );
        SimTKOpenMMUtilities::freeOneDRealOpenMMArray( _distanceTolerance, "distanceTolerance" );
        SimTKOpenMMUtilities::freeOneDRealOpenMMArray( _reducedMasses, "reducedMasses" );
@@ -348,7 +344,31 @@ void ReferenceCCMAAlgorithm::setTolerance( RealOpenMM tolerance ){
 int ReferenceCCMAAlgorithm::apply( int numberOfAtoms, vector<RealVec>& atomCoordinates,
                                          vector<RealVec>& atomCoordinatesP,
                                          vector<RealOpenMM>& inverseMasses ){
+    applyConstraints(numberOfAtoms, atomCoordinates, atomCoordinatesP, inverseMasses, false);
+}
+
+/**---------------------------------------------------------------------------------------
+
+   Apply constraint algorithm to velocities.
+
+   @param numberOfAtoms    number of atoms
+   @param atomCoordinates  atom coordinates
+   @param velocities       atom velocities
+   @param inverseMasses    1/mass
+
+   @return SimTKOpenMMCommon::DefaultReturn if converge; else
+    return SimTKOpenMMCommon::ErrorReturn
+
+   --------------------------------------------------------------------------------------- */
 
+int ReferenceCCMAAlgorithm::applyToVelocities(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
+               std::vector<OpenMM::RealVec>& velocities, std::vector<RealOpenMM>& inverseMasses) {
+    applyConstraints(numberOfAtoms, atomCoordinates, velocities, inverseMasses, true);
+}
+
+int ReferenceCCMAAlgorithm::applyConstraints(int numberOfAtoms, vector<RealVec>& atomCoordinates,
+                                         vector<RealVec>& atomCoordinatesP,
+                                         vector<RealOpenMM>& inverseMasses, bool constrainingVelocities){
    // ---------------------------------------------------------------------------------------
 
    static const char* methodName = "\nReferenceCCMAAlgorithm::apply";
@@ -364,9 +384,8 @@ int ReferenceCCMAAlgorithm::apply( int numberOfAtoms, vector<RealVec>& atomCoord
 
    // temp arrays
 
-   RealOpenMM** r_ij                   = _r_ij;
+   vector<RealVec>& r_ij               = _r_ij;
    RealOpenMM* d_ij2                   = _d_ij2;
-   RealOpenMM* distanceTolerance       = _distanceTolerance;
    RealOpenMM* reducedMasses           = _reducedMasses;
 
    // calculate reduced masses on 1st pass
@@ -385,17 +404,10 @@ int ReferenceCCMAAlgorithm::apply( int numberOfAtoms, vector<RealVec>& atomCoord
    RealOpenMM tolerance     = getTolerance();
               tolerance    *= two;
    for( int ii = 0; ii < _numberOfConstraints; ii++ ){
-
       int atomI   = _atomIndices[ii][0];
       int atomJ   = _atomIndices[ii][1];
-      for( int jj = 0; jj < 3; jj++ ){
-         r_ij[ii][jj] = atomCoordinates[atomI][jj] - atomCoordinates[atomJ][jj];
-      }
-      d_ij2[ii]              = DOT3( r_ij[ii], r_ij[ii] );
-      distanceTolerance[ii]  = d_ij2[ii]*tolerance;
-      if( distanceTolerance[ii] > zero ){
-         distanceTolerance[ii] = one/distanceTolerance[ii];
-      }
+      r_ij[ii] = atomCoordinates[atomI] - atomCoordinates[atomJ];
+      d_ij2[ii] = r_ij[ii].dot(r_ij[ii]);
    }
    RealOpenMM lowerTol = one-two*getTolerance()+getTolerance()*getTolerance();
    RealOpenMM upperTol = one+two*getTolerance()+getTolerance()*getTolerance();
@@ -413,24 +425,28 @@ int ReferenceCCMAAlgorithm::apply( int numberOfAtoms, vector<RealVec>& atomCoord
          int atomI   = _atomIndices[ii][0];
          int atomJ   = _atomIndices[ii][1];
 
-         RealOpenMM rp_ij[3];
-         for( int jj = 0; jj < 3; jj++ ){
-            rp_ij[jj] = atomCoordinatesP[atomI][jj] - atomCoordinatesP[atomJ][jj];
-         }
-         RealOpenMM rp2  = DOT3( rp_ij, rp_ij );
-         RealOpenMM dist2= _distance[ii]*_distance[ii];
-         RealOpenMM diff = dist2 - rp2;
-         constraintDelta[ii] = zero;
-         RealOpenMM rrpr  = DOT3(  rp_ij, r_ij[ii] );
-         if( rrpr <  d_ij2[ii]*epsilon6 ){
-             std::stringstream message;
-             message << iterations <<" "<<atomI<<" "<<atomJ<< " Error: sign of rrpr < 0?\n";
-             SimTKOpenMMLog::printMessage( message );
-         } else {
-             constraintDelta[ii] = reducedMasses[ii]*diff/rrpr;
+         RealVec rp_ij = atomCoordinatesP[atomI] - atomCoordinatesP[atomJ];
+         if (constrainingVelocities) {
+             RealOpenMM rrpr = rp_ij.dot(r_ij[ii]);
+             constraintDelta[ii] = -2*reducedMasses[ii]*rrpr/d_ij2[ii];
+             if (constraintDelta[ii] <= getTolerance())
+                numberConverged++;
          }
-         if (rp2 >= lowerTol*dist2 && rp2 <= upperTol*dist2) {
-            numberConverged++;
+         else {
+             RealOpenMM rp2  = rp_ij.dot(rp_ij);
+             RealOpenMM dist2 = _distance[ii]*_distance[ii];
+             RealOpenMM diff = dist2 - rp2;
+             constraintDelta[ii] = zero;
+             RealOpenMM rrpr = DOT3(  rp_ij, r_ij[ii] );
+             if( rrpr <  d_ij2[ii]*epsilon6 ){
+                 std::stringstream message;
+                 message << iterations <<" "<<atomI<<" "<<atomJ<< " Error: sign of rrpr < 0?\n";
+                 SimTKOpenMMLog::printMessage( message );
+             } else {
+                 constraintDelta[ii] = reducedMasses[ii]*diff/rrpr;
+             }
+             if (rp2 >= lowerTol*dist2 && rp2 <= upperTol*dist2)
+                numberConverged++;
          }
       }
       if( numberConverged == _numberOfConstraints )
@@ -452,11 +468,9 @@ int ReferenceCCMAAlgorithm::apply( int numberOfAtoms, vector<RealVec>& atomCoord
 
          int atomI   = _atomIndices[ii][0];
          int atomJ   = _atomIndices[ii][1];
-         for( int jj = 0; jj < 3; jj++ ){
-            RealOpenMM dr                = constraintDelta[ii]*r_ij[ii][jj];
-            atomCoordinatesP[atomI][jj] += inverseMasses[atomI]*dr;
-            atomCoordinatesP[atomJ][jj] -= inverseMasses[atomJ]*dr;
-         }
+         RealVec dr = r_ij[ii]*constraintDelta[ii];
+         atomCoordinatesP[atomI] += dr*inverseMasses[atomI];
+         atomCoordinatesP[atomJ] -= dr*inverseMasses[atomJ];
       }
    }
 

platforms/reference/src/SimTKReference/ReferenceCCMAAlgorithm.h

@@ -43,13 +43,18 @@ class OPENMM_EXPORT ReferenceCCMAAlgorithm : public ReferenceConstraintAlgorithm
       int** _atomIndices;
       RealOpenMM* _distance;
 
-      RealOpenMM** _r_ij;
+      std::vector<OpenMM::RealVec> _r_ij;
       RealOpenMM* _d_ij2;
       RealOpenMM* _distanceTolerance;
       RealOpenMM* _reducedMasses;
       bool _hasInitializedMasses;
       std::vector<std::vector<std::pair<int, RealOpenMM> > > _matrix;
 
+   private:
+
+      int applyConstraints(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
+                       std::vector<OpenMM::RealVec>& atomCoordinatesP, std::vector<RealOpenMM>& inverseMasses, bool constrainingVelocities);
+          
    public:
       class AngleInfo;
 
@@ -144,6 +149,23 @@ class OPENMM_EXPORT ReferenceCCMAAlgorithm : public ReferenceConstraintAlgorithm
       int apply( int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
                        std::vector<OpenMM::RealVec>& atomCoordinatesP, std::vector<RealOpenMM>& inverseMasses );
 
+      /**---------------------------------------------------------------------------------------
+
+         Apply constraint algorithm to velocities.
+
+         @param numberOfAtoms    number of atoms
+         @param atomCoordinates  atom coordinates
+         @param velocities       atom velocities
+         @param inverseMasses    1/mass
+
+         @return SimTKOpenMMCommon::DefaultReturn if converge; else
+          return SimTKOpenMMCommon::ErrorReturn
+
+         --------------------------------------------------------------------------------------- */
+
+      int applyToVelocities(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
+                     std::vector<OpenMM::RealVec>& velocities, std::vector<RealOpenMM>& inverseMasses);
+
       /**---------------------------------------------------------------------------------------
 
          Report any violated constraints

platforms/reference/src/SimTKReference/ReferenceConstraintAlgorithm.h

@@ -68,6 +68,23 @@ public:
 
     virtual int apply(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
                      std::vector<OpenMM::RealVec>& atomCoordinatesP, std::vector<RealOpenMM>& inverseMasses) = 0;
+
+      /**---------------------------------------------------------------------------------------
+
+         Apply constraint algorithm to velocities.
+
+         @param numberOfAtoms    number of atoms
+         @param atomCoordinates  atom coordinates
+         @param velocities       atom velocities
+         @param inverseMasses    1/mass
+
+         @return SimTKOpenMMCommon::DefaultReturn if converge; else
+          return SimTKOpenMMCommon::ErrorReturn
+
+         --------------------------------------------------------------------------------------- */
+
+    virtual int applyToVelocities(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
+                     std::vector<OpenMM::RealVec>& velocities, std::vector<RealOpenMM>& inverseMasses) = 0;
 };
 
 // ---------------------------------------------------------------------------------------

platforms/reference/src/SimTKReference/ReferenceCustomDynamics.cpp

@@ -194,6 +194,11 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
             }
             case CustomIntegrator::ConstrainPositions: {
                 getReferenceConstraintAlgorithm()->apply(numberOfAtoms, atomCoordinates, atomCoordinates, inverseMasses);
+                break;
+            }
+            case CustomIntegrator::ConstrainVelocities: {
+                getReferenceConstraintAlgorithm()->applyToVelocities(numberOfAtoms, atomCoordinates, velocities, inverseMasses);
+                break;
             }
             case CustomIntegrator::UpdateContextState: {
                 recordChangedParameters(context, globals);

platforms/reference/src/SimTKReference/ReferenceLincsAlgorithm.cpp

@@ -29,6 +29,7 @@
 #include "../SimTKUtilities/SimTKOpenMMLog.h"
 #include "ReferenceLincsAlgorithm.h"
 #include "ReferenceDynamics.h"
+#include "openmm/OpenMMException.h"
 
 using std::vector;
 using OpenMM::RealVec;
@@ -292,3 +293,21 @@ int ReferenceLincsAlgorithm::apply( int numberOfAtoms, vector<RealVec>& atomCoor
 
 }
 
+/**---------------------------------------------------------------------------------------
+
+   Apply constraint algorithm to velocities.
+
+   @param numberOfAtoms    number of atoms
+   @param atomCoordinates  atom coordinates
+   @param velocities       atom velocities
+   @param inverseMasses    1/mass
+
+   @return SimTKOpenMMCommon::DefaultReturn if converge; else
+    return SimTKOpenMMCommon::ErrorReturn
+
+   --------------------------------------------------------------------------------------- */
+
+int ReferenceLincsAlgorithm::applyToVelocities(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
+               std::vector<OpenMM::RealVec>& velocities, std::vector<RealOpenMM>& inverseMasses) {
+    throw OpenMM::OpenMMException("applyToVelocities is not implemented");
+}

platforms/reference/src/SimTKReference/ReferenceLincsAlgorithm.h

@@ -138,6 +138,23 @@ class ReferenceLincsAlgorithm : public ReferenceConstraintAlgorithm {
 
       int apply( int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
                        std::vector<OpenMM::RealVec>& atomCoordinatesP, std::vector<RealOpenMM>& inverseMasses );
+
+      /**---------------------------------------------------------------------------------------
+
+         Apply constraint algorithm to velocities.
+
+         @param numberOfAtoms    number of atoms
+         @param atomCoordinates  atom coordinates
+         @param velocities       atom velocities
+         @param inverseMasses    1/mass
+
+         @return SimTKOpenMMCommon::DefaultReturn if converge; else
+          return SimTKOpenMMCommon::ErrorReturn
+
+         --------------------------------------------------------------------------------------- */
+
+      int applyToVelocities(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
+                     std::vector<OpenMM::RealVec>& velocities, std::vector<RealOpenMM>& inverseMasses);
 };
 
 // ---------------------------------------------------------------------------------------

platforms/reference/src/SimTKReference/ReferenceShakeAlgorithm.cpp

@@ -30,6 +30,7 @@
 #include "../SimTKUtilities/SimTKOpenMMLog.h"
 #include "ReferenceShakeAlgorithm.h"
 #include "ReferenceDynamics.h"
+#include "openmm/OpenMMException.h"
 
 using std::vector;
 using OpenMM::RealVec;
@@ -323,6 +324,25 @@ int ReferenceShakeAlgorithm::apply( int numberOfAtoms, vector<RealVec>& atomCoor
 
 }
 
+/**---------------------------------------------------------------------------------------
+
+   Apply constraint algorithm to velocities.
+
+   @param numberOfAtoms    number of atoms
+   @param atomCoordinates  atom coordinates
+   @param velocities       atom velocities
+   @param inverseMasses    1/mass
+
+   @return SimTKOpenMMCommon::DefaultReturn if converge; else
+    return SimTKOpenMMCommon::ErrorReturn
+
+   --------------------------------------------------------------------------------------- */
+
+int ReferenceShakeAlgorithm::applyToVelocities(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
+               std::vector<OpenMM::RealVec>& velocities, std::vector<RealOpenMM>& inverseMasses) {
+    throw OpenMM::OpenMMException("applyToVelocities is not implemented");
+}
+
 /**---------------------------------------------------------------------------------------
 
    Report any violated constriants

platforms/reference/src/SimTKReference/ReferenceShakeAlgorithm.h

@@ -135,6 +135,23 @@ class ReferenceShakeAlgorithm : public ReferenceConstraintAlgorithm {
       
       int apply( int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
                        std::vector<OpenMM::RealVec>& atomCoordinatesP, std::vector<RealOpenMM>& inverseMasses );
+
+      /**---------------------------------------------------------------------------------------
+
+         Apply constraint algorithm to velocities.
+
+         @param numberOfAtoms    number of atoms
+         @param atomCoordinates  atom coordinates
+         @param velocities       atom velocities
+         @param inverseMasses    1/mass
+
+         @return SimTKOpenMMCommon::DefaultReturn if converge; else
+          return SimTKOpenMMCommon::ErrorReturn
+
+         --------------------------------------------------------------------------------------- */
+
+      int applyToVelocities(int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
+                     std::vector<OpenMM::RealVec>& velocities, std::vector<RealOpenMM>& inverseMasses);
       
       /**---------------------------------------------------------------------------------------
       

platforms/reference/tests/TestReferenceCustomIntegrator.cpp

@@ -151,6 +151,65 @@ void testConstraints() {
     }
 }
 
+/**
+ * Test an integrator that applies constraints directly to velocities.
+ */
+void testVelocityConstraints() {
+    const int numParticles = 8;
+    const double temp = 500.0;
+    ReferencePlatform platform;
+    System system;
+    CustomIntegrator integrator(0.002);
+    integrator.addComputePerDof("v", "v+0.5*dt*f/m");
+    integrator.addComputePerDof("x", "x+dt*v");
+    integrator.addConstrainPositions();
+    integrator.addComputePerDof("v", "v+0.5*dt*f/m");
+    integrator.addConstrainVelocities();
+    integrator.setConstraintTolerance(1e-5);
+    NonbondedForce* forceField = new NonbondedForce();
+    for (int i = 0; i < numParticles; ++i) {
+        system.addParticle(i%2 == 0 ? 5.0 : 10.0);
+        forceField->addParticle((i%2 == 0 ? 0.2 : -0.2), 0.5, 5.0);
+    }
+    for (int i = 0; i < numParticles-1; ++i)
+        system.addConstraint(i, i+1, 1.0);
+    system.addForce(forceField);
+    Context context(system, integrator, platform);
+    vector<Vec3> positions(numParticles);
+    vector<Vec3> velocities(numParticles);
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+
+    for (int i = 0; i < numParticles; ++i) {
+        positions[i] = Vec3(i/2, (i+1)/2, 0);
+        velocities[i] = Vec3(genrand_real2(sfmt)-0.5, genrand_real2(sfmt)-0.5, genrand_real2(sfmt)-0.5);
+    }
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+    
+    // Simulate it and see whether the constraints remain satisfied.
+    
+    double initialEnergy = 0.0;
+    for (int i = 0; i < 1000; ++i) {
+        State state = context.getState(State::Positions | State::Energy);
+        for (int j = 0; j < system.getNumConstraints(); ++j) {
+            int particle1, particle2;
+            double distance;
+            system.getConstraintParameters(j, particle1, particle2, distance);
+            Vec3 p1 = state.getPositions()[particle1];
+            Vec3 p2 = state.getPositions()[particle2];
+            double dist = std::sqrt((p1[0]-p2[0])*(p1[0]-p2[0])+(p1[1]-p2[1])*(p1[1]-p2[1])+(p1[2]-p2[2])*(p1[2]-p2[2]));
+            ASSERT_EQUAL_TOL(distance, dist, 2e-5);
+        }
+        double energy = state.getKineticEnergy()+state.getPotentialEnergy();
+        if (i == 1)
+            initialEnergy = energy;
+        else if (i > 1)
+            ASSERT_EQUAL_TOL(initialEnergy, energy, 0.05);
+        integrator.step(2);
+    }
+}
+
 /**
  * Test an integrator with an AndersenThermostat to see if updateContextState()
  * is being handled correctly.
@@ -254,6 +313,7 @@ int main() {
     try {
         testSingleBond();
         testConstraints();
+        testVelocityConstraints();
         testWithThermostat();
         testMonteCarlo();
     }