Context.reinitialize() can optionally preserve state information

openmmapi/include/openmm/Context.h

@@ -217,15 +217,25 @@ public:
      */
     void computeVirtualSites();
     /**
-     * When a Context is created, it may cache information about the System being simulated
+     * When a Context is created, it caches information about the System being simulated
      * and the Force objects contained in it.  This means that, if the System or Forces are then
-     * modified, the Context might not see all of the changes.  Call reinitialize() to force
+     * modified, the Context does not see the changes.  Call reinitialize() to force
      * the Context to rebuild its internal representation of the System and pick up any changes
      * that have been made.
      * 
      * This is an expensive operation, so you should try to avoid calling it too frequently.
+     * Most Force classes have an updateParametersInContext() method that provides a less expensive
+     * way of updating certain types of information.  However, this method is the only way to
+     * make some types of changes, so it is sometimes necessary to call it.
+     * 
+     * By default, reinitializing a Context causes all state information (positions, velocities,
+     * etc.) to be discarded.  You can optionally tell it to try to preserve state information.
+     * It does this by internally creating a checkpoint, then reinitializing the Context, then
+     * loading the checkpoint.  Be aware that if the System has changed in a way that prevents
+     * the checkpoint from being loaded (such as changing the number of particles), this will
+     * throw an exception and the state information will be lost.
      */
-    void reinitialize();
+    void reinitialize(bool preserveState=false);
     /**
      * Create a checkpoint recording the current state of the Context.  This should be treated
      * as an opaque block of binary data.  See loadCheckpoint() for more details.

openmmapi/src/Context.cpp

@@ -36,6 +36,8 @@
 #include "SimTKOpenMMRealType.h"
 #include "sfmt/SFMT.h"
 #include <cmath>
+#include <iostream>
+#include <sstream>
 
 using namespace OpenMM;
 using namespace std;
@@ -242,14 +244,19 @@ void Context::computeVirtualSites() {
     impl->computeVirtualSites();
 }
 
-void Context::reinitialize() {
+void Context::reinitialize(bool preserveState) {
     const System& system = impl->getSystem();
     Integrator& integrator = impl->getIntegrator();
     Platform& platform = impl->getPlatform();
+    stringstream checkpoint(ios_base::out | ios_base::in | ios_base::binary);
+    if (preserveState)
+        createCheckpoint(checkpoint);
     integrator.cleanup();
     delete impl;
     impl = new ContextImpl(*this, system, integrator, &platform, properties);
     impl->initialize();
+    if (preserveState)
+        loadCheckpoint(checkpoint);
 }
 
 void Context::createCheckpoint(ostream& stream) {

tests/TestNonbondedForce.h

@@ -261,6 +261,7 @@ void testCutoff14() {
     positions[2] = Vec3(2, 0, 0);
     positions[3] = Vec3(3, 0, 0);
     positions[4] = Vec3(4, 0, 0);
+    context.setPositions(positions);
     for (int i = 1; i < 5; ++i) {
 
         // Test LJ forces
@@ -271,8 +272,7 @@ void testCutoff14() {
         nonbonded->setParticleParameters(i, 0, 1.5, 1);
         nonbonded->setExceptionParameters(first14, 0, 3, 0, 1.5, i == 3 ? 0.5 : 0.0);
         nonbonded->setExceptionParameters(second14, 1, 4, 0, 1.5, 0.0);
-        context.reinitialize();
-        context.setPositions(positions);
+        context.reinitialize(true);
         State state = context.getState(State::Forces | State::Energy);
         const vector<Vec3>& forces = state.getForces();
         double r = positions[i][0];
@@ -299,8 +299,7 @@ void testCutoff14() {
         nonbonded->setParticleParameters(i, q, 1.5, 0);
         nonbonded->setExceptionParameters(first14, 0, 3, i == 3 ? q*q/1.2 : 0, 1.5, 0);
         nonbonded->setExceptionParameters(second14, 1, 4, 0, 1.5, 0);
-        context.reinitialize();
-        context.setPositions(positions);
+        context.reinitialize(true);
         state = context.getState(State::Forces | State::Energy);
         const vector<Vec3>& forces2 = state.getForces();
         force = ONE_4PI_EPS0*q*q/(r*r);