Merge pull request #21 from OpenBioSim/fix_20

Fix 20
OpenBioSim · Feb 16, 2024 · 266e951 · 266e951
2 parents 83cfd45 + 2b7c280
commit 266e951
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 7 deletions.
diff --git a/src/somd2/runner/_dynamics.py b/src/somd2/runner/_dynamics.py
@@ -166,19 +166,21 @@ def _setup_dynamics(self, equilibration=False):
         self._dyn = self._system.dynamics(
             temperature=self._config.temperature,
             pressure=pressure,
-            timestep=self._config.equilibration_timestep
-            if equilibration
-            else self._config.timestep,
+            timestep=(
+                self._config.equilibration_timestep
+                if equilibration
+                else self._config.timestep
+            ),
             lambda_value=self._lambda_val,
             cutoff_type=self._config.cutoff_type,
             cutoff=self._config.cutoff,
             schedule=self._config.lambda_schedule,
             platform=self._config.platform,
             device=self._device,
             constraint="none" if equilibration else self._config.constraint,
-            perturbable_constraint="none"
-            if equilibration
-            else self._config.perturbable_constraint,
+            perturbable_constraint=(
+                "none" if equilibration else self._config.perturbable_constraint
+            ),
             include_constrained_energies=self._config.include_constrained_energies,
             dynamic_constraints=self._config.dynamic_constraints,
             swap_end_states=self._config.swap_end_states,
@@ -411,8 +413,14 @@ def generate_lam_vals(lambda_base, increment):
             from sire import save as _save
 
             _save(self._system.trajectory(), traj_filename, format=["DCD"])
+        # Add config and lambda value to the system properties
+        self._system.add_shared_property(
+            "config", self._config.as_dict(sire_compatible=True)
+        )
+        self._system.add_shared_property("lambda", self._lambda_val)
         # dump final system to checkpoint file
         _stream.save(self._system, sire_checkpoint_name)
+        _logger.debug(f"Properties on system: {self._system.property_keys()}")
         df = self._system.energy_trajectory(to_alchemlyb=True, energy_unit="kT")
         return df
 

diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py
@@ -772,6 +772,7 @@ def _run(sim, is_restart=False):
                     )
 
         if self._config.restart:
+            _logger.debug(f"Restarting {_lam_sym} = {lambda_value} from file")
             try:
                 system = _stream.load(
                     str(
@@ -812,7 +813,6 @@ def _run(sim, is_restart=False):
                             f"Lambda value from checkpoint file {fname} ({lambda_encoded}) does not match expected value ({lambda_value})."
                         )
                 is_restart = True
-
         else:
             system = self._system.clone()
             is_restart = False

diff --git a/tests/runner/test_restart.py b/tests/runner/test_restart.py
@@ -44,6 +44,13 @@ def test_restart(mols, request):
 
         num_entries = len(energy_traj_1.index)
 
+        # Check that both config and lambda have been written
+        # as properties to the streamed checkpoint file.
+        checkpoint = sr.stream.load(str(Path(tmpdir) / "checkpoint_0.s3"))
+        props = checkpoint.property_keys()
+        assert "config" in props
+        assert "lambda" in props
+
         del runner
 
         config_new = {