fix: track written frames for traj2box step indices

PQAnalysis/io/box_writer.py

@@ -58,6 +58,7 @@ def __init__(
 
         super().__init__(filename, FileWritingMode(mode))
         self.output_format = BoxFileFormat(output_format)
+        self._n_frames_written = 0
 
     @runtime_type_checking
     def write(self, traj: Trajectory, reset_counter: bool = True) -> None:
@@ -117,7 +118,7 @@ def write_vmd(self, traj: Trajectory) -> None:
     def write_box_file(
         self,
         traj: Trajectory,
-        reset_counter: bool = True  # pylint: disable=unused-argument # is needed for the decorator
+        reset_counter: bool = True
     ) -> None:
         """
         Writes the given trajectory to the file in data file format.
@@ -141,8 +142,10 @@ def write_box_file(
         """
         self.__check_pbc__(traj)
 
-        counter = self.counter[BoxWriter.write_box_file.__name__]  # pylint: disable=no-member # is added via decorator
-        counter = len(traj) * (counter - 1)
+        if reset_counter:
+            self._n_frames_written = 0
+
+        counter = self._n_frames_written
 
         for i, frame in enumerate(traj):
             cell = frame.cell
@@ -155,6 +158,8 @@ def write_box_file(
                 file=self.file
             )
 
+        self._n_frames_written += len(traj)
+
     def __check_pbc__(self, traj: Trajectory) -> None:
         """
         Checks if the cell of the trajectory is not None.

tests/io/test_boxWriter.py

@@ -92,6 +92,30 @@ def test_write_box_file(self, capsys: CaptureFixture):
 
         assert captured.out == "1 10 10 10 90 90 90\n2 10 10 11 90 90 120\n"
 
+    def test_write_box_file_unequal_trajectories(self, capsys: CaptureFixture):
+        writer = BoxWriter()
+
+        cell = Cell(10, 10, 10, 90, 90, 90)
+        frame = AtomicSystem(atoms=self.atoms1, pos=self.pos1, cell=cell)
+
+        traj3 = Trajectory([frame, frame, frame])
+        traj2 = Trajectory([frame, frame])
+
+        writer.write_box_file(traj3, reset_counter=False)
+        writer.write_box_file(traj2, reset_counter=False)
+
+        captured = capsys.readouterr()
+        steps = [int(line.split()[0]) for line in captured.out.splitlines()]
+
+        assert steps == [1, 2, 3, 4, 5]
+
+        writer.write_box_file(traj2, reset_counter=True)
+
+        captured = capsys.readouterr()
+        steps = [int(line.split()[0]) for line in captured.out.splitlines()]
+
+        assert steps == [1, 2]
+
     def test_write_vmd(self, capsys: CaptureFixture):
         writer = BoxWriter()