Fix num_of_atomic_wfc parsing + add test

mbercx · Jul 24, 2024 · 7c4fbb1 · 7c4fbb1
1 parent e831ac9
commit 7c4fbb1
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 10 deletions.
diff --git a/src/aiida_quantumespresso/parsers/parse_xml/parse.py b/src/aiida_quantumespresso/parsers/parse_xml/parse.py
@@ -295,7 +295,9 @@ def parse_xml_post_6_2(xml):
 
         num_k_points = band_structure['nks']
         num_electrons = band_structure['nelec']
-        num_atomic_wfc = band_structure['num_of_atomic_wfc']
+
+        # In schema v240411 (QE v7.3.1), the `number_of_atomic_wfc` is moved to the `atomic_structure` tag as an attribute
+        num_atomic_wfc = band_structure.get('num_of_atomic_wfc', None) or outputs['atomic_structure']['@num_of_atomic_wfc']
         num_bands = band_structure.get('nbnd', None)
         num_bands_up = band_structure.get('nbnd_up', None)
         num_bands_down = band_structure.get('nbnd_dw', None)

diff --git a/tests/parsers/test_pw.py b/tests/parsers/test_pw.py
@@ -97,15 +97,18 @@ def test_pw_default_no_xml(
     })
 
 
-@pytest.mark.parametrize('xml_format', [
-    '190304',
-    '191206',
-    '200420',
-    '210716',
-    '211101',
-    '220603',
-    '230310',
-])
+@pytest.mark.parametrize(
+    'xml_format', [
+        '190304',
+        '191206',
+        '200420',
+        '210716',
+        '211101',
+        '220603',
+        '230310',
+        '240411',
+    ]
+)
 def test_pw_default_xml(
     fixture_localhost, generate_calc_job_node, generate_parser, generate_inputs, data_regression, xml_format
 ):