Fix frame clearing for trajectory chunking.

OpenBioSim · Oct 16, 2024 · def9309 · def9309
1 parent 4665bcb
commit def9309
Showing 1 changed file with 10 additions and 11 deletions.
diff --git a/src/somd2/runner/_dynamics.py b/src/somd2/runner/_dynamics.py
@@ -126,6 +126,7 @@ def __init__(
             self._config.restart,
         )
 
+        self._num_frames = 0
         self._nrg_sample = 0
         self._nrg_file = "energy_components.txt"
 
@@ -489,11 +490,11 @@ def generate_lam_vals(lambda_base, increment):
                             + f"{x}.dcd"
                         )
                         sr.save(
-                            self._system.trajectory(), traj_filename, format=["DCD"]
+                            self._system.trajectory()[self._num_frames :],
+                            traj_filename,
+                            format=["DCD"],
                         )
-
-                        # Delete the trajectory from memory.
-                        self._system.delete_all_frames()
+                        self._num_frames = self._system.num_frames()
 
                     # Stream the checkpoint to file.
                     sr.stream.save(self._system, str(checkpoint_file))
@@ -557,11 +558,11 @@ def generate_lam_vals(lambda_base, increment):
                             + f"{x}.dcd"
                         )
                         sr.save(
-                            self._system.trajectory(), traj_filename, format=["DCD"]
+                            self._system.trajectory()[self._num_frames :],
+                            traj_filename,
+                            format=["DCD"],
                         )
-
-                        # Delete the trajectory from memory.
-                        self._system.delete_all_frames()
+                        self._num_frames = self._system.num_frames()
 
                     _logger.info(
                         f"Finished block {x+1} of {self._current_block + num_blocks + int(rem > 0)} "
@@ -622,9 +623,7 @@ def generate_lam_vals(lambda_base, increment):
                 self._config.output_directory / self._filenames["trajectory"]
             )
             sr.save(system.trajectory(), traj_filename, format=["DCD"])
-
-            # Delete the trajectory from memory.
-            self._system.delete_all_frames()
+            self._num_frames = system.num_frames()
 
             # Now remove the chunked trajectory files.
             for chunk in traj_chunks: