Allow querying current step count (#3248)

* Allow querying current step count

* Fixed error building Python wrapper

serialization/src/SerializationNode.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 Stanford University and the Authors.           *
+ * Portions copyright (c) 2010-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -121,6 +121,30 @@ SerializationNode& SerializationNode::setIntProperty(const string& name, int val
     return *this;
 }
 
+long long SerializationNode::getLongProperty(const string& name) const {
+    map<string, string>::const_iterator iter = properties.find(name);
+    if (iter == properties.end())
+        throw OpenMMException("Unknown property '"+name+"' in node '"+getName()+"'");
+    long long value;
+    stringstream(iter->second) >> value;
+    return value;
+}
+
+long long SerializationNode::getLongProperty(const string& name, long long defaultValue) const {
+    map<string, string>::const_iterator iter = properties.find(name);
+    if (iter == properties.end())
+        return defaultValue;
+    long long value;
+    stringstream(iter->second) >> value;
+    return value;
+}
+
+SerializationNode& SerializationNode::setLongProperty(const string& name, long long value) {
+    stringstream s;
+    s << value;
+    properties[name] = s.str();
+    return *this;
+}
 
 bool SerializationNode::getBoolProperty(const string& name) const {
     map<string, string>::const_iterator iter = properties.find(name);

serialization/src/StateProxy.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-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -47,6 +47,7 @@ void StateProxy::serialize(const void* object, SerializationNode& node) const {
     node.setStringProperty("openmmVersion", Platform::getOpenMMVersion());
     const State& s = *reinterpret_cast<const State*>(object);
     node.setDoubleProperty("time", s.getTime());
+    node.setLongProperty("stepCount", s.getStepCount());
     Vec3 a,b,c;
     s.getPeriodicBoxVectors(a,b,c);
     SerializationNode& boxVectorsNode = node.createChildNode("PeriodicBoxVectors");
@@ -98,6 +99,7 @@ void* StateProxy::deserialize(const SerializationNode& node) const {
     if (node.getIntProperty("version") != 1)
         throw OpenMMException("Unsupported version number");
     double outTime = node.getDoubleProperty("time");
+    long long outStepCount = node.getLongProperty("stepCount", 0);
     const SerializationNode& boxVectorsNode = node.getChildNode("PeriodicBoxVectors");
     const SerializationNode& AVec = boxVectorsNode.getChildNode("A");
     Vec3 outAVec(AVec.getDoubleProperty("x"),AVec.getDoubleProperty("y"),AVec.getDoubleProperty("z"));
@@ -107,7 +109,7 @@ void* StateProxy::deserialize(const SerializationNode& node) const {
     Vec3 outCVec(CVec.getDoubleProperty("x"),CVec.getDoubleProperty("y"),CVec.getDoubleProperty("z"));
     int types = 0;
     vector<int> arraySizes;
-    State::StateBuilder builder(outTime);
+    State::StateBuilder builder(outTime, outStepCount);
     for (auto& child : node.getChildren()) {
         if (child.getName() == "Parameters") {
             map<string, double> outStateParams;

serialization/tests/TestSerializeState.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 Stanford University and the Authors.           *
+ * Portions copyright (c) 2010-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -82,6 +82,8 @@ void testSerialization() {
 
     context.setPositions(positions);
     context.setVelocities(velocities);
+    context.setTime(2.5);
+    context.setStepCount(100);
 
     // Serialize and then deserialize it.
     State s1 = context.getState(State::Positions | State::Velocities | State::Forces | State::Energy | State::Parameters);
@@ -121,6 +123,7 @@ void testSerialization() {
     ASSERT_EQUAL(s1.getPotentialEnergy(), s2.getPotentialEnergy());
     ASSERT_EQUAL(s1.getKineticEnergy(), s2.getKineticEnergy());
     ASSERT_EQUAL(s1.getTime(), s2.getTime());
+    ASSERT_EQUAL(s1.getStepCount(), s2.getStepCount());
 
     map<string, double> p1 = s1.getParameters();
     map<string, double> p2 = s2.getParameters();

wrappers/python/openmm/app/simulation.py

@@ -92,8 +92,6 @@ class Simulation(object):
                 self.integrator = mm.XmlSerializer.deserialize(f.read())
         else:
             self.integrator = integrator
-        ## The index of the current time step
-        self.currentStep = 0
         ## A list of reporters to invoke during the simulation
         self.reporters = []
         if platform is None:
@@ -113,6 +111,15 @@ class Simulation(object):
         except Exception: # OpenMM just raises Exception if it's not implemented everywhere
             self._usesPBC = topology.getUnitCellDimensions() is not None
 
+    @property
+    def currentStep(self):
+        """The index of the current time step."""
+        return self.context.getStepCount()
+
+    @currentStep.setter
+    def currentStep(self, step):
+        self.context.setStepCount(step)
+
     def minimizeEnergy(self, tolerance=10*unit.kilojoules_per_mole/unit.nanometer, maxIterations=0):
         """Perform a local energy minimization on the system.
 
@@ -198,11 +205,9 @@ class Simulation(object):
             while stepsToGo > 10:
                 self.integrator.step(10) # Only take 10 steps at a time, to give Python more chances to respond to a control-c.
                 stepsToGo -= 10
-                self.currentStep += 10
                 if endTime is not None and datetime.now() >= endTime:
                     return
             self.integrator.step(stepsToGo)
-            self.currentStep += stepsToGo
             if anyReport:
                 # One or more reporters are ready to generate reports.  Organize them into three
                 # groups: ones that want wrapped positions, ones that want unwrapped positions,

wrappers/python/src/swig_doxygen/swigInputConfig.py

@@ -234,6 +234,7 @@ UNITS = {
 ("SerializationNode", "getProperties") : (None, ()),
 ("SerializationNode", "getStringProperty") : (None, ()),
 ("SerializationNode", "getIntProperty") : (None, ()),
+("SerializationNode", "getLongProperty") : (None, ()),
 ("SerializationNode", "getDoubleProperty") : (None, ()),
 ("SerializationProxy", "getProxy") : (None, ()),
 ("SerializationProxy", "getTypeName") : (None, ()),
@@ -324,7 +325,10 @@ UNITS = {
 ("Context", "getState") : (None, (None, None, None)),
 ("Context", "setPeriodicBoxVectors") : (None, ("unit.nanometer", "unit.nanometer", "unit.nanometer")),
 ("Context", "setPositions") : (None, ("unit.nanometer",)),
+("Context", "getTime") : ("unit.picosecond", ()),
 ("Context", "setTime") : (None, ("unit.picosecond",)),
+("Context", "getStepCount") : (None, ()),
+("Context", "setStepCount") : (None, (None,)),
 ("Context", "setVelocities") : (None, ("unit.nanometer/unit.picosecond",)),
 ("CMAPTorsionForce", "getMapParameters") : (None, (None, "unit.kilojoule_per_mole")),
 ("CMAPTorsionForce", "setMapParameters") : (None, (None, None, "unit.kilojoule_per_mole")),
@@ -439,6 +443,7 @@ UNITS = {
 ("RBTorsionForce", "getTorsionParameters") : (None, (None, None, None, None, "unit.kilojoules_per_mole", "unit.kilojoules_per_mole", "unit.kilojoules_per_mole", "unit.kilojoules_per_mole", "unit.kilojoules_per_mole", "unit.kilojoules_per_mole")),
 ("RBTorsionForce", "setTorsionParameters") : (None, (None, None, None, None, None, "unit.kilojoules_per_mole", "unit.kilojoules_per_mole", "unit.kilojoules_per_mole", "unit.kilojoules_per_mole", "unit.kilojoules_per_mole", "unit.kilojoules_per_mole")),
 ("State", "getTime") : ("unit.picosecond", ()),
+("State", "getStepCount") : (None, ()),
 ("State", "getKineticEnergy") : ("unit.kilojoules_per_mole", ()),
 ("State", "getPotentialEnergy") : ("unit.kilojoules_per_mole", ()),
 ("State", "getPeriodicBoxVolume") : ("unit.nanometers**3", ()),

wrappers/python/tests/TestSimulation.py

@@ -112,11 +112,13 @@ class TestSimulation(unittest.TestCase):
         simulation.context.setVelocitiesToTemperature(300*kelvin)
         self.assertEqual(0, simulation.currentStep)
         self.assertEqual(0*picoseconds, simulation.context.getState().getTime())
+        simulation.currentStep = 5
+        self.assertEqual(5, simulation.currentStep)
 
         # Take some steps and verify the simulation has advanced by the correct amount.
 
         simulation.step(23)
-        self.assertEqual(23, simulation.currentStep)
+        self.assertEqual(28, simulation.currentStep)
         self.assertAlmostEqual(0.023, simulation.context.getState().getTime().value_in_unit(picoseconds))
 
     def testRunForClockTime(self):