fix: read the cells of extended xyz trajectories

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

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 @@
         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 @@
 
                     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,

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__,
+        )