fix(qe): honor angstrom cells without cel files

dpdata/formats/qe/traj.py

@@ -12,6 +12,7 @@
     from dpdata.utils import FileType
 
 import os
+import re
 
 from ...unit import (
     EnergyConversion,
@@ -28,6 +29,13 @@
 energy_convert = EnergyConversion("hartree", "eV").value()
 force_convert = ForceConversion("hartree/bohr", "eV/angstrom").value()
 
+_QE_FLOAT_PATTERN = r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[EeDd][-+]?\d+)?"
+_CELL_PARAMETERS_PATTERN = re.compile(
+    r"^CELL_PARAMETERS(?:\s*(?:\{\s*([A-Za-z]+)\s*\}"
+    r"|\(\s*([A-Za-z]+)\s*\)|([A-Za-z]+)))?$",
+    re.IGNORECASE,
+)
+
 
 def load_key(lines, key):
     for ii in lines:
@@ -65,12 +73,53 @@ def convert_celldm(ibrav, celldm):
         # raise RuntimeError('unsupported ibrav ' + str(ibrav))
 
 
-def load_cell_parameters(lines):
-    blk = load_block(lines, "CELL_PARAMETERS", 3)
-    ret = []
-    for ii in blk:
-        ret.append([float(jj) for jj in ii.split()[0:3]])
-    return np.array(ret)
+def load_cell_parameters(lines, lattice_parameter=None):
+    """Load ``CELL_PARAMETERS`` and convert its vectors to angstrom.
+
+    CP trajectory ``.cel`` files always use atomic units, but the fallback
+    cell in the QE input file follows the unit declared on the
+    ``CELL_PARAMETERS`` card.  Keeping that distinction here prevents an
+    angstrom input cell from being converted a second time when no ``.cel``
+    file is available. ``lattice_parameter`` is the QE ``alat`` value in
+    angstrom, derived from either ``celldm(1)`` or ``A``.
+    """
+    for idx, line in enumerate(lines):
+        # Ignore commented examples such as ``!CELL_PARAMETERS {bohr}``, which
+        # are common in QE inputs and must not shadow the active card below.
+        card = line.split("!", 1)[0].strip()
+        if card.upper().startswith("CELL_PARAMETERS"):
+            blk = lines[idx + 1 : idx + 4]
+            break
+    else:
+        raise ValueError("CELL_PARAMETERS is required when ibrav is 0")
+
+    matched = _CELL_PARAMETERS_PATTERN.fullmatch(card)
+    if matched is None:
+        raise ValueError(f"ambiguous CELL_PARAMETERS unit in {card!r}")
+    unit = next((value for value in matched.groups() if value is not None), None)
+    if unit is None:
+        if lattice_parameter is None:
+            raise ValueError(
+                "CELL_PARAMETERS without a unit requires celldm(1) or A to define alat"
+            )
+        unit = "alat"
+    unit = unit.lower()
+
+    if unit == "angstrom":
+        scale = 1.0
+    elif unit == "bohr":
+        scale = length_convert
+    elif unit == "alat":
+        if lattice_parameter is None:
+            raise ValueError(
+                "CELL_PARAMETERS {alat} requires celldm(1) or A to define alat"
+            )
+        scale = lattice_parameter
+    else:
+        raise ValueError(f"unsupported CELL_PARAMETERS unit {unit!r}")
+
+    cell = np.array([[float(value) for value in row.split()[:3]] for row in blk])
+    return cell * scale
 
 
 def load_atom_names(lines, ntypes):
@@ -88,6 +137,32 @@ def load_celldm(lines):
     return celldm
 
 
+def load_lattice_parameter(lines, celldm):
+    """Return QE's ``alat`` in angstrom, rejecting conflicting definitions."""
+    a_value = None
+    in_system = False
+    a_pattern = re.compile(rf"\bA\s*=\s*({_QE_FLOAT_PATTERN})", re.IGNORECASE)
+    for raw_line in lines:
+        line = raw_line.split("!", 1)[0]
+        if re.search(r"&SYSTEM\b", line, re.IGNORECASE):
+            in_system = True
+        if in_system:
+            matched = a_pattern.search(line)
+            if matched is not None:
+                a_value = float(matched.group(1).replace("d", "e").replace("D", "E"))
+            if "/" in line:
+                break
+
+    celldm_value = celldm[0] if celldm[0] != 0 else None
+    if a_value is not None and celldm_value is not None:
+        raise ValueError("both A and celldm(1) define the QE lattice parameter")
+    if a_value is not None:
+        return a_value
+    if celldm_value is not None:
+        return celldm_value * length_convert
+    return None
+
+
 def load_atom_types(lines, natoms, atom_names):
     blk = load_block(lines, "ATOMIC_POSITIONS", natoms)
     ret = []
@@ -109,10 +184,11 @@ def load_param_file(fname: FileType):
     ibrav = int(load_key(lines, "ibrav"))
     celldm = load_celldm(lines)
     if ibrav == 0:
-        cell = load_cell_parameters(lines)
+        lattice_parameter = load_lattice_parameter(lines, celldm)
+        cell = load_cell_parameters(lines, lattice_parameter)
     else:
-        cell = convert_celldm(ibrav, celldm)
-    cell = cell * length_convert
+        # celldm and cells reconstructed from it are expressed in Bohr.
+        cell = convert_celldm(ibrav, celldm) * length_convert
     # print(atom_names)
     # print(atom_numbs)
     # print(atom_types)

tests/qe.traj/a_no_cel/cp.in

@@ -0,0 +1,21 @@
+&CONTROL
+  calculation = 'cp',
+  prefix = 'cp',
+/
+&SYSTEM
+  ibrav = 0,
+  nat = 1,
+  ntyp = 1,
+  A = 5.5,
+/
+
+CELL_PARAMETERS
+  1.0  0.0  0.0
+  0.0  1.0  0.0
+  0.0  0.0  1.0
+
+ATOMIC_SPECIES
+H  1.00794  H.UPF
+
+ATOMIC_POSITIONS { bohr }
+H  0.0  0.0  0.0

tests/qe.traj/a_no_cel/cp.pos

@@ -0,0 +1,2 @@
+0 0.0
+0.0 0.0 0.0

tests/qe.traj/alat_no_cel/cp.in

@@ -0,0 +1,21 @@
+&CONTROL
+  calculation = 'cp',
+  prefix = 'cp',
+/
+&SYSTEM
+  ibrav = 0,
+  nat = 1,
+  ntyp = 1,
+  celldm(1) = 10.0,
+/
+
+CELL_PARAMETERS { alat }
+  1.0  0.0  0.0
+  0.0  1.0  0.0
+  0.0  0.0  1.0
+
+ATOMIC_SPECIES
+H  1.00794  H.UPF
+
+ATOMIC_POSITIONS { bohr }
+H  0.0  0.0  0.0

tests/qe.traj/alat_no_cel/cp.pos

@@ -0,0 +1,2 @@
+0 0.0
+0.0 0.0 0.0

tests/qe.traj/angstrom_no_cel/cp.in

@@ -0,0 +1,20 @@
+&CONTROL
+  calculation = 'cp',
+  prefix = 'cp',
+/
+&SYSTEM
+  ibrav = 0,
+  nat = 1,
+  ntyp = 1,
+/
+
+CELL_PARAMETERS { angstrom }
+  19.7299995422   0.0000000000   0.0000000000
+   0.0000000000  19.7299995422   0.0000000000
+   0.0000000000   0.0000000000  19.7299995422
+
+ATOMIC_SPECIES
+H  1.00794  H.UPF
+
+ATOMIC_POSITIONS { bohr }
+H  0.0  0.0  0.0

tests/qe.traj/angstrom_no_cel/cp.pos

@@ -0,0 +1,2 @@
+0 0.0
+0.0 0.0 0.0

tests/qe.traj/bohr_no_cel/cp.in

@@ -0,0 +1,20 @@
+&CONTROL
+  calculation = 'cp',
+  prefix = 'cp',
+/
+&SYSTEM
+  ibrav = 0,
+  nat = 1,
+  ntyp = 1,
+/
+
+CELL_PARAMETERS { bohr }
+  2.0  0.0  0.0
+  0.0  2.0  0.0
+  0.0  0.0  2.0
+
+ATOMIC_SPECIES
+H  1.00794  H.UPF
+
+ATOMIC_POSITIONS { bohr }
+H  0.0  0.0  0.0

tests/qe.traj/bohr_no_cel/cp.pos

@@ -0,0 +1,2 @@
+0 0.0
+0.0 0.0 0.0

tests/qe.traj/missing_cell_no_cel/cp.in

@@ -0,0 +1,15 @@
+&CONTROL
+  calculation = 'cp',
+  prefix = 'cp',
+/
+&SYSTEM
+  ibrav = 0,
+  nat = 1,
+  ntyp = 1,
+/
+
+ATOMIC_SPECIES
+H  1.00794  H.UPF
+
+ATOMIC_POSITIONS { bohr }
+H  0.0  0.0  0.0

tests/test_qe_cp_traj.py

@@ -5,7 +5,7 @@
 import numpy as np
 from context import dpdata
 
-from dpdata.formats.qe.traj import convert_celldm
+from dpdata.formats.qe.traj import convert_celldm, load_cell_parameters
 
 bohr2ang = dpdata.unit.LengthConversion("bohr", "angstrom").value()
 
@@ -61,6 +61,82 @@ def setUp(self):
         self.system = dpdata.LabeledSystem("qe.traj/oh-md", fmt="qe/cp/traj")
 
 
+class TestCPTRAJInputCellUnits(unittest.TestCase):
+    def test_angstrom_cell_without_cel_trajectory(self):
+        # Earlier tests did not expose this regression: the only missing-.cel
+        # fixture used ibrav/celldm, whose cell is correctly expressed in Bohr.
+        # Exercise the distinct fallback path where CELL_PARAMETERS already
+        # stores angstrom values and therefore must not be converted again.
+        system = dpdata.System(
+            "qe.traj/angstrom_no_cel/cp",
+            fmt="qe/cp/traj",
+        )
+
+        np.testing.assert_allclose(
+            system["cells"][0],
+            np.eye(3) * 19.7299995422,
+        )
+
+    def test_bohr_cell_without_cel_trajectory(self):
+        system = dpdata.System(
+            "qe.traj/bohr_no_cel/cp",
+            fmt="qe/cp/traj",
+        )
+
+        np.testing.assert_allclose(
+            system["cells"][0],
+            np.eye(3) * 2.0 * bohr2ang,
+        )
+
+    def test_alat_cell_without_cel_trajectory(self):
+        system = dpdata.System(
+            "qe.traj/alat_no_cel/cp",
+            fmt="qe/cp/traj",
+        )
+
+        np.testing.assert_allclose(
+            system["cells"][0],
+            np.eye(3) * 10.0 * bohr2ang,
+        )
+
+    def test_omitted_unit_uses_a_lattice_parameter(self):
+        system = dpdata.System(
+            "qe.traj/a_no_cel/cp",
+            fmt="qe/cp/traj",
+        )
+        np.testing.assert_allclose(system["cells"][0], np.eye(3) * 5.5)
+
+    def test_alat_without_lattice_parameter_raises(self):
+        with self.assertRaisesRegex(ValueError, "requires celldm\\(1\\) or A"):
+            load_cell_parameters(["CELL_PARAMETERS {alat}", "1 0 0", "0 1 0", "0 0 1"])
+
+    def test_omitted_unit_without_lattice_parameter_raises(self):
+        with self.assertRaisesRegex(ValueError, "without a unit requires"):
+            load_cell_parameters(["CELL_PARAMETERS", "1 0 0", "0 1 0", "0 0 1"])
+
+    def test_unsupported_cell_unit_raises(self):
+        with self.assertRaisesRegex(ValueError, "unsupported CELL_PARAMETERS unit"):
+            load_cell_parameters(
+                ["CELL_PARAMETERS {crystal}", "1 0 0", "0 1 0", "0 0 1"]
+            )
+
+    def test_ambiguous_cell_unit_raises(self):
+        with self.assertRaisesRegex(ValueError, "ambiguous CELL_PARAMETERS unit"):
+            load_cell_parameters(
+                ["CELL_PARAMETERS {alat} bohr", "1 0 0", "0 1 0", "0 0 1"],
+                lattice_parameter=5.0,
+            )
+
+    def test_missing_cell_parameters_for_ibrav_zero_raises(self):
+        with self.assertRaisesRegex(
+            ValueError, "CELL_PARAMETERS is required when ibrav is 0"
+        ):
+            dpdata.System(
+                "qe.traj/missing_cell_no_cel/cp",
+                fmt="qe/cp/traj",
+            )
+
+
 class TestConverCellDim(unittest.TestCase):
     def test_case_null(self):
         cell = convert_celldm(8, [1, 1, 1])