From 8bd89dae5334a8fe8b4de9cea4cbcce9699bac3a Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 9 Oct 2026 14:12:47 +0000 Subject: [PATCH 1/6] feat(qe): support pw.x AIMD output --- dpdata/formats/qe/__init__.py | 4 +- dpdata/formats/qe/pwmd.py | 14 + dpdata/formats/qe/scf.py | 468 ++++++++++++++++++++++++++++++---- dpdata/plugins/qe.py | 30 ++- dpdata/system.py | 8 +- tests/test_qe_pw_md.py | 101 ++++++++ 6 files changed, 564 insertions(+), 61 deletions(-) create mode 100644 dpdata/formats/qe/pwmd.py create mode 100644 tests/test_qe_pw_md.py diff --git a/dpdata/formats/qe/__init__.py b/dpdata/formats/qe/__init__.py index 35d20788f..d4109559e 100644 --- a/dpdata/formats/qe/__init__.py +++ b/dpdata/formats/qe/__init__.py @@ -1,5 +1,5 @@ from __future__ import annotations -from . import scf, traj +from . import pwmd, scf, traj -__all__ = ["scf", "traj"] +__all__ = ["pwmd", "scf", "traj"] diff --git a/dpdata/formats/qe/pwmd.py b/dpdata/formats/qe/pwmd.py new file mode 100644 index 000000000..fa962a323 --- /dev/null +++ b/dpdata/formats/qe/pwmd.py @@ -0,0 +1,14 @@ +"""Quantum ESPRESSO ``pw.x`` molecular-dynamics output reader. + +The implementation is shared with the ``qe/pw/scf`` reader because a PWscf +output can contain either one labeled structure or a sequence of MD frames. +""" + +from __future__ import annotations + +from .scf import get_frames + + +def to_system_data(fname, begin=0, step=1): + """Read PWscf AIMD frames from an output file and its matching input.""" + return get_frames(fname, begin=begin, step=step) diff --git a/dpdata/formats/qe/scf.py b/dpdata/formats/qe/scf.py index 341261d22..abb43d51d 100755 --- a/dpdata/formats/qe/scf.py +++ b/dpdata/formats/qe/scf.py @@ -2,12 +2,14 @@ from __future__ import annotations import os +import re import numpy as np from dpdata.utils import open_file from .traj import ( + convert_celldm, kbar2evperang3, ry2ev, ) @@ -49,63 +51,119 @@ def get_block(lines, keyword, skip=0): def get_cell(lines): ret = [] - for idx, ii in enumerate(lines): - if "ibrav" in ii: + idx = None + match = None + for line_idx, ii in enumerate(lines): + match = re.search(r"ibrav\s*=\s*([+-]?\d+)", ii, re.IGNORECASE) + if match is not None: + idx = line_idx break - blk = lines[idx : idx + 2] - ibrav = int(blk[0].replace(",", "").split("=")[-1]) + if idx is None or match is None: + raise RuntimeError("ibrav cannot be found in QE input") + ibrav = int(match.group(1)) if ibrav == 0: - for iline in lines: - if "CELL_PARAMETERS" in iline and "angstrom" not in iline.lower(): + for line_idx, line in enumerate(lines): + if "CELL_PARAMETERS" not in line.upper(): + continue + block = lines[line_idx + 1 : line_idx + 4] + if len(block) != 3: + raise RuntimeError("CELL_PARAMETERS block must contain three rows") + ret = np.array( + [ + [ + float(value.replace("D", "E").replace("d", "e")) + for value in row.split()[:3] + ] + for row in block + ] + ) + unit = _cell_unit(line) + if unit == "bohr": + ret *= bohr2ang + elif unit == "alat": + ret *= (_card_alat(line) or _input_alat(lines)) * bohr2ang + elif unit not in (None, "angstrom"): raise RuntimeError( - "CELL_PARAMETERS must be written in Angstrom. Other units are not supported yet." + "CELL_PARAMETERS must be written in Angstrom, bohr, or alat." ) - blk = get_block(lines, "CELL_PARAMETERS") - for ii in blk: - ret.append([float(jj) for jj in ii.split()[0:3]]) - ret = np.array(ret) + break + else: + raise RuntimeError("CELL_PARAMETERS block cannot be found") elif ibrav == 1: a = None for iline in lines: line = iline.replace("=", " ").replace(",", "").split() if len(line) >= 2 and "a" == line[0]: # print("line = ", line) - a = float(line[1]) + a = float(line[1].replace("D", "E").replace("d", "e")) if len(line) >= 2 and "celldm(1)" == line[0]: - a = float(line[1]) * bohr2ang + a = float(line[1].replace("D", "E").replace("d", "e")) * bohr2ang # print("a = ", a) if not a: raise RuntimeError("parameter 'a' or 'celldm(1)' cannot be found.") ret = np.array([[a, 0.0, 0.0], [0.0, a, 0.0], [0.0, 0.0, a]]) + elif ibrav in (2, 3, -3): + ret = convert_celldm(ibrav, [_input_alat(lines)]) * bohr2ang else: raise RuntimeError("ibrav > 1 not supported yet.") return ret +def _cell_unit(line): + match = re.search(r"[({]\s*([^})]+?)\s*[})]", line, re.IGNORECASE) + if match is None: + return None + unit = match.group(1).strip().lower() + if unit.startswith("alat"): + return "alat" + if unit.startswith("angstrom"): + return "angstrom" + if unit.startswith("bohr"): + return "bohr" + return unit + + +def _input_alat(lines): + for line in lines: + words = line.replace("=", " ").replace(",", " ").split() + if len(words) >= 2 and words[0].lower() == "celldm(1)": + return float(words[1].replace("D", "E").replace("d", "e")) + if len(words) >= 2 and words[0].lower() == "a": + return float(words[1].replace("D", "E").replace("d", "e")) + raise RuntimeError("parameter 'a' or 'celldm(1)' cannot be found") + + def get_coords(lines, cell): - coord = [] - atom_symbol_list = [] - for iline in lines: - if "ATOMIC_POSITIONS" in iline and ( - "angstrom" not in iline.lower() and "crystal" not in iline.lower() - ): - raise RuntimeError( - "ATOMIC_POSITIONS must be written in Angstrom or crystal. Other units are not supported yet." - ) - if "ATOMIC_POSITIONS" in iline and "angstrom" in iline.lower(): - blk = get_block(lines, "ATOMIC_POSITIONS") - for ii in blk: - coord.append([float(jj) for jj in ii.split()[1:4]]) - atom_symbol_list.append(ii.split()[0]) - coord = np.array(coord) - elif "ATOMIC_POSITIONS" in iline and "crystal" in iline.lower(): - blk = get_block(lines, "ATOMIC_POSITIONS") - for ii in blk: - coord.append([float(jj) for jj in ii.split()[1:4]]) - atom_symbol_list.append(ii.split()[0]) - coord = np.array(coord) - coord = np.matmul(coord, cell) - atom_symbol_list = np.array(atom_symbol_list) + """Read the first ``ATOMIC_POSITIONS`` block from a QE input file. + + PWscf prints an ``ATOMIC_POSITIONS`` block for every MD step. The old + implementation walked over every matching line and converted the list to + an array inside the loop, so a second block attempted ``append`` on an + ndarray. Reading one block here also makes this helper useful for the + static input metadata needed by the AIMD reader. + """ + for line_idx, line in enumerate(lines): + if "ATOMIC_POSITIONS" not in line.upper(): + continue + unit = _position_unit(line) + block = _read_atom_block(lines, line_idx + 1) + if not block: + continue + atom_symbol_list = np.array([parts[0] for parts in block]) + coord = np.array( + [ + [ + float(value.replace("D", "E").replace("d", "e")) + for value in parts[1:4] + ] + for parts in block + ] + ) + coord = _convert_positions(coord, unit, cell, alat=_card_alat(line)) + break + else: + raise RuntimeError("ATOMIC_POSITIONS block cannot be found") + tmp_names, symbol_idx = np.unique(atom_symbol_list, return_index=True) atom_types = [] atom_numbs = [] @@ -122,6 +180,73 @@ def get_coords(lines, cell): return list(atom_names), atom_numbs, atom_types, coord +def _position_unit(line): + """Return the coordinate unit in an ``ATOMIC_POSITIONS`` header.""" + match = re.search(r"[({]\s*([^})]+?)\s*[})]", line, re.IGNORECASE) + if match is None: + return None + unit = match.group(1).strip().lower() + if unit.startswith("alat"): + return "alat" + if unit.startswith("crystal"): + return "crystal" + if unit.startswith("angstrom"): + return "angstrom" + if unit.startswith("bohr"): + return "bohr" + return unit + + +def _card_alat(line): + """Return an explicit ``alat`` value (in bohr) from a QE card header.""" + match = re.search(r"alat\s*=\s*([-+0-9.eEdD]+)", line, re.IGNORECASE) + if match is None: + return None + return float(match.group(1).replace("D", "E").replace("d", "e")) + + +def _read_atom_block(lines, start, natoms=None): + """Read atom rows following an ``ATOMIC_POSITIONS`` header.""" + block = [] + idx = start + while idx < len(lines) and (natoms is None or len(block) < natoms): + words = lines[idx].split() + if not words: + if block: + break + idx += 1 + continue + if words[0].upper() in _QE_BLOCK_KEYWORDS: + break + if len(words) < 4: + break + try: + [float(value.replace("D", "E").replace("d", "e")) for value in words[1:4]] + except ValueError: + break + block.append(words) + idx += 1 + return block + + +def _convert_positions(coord, unit, cell, alat=None): + """Convert QE positions to Cartesian Angstrom coordinates.""" + if unit in (None, "angstrom"): + return coord + if unit == "bohr": + return coord * bohr2ang + if unit == "crystal": + return np.matmul(coord, cell) + if unit == "alat": + if alat is None: + alat = np.linalg.norm(cell[0]) / bohr2ang + return coord * alat * bohr2ang + raise RuntimeError( + "ATOMIC_POSITIONS must be written in Angstrom, bohr, alat, or crystal. " + f"Unit {unit!r} is not supported." + ) + + def get_energy(lines): energy = None for ii in lines: @@ -153,36 +278,269 @@ def get_stress(lines): return ret -def get_frame(fname): +def _output_alat(lines, fallback_cell): + """Get the QE ``alat`` value in bohr from an output header.""" + for line in lines: + if "lattice parameter (alat)" not in line.lower(): + continue + match = re.search(r"=\s*([-+0-9.eEdD]+)", line) + if match: + return float(match.group(1).replace("D", "E").replace("d", "e")) + return np.linalg.norm(fallback_cell[0]) / bohr2ang + + +def _output_position_frames(lines, natoms, cell): + """Read all Cartesian position frames printed by ``pw.x``.""" + frames = [] + units = [] + symbols = [] + for line_idx, line in enumerate(lines): + if "ATOMIC_POSITIONS" not in line.upper(): + continue + block = _read_atom_block(lines, line_idx + 1, natoms) + if len(block) != natoms: + continue + unit = _position_unit(line) + symbols_in_block = [parts[0] for parts in block] + coords = np.array( + [ + [ + float(value.replace("D", "E").replace("d", "e")) + for value in parts[1:4] + ] + for parts in block + ] + ) + frames.append(coords) + units.append((unit, _card_alat(line))) + if not symbols: + symbols = symbols_in_block + return np.array(frames), symbols, units + + +def _output_cell_frames(lines, nframes, fallback_cell): + """Read all ``CELL_PARAMETERS`` blocks from a ``pw.x`` output.""" + cells = [] + alat = _output_alat(lines, fallback_cell) + for line_idx, line in enumerate(lines): + if "CELL_PARAMETERS" not in line.upper(): + continue + block = lines[line_idx + 1 : line_idx + 4] + if len(block) != 3: + continue + try: + current = np.array( + [ + [ + float(value.replace("D", "E").replace("d", "e")) + for value in row.split()[:3] + ] + for row in block + ] + ) + except (ValueError, IndexError): + continue + unit = _cell_unit(line) + card_alat = _card_alat(line) + if unit == "bohr": + current *= bohr2ang + elif unit == "alat": + current *= (card_alat or alat) * bohr2ang + elif unit not in (None, "angstrom"): + continue + cells.append(current) + if not cells: + return np.tile(fallback_cell, (nframes, 1, 1)) + if len(cells) == 1: + return np.tile(cells[0], (nframes, 1, 1)) + if len(cells) != nframes: + raise RuntimeError( + "the number of CELL_PARAMETERS blocks does not match " + f"the number of coordinate frames ({len(cells)} != {nframes})" + ) + return np.array(cells) + + +def _output_energies(lines): + energies = [] + for line in lines: + if re.search(r"!\s*total energy", line, re.IGNORECASE) is None: + continue + try: + value = line.split("=", 1)[1].split()[0] + energies.append(ry2ev * float(value.replace("D", "E").replace("d", "e"))) + except (IndexError, ValueError): + continue + return np.array(energies) + + +def _output_forces(lines, natoms): + forces = [] + for line_idx, line in enumerate(lines): + if "forces acting on atoms" not in line.lower(): + continue + frame = [] + for candidate in lines[line_idx + 1 :]: + if "force =" not in candidate.lower(): + if frame and ( + "total force" in candidate.lower() + or "atomic positions" in candidate.lower() + ): + break + continue + try: + values = candidate.lower().split("force =", 1)[1].split()[:3] + frame.append([float(value.replace("d", "e")) for value in values]) + except (IndexError, ValueError): + continue + if len(frame) == natoms: + break + if len(frame) == natoms: + forces.append(frame) + if not forces: + return np.empty((0, natoms, 3)) + return np.asarray(forces) * (ry2ev / bohr2ang) + + +def _output_stresses(lines): + stresses = [] + for line_idx, line in enumerate(lines): + if re.search(r"total\s+stress", line, re.IGNORECASE) is None: + continue + frame = [] + for candidate in lines[line_idx + 1 :]: + words = candidate.split() + if len(words) < 6: + if frame: + break + continue + try: + # QE prints the kbar values in columns 4--6. + frame.append( + [ + float(value.replace("D", "E").replace("d", "e")) + for value in words[3:6] + ] + ) + except ValueError: + continue + if len(frame) == 3: + break + if len(frame) == 3: + stresses.append(frame) + if not stresses: + return None + return np.asarray(stresses) * kbar2evperang3 + + +def _frame_count(*arrays): + counts = [len(array) for array in arrays if len(array)] + return min(counts) if counts else 0 + + +def get_frames(fname, begin=0, step=1): + """Read one or more PWscf single-point or AIMD frames. + + ``pw.x`` writes an ``ATOMIC_POSITIONS`` block for each MD step while the + historical ``qe/pw/scf`` reader assumed a single block. This reader keeps + the single-point behavior and additionally aligns repeated positions, + energies, forces, cells, and stresses into frame-major arrays. + """ + if begin < 0 or step <= 0: + raise ValueError("begin must be non-negative and step must be positive") if isinstance(fname, str): path_out = fname outname = os.path.basename(path_out) - # the name of the input file is assumed to be different from the output by 'in' and 'out' + # The input file is conventionally obtained by replacing ``out`` by ``in``. inname = outname.replace("out", "in") path_in = os.path.join(os.path.dirname(path_out), inname) elif isinstance(fname, list) and len(fname) == 2: - path_in = fname[0] - path_out = fname[1] + path_in, path_out = fname else: raise RuntimeError("invalid input") with open_file(path_out) as fp: - outlines = fp.read().split("\n") + outlines = fp.read().splitlines() + # Restarted pw.x jobs may append a second run to the same output file. + # Parse only the latest run, matching the trajectory files written by QE. + program_headers = [ + idx for idx, line in enumerate(outlines) if "program pwscf" in line.lower() + ] + if program_headers: + outlines = outlines[program_headers[-1] :] with open_file(path_in) as fp: - inlines = fp.read().split("\n") + inlines = fp.read().splitlines() + cell = get_cell(inlines) - atom_names, natoms, types, coords = get_coords(inlines, cell) - energy = get_energy(outlines) - force = get_force(outlines, natoms) - stress = get_stress(outlines) - if stress is not None: - stress = (stress * np.linalg.det(cell))[np.newaxis, :, :] + atom_names, natoms, types, input_coords = get_coords(inlines, cell) + natom = sum(natoms) + output_coords, output_symbols, output_units = _output_position_frames( + outlines, natom, cell + ) + energies = _output_energies(outlines) + forces = _output_forces(outlines, natom) + stresses = _output_stresses(outlines) + + # A regular SCF output has no positions block. Use the input structure in + # that case; this preserves the existing one-frame ``qe/pw/scf`` behavior. + if len(output_coords) == 0: + output_coords = input_coords[np.newaxis, :, :] + # A static pw.x run prints one ``total energy`` line for every SCF + # iteration but only one force frame. Keep the converged (last) energy in + # that case; AIMD output normally has one final energy per position frame. + if len(energies) > len(output_coords): + energies = energies[-len(output_coords) :] + nframes = _frame_count(output_coords, energies, forces) + if nframes == 0: + raise RuntimeError("QE output does not contain a complete labeled frame") + if output_symbols: + output_names = list(dict.fromkeys(output_symbols)) + if output_names != atom_names: + raise RuntimeError("atom order in QE output does not match the input") + output_coords = output_coords[:nframes] + energies = energies[:nframes] + forces = forces[:nframes] + cell_frame_count = len(output_coords) if output_units else nframes + cells = _output_cell_frames(outlines, cell_frame_count, cell)[:nframes] + if output_units: + alat = _output_alat(outlines, cell) + output_coords = np.array( + [ + _convert_positions(coords, unit, frame_cell, alat=card_alat or alat) + for coords, (unit, card_alat), frame_cell in zip( + output_coords, output_units[:nframes], cells + ) + ] + ) + virials = None + if stresses is not None: + if len(stresses) < nframes: + raise RuntimeError( + "the number of stress blocks does not match the number of frames" + ) + virials = stresses[:nframes] * np.linalg.det(cells)[:, None, None] + selection = slice(begin, None, step) + output_coords = output_coords[selection] + energies = energies[selection] + forces = forces[selection] + cells = cells[selection] + if virials is not None: + virials = virials[selection] + return atom_names, natoms, types, cells, output_coords, energies, forces, virials + + +def get_frame(fname): + """Read the first frame from a PWscf output for backward compatibility.""" + data = get_frames(fname) + atom_names, atom_numbs, atom_types, cells, coords, energies, forces, virials = data + if virials is not None: + virials = virials[:1] return ( atom_names, - natoms, - types, - cell[np.newaxis, :, :], - coords[np.newaxis, :, :], - np.array(energy)[np.newaxis], - force[np.newaxis, :, :], - stress, + atom_numbs, + atom_types, + cells[:1], + coords[:1], + energies[:1], + forces[:1], + virials, ) diff --git a/dpdata/plugins/qe.py b/dpdata/plugins/qe.py index 975a8cd46..72ba74a69 100644 --- a/dpdata/plugins/qe.py +++ b/dpdata/plugins/qe.py @@ -1,5 +1,6 @@ from __future__ import annotations +import dpdata.formats.qe.pwmd import dpdata.formats.qe.scf import dpdata.formats.qe.traj import dpdata.md.pbc @@ -82,6 +83,8 @@ def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): return data +@Format.register("qe/pw/aimd") +@Format.register("qe/pw/aimd_output") @Format.register("qe/pw/scf") class QECPPWSCFFormat(Format): """Quantum ESPRESSO PWscf self-consistent-field output. @@ -94,7 +97,7 @@ class QECPPWSCFFormat(Format): """ @Format.post("rot_lower_triangular") - def from_labeled_system(self, file_name, **kwargs): + def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): """Load a labeled Quantum ESPRESSO PWscf calculation. Parameters @@ -121,7 +124,30 @@ def from_labeled_system(self, file_name, **kwargs): data["energies"], data["forces"], tmp_virial, - ) = dpdata.formats.qe.scf.get_frame(file_name) + ) = dpdata.formats.qe.scf.get_frames(file_name, begin=begin, step=step) + if tmp_virial is not None: + data["virials"] = tmp_virial + return data + + +@Format.register("qe/pw/md") +class QECPPWMDFormat(Format): + """Quantum ESPRESSO ``pw.x`` molecular-dynamics output.""" + + @Format.post("rot_lower_triangular") + def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): + data = {} + ( + data["atom_names"], + data["atom_numbs"], + data["atom_types"], + data["cells"], + data["coords"], + data["energies"], + data["forces"], + tmp_virial, + ) = dpdata.formats.qe.pwmd.to_system_data(file_name, begin=begin, step=step) if tmp_virial is not None: data["virials"] = tmp_virial + data["coords"] = dpdata.md.pbc.apply_pbc(data["coords"], data["cells"]) return data diff --git a/dpdata/system.py b/dpdata/system.py index a7607bd9b..4d10158bd 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -130,7 +130,9 @@ def __init__( - ``vasp/outcar``: vasp outcar - ``vasp/xml``: vasp xml - ``qe/cp/traj``: Quantum Espresso CP trajectory files. should have: file_name+'.in' and file_name+'.pos' - - ``qe/pw/scf``: Quantum Espresso PW single point calculations. Both input and output files are required. If file_name is a string, it denotes the output file name. Input file name is obtained by replacing 'out' by 'in' from file_name. Or file_name is a list, with the first element being the input file name and the second element being the output filename. + - ``qe/pw/scf``: Quantum Espresso PW single point or AIMD calculations. Both input and output files are required. If file_name is a string, it denotes the output file name. Input file name is obtained by replacing 'out' by 'in' from file_name. Or file_name is a list, with the first element being the input file name and the second element being the output filename. + - ``qe/pw/aimd``: Quantum Espresso PW AIMD output (alias of ``qe/pw/scf``). + - ``qe/pw/md``: Quantum Espresso PW AIMD output. - ``abacus/scf``: ABACUS pw/lcao scf. The directory containing INPUT file is required. - ``abacus/md``: ABACUS pw/lcao MD. The directory containing INPUT file is required. - ``abacus/relax``: ABACUS pw/lcao relax or cell-relax. The directory containing INPUT file is required. @@ -1216,7 +1218,9 @@ class LabeledSystem(System): - ``deepmd/hdf5``: deepmd hdf5 - ``deepmd/hdf5/mixed``: deepmd mixed type hdf5 - ``qe/cp/traj``: Quantum Espresso CP trajectory files. should have: file_name+'.in', file_name+'.pos', file_name+'.evp' and file_name+'.for' - - ``qe/pw/scf``: Quantum Espresso PW single point calculations. Both input and output files are required. If file_name is a string, it denotes the output file name. Input file name is obtained by replacing 'out' by 'in' from file_name. Or file_name is a list, with the first element being the input file name and the second element being the output filename. + - ``qe/pw/scf``: Quantum Espresso PW single point or AIMD calculations. Both input and output files are required. If file_name is a string, it denotes the output file name. Input file name is obtained by replacing 'out' by 'in' from file_name. Or file_name is a list, with the first element being the input file name and the second element being the output filename. + - ``qe/pw/aimd``: Quantum Espresso PW AIMD output (alias of ``qe/pw/scf``). + - ``qe/pw/md``: Quantum Espresso PW AIMD output. - ``siesta/output``: siesta SCF output file - ``siesta/aimd_output``: siesta aimd output file - ``gaussian/log``: gaussian logs diff --git a/tests/test_qe_pw_md.py b/tests/test_qe_pw_md.py new file mode 100644 index 000000000..bc970d1da --- /dev/null +++ b/tests/test_qe_pw_md.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import os +import tempfile +import unittest + +import numpy as np +from context import dpdata + + +class TestQEPWMD(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.prefix = os.path.join(self.tmpdir.name, "md") + with open(self.prefix + ".in", "w") as fp: + fp.write( + """&CONTROL + calculation = 'md' +/ +&SYSTEM + ibrav = 0, nat = 2, ntyp = 1 +/ +CELL_PARAMETERS (angstrom) + 4.0 0.0 0.0 + 0.0 4.0 0.0 + 0.0 0.0 4.0 +ATOMIC_SPECIES +H 1.0 H.UPF +ATOMIC_POSITIONS (angstrom) +H 0.0 0.0 0.0 +H 1.0 1.0 1.0 +""" + ) + with open(self.prefix + ".out", "w") as fp: + fp.write( + """ Program PWSCF +ATOMIC_POSITIONS (angstrom) +H 0.0 0.0 0.0 +H 1.0 1.0 1.0 + +! total energy = -1.00000000 Ry + Forces acting on atoms (Ry/au): + atom 1 type 1 force = 0.10000000 0.00000000 0.00000000 + atom 2 type 1 force = 0.00000000 0.10000000 0.00000000 + +ATOMIC_POSITIONS (bohr) +H 0.5 0.5 0.5 +H 1.5 1.5 1.5 + +! total energy = -2.00000000 Ry + Forces acting on atoms (Ry/au): + atom 1 type 1 force = 0.20000000 0.00000000 0.00000000 + atom 2 type 1 force = 0.00000000 0.20000000 0.00000000 +""" + ) + + def tearDown(self): + self.tmpdir.cleanup() + + def test_pw_md_frames_and_sampling(self): + system = dpdata.LabeledSystem( + self.prefix + ".out", fmt="qe/pw/md", begin=1, step=1 + ) + self.assertEqual(system.get_nframes(), 1) + self.assertEqual(system["coords"].shape, (1, 2, 3)) + self.assertAlmostEqual( + system["energies"][0], -2.0 * dpdata.formats.qe.scf.ry2ev + ) + np.testing.assert_allclose( + system["coords"][0, 0], [0.5 * dpdata.formats.qe.scf.bohr2ang] * 3 + ) + self.assertNotIn("virials", system.data) + npy_dir = os.path.join(self.tmpdir.name, "npy") + system.to("deepmd/npy", npy_dir) + self.assertTrue(os.path.isfile(os.path.join(npy_dir, "set.000", "coord.npy"))) + + def test_variable_cell_units(self): + output = self.prefix + ".vc.out" + with open(self.prefix + ".vc.in", "w") as fp: + fp.write(open(self.prefix + ".in").read()) + with open(output, "w") as fp: + fp.write( + """ Program PWSCF +CELL_PARAMETERS (bohr) + 5.0 0.0 0.0 + 0.0 5.0 0.0 + 0.0 0.0 5.0 +ATOMIC_POSITIONS (crystal) +H 0.0 0.0 0.0 +H 0.5 0.5 0.5 + +! total energy = -1.00000000 Ry + Forces acting on atoms (Ry/au): + atom 1 type 1 force = 0.0 0.0 0.0 + atom 2 type 1 force = 0.0 0.0 0.0 +""" + ) + system = dpdata.LabeledSystem(output, fmt="qe/pw/md") + cell = 5.0 * dpdata.formats.qe.scf.bohr2ang + np.testing.assert_allclose(system["cells"][0], np.eye(3) * cell) + np.testing.assert_allclose(system["coords"][0, 1], [cell / 2] * 3) From ec9c260a585301962d7f23c606d91ec0cbaa5f85 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 9 Oct 2026 14:29:34 +0000 Subject: [PATCH 2/6] test(qe): close temporary input fixture --- tests/test_qe_pw_md.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_qe_pw_md.py b/tests/test_qe_pw_md.py index bc970d1da..cd64a5f4a 100644 --- a/tests/test_qe_pw_md.py +++ b/tests/test_qe_pw_md.py @@ -77,7 +77,8 @@ def test_pw_md_frames_and_sampling(self): def test_variable_cell_units(self): output = self.prefix + ".vc.out" with open(self.prefix + ".vc.in", "w") as fp: - fp.write(open(self.prefix + ".in").read()) + with open(self.prefix + ".in") as source: + fp.write(source.read()) with open(output, "w") as fp: fp.write( """ Program PWSCF From a159f0cc52f0b9193cb20a89f90d4b0912141e4a Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 9 Oct 2026 14:40:17 +0000 Subject: [PATCH 3/6] test(qe): cover pw scf AIMD loading --- tests/test_qe_pw_md.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_qe_pw_md.py b/tests/test_qe_pw_md.py index cd64a5f4a..24d4c9da5 100644 --- a/tests/test_qe_pw_md.py +++ b/tests/test_qe_pw_md.py @@ -74,6 +74,11 @@ def test_pw_md_frames_and_sampling(self): system.to("deepmd/npy", npy_dir) self.assertTrue(os.path.isfile(os.path.join(npy_dir, "set.000", "coord.npy"))) + # The original issue used the existing qe/pw/scf entry point for a + # pw.x AIMD output. It should now retain every labeled frame too. + scf_system = dpdata.LabeledSystem(self.prefix + ".out", fmt="qe/pw/scf") + self.assertEqual(scf_system.get_nframes(), 2) + def test_variable_cell_units(self): output = self.prefix + ".vc.out" with open(self.prefix + ".vc.in", "w") as fp: From 744a0d4cd654e5a7ccda2411735f25cc86a34c93 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 9 Oct 2026 14:54:37 +0000 Subject: [PATCH 4/6] docs(qe): document AIMD reader options --- dpdata/plugins/qe.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/dpdata/plugins/qe.py b/dpdata/plugins/qe.py index 72ba74a69..47ac9ea5d 100644 --- a/dpdata/plugins/qe.py +++ b/dpdata/plugins/qe.py @@ -106,6 +106,10 @@ def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): Quantum ESPRESSO ``pw.x`` output file. The matching input file is inferred by replacing ``out`` with ``in`` in the base name; pass ``[input_file, output_file]`` to give both paths explicitly. + begin : int + Index of the first frame to return. + step : int + Stride between returned frames. **kwargs : dict Additional format arguments accepted for API compatibility. @@ -136,6 +140,25 @@ class QECPPWMDFormat(Format): @Format.post("rot_lower_triangular") def from_labeled_system(self, file_name, begin=0, step=1, **kwargs): + """Load labeled frames from a Quantum ESPRESSO PW MD output. + + Parameters + ---------- + file_name : str or list[str] + Quantum ESPRESSO ``pw.x`` output file and, optionally, its input + file as ``[input_file, output_file]``. + begin : int + Index of the first frame to return. + step : int + Stride between returned frames. + **kwargs : dict + Additional format arguments accepted for API compatibility. + + Returns + ------- + dict + Labeled system data for the selected trajectory frames. + """ data = {} ( data["atom_names"], From a6d0ea149537239921d26032a4ea6569bef61c64 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 9 Oct 2026 15:08:55 +0000 Subject: [PATCH 5/6] fix(qe): align AIMD labels with input geometry --- dpdata/formats/qe/scf.py | 8 ++++++++ tests/test_qe_pw_md.py | 29 +++++++++++++++++++++++------ 2 files changed, 31 insertions(+), 6 deletions(-) diff --git a/dpdata/formats/qe/scf.py b/dpdata/formats/qe/scf.py index abb43d51d..35da26955 100755 --- a/dpdata/formats/qe/scf.py +++ b/dpdata/formats/qe/scf.py @@ -511,6 +511,14 @@ def get_frames(fname, begin=0, step=1): ) ] ) + # In a pw.x MD output, the energy and forces are printed for the + # current geometry before the updated positions for the next step. + # Associate the first labels with the input geometry and discard the + # final unlabelled update. + output_coords = np.concatenate( + (input_coords[np.newaxis, :, :], output_coords[:-1]), axis=0 + ) + cells = np.concatenate((cell[np.newaxis, :, :], cells[:-1]), axis=0) virials = None if stresses is not None: if len(stresses) < nframes: diff --git a/tests/test_qe_pw_md.py b/tests/test_qe_pw_md.py index 24d4c9da5..9fe9c7868 100644 --- a/tests/test_qe_pw_md.py +++ b/tests/test_qe_pw_md.py @@ -34,10 +34,6 @@ def setUp(self): with open(self.prefix + ".out", "w") as fp: fp.write( """ Program PWSCF -ATOMIC_POSITIONS (angstrom) -H 0.0 0.0 0.0 -H 1.0 1.0 1.0 - ! total energy = -1.00000000 Ry Forces acting on atoms (Ry/au): atom 1 type 1 force = 0.10000000 0.00000000 0.00000000 @@ -51,6 +47,10 @@ def setUp(self): Forces acting on atoms (Ry/au): atom 1 type 1 force = 0.20000000 0.00000000 0.00000000 atom 2 type 1 force = 0.00000000 0.20000000 0.00000000 + +ATOMIC_POSITIONS (bohr) +H 1.0 1.0 1.0 +H 2.0 2.0 2.0 """ ) @@ -78,6 +78,10 @@ def test_pw_md_frames_and_sampling(self): # pw.x AIMD output. It should now retain every labeled frame too. scf_system = dpdata.LabeledSystem(self.prefix + ".out", fmt="qe/pw/scf") self.assertEqual(scf_system.get_nframes(), 2) + np.testing.assert_allclose(scf_system["coords"][0, 0], [0.0, 0.0, 0.0]) + self.assertAlmostEqual( + scf_system["energies"][0], -1.0 * dpdata.formats.qe.scf.ry2ev + ) def test_variable_cell_units(self): output = self.prefix + ".vc.out" @@ -87,6 +91,11 @@ def test_variable_cell_units(self): with open(output, "w") as fp: fp.write( """ Program PWSCF +! total energy = -1.00000000 Ry + Forces acting on atoms (Ry/au): + atom 1 type 1 force = 0.0 0.0 0.0 + atom 2 type 1 force = 0.0 0.0 0.0 + CELL_PARAMETERS (bohr) 5.0 0.0 0.0 0.0 5.0 0.0 @@ -95,13 +104,21 @@ def test_variable_cell_units(self): H 0.0 0.0 0.0 H 0.5 0.5 0.5 -! total energy = -1.00000000 Ry +! total energy = -2.00000000 Ry Forces acting on atoms (Ry/au): atom 1 type 1 force = 0.0 0.0 0.0 atom 2 type 1 force = 0.0 0.0 0.0 + +CELL_PARAMETERS (bohr) + 6.0 0.0 0.0 + 0.0 6.0 0.0 + 0.0 0.0 6.0 +ATOMIC_POSITIONS (crystal) +H 0.0 0.0 0.0 +H 0.25 0.25 0.25 """ ) - system = dpdata.LabeledSystem(output, fmt="qe/pw/md") + system = dpdata.LabeledSystem(output, fmt="qe/pw/md", begin=1) cell = 5.0 * dpdata.formats.qe.scf.bohr2ang np.testing.assert_allclose(system["cells"][0], np.eye(3) * cell) np.testing.assert_allclose(system["coords"][0, 1], [cell / 2] * 3) From 005d22996388ff5ba614b38bfb05c40ca920d944 Mon Sep 17 00:00:00 2001 From: SchrodingersCattt Date: Fri, 9 Oct 2026 16:35:51 +0000 Subject: [PATCH 6/6] fix(qe): reject incomplete AIMD position blocks --- dpdata/formats/qe/scf.py | 2 +- tests/test_qe_pw_md.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/dpdata/formats/qe/scf.py b/dpdata/formats/qe/scf.py index 35da26955..2da17eac5 100755 --- a/dpdata/formats/qe/scf.py +++ b/dpdata/formats/qe/scf.py @@ -299,7 +299,7 @@ def _output_position_frames(lines, natoms, cell): continue block = _read_atom_block(lines, line_idx + 1, natoms) if len(block) != natoms: - continue + raise RuntimeError("incomplete ATOMIC_POSITIONS block in QE output") unit = _position_unit(line) symbols_in_block = [parts[0] for parts in block] coords = np.array( diff --git a/tests/test_qe_pw_md.py b/tests/test_qe_pw_md.py index 9fe9c7868..14a554dd2 100644 --- a/tests/test_qe_pw_md.py +++ b/tests/test_qe_pw_md.py @@ -122,3 +122,11 @@ def test_variable_cell_units(self): cell = 5.0 * dpdata.formats.qe.scf.bohr2ang np.testing.assert_allclose(system["cells"][0], np.eye(3) * cell) np.testing.assert_allclose(system["coords"][0, 1], [cell / 2] * 3) + + def test_incomplete_output_positions_are_rejected(self): + with self.assertRaisesRegex(RuntimeError, "incomplete ATOMIC_POSITIONS"): + dpdata.formats.qe.scf._output_position_frames( + ["ATOMIC_POSITIONS (angstrom)", "H 0.0 0.0 0.0", ""], + natoms=2, + cell=np.eye(3), + )