[HREMD] Fix storing wrongly indexed coords to checkpoint (#30)

docs/migration.md

@@ -22,9 +22,8 @@ field in OpenFF, Amber, and OpenMM FFXML formats. This re-write also introduced
 -   Ligand force field parameters no longer need to be provided. The [femto.md.config.Prepare][] configuration now
     exposes a `default_ligand_ff` field that can be used to automatically parameterize ligands with an OpenFF based
     force field.
--   The labelling of atoms in Boresch restraints has been updated to be more consistent with the literature, and also
-    properly documented. See the [femto.md.restraints.create_boresch_restraint][] for more information on the exact
-    definition of the distances, angles, and dihedrals that will be restrained.
+-   HREMD now correctly stores coordinates as ``coords[i] = replica_i_coords`` rather than ``coords[i] = state_i_coords``.
+    Checkpoints from previous versions will likely be incorrect.
 
 #### FE
 

femto/md/hremd.py

@@ -509,7 +509,7 @@ def _store_checkpoint(
     if mpi_comm.rank != 0:
         return
 
-    coords = [coords_dict[replica_to_state_idx[i]] for i in range(len(u_kn))]
+    coords = [coords_dict[i] for i in range(len(u_kn))]
 
     path.parent.mkdir(exist_ok=True, parents=True)
 

femto/md/tests/test_hremd.py

@@ -361,8 +361,6 @@ def test_hremd_sampling_checkpoint(harmonic_test_case, tmp_cwd, mocker):
         return_value=1.0,
     )
 
-    spied_load_checkpoint = mocker.spy(femto.md.hremd, "_load_checkpoint")
-
     simulation, temperature, states, expected_delta_f_ij = harmonic_test_case
     top = mdtraj.Topology.from_openmm(simulation.topology)
 
@@ -401,22 +399,6 @@ def test_hremd_sampling_checkpoint(harmonic_test_case, tmp_cwd, mocker):
         simulation=simulation, states=states, config=config_2, output_dir=tmp_cwd
     )
 
-    expected_init_coords = [coord.getPositions(asNumpy=True) for coord in coords_1]
-    actual_init_coords = [
-        coord.getPositions(asNumpy=True)
-        for coord in spied_load_checkpoint.spy_return[1]
-    ]
-
-    assert all(
-        numpy.allclose(
-            actual_coord.value_in_unit(openmm.unit.angstrom),
-            expected_coord.value_in_unit(openmm.unit.angstrom),
-        )
-        for actual_coord, expected_coord in zip(
-            actual_init_coords, expected_init_coords, strict=True
-        )
-    )
-
     traj_2 = mdtraj.load_dcd(str(tmp_cwd / "trajectories/r1.dcd"), top=top)
     assert len(traj_2) == n_steps_2