Augment checkpoint storage to save velocities

[zhang-ivy]

I am seeing dips in potential energy when I resume a repex simulation:

I next decided to investigate the first dip in potential energy (at iteration 791). Here are histograms of replica energies for each iteration, colored by iteration number:

We figured out that these dips are caused by resuming simulations from a checkpoint file. Since we are not saving velocities in the checkpoint, the velocities start at 0 when resuming a simulation.

The short term solution for this is to set reassign_velocities=True in the LangevinSplittingDynamicsMove() when I instantiate it.
The long term solution is to save velocities in the checkpoint.

[zhang-ivy]

@ijpulidos : I just checked to see if this issue was resolved, but I am still seeing the same issue, even when using openmmtools >= 0.21.3 (the PR for saving velocities in the checkpoint file was introduced in 0.21.3).

I ran repex on ala->thr dipeptide in solvent with reassign_velocities=False:

move = mcmc.LangevinSplittingDynamicsMove(timestep= 4.0 * unit.femtoseconds,
                                          collision_rate=1.0 / unit.picosecond,
                                          n_steps=250,
                                          reassign_velocities=False,
                                          n_restart_attempts=20,
                                          splitting="V R R R O R R R V",
                                          constraint_tolerance=1e-06)

Here is the potential energy vs time plot:

You can see dips in the energy when I resumed it (2 different times: once at 1000 iterations, and once shortly thereafter).

The problem goes away when I set reassign_velocities=True, which was the short term fix I mentioned in the original post of this issue.

Attached are files to reproduce this behavior:

reproduce_velocity_problem.zip

[mikemhenry]

I guess the next step is to confirm that we really are using the velocities that we save. Another thing that could be weird is, depending on the integrator, it might either need the forces saved or some sort of flag to calculate the forces before doing a step.

[jchodera]

@mikemhenry : We're missing test coverage of velocity restoration, so I'd start there.

For example:

• This test only checks positions
• Here too
• Here as well

and so on