Backport fixes from PR #244

src/somd2/runner/_base.py

@@ -905,15 +905,21 @@ def __init__(self, system, config):
         # Create the lock file name.
         self._lock_file = str(self._config.output_directory / "somd2.lock")
 
-        # Write the end-state topologies to the output directory.
+        # Write the end-state topologies. A GCMC restart keeps those already
+        # written, since they include the additional waters.
         if isinstance(self._system, list):
             mols = self._system[0]
         else:
             mols = self._system
-        mols0 = _sr.morph.link_to_reference(mols)
-        mols1 = _sr.morph.link_to_perturbed(mols)
-        _sr.save(mols0, self._filenames["topology0"])
-        _sr.save(mols1, self._filenames["topology1"])
+        topologies = (self._filenames["topology0"], self._filenames["topology1"])
+        if not (
+            self._config.restart
+            and self._config.gcmc
+            and all(_Path(t).exists() for t in topologies)
+        ):
+            _sr.save(_sr.morph.link_to_reference(mols), topologies[0])
+            _sr.save(_sr.morph.link_to_perturbed(mols), topologies[1])
+        self._gcmc_topology_saved = False
 
         # Update the tajectory page size.
         if self._config.page_size is not None:
@@ -2450,28 +2456,13 @@ def _checkpoint(
                         f"{self._standard_state_correction:.6f}"
                     )
 
-            if is_final_block:
-                # Save the end-state GCMC topologies for trajectory analysis and visualisation.
-                # This topology contains additional water molecules that are used for GCMC
-                # insertion moves.
-                if self._config.gcmc:
-                    mols0 = _sr.morph.link_to_reference(system)
-                    mols1 = _sr.morph.link_to_perturbed(system)
-
-                    # Save to AMBER format.
-                    _sr.save(mols0, self._filenames["topology0"])
-                    _sr.save(mols1, self._filenames["topology1"])
-
-                    # Save to PDB format.
-                    _sr.save(
-                        mols0,
-                        self._filenames["topology0"].replace(".prm7", ".pdb"),
-                    )
-                    _sr.save(
-                        mols1,
-                        self._filenames["topology1"].replace(".prm7", ".pdb"),
-                    )
+            # Saved at the first checkpoint of each run, and the final one for the
+            # final coordinates in the PDB files.
+            if self._config.gcmc and (is_final_block or not self._gcmc_topology_saved):
+                self._save_gcmc_topology(system)
+                self._gcmc_topology_saved = True
 
+            if is_final_block:
                 # Assemble and save the final trajectory.
                 if self._config.save_trajectories:
                     # Save the final trajectory chunk to file.
@@ -2825,6 +2816,27 @@ def _flush_energy_components(self, index):
             )
             _pq_local.write_table(table, filepath)
 
+    def _save_gcmc_topology(self, system):
+        """
+        Save the end-state topologies, including the additional waters used for
+        GCMC insertion moves, in AMBER and PDB format.
+
+        Parameters
+        ----------
+
+        system : :class: `System <sr.system.System>`
+            The system being simulated.
+        """
+        mols0 = _sr.morph.link_to_reference(system)
+        mols1 = _sr.morph.link_to_perturbed(system)
+
+        for mols, topology in (
+            (mols0, self._filenames["topology0"]),
+            (mols1, self._filenames["topology1"]),
+        ):
+            _sr.save(mols, topology)
+            _sr.save(mols, topology.replace(".prm7", ".pdb"))
+
     def _save_ghost_residues(self, index, gcmc_sampler):
         """
         Record the current GCMC ghost residue indices for a window. This must