dwhswenson · Oct 2, 2024 · Feb 3, 2022 · Feb 3, 2022 · Feb 3, 2022 · Feb 3, 2022
diff --git a/examples/gromacs/AD_tps_2a_run_flex.ipynb b/examples/gromacs/AD_tps_2a_run_flex.ipynb
@@ -220,11 +220,83 @@
    "source": [
     "storage.close()"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Extra: Restarting a simulation with the gromacs engine\n",
+    "\n",
+    "Because engines are only stored once at the start, some extra info (like the last filename) is needed after a restart to prevent old trajectories from being overwritten.\n",
+    "\n",
+    "Start by opening the storage in append mode:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "storage = paths.Storage(\"alanine_dipeptide_tps.nc\", mode=\"a\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Then for every Gromacs engine in the storage, run\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for eng in storage.engines:\n",
+    "    eng.restart_at_step(storage.steps[-1])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "That's it, we now can safely continue the simulation (after also restarting that object)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "DONE! Completed 110 Monte Carlo cycles.\n"
+     ]
+    }
+   ],
+   "source": [
+    "sampler = storage.pathsimulators[-1]\n",
+    "sampler.restart_at_step(storage.steps[-1])\n",
+    "sampler.run(10)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "storage.close()"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -238,7 +310,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.3"
+   "version": "3.7.12"
   },
   "toc": {
    "base_numbering": 1,

diff --git a/openpathsampling/engines/external_engine.py b/openpathsampling/engines/external_engine.py
@@ -333,3 +333,31 @@ def set_filenames(self, number):
     def engine_command(self):
         """Generates a string for the command to run the engine."""
         raise NotImplementedError()
+
+    def restart_at_step(self, step):
+        """Restart the filecounter (if used) to the right number."""
+        if not type(self.filename_setter) == FilenameSetter:
+            # Nothing to do
+            # TODO OPS 2.0: this can be turned into isinstance after
+            # inheritance is swapped
+            return
+        # Grab the last trajectory that was written to the storage
+        last_trajectory = step.simulation.storage.trajectories[-1]
+        # As this trajectory can be made from different files, querry all files
+        possible_last_trajs = set(frame.file_name for frame in last_trajectory)
+
+        # Split off extensions and directories,
+        # loop backwards untill we find a non-number, and cast to int
+        possible_numbers = []
+        numbers = '0123456789'
+        for traj in possible_last_trajs:
+            traj = traj.rsplit('.', 1)[0]  # split off extension
+            traj = traj.replace('\\', '/').split('/')[-1]  # split directories
+            string = ""
+            for c in traj[::-1]:
+                if c not in numbers:
+                    break
+                string += c
+            possible_numbers.append(int(string[::-1]))
+        last_number = max(possible_numbers)
+        self.filename_setter.count = last_number + 1
diff --git a/openpathsampling/tests/test_external_engine.py b/openpathsampling/tests/test_external_engine.py
@@ -26,6 +26,21 @@
                           "external_engine")
 
 
+class FakeStep(object):
+    """Only here to allow restart_at_step to be tested"""
+    def __init__(self):
+        # Somehow object()'s can't have attributes but lambdas can ....
+        self.simulation = lambda: None
+        self.simulation.storage = lambda: None
+        self.simulation.storage.trajectories = [[]]
+
+
+class FakeExternalFrame(object):
+    """A frame that only has a filename"""
+    def __init__(self, file_name):
+        self.file_name = file_name
+
+
 class ExampleExternalEngine(peng.ExternalEngine):
     """Trivial external engine for engine.c in the tests.
     """
@@ -235,6 +250,64 @@ def test_in_shooting_move(self):
         for testfile in glob.glob("test*out") + glob.glob("test*inp"):
             os.remove(testfile)
 
+    def test_restart_at_step_no_op(self):
+        options = {'n_frames_max': 10000,
+                   'engine_sleep': 100,
+                   'name_prefix': "test",
+                   'engine_directory': engine_dir,
+                   'filename_setter': RandomStringFilenames()}
+        eng = ExampleExternalEngine(options,
+                                    self.descriptor,
+                                    self.template)
+        eng.restart_at_step('foo')
+
+    def test_restart_at_step(self):
+        step = FakeStep()
+        frame = FakeExternalFrame('test123.fake')
+        step.simulation.storage.trajectories[-1] = [frame]
+        self.slow_engine.restart_at_step(step)
+        assert self.slow_engine.filename_setter() == 124
+
+    def test_restart_at_step_numbers_in_file_name(self):
+        step = FakeStep()
+        frame = FakeExternalFrame('simulation3-test12.fake')
+        step.simulation.storage.trajectories[-1] = [frame]
+        self.slow_engine.restart_at_step(step)
+        assert self.slow_engine.filename_setter() == 13
+
+    def test_restart_at_step_numbers_in_extensions(self):
+        step = FakeStep()
+        frame = FakeExternalFrame('test135.h5')
+        step.simulation.storage.trajectories[-1] = [frame]
+        self.slow_engine.restart_at_step(step)
+        assert self.slow_engine.filename_setter() == 136
+
+    def test_restart_at_step_directories_in_file_name(self):
+        step = FakeStep()
+        frame = FakeExternalFrame('/simulations/3/test13.fake')
+        step.simulation.storage.trajectories[-1] = [frame]
+        self.slow_engine.restart_at_step(step)
+        assert self.slow_engine.filename_setter() == 14
+
+    def test_restart_at_step_multiple_files_in_traj(self):
+        step = FakeStep()
+        frame1 = FakeExternalFrame('test1.fake')
+        frame2 = FakeExternalFrame('test1001.fake')
+        frame3 = FakeExternalFrame('test12.fake')
+        step.simulation.storage.trajectories[-1] = [frame1, frame2, frame3]
+        self.slow_engine.restart_at_step(step)
+        assert self.slow_engine.filename_setter() == 1002
+
+    def test_restart_at_step_multiple_trajectories(self):
+        step = FakeStep()
+        frame1 = FakeExternalFrame('test1.fake')
+        frame2 = FakeExternalFrame('test1001.fake')
+        frame3 = FakeExternalFrame('test12.fake')
+        for frame in [frame1, frame2, frame3]:
+            step.simulation.storage.trajectories.append([frame])
+        self.slow_engine.restart_at_step(step)
+        assert self.slow_engine.filename_setter() == 13
+
 
 class TestFilenameSetter(object):
     def test_default_setter(self):