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