diff --git a/PQAnalysis/io/box_writer.py b/PQAnalysis/io/box_writer.py index d6281f2f..fd6c7793 100644 --- a/PQAnalysis/io/box_writer.py +++ b/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. diff --git a/tests/io/test_boxWriter.py b/tests/io/test_boxWriter.py index 2844104b..e7c98d41 100644 --- a/tests/io/test_boxWriter.py +++ b/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()