Returning velocities in sampler state when propagating replicas (#602)

* Enabling velocities when returning sampler state when propagating replicas.

* Testing velocities match after resuming simulation.

* Better naming for test.

* Adding notes to the changelog.

* Velocities NOT ignored when updating sampler states in replica propagation.

docs/releasehistory.rst

@@ -1,6 +1,13 @@
 Release History
 ***************
 
+0.21.5 - Bugfix release
+=======================
+
+Bugfixes
+--------
+- Bug in returning velocities when propagating replicas. Fixed by using ``ignore_velocities=False`` in ``_propagate_replica``. Issue `#531 <https://github.com/choderalab/openmmtools/issues/531>`_ (PR `#602 <https://github.com/choderalab/openmmtools/pull/602>`_).
+
 0.21.4 - Bugfix release
 =======================
 

openmmtools/multistate/multistatereporter.py

@@ -1664,6 +1664,8 @@ def _write_sampler_states_to_given_file(self, sampler_states: list, iteration: i
                     x = sampler_state.velocities / (unit.nanometer/unit.picoseconds) # _unitless_velocities
                     velocities[replica_index, :, :] = x[:, :]
              # Store velocites
+            # TODO: This stores velocities as zeros if no velocities are present in the sampler state. Making restored
+            #  sampler_state different from origin.
             storage.variables['velocities'][write_iteration, :, :, :] = velocities
 
             if is_periodic:

openmmtools/multistate/multistatesampler.py

@@ -1298,7 +1298,7 @@ def _propagate_replicas(self):
         # Update all sampler states. For non-0 nodes, this will update only the
         # sampler states associated to the replicas propagated by this node.
         for replica_id, propagated_state in zip(replica_ids, propagated_states):
-            self._sampler_states[replica_id].__setstate__(propagated_state, ignore_velocities=True)
+            self._sampler_states[replica_id].__setstate__(propagated_state, ignore_velocities=False)
 
         # Gather all MCMCMoves statistics. All nodes must have these up-to-date
         # since they are tied to the ThermodynamicState, not the replica.
@@ -1332,8 +1332,8 @@ def _propagate_replica(self, replica_id):
             logger.critical(message)
             raise SimulationNaNError(message)
 
-        # Send the new state to the root node. We can ignore velocities as we're not saving them.
-        return sampler_state.__getstate__(ignore_velocities=True)
+        # Send the new state to the root node.
+        return sampler_state.__getstate__(ignore_velocities=False)
 
     def _get_replica_move_statistics(self, replica_id):
         """Return the statistics of the MCMCMove currently associated to this replica."""

openmmtools/tests/test_sampling.py

@@ -360,23 +360,28 @@ def test_write_sampler_states(self):
             restored_sampler_states = reporter.read_sampler_states(iteration=0)
             for state, restored_state in zip(sampler_states, restored_sampler_states):
                 assert np.allclose(state.positions, restored_state.positions)
+                # By default stored velocities are zeros if not present in origin sampler_state
+                assert np.allclose(np.zeros(state.positions.shape), restored_state.velocities)
                 assert np.allclose(state.box_vectors / unit.nanometer, restored_state.box_vectors / unit.nanometer)
             # Check that the analysis particles are written off checkpoint whereas full trajectory is not
             restored_analysis_states = reporter.read_sampler_states(iteration=1, analysis_particles_only=True)
             restored_checkpoint_states = reporter.read_sampler_states(iteration=1)
             assert type(restored_analysis_states) is list
             for state in restored_analysis_states:
                 assert state.positions.shape == (len(analysis_particles), 3)
+                assert state.velocities.shape == (len(analysis_particles), 3)
             assert restored_checkpoint_states is None
             # Check that the analysis particles are written separate from the checkpoint particles
             restored_analysis_states = reporter.read_sampler_states(iteration=2, analysis_particles_only=True)
             restored_checkpoint_states = reporter.read_sampler_states(iteration=2)
             assert len(restored_analysis_states) == len(restored_checkpoint_states)
             for analysis_state, checkpoint_state in zip(restored_analysis_states, restored_checkpoint_states):
-                # This assert is dual purpose: Positions are identical; Analysis shape is correct
+                # This assert is multiple purpose: Positions are identical; Velocities are indetical and zeros
+                # (since unspecified); Analysis shape is correct
                 # Will raise a ValueError for np.allclose(x,y) if x.shape != y.shape
                 # Will raise AssertionError if the values are not allclose
                 assert np.allclose(analysis_state.positions, checkpoint_state.positions[analysis_particles, :])
+                assert np.allclose(analysis_state.velocities, checkpoint_state.velocities[analysis_particles, :])
                 assert np.allclose(analysis_state.box_vectors / unit.nanometer,
                                    checkpoint_state.box_vectors / unit.nanometer)
 
@@ -1268,6 +1273,41 @@ def test_checkpointing(self):
                 else:
                     assert states is None
 
+    def test_resume_positions_velocities_from_storage(self):
+        """Test that positions and velocities are the same when resuming a simulation from reporter storage file."""
+        # TODO: Find a way to extend this test to use MPI since resuming velocities has a problem there.
+        test_cases = [self.alanine_test, self.hostguest_test]
+
+        for test_case in test_cases:
+            thermodynamic_states, sampler_states, unsampled_states = copy.deepcopy(test_case)
+
+            with self.temporary_storage_path() as storage_path:
+                moves = mmtools.mcmc.SequenceMove([
+                    mmtools.mcmc.LangevinDynamicsMove(n_steps=1),
+                    mmtools.mcmc.MCRotationMove(),
+                    mmtools.mcmc.GHMCMove(n_steps=1)
+                ])
+
+                sampler = self.SAMPLER(mcmc_moves=moves, number_of_iterations=3)
+                reporter = self.REPORTER(storage_path, checkpoint_interval=1)
+                self.call_sampler_create(sampler, reporter,
+                                         thermodynamic_states, sampler_states,
+                                         unsampled_states)
+                # Run 3 iterations
+                sampler.run(n_iterations=3)
+                # store a copy of the original states
+                original_states = sampler.sampler_states
+                # Unallocate current objects and close reporter
+                del sampler
+                reporter.close()
+                # recreate sampler from storage
+                sampler = self.SAMPLER.from_storage(reporter)
+                restored_states = sampler.sampler_states
+                for original_state, restored_state in zip(original_states, restored_states):
+                    assert np.allclose(original_state.positions, restored_state.positions)
+                    assert np.allclose(original_state.velocities, restored_state.velocities)
+
+
     def test_last_iteration_functions(self):
         """Test that the last_iteration functions work right"""
         thermodynamic_states, sampler_states, unsampled_states = copy.deepcopy(self.alanine_test)