From 4c02f93cf0ffbb839a488250d239b03316f43d86 Mon Sep 17 00:00:00 2001 From: "Josef M. Gallmetzer" <64498081+galjos@users.noreply.github.com> Date: Sun, 23 Aug 2026 21:32:30 +0200 Subject: [PATCH] fix: read cells of extended xyz trajectories The cell generator behind TrajectoryReader.cells parsed every header line as a PQ xyz header of the form "n_atoms a b c alpha beta gamma", regardless of the trajectory format. An extended xyz frame stores the box as Lattice metadata in the comment line and puts only the atom count in the header line, so every frame matched the single-field case and yielded a vacuum cell while the comment line was skipped together with the atoms. Reading the same file with read() returned the correct box, so the two entry points disagreed. Any analysis that takes the cells from the reader instead of from the frames therefore worked with an infinite volume: an RDF over an extended xyz trajectory produced an average volume of inf and wrote a column of nan and inf, with an exit status of zero and nothing but a numpy warning to indicate that the result was meaningless. The generator now dispatches on the trajectory format and reads the Lattice metadata for extended xyz files, carrying the previous cell over to frames that define no Lattice, exactly as the frame generator does. A format for which no cell layout is known raises instead of silently reporting vacuum. The xyz branch also takes the atom count from the header of each frame rather than from the first line of the first file, so that a set of files with differing atom counts no longer runs out of step. --- PQAnalysis/io/traj_file/frame_reader.py | 6 + PQAnalysis/io/traj_file/trajectory_reader.py | 70 ++++++++++- tests/io/test_trajectoryReader.py | 117 +++++++++++++++++++ 3 files changed, 189 insertions(+), 4 deletions(-) diff --git a/PQAnalysis/io/traj_file/frame_reader.py b/PQAnalysis/io/traj_file/frame_reader.py index 19005503..dcaa8b53 100644 --- a/PQAnalysis/io/traj_file/frame_reader.py +++ b/PQAnalysis/io/traj_file/frame_reader.py @@ -505,6 +505,12 @@ def read( cell=self._read_cell(metadata), ) + def read_cell(self, comment_line: str) -> Cell: + """ + Reads the cell from the comment line of an extended xyz frame. + """ + return self._read_cell(self._read_metadata(comment_line)) + def _read_atom_count( # pylint: disable=inconsistent-return-statements self, line: str diff --git a/PQAnalysis/io/traj_file/trajectory_reader.py b/PQAnalysis/io/traj_file/trajectory_reader.py index 6a7a18e0..14b32ffe 100644 --- a/PQAnalysis/io/traj_file/trajectory_reader.py +++ b/PQAnalysis/io/traj_file/trajectory_reader.py @@ -21,7 +21,7 @@ # Local relative modules from .exceptions import TrajectoryReaderError -from .frame_reader import get_frame_reader +from .frame_reader import get_frame_reader, XYZ_FRAME_READER_TRAJ_FORMATS @@ -597,10 +597,20 @@ def _cell_generator(self) -> Generator[List[Cell]]: list of Cell The list of cells read from the trajectory. """ + if self.traj_format == TrajectoryFormat.EXTXYZ: + yield from self._extxyz_cell_generator() + return + + if self.traj_format not in XYZ_FRAME_READER_TRAJ_FORMATS: + self.logger.error( + ( + "Reading the cells is not implemented for the " + f"trajectory format {self.traj_format}." + ), + exception=TrajectoryReaderError, + ) + last_cell = None - with open(self.filenames[0], "r", encoding="utf-8") as f: - line = f.readline() - n_atoms = int(line.split()[0]) for filename in self.filenames: line_number = 0 @@ -654,9 +664,61 @@ def _cell_generator(self) -> Generator[List[Cell]]: last_cell = cell + n_atoms = int(splitted_line[0]) + for _ in range(n_atoms + 1): next(f, None) # Skip the next n_atoms+1 lines + def _extxyz_cell_generator(self) -> Generator[Cell]: + """ + A generator that yields the cells of an extended xyz trajectory. + + The cell of an extended xyz frame is stored as Lattice + metadata in the comment line of the frame and not in the + atom count line. If a frame does not define a Lattice, the + cell of the last frame is used. + + Yields + ------ + Cell + The cells read from the trajectory. + """ + last_cell = None + + for filename in self.filenames: + with open(filename, "r", encoding="utf-8") as f: + while True: + atom_count_line = f.readline() + if atom_count_line == "": + break + + if atom_count_line.strip() == "": + continue + + n_atoms = 0 + try: + n_atoms = int(atom_count_line.split()[0]) + except (ValueError, IndexError): + self.logger.error( + ( + "Invalid number of atoms encountered " + f"in file {filename}." + ), + exception=TrajectoryReaderError, + ) + + cell = self.frame_reader.read_cell(f.readline()) + + if cell.is_vacuum and last_cell is not None: + cell = last_cell + + last_cell = cell + + for _ in range(n_atoms): + next(f, None) # Skip the atom lines of the frame + + yield cell + def _read_single_frame( self, frame_string: str, diff --git a/tests/io/test_trajectoryReader.py b/tests/io/test_trajectoryReader.py index 76e0b8f8..c7aeaee5 100644 --- a/tests/io/test_trajectoryReader.py +++ b/tests/io/test_trajectoryReader.py @@ -878,3 +878,120 @@ def test_cells(self, caplog): ), function=reader._cell_generator().__next__, ) + + # -------------------------------------------------------------------------------- # + @pytest.mark.usefixtures("tmpdir") + def test_cells_extxyz(self): + file = open("tmp.extxyz", "w") + print("2", file=file) + print( + 'Lattice="1.0 0.0 0.0 0.0 2.0 0.0 0.0 0.0 3.0" ' + 'Properties=species:S:1:pos:R:3', + file=file + ) + print("h 0.0 0.0 0.0", file=file) + print("o 0.0 1.0 0.0", file=file) + print("2", file=file) + print( + 'Lattice="4.0 0.0 0.0 0.0 5.0 0.0 0.0 0.0 6.0" ' + 'Properties=species:S:1:pos:R:3', + file=file + ) + print("h 1.0 0.0 0.0", file=file) + print("o 1.0 1.0 0.0", file=file) + file.close() + + reader = TrajectoryReader("tmp.extxyz") + assert reader.traj_format is TrajectoryFormat.EXTXYZ + + assert reader.cells == [Cell(1.0, 2.0, 3.0), Cell(4.0, 5.0, 6.0)] + assert reader.cells == [frame.cell for frame in reader.read()] + + # -------------------------------------------------------------------------------- # + @pytest.mark.usefixtures("tmpdir") + def test_cells_variable_size_extxyz(self): + file = open("tmp.extended.xyz", "w") + print("2", file=file) + print( + 'Lattice="1.0 0.0 0.0 0.0 2.0 0.0 0.0 0.0 3.0" ' + 'Properties=species:S:1:pos:R:3', + file=file + ) + print("h 0.0 0.0 0.0", file=file) + print("o 0.0 1.0 0.0", file=file) + print("1", file=file) + print( + 'Lattice="4.0 0.0 0.0 0.0 5.0 0.0 0.0 0.0 6.0" ' + 'Properties=species:S:1:pos:R:3', + file=file + ) + print("c 1.0 0.0 0.0", file=file) + print("3", file=file) + print("Properties=species:S:1:pos:R:3", file=file) + print("h 0.0 0.0 0.0", file=file) + print("o 0.0 1.0 0.0", file=file) + print("c 1.0 0.0 0.0", file=file) + file.close() + + reader = TrajectoryReader("tmp.extended.xyz") + + # the last frame does not define a Lattice and + # therefore reuses the cell of the previous frame + assert reader.cells == [ + Cell(1.0, 2.0, 3.0), + Cell(4.0, 5.0, 6.0), + Cell(4.0, 5.0, 6.0), + ] + assert reader.cells == [frame.cell for frame in reader.read().frames] + + # -------------------------------------------------------------------------------- # + @pytest.mark.usefixtures("tmpdir") + def test_cells_variable_size_xyz_files(self): + file = open("tmp1.xyz", "w") + print("2 1.0 1.0 1.0", file=file) + print("", file=file) + print("h 0.0 0.0 0.0", file=file) + print("o 0.0 1.0 0.0", file=file) + file.close() + + file = open("tmp2.xyz", "w") + print("1 2.0 2.0 2.0", file=file) + print("", file=file) + print("c 1.0 0.0 0.0", file=file) + print("1 3.0 3.0 3.0", file=file) + print("", file=file) + print("c 2.0 0.0 0.0", file=file) + file.close() + + reader = TrajectoryReader(["tmp1.xyz", "tmp2.xyz"]) + + assert reader.cells == [ + Cell(1.0, 1.0, 1.0), + Cell(2.0, 2.0, 2.0), + Cell(3.0, 3.0, 3.0), + ] + + # -------------------------------------------------------------------------------- # + @pytest.mark.usefixtures("tmpdir") + def test_cells_unsupported_format(self, caplog): + file = open("tmp.xyz", "w") + print("2 1.0 1.0 1.0", file=file) + print("", file=file) + print("h 0.0 0.0 0.0", file=file) + print("o 0.0 1.0 0.0", file=file) + file.close() + + reader = TrajectoryReader("tmp.xyz") + reader.traj_format = TrajectoryFormat.AUTO + + assert_logging_with_exception( + caplog, + TrajectoryReader.__qualname__, + exception=TrajectoryReaderError, + logging_level="ERROR", + message_to_test=( + "Reading the cells is not implemented for the " + "trajectory format TrajectoryFormat.AUTO." + ), + function=reader._cell_generator().__next__, + )