Skip to content

Commit eee737b

Browse files
Copilotnjzjz
Copilot
and
njzjz
committed
feat(qe): Add support for ibrav > 1 in QE SCF parser
Co-authored-by: njzjz <[email protected]>
1 parent 94f69b9 commit eee737b

File tree

2 files changed

+340
-14
lines changed

2 files changed

+340
-14
lines changed

‎dpdata/qe/scf.py‎

Lines changed: 204 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -47,13 +47,213 @@ def get_block(lines, keyword, skip=0):
4747
return ret
4848

4949

50+
def _parse_lattice_parameters(lines):
51+
"""Parse lattice parameters from QE input lines."""
52+
params = {}
53+
for iline in lines:
54+
line = iline.replace("=", " ").replace(",", "").split()
55+
if len(line) >= 2:
56+
if line[0] == "a":
57+
params["a"] = float(line[1])
58+
elif line[0] == "b":
59+
params["b"] = float(line[1])
60+
elif line[0] == "c":
61+
params["c"] = float(line[1])
62+
elif line[0] == "cosab":
63+
params["cosab"] = float(line[1])
64+
elif line[0] == "cosac":
65+
params["cosac"] = float(line[1])
66+
elif line[0] == "cosbc":
67+
params["cosbc"] = float(line[1])
68+
elif line[0].startswith("celldm("):
69+
# Extract celldm index from celldm(1), celldm(2), etc.
70+
idx = int(line[0][7:-1]) # Extract number from celldm(n)
71+
if "celldm" not in params:
72+
params["celldm"] = {}
73+
params["celldm"][idx] = float(line[1])
74+
return params
75+
76+
77+
def _convert_ibrav_to_cell(ibrav, params):
78+
"""Convert ibrav and lattice parameters to cell matrix."""
79+
# Extract parameters
80+
a = params.get("a")
81+
b = params.get("b")
82+
c = params.get("c")
83+
cosab = params.get("cosab", 0.0)
84+
cosac = params.get("cosac", 0.0)
85+
cosbc = params.get("cosbc", 0.0)
86+
celldm = params.get("celldm", {})
87+
88+
# Convert celldm parameters if present (celldm is in atomic units)
89+
if 1 in celldm:
90+
a = celldm[1] * bohr2ang if a is None else a
91+
if 2 in celldm and a is not None:
92+
b = celldm[2] * a if b is None else b
93+
if 3 in celldm and a is not None:
94+
c = celldm[3] * a if c is None else c
95+
if 4 in celldm:
96+
cosab = celldm[4] if cosab == 0.0 else cosab
97+
if 5 in celldm:
98+
cosac = celldm[5] if cosac == 0.0 else cosac
99+
if 6 in celldm:
100+
cosbc = celldm[6] if cosbc == 0.0 else cosbc
101+
102+
# Set defaults only for ibrav types that don't require explicit b,c
103+
if ibrav in [1, 2, 3, -3]: # Cubic lattices
104+
if b is None:
105+
b = a
106+
if c is None:
107+
c = a
108+
elif ibrav in [6, 7]: # Tetragonal lattices
109+
if b is None:
110+
b = a
111+
# c must be specified explicitly
112+
113+
# Validate required parameters
114+
if a is None:
115+
raise RuntimeError("parameter 'a' or 'celldm(1)' cannot be found.")
116+
117+
# Generate cell matrix based on ibrav
118+
if ibrav == 1: # simple cubic
119+
return np.array([[a, 0.0, 0.0], [0.0, a, 0.0], [0.0, 0.0, a]])
120+
121+
elif ibrav == 2: # face-centered cubic
122+
return a * 0.5 * np.array([[-1, 0, 1], [0, 1, 1], [-1, 1, 0]])
123+
124+
elif ibrav == 3: # body-centered cubic
125+
return a * 0.5 * np.array([[1, 1, 1], [-1, 1, 1], [-1, -1, 1]])
126+
127+
elif ibrav == -3: # reverse body-centered cubic
128+
return a * 0.5 * np.array([[-1, 1, 1], [1, -1, 1], [1, 1, -1]])
129+
130+
elif ibrav == 4: # hexagonal
131+
if c is None:
132+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=4")
133+
return np.array([[a, 0, 0],
134+
[-a/2, a*np.sqrt(3)/2, 0],
135+
[0, 0, c]])
136+
137+
elif ibrav == 6: # simple tetragonal
138+
if c is None:
139+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=6")
140+
return np.array([[a, 0, 0], [0, a, 0], [0, 0, c]])
141+
142+
elif ibrav == 7: # body-centered tetragonal
143+
if c is None:
144+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=7")
145+
return np.array([[a/2, -a/2, c/2],
146+
[a/2, a/2, c/2],
147+
[-a/2, -a/2, c/2]])
148+
149+
elif ibrav == 8: # simple orthorhombic
150+
if b is None:
151+
raise RuntimeError("parameter 'b' or 'celldm(2)' required for ibrav=8")
152+
if c is None:
153+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=8")
154+
return np.array([[a, 0, 0], [0, b, 0], [0, 0, c]])
155+
156+
elif ibrav == 9: # base-centered orthorhombic
157+
if b is None:
158+
raise RuntimeError("parameter 'b' or 'celldm(2)' required for ibrav=9")
159+
if c is None:
160+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=9")
161+
return np.array([[a/2, b/2, 0],
162+
[-a/2, b/2, 0],
163+
[0, 0, c]])
164+
165+
elif ibrav == -9: # reverse base-centered orthorhombic
166+
if b is None:
167+
raise RuntimeError("parameter 'b' or 'celldm(2)' required for ibrav=-9")
168+
if c is None:
169+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=-9")
170+
return np.array([[a/2, -b/2, 0],
171+
[a/2, b/2, 0],
172+
[0, 0, c]])
173+
174+
elif ibrav == 10: # face-centered orthorhombic
175+
if b is None:
176+
raise RuntimeError("parameter 'b' or 'celldm(2)' required for ibrav=10")
177+
if c is None:
178+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=10")
179+
return np.array([[a/2, 0, c/2],
180+
[a/2, b/2, 0],
181+
[0, b/2, c/2]])
182+
183+
elif ibrav == 11: # body-centered orthorhombic
184+
if b is None:
185+
raise RuntimeError("parameter 'b' or 'celldm(2)' required for ibrav=11")
186+
if c is None:
187+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=11")
188+
return np.array([[a/2, b/2, c/2],
189+
[-a/2, b/2, c/2],
190+
[-a/2, -b/2, c/2]])
191+
192+
elif ibrav == 12: # simple monoclinic
193+
if b is None:
194+
raise RuntimeError("parameter 'b' or 'celldm(2)' required for ibrav=12")
195+
if c is None:
196+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=12")
197+
sinab = np.sqrt(1 - cosab**2)
198+
return np.array([[a, 0, 0],
199+
[b*cosab, b*sinab, 0],
200+
[0, 0, c]])
201+
202+
elif ibrav == -12: # reverse monoclinic
203+
if b is None:
204+
raise RuntimeError("parameter 'b' or 'celldm(2)' required for ibrav=-12")
205+
if c is None:
206+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=-12")
207+
sinac = np.sqrt(1 - cosac**2)
208+
return np.array([[a, 0, 0],
209+
[0, b, 0],
210+
[c*cosac, 0, c*sinac]])
211+
212+
elif ibrav == 13: # base-centered monoclinic
213+
if b is None:
214+
raise RuntimeError("parameter 'b' or 'celldm(2)' required for ibrav=13")
215+
if c is None:
216+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=13")
217+
sinab = np.sqrt(1 - cosab**2)
218+
return np.array([[a/2, 0, -c/2],
219+
[b*cosab, b*sinab, 0],
220+
[a/2, 0, c/2]])
221+
222+
elif ibrav == -13: # reverse base-centered monoclinic
223+
if b is None:
224+
raise RuntimeError("parameter 'b' or 'celldm(2)' required for ibrav=-13")
225+
if c is None:
226+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=-13")
227+
sinac = np.sqrt(1 - cosac**2)
228+
return np.array([[a/2, -b/2, 0],
229+
[a/2, b/2, 0],
230+
[c*cosac, 0, c*sinac]])
231+
232+
elif ibrav == 14: # triclinic
233+
if b is None:
234+
raise RuntimeError("parameter 'b' or 'celldm(2)' required for ibrav=14")
235+
if c is None:
236+
raise RuntimeError("parameter 'c' or 'celldm(3)' required for ibrav=14")
237+
sinab = np.sqrt(1 - cosab**2)
238+
sinac = np.sqrt(1 - cosac**2)
239+
cosbc_prime = (cosbc - cosab*cosac) / (sinab*sinac)
240+
sinbc_prime = np.sqrt(1 - cosbc_prime**2)
241+
return np.array([[a, 0, 0],
242+
[b*cosab, b*sinab, 0],
243+
[c*cosac, c*sinac*cosbc_prime, c*sinac*sinbc_prime]])
244+
245+
else:
246+
raise RuntimeError(f"ibrav = {ibrav} is not supported yet.")
247+
248+
50249
def get_cell(lines):
51250
ret = []
52251
for idx, ii in enumerate(lines):
53252
if "ibrav" in ii:
54253
break
55254
blk = lines[idx : idx + 2]
56255
ibrav = int(blk[0].replace(",", "").split("=")[-1])
256+
57257
if ibrav == 0:
58258
for iline in lines:
59259
if "CELL_PARAMETERS" in iline and "angstrom" not in iline.lower():
@@ -64,21 +264,11 @@ def get_cell(lines):
64264
for ii in blk:
65265
ret.append([float(jj) for jj in ii.split()[0:3]])
66266
ret = np.array(ret)
67-
elif ibrav == 1:
68-
a = None
69-
for iline in lines:
70-
line = iline.replace("=", " ").replace(",", "").split()
71-
if len(line) >= 2 and "a" == line[0]:
72-
# print("line = ", line)
73-
a = float(line[1])
74-
if len(line) >= 2 and "celldm(1)" == line[0]:
75-
a = float(line[1]) * bohr2ang
76-
# print("a = ", a)
77-
if not a:
78-
raise RuntimeError("parameter 'a' or 'celldm(1)' cannot be found.")
79-
ret = np.array([[a, 0.0, 0.0], [0.0, a, 0.0], [0.0, 0.0, a]])
80267
else:
81-
raise RuntimeError("ibrav > 1 not supported yet.")
268+
# Parse lattice parameters and convert based on ibrav
269+
params = _parse_lattice_parameters(lines)
270+
ret = _convert_ibrav_to_cell(ibrav, params)
271+
82272
return ret
83273

84274

‎tests/test_qe_ibrav_support.py‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
from __future__ import annotations
2+
3+
import os
4+
import tempfile
5+
import unittest
6+
7+
import numpy as np
8+
from context import dpdata
9+
10+
11+
class TestQEIbravSupport(unittest.TestCase):
12+
"""Test support for various ibrav values in Quantum Espresso files."""
13+
14+
def setUp(self):
15+
# Create temporary directory for test files
16+
self.temp_dir = tempfile.mkdtemp()
17+
18+
def tearDown(self):
19+
# Clean up temporary files
20+
import shutil
21+
shutil.rmtree(self.temp_dir)
22+
23+
def _create_qe_files(self, ibrav, params_str, expected_cell):
24+
"""Create QE input and output files for testing."""
25+
input_content = f"""&control
26+
calculation='scf'
27+
/
28+
&system
29+
ibrav={ibrav}
30+
{params_str}
31+
nat=1
32+
ntyp=1
33+
/
34+
ATOMIC_SPECIES
35+
H 1.0 H.pbe.UPF
36+
ATOMIC_POSITIONS {{angstrom}}
37+
H 0.0 0.0 0.0
38+
"""
39+
40+
output_content = """&control
41+
calculation='scf'
42+
/
43+
44+
! total energy = -10.0 Ry
45+
46+
Forces acting on atoms (cartesian axes, Ry/au):
47+
48+
atom 1 type 1 force = 0.00000000 0.00000000 0.00000000
49+
"""
50+
51+
input_file = os.path.join(self.temp_dir, "test.in")
52+
output_file = os.path.join(self.temp_dir, "test.out")
53+
54+
with open(input_file, 'w') as f:
55+
f.write(input_content)
56+
with open(output_file, 'w') as f:
57+
f.write(output_content)
58+
59+
return input_file, output_file
60+
61+
def test_ibrav_1_cubic(self):
62+
"""Test ibrav=1 (simple cubic) - existing functionality."""
63+
_, output_file = self._create_qe_files(1, "a=10", np.eye(3) * 10)
64+
65+
sys = dpdata.LabeledSystem(output_file, fmt='qe/pw/scf')
66+
expected_cell = np.array([[10, 0, 0], [0, 10, 0], [0, 0, 10]])
67+
np.testing.assert_allclose(sys.data['cells'][0], expected_cell, rtol=1e-10)
68+
69+
def test_ibrav_8_simple_orthorhombic(self):
70+
"""Test ibrav=8 (simple orthorhombic) - main issue case."""
71+
_, output_file = self._create_qe_files(8, "a=10\nb=12\nc=8", None)
72+
73+
sys = dpdata.LabeledSystem(output_file, fmt='qe/pw/scf')
74+
expected_cell = np.array([[10, 0, 0], [0, 12, 0], [0, 0, 8]])
75+
np.testing.assert_allclose(sys.data['cells'][0], expected_cell, rtol=1e-10)
76+
77+
def test_ibrav_8_with_celldm(self):
78+
"""Test ibrav=8 with celldm parameters."""
79+
# celldm in atomic units (bohr)
80+
bohr2ang = dpdata.qe.scf.bohr2ang
81+
a_bohr = 10 / bohr2ang
82+
_, output_file = self._create_qe_files(8, f"celldm(1)={a_bohr}\ncelldm(2)=1.2\ncelldm(3)=0.8", None)
83+
84+
sys = dpdata.LabeledSystem(output_file, fmt='qe/pw/scf')
85+
expected_cell = np.array([[10, 0, 0], [0, 12, 0], [0, 0, 8]])
86+
np.testing.assert_allclose(sys.data['cells'][0], expected_cell, rtol=1e-6)
87+
88+
def test_ibrav_2_fcc(self):
89+
"""Test ibrav=2 (face-centered cubic)."""
90+
_, output_file = self._create_qe_files(2, "a=10", None)
91+
92+
sys = dpdata.LabeledSystem(output_file, fmt='qe/pw/scf')
93+
# After rot_lower_triangular transformation
94+
expected_cell = np.array([[7.071068e+00, 0.000000e+00, 0.000000e+00],
95+
[3.535534e+00, 6.123724e+00, 0.000000e+00],
96+
[3.535534e+00, 2.041241e+00, 5.773503e+00]])
97+
np.testing.assert_allclose(sys.data['cells'][0], expected_cell, atol=1e-6)
98+
99+
def test_ibrav_3_bcc(self):
100+
"""Test ibrav=3 (body-centered cubic)."""
101+
_, output_file = self._create_qe_files(3, "a=10", None)
102+
103+
sys = dpdata.LabeledSystem(output_file, fmt='qe/pw/scf')
104+
# After rot_lower_triangular transformation
105+
expected_cell = np.array([[8.660254e+00, 0.000000e+00, 0.000000e+00],
106+
[2.886751e+00, 8.164966e+00, 0.000000e+00],
107+
[-2.886751e+00, 4.082483e+00, 7.071068e+00]])
108+
np.testing.assert_allclose(sys.data['cells'][0], expected_cell, atol=1e-6)
109+
110+
def test_ibrav_6_tetragonal(self):
111+
"""Test ibrav=6 (simple tetragonal)."""
112+
_, output_file = self._create_qe_files(6, "a=10\nc=8", None)
113+
114+
sys = dpdata.LabeledSystem(output_file, fmt='qe/pw/scf')
115+
expected_cell = np.array([[10, 0, 0], [0, 10, 0], [0, 0, 8]])
116+
np.testing.assert_allclose(sys.data['cells'][0], expected_cell, rtol=1e-10)
117+
118+
def test_ibrav_missing_parameters(self):
119+
"""Test error handling for missing required parameters."""
120+
_, output_file = self._create_qe_files(8, "a=10", None) # Missing b and c
121+
122+
with self.assertRaises(RuntimeError) as cm:
123+
dpdata.LabeledSystem(output_file, fmt='qe/pw/scf')
124+
self.assertIn("parameter 'b'", str(cm.exception))
125+
126+
def test_unsupported_ibrav(self):
127+
"""Test error handling for unsupported ibrav values."""
128+
_, output_file = self._create_qe_files(99, "a=10", None)
129+
130+
with self.assertRaises(RuntimeError) as cm:
131+
dpdata.LabeledSystem(output_file, fmt='qe/pw/scf')
132+
self.assertIn("ibrav = 99 is not supported", str(cm.exception))
133+
134+
135+
if __name__ == '__main__':
136+
unittest.main()

0 commit comments

Comments
 (0)