Feature: Enable appending to trajectory when using ts.optimize/ts.integrate

tests/test_trajectory.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import tempfile
 from collections.abc import Callable, Generator
 from pathlib import Path
 
@@ -834,3 +835,99 @@ def test_write_ase_trajectory_importerror(
     with pytest.raises(ImportError, match="ASE is required to convert to ASE trajectory"):
         traj.write_ase_trajectory(tmp_path / "dummy.traj")
     traj.close()
+
+
+def test_optimize_append_to_trajectory(
+    si_double_sim_state: SimState, lj_model: LennardJonesModel
+) -> None:
+    """Test appending to an existing trajectory when running ts.optimize."""
+
+    # Create a temporary trajectory file
+    with tempfile.TemporaryDirectory() as temp_dir:
+        traj_files = [f"{temp_dir}/optimize_trajectory_{idx}.h5" for idx in range(2)]
+
+        # Initialize model and state
+        trajectory_reporter = ts.TrajectoryReporter(
+            traj_files,
+            state_frequency=1,
+        )
+
+        # First optimization run
+        opt_state = ts.optimize(
+            system=si_double_sim_state,
+            model=lj_model,
+            max_steps=5,
+            optimizer=ts.Optimizer.fire,
+            trajectory_reporter=trajectory_reporter,
+            steps_between_swaps=100,
+        )
+
+        for traj in trajectory_reporter.trajectories:
+            with TorchSimTrajectory(traj._file.filename, mode="r") as traj:
+                # Check that the trajectory file has 5 frames
+                np.testing.assert_allclose(traj.get_steps("positions"), range(1, 6))
+
+        trajectory_reporter_2 = ts.TrajectoryReporter(
+            traj_files, state_frequency=1, trajectory_kwargs=dict(mode="a")
+        )
+        _ = ts.optimize(
+            system=opt_state,
+            model=lj_model,
+            max_steps=7,
+            optimizer=ts.Optimizer.fire,
+            trajectory_reporter=trajectory_reporter_2,
+            steps_between_swaps=100,
+        )
+        for traj in trajectory_reporter_2.trajectories:
+            with TorchSimTrajectory(traj._file.filename, mode="r") as traj:
+                # Check that the trajectory file now has 12 frames
+                np.testing.assert_allclose(traj.get_steps("positions"), range(1, 13))
+
+
+def test_integrate_append_to_trajectory(
+    si_double_sim_state: SimState, lj_model: LennardJonesModel
+) -> None:
+    """Test appending to an existing trajectory when running ts.integrate."""
+
+    # Create a temporary trajectory file
+    with tempfile.TemporaryDirectory() as temp_dir:
+        traj_files = [f"{temp_dir}/integrate_trajectory_{idx}.h5" for idx in range(2)]
+
+        # Initialize model and state
+        trajectory_reporter = ts.TrajectoryReporter(
+            traj_files,
+            state_frequency=1,
+        )
+
+        # First integration run
+        int_state = ts.integrate(
+            system=si_double_sim_state,
+            model=lj_model,
+            timestep=0.1,
+            n_steps=5,
+            temperature=300.0,
+            integrator=ts.Integrator.nvt_langevin,
+            trajectory_reporter=trajectory_reporter,
+        )
+
+        for traj in trajectory_reporter.trajectories:
+            with TorchSimTrajectory(traj._file.filename, mode="r") as traj:
+                # Check that the trajectory file has 5 frames
+                np.testing.assert_allclose(traj.get_steps("positions"), range(1, 6))
+
+        trajectory_reporter_2 = ts.TrajectoryReporter(
+            traj_files, state_frequency=1, trajectory_kwargs=dict(mode="a")
+        )
+        _ = ts.integrate(
+            system=int_state,
+            model=lj_model,
+            timestep=0.1,
+            temperature=300.0,
+            n_steps=7,
+            integrator=ts.Integrator.nvt_langevin,
+            trajectory_reporter=trajectory_reporter_2,
+        )
+        for traj in trajectory_reporter_2.trajectories:
+            with TorchSimTrajectory(traj._file.filename, mode="r") as traj:
+                # Check that the trajectory file now has 12 frames
+                np.testing.assert_allclose(traj.get_steps("positions"), range(1, 13))

torch_sim/runners.py

@@ -175,7 +175,12 @@ def integrate[T: SimState](  # noqa: C901
             trajectory_reporter,
             properties=["kinetic_energy", "potential_energy", "temperature"],
         )
-
+    # Auto-detect initial step from trajectory files for resuming integration
+    initial_step: int = 1
+    if trajectory_reporter is not None and trajectory_reporter.mode == "a":
+        last_step = trajectory_reporter.last_step
+        if last_step > 0:
+            initial_step = last_step + 1
     final_states: list[T] = []
     og_filenames = trajectory_reporter.filenames if trajectory_reporter else None
 
@@ -199,9 +204,13 @@ def integrate[T: SimState](  # noqa: C901
             )
 
         # run the simulation
-        for step in range(1, n_steps + 1):
+        for step in range(initial_step, initial_step + n_steps):
             state = step_func(
-                state=state, model=model, dt=dt, kT=kTs[step - 1], **integrator_kwargs
+                state=state,
+                model=model,
+                dt=dt,
+                kT=kTs[step - initial_step],
+                **integrator_kwargs,
             )
 
             if trajectory_reporter:
@@ -457,7 +466,14 @@ def optimize[T: OptimState](  # noqa: C901, PLR0915
             trajectory_reporter, properties=["potential_energy"]
         )
 
-    step: int = 1
+    # Auto-detect initial step from trajectory files for resuming optimizations
+    initial_step: int = 1
+    if trajectory_reporter is not None and trajectory_reporter.mode == "a":
+        last_step = trajectory_reporter.last_step
+        if last_step > 0:
+            initial_step = last_step + 1
+    step: int = initial_step
+
     last_energy = None
     all_converged_states: list[T] = []
     convergence_tensor = None
@@ -485,9 +501,13 @@ def optimize[T: OptimState](  # noqa: C901, PLR0915
             and og_filenames is not None
             and (step == 1 or len(converged_states) > 0)
         ):
+            mode_before = trajectory_reporter.trajectory_kwargs["mode"]
+            # temporarily set to "append" mode to avoid overwriting existing files
+            trajectory_reporter.trajectory_kwargs["mode"] = "a"
             trajectory_reporter.load_new_trajectories(
                 filenames=[og_filenames[i] for i in autobatcher.current_idx]
             )
+            trajectory_reporter.trajectory_kwargs["mode"] = mode_before
 
         for _step in range(steps_between_swaps):
             if hasattr(state, "energy"):
@@ -498,7 +518,7 @@ def optimize[T: OptimState](  # noqa: C901, PLR0915
             if trajectory_reporter:
                 trajectory_reporter.report(state, step, model=model)
             step += 1
-            if step > max_steps:
+            if step > max_steps + initial_step - 1:
                 # TODO: max steps should be tracked for each structure in the batch
                 warnings.warn(f"Optimize has reached max steps: {step}", stacklevel=2)
                 break

torch_sim/trajectory.py

@@ -41,7 +41,6 @@
 from torch_sim.models.interface import ModelInterface
 from torch_sim.state import SimState
 
-
 if TYPE_CHECKING:
     from ase import Atoms
     from ase.io.trajectory import TrajectoryReader
@@ -302,6 +301,43 @@ def close(self) -> None:
         for trajectory in self.trajectories:
             trajectory.close()
 
+    @property
+    def mode(self) -> Literal["r", "w", "a"]:
+        """Get the mode of the first trajectory file.
+
+        Returns:
+            "r" | "w" | "a": Mode from the trajectory_kwargs used during initialization.
+        """
+        if not self.trajectories:
+            raise ValueError("No trajectories loaded.")
+        # Key is guaranteed to exist because we set it during initialization.
+        return self.trajectory_kwargs["mode"]
+
+    @property
+    def last_step(self) -> int:
+        """Get the maximum last step across all trajectory files.
+
+        Returns the highest step number found across all trajectory files.
+        This is useful for resuming optimizations from where they left off.
+
+        Returns:
+            int: The maximum last step number across all trajectories, or 0 if
+                no trajectories exist or all are empty
+        """
+        if not self.trajectories:
+            return 0
+
+        max_step = 0
+        for trajectory in self.trajectories:
+            if trajectory._file.isopen:
+                last_step = trajectory.last_step
+            else:
+                with TorchSimTrajectory(trajectory._file.filename, mode="r") as traj:
+                    last_step = traj.last_step
+            max_step = max(max_step, last_step)
+
+        return max_step
+
     def __enter__(self) -> "TrajectoryReporter":
         """Support the context manager protocol.
 
@@ -594,7 +630,7 @@ def _validate_array(self, name: str, data: np.ndarray, steps: list[int]) -> None
             )
 
         # Validate step is monotonically increasing by checking HDF5 file directly
-        steps_node = self._file.get_node("/steps/", name=name)
+        steps_node = self.get_steps(name)
         if len(steps_node) > 0:
             last_step = steps_node[-1]  # Get the last recorded step
             if steps[0] <= last_step:
@@ -658,9 +694,6 @@ def get_array(
     def get_steps(
         self,
         name: str,
-        start: int | None = None,
-        stop: int | None = None,
-        step: int = 1,
     ) -> np.ndarray:
         """Get the steps for an array.
 
@@ -675,9 +708,29 @@ def get_steps(
         Returns:
             np.ndarray: Array of step numbers with shape [n_selected_frames]
         """
-        return self._file.root.steps.__getitem__(name).read(
-            start=start, stop=stop, step=step
-        )
+        return self._file.get_node("/steps/", name=name).read()
+
+    @property
+    def last_step(self) -> int:
+        """Get the last step number from the trajectory.
+
+        Retrieves the maximum step number across all arrays in the trajectory.
+        If the trajectory is empty or has no arrays, returns 0.
+
+        Returns:
+            int: The last (maximum) step number in the trajectory, or 0 if empty
+        """
+        if not self.array_registry:
+            return 0
+
+        max_step = 0
+        for name in self.array_registry:
+            steps_node = self.get_steps(name)
+            if len(steps_node) > 0:
+                last_step = int(steps_node[-1])
+                max_step = max(max_step, last_step)
+
+        return max_step
 
     def __str__(self) -> str:
         """Get a string representation of the trajectory.