Skip to content

Commit 0f4318b

Browse files
alberto-cartaalberto-carta
authored
✨ PwTools: add method to parse and return occupation matrices (#1215)
When performing Hubbard-corrected DFT calculations, Quantum ESPRESSO calculates the occupations matrices related to the target manifolds. The final matrices are saved in the xml file, which it is always retrieved. To access these matrices, instead of storing them and duplicate the data, we here add an auxiliary method to parse these matrices and return them in a format compatible with pandas, so that they can be used for post-processing purposes (e.g., calculation of oxidation states, magnetic analysis, ...). This has been implemented as a "tool" method, such that one can access the data as follows: ```python occupations = load_node(<PwCalculation PK/ID>).tools.get_occupations() ```
1 parent 1341c41 commit 0f4318b

11 files changed

+817
-0
lines changed

‎src/aiida_quantumespresso/tools/calculations/pw.py‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,16 @@
11
"""Tools for nodes created by running the `PwCalculation` class."""
2+
import numpy as np
23

34
from aiida.common import AttributeDict, exceptions
45
from aiida.tools.calculations.base import CalculationTools
6+
from aiida_quantumespresso.data.hubbard_structure import HubbardStructureData
57

68
from aiida_quantumespresso.calculations.functions.create_magnetic_configuration import create_magnetic_configuration
79

10+
from xmlschema import XMLSchema
11+
from xml.etree import ElementTree
12+
from aiida_quantumespresso.parsers.parse_xml.parse import get_schema_filepath
13+
814

915
class PwCalculationTools(CalculationTools):
1016
"""Calculation tools for `PwCalculation`.
@@ -92,3 +98,141 @@ def get_magnetic_configuration(self, atol: float = 0.5, ztol: float = 0.05) -> A
9298
return AttributeDict(
9399
{'structure': structure, 'magnetic_moments': None if non_magnetic else results['magnetic_moments']}
94100
)
101+
102+
def get_occupations(self, reshape=False) -> list[dict]:
103+
"""Return the occupations for a PwCalculation/PwBaseWorkChain node as a list of python dictionaries.
104+
Output depends on the type of calculation (unpolarized, collinear spin-polarized, non-collinear spin-polarized).
105+
106+
Example of an entry in the returned list:
107+
108+
Unpolarized:
109+
{
110+
'atom_index': 1,
111+
'kind_name': 'Fe',
112+
'manifold': 'd',
113+
'occupations': {
114+
'up-down': np.array([[...], [...], ...])
115+
}
116+
}
117+
Collinear spin-polarized:
118+
{
119+
'atom_index': 1,
120+
'kind_name': 'Fe',
121+
'manifold': 'd',
122+
'occupations': {
123+
'up': np.array([[...], [...], ...]),
124+
'down': np.array([[...], [...], ...])
125+
}
126+
}
127+
Non-collinear spin-polarized:
128+
{
129+
'atom_index': 1,
130+
'kind_name': 'Fe',
131+
'manifold': 'd',
132+
'occupations': {
133+
'up-down': np.array([[...], [...], ...])
134+
}
135+
}
136+
137+
"""
138+
139+
# assert first that this is a Hubbard calculation
140+
# so for instance if you try to call this method on a normal pw calculation it raises a ValueError
141+
if not isinstance(self._node.inputs.structure, HubbardStructureData):
142+
raise ValueError('The input structure is not a HubbardStructureData node.')
143+
144+
145+
with self._node.outputs.retrieved.base.repository.open('data-file-schema.xml') as xml_file:
146+
xml_parsed = ElementTree.parse(xml_file)
147+
148+
schema_filepath = get_schema_filepath(xml_parsed)
149+
xsd = XMLSchema(schema_filepath)
150+
xml_dictionary = xsd.to_dict(xml_parsed, validation='skip')
151+
152+
dft_u = xml_dictionary['output']['dft']['dftU']
153+
154+
# We aggregate data by atom index
155+
aggregated_data = {}
156+
157+
# Helper to standardize processing.
158+
# We treat the two lists differently based on your specifications.
159+
# 1. Hubbard_ns_nc: Non-collinear (Single matrix, ignore @spin)
160+
# 2. Hubbard_ns: Collinear (Split by @spin) OR No-Polarization (Single matrix, no @spin)
161+
162+
source_lists = [
163+
(dft_u.get('Hubbard_ns_nc', []), True), # (List, is_non_collinear)
164+
(dft_u.get('Hubbard_ns', []), False) # (List, is_standard_ns)
165+
]
166+
167+
# this needs to be verified together with the QE developers, but it seems that in the case of non-collinear
168+
# calculations, what is reported as "occupation_matrix" is actually the absolute value,
169+
# since it should have both real and imaginary parts
170+
171+
for entries, is_non_collinear in source_lists:
172+
for entry in entries:
173+
atom_index = entry['@index']
174+
175+
# The quantum espresso matrices are always as big as the biggest shell
176+
# so if we have a d and a p orbital, the p shell matrix will also be 5x5 with zeros padding
177+
# same, if we have f and d, the d will be 7x7 with zeros padding
178+
# the @dims attribute for now does not reflect this
179+
# so we can for now extract a dimension based on the name of the manifold
180+
#
181+
#
182+
# this will need to be done more intelligently (if we have Wannier orbitals there is no
183+
# guarantee for the dimension of the underlying shell) or QE should made to print the
184+
# correct dimensions in the @dims attribute
185+
actual_dim_map = {'s': 1, 'p': 3, 'd': 5, 'f': 7}
186+
187+
188+
# Parse matrix
189+
occ_matrix = np.array(entry['$'])
190+
191+
# Initialize container if new
192+
if atom_index not in aggregated_data:
193+
aggregated_data[atom_index] = {
194+
'atom_index': atom_index,
195+
'kind_name': entry['@specie'],
196+
'manifold': entry['@label'],
197+
'occupations': {}
198+
}
199+
# strip the number in the manifold label to get the actual dimension
200+
manifold_type = ''.join(filter(str.isalpha, entry['@label']))
201+
if manifold_type in actual_dim_map.keys():
202+
actual_dim = actual_dim_map[manifold_type]
203+
if is_non_collinear:
204+
205+
actual_dim *= 2 # Non-collinear has double size
206+
# print(f"Non-collinear manifold detected, doubling actual_dim to {actual_dim}")
207+
else:
208+
actual_dim = entry['@dims'][0] # Fallback to reported dimension
209+
210+
# reshape to @dims get the actual_dim x actual_dim block and reshape back to linear
211+
occ_matrix = occ_matrix.reshape(entry['@dims'])[:actual_dim, :actual_dim].reshape(-1)
212+
213+
if reshape:
214+
occ_matrix = occ_matrix.reshape((actual_dim, actual_dim))
215+
216+
# --- LOGIC FOR THE 3 CASES ---
217+
218+
# Case 3: Hubbard_ns_nc
219+
# (@spin is always 1, but we treat it as a single full block)
220+
if is_non_collinear:
221+
aggregated_data[atom_index]['occupations']['up-down'] = occ_matrix
222+
223+
# Case 1 & 2: Hubbard_ns
224+
else:
225+
spin_val = entry.get('@spin')
226+
227+
if spin_val is None:
228+
# Case 1: Hubbard_ns with no polarization (No @spin)
229+
# Treat as single matrix
230+
aggregated_data[atom_index]['occupations']['up-down'] = occ_matrix
231+
else:
232+
# Case 2: Hubbard_ns with polarization (@spin is 1 or 2)
233+
# Treat as dictionary with 'up'/'down'
234+
spin_label = 'up' if spin_val == 1 else 'down'
235+
236+
aggregated_data[atom_index]['occupations'][spin_label] = occ_matrix
237+
238+
return list(aggregated_data.values())

‎tests/tools/calculations/collinear-data-file-schema.xml‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<qes:espresso xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xmlns:qes="http://www.quantum-espresso.org/ns/qes/qes-1.0" xsi:schemaLocation="http://www.quantum-espresso.org/ns/qes/qes-1.0 http://www.quantum-espresso.org/ns/qes/qes_250521.xsd" Units="Hartree atomic units">
3+
<!-- All quantities are in Hartree atomic units unless otherwise specified -->
4+
<general_info>
5+
<xml_format NAME="QEXSD" VERSION="24.11.04">QEXSD_24.11.04</xml_format>
6+
<creator NAME="PWSCF" VERSION="7.4.1">XML file generated by PWSCF</creator>
7+
<created DATE="30Jul2025" TIME="20:58:46">This run was terminated on: 20:58:46 30 Jul 2025</created>
8+
<job></job>
9+
</general_info>
10+
<parallel_info>
11+
<nprocs>1</nprocs>
12+
<nthreads>1</nthreads>
13+
<ntasks>1</ntasks>
14+
<nbgrp>1</nbgrp>
15+
<npool>1</npool>
16+
<ndiag>1</ndiag>
17+
</parallel_info>
18+
<output>
19+
<atomic_species ntyp="3" pseudo_dir="./pseudo/">
20+
<species name="Fe1">
21+
<mass>5.584500000000000E+001</mass>
22+
<pseudo_file>Fe.pbesol-spn-kjpaw_psl.0.2.1.UPF</pseudo_file>
23+
<starting_magnetization>1.000000000000000E-009</starting_magnetization>
24+
</species>
25+
<species name="Fe2">
26+
<mass>5.584500000000000E+001</mass>
27+
<pseudo_file>Fe.pbesol-spn-kjpaw_psl.0.2.1.UPF</pseudo_file>
28+
<starting_magnetization>1.000000000000000E-009</starting_magnetization>
29+
</species>
30+
<species name="O1">
31+
<mass>1.599940000000000E+001</mass>
32+
<pseudo_file>O.pbesol-n-kjpaw_psl.0.1.UPF</pseudo_file>
33+
<starting_magnetization>0.000000000000000E+000</starting_magnetization>
34+
</species>
35+
</atomic_species>
36+
<atomic_structure nat="4" num_of_atomic_wfc="34" alat="10.03066049663022">
37+
<atomic_positions>
38+
<atom name="Fe1" index="1">
39+
0.000000000000000E+000 0.000000000000000E+000 0.000000000000000E+000
40+
</atom>
41+
<atom name="Fe2" index="2">8.190000000008371E+000 8.190000000008371E+000 8.190000000008371E+000</atom>
42+
<atom name="O1" index="3">4.094999999909700E+000 4.094999999909700E+000 4.094999999909700E+000</atom>
43+
<atom name="O1" index="4">1.228499999991807E+001 1.228499999991807E+001 1.228499999991807E+001</atom>
44+
</atomic_positions>
45+
<cell>
46+
<a1>4.094999999909700E+000 4.094999999909700E+000 8.190000000008371E+000</a1>
47+
<a2>4.094999999909700E+000 8.190000000008371E+000 4.094999999909700E+000</a2>
48+
<a3>8.190000000008371E+000 4.094999999909700E+000 4.094999999909700E+000</a3>
49+
</cell>
50+
</atomic_structure>
51+
<dft>
52+
<functional>PBESOL</functional>
53+
<dftU new_format="true">
54+
<lda_plus_u_kind>2</lda_plus_u_kind>
55+
<Hubbard_Occ channels="1" specie="Fe1">
56+
<channel_occ specie="Fe1" label="3d" index="1">6.000000000000000E+000</channel_occ>
57+
</Hubbard_Occ>
58+
<Hubbard_Occ channels="1" specie="Fe2">
59+
<channel_occ specie="Fe2" label="3d" index="1">6.000000000000000E+000</channel_occ>
60+
</Hubbard_Occ>
61+
<Hubbard_Occ channels="1" specie="O1">
62+
<channel_occ specie="O1" label="2p" index="1">4.000000000000000E+000</channel_occ>
63+
</Hubbard_Occ>
64+
<Hubbard_V specie1="Fe1" index1="1" label1="3d" specie2="Fe1" index2="1" label2="3d">1.837466108782747E-001</Hubbard_V>
65+
<Hubbard_V specie1="Fe1" index1="1" label1="3d" specie2="O1" index2="12" label2="2p">3.674932217565494E-002</Hubbard_V>
66+
<Hubbard_V specie1="Fe1" index1="1" label1="3d" specie2="O1" index2="20" label2="2p">3.674932217565494E-002</Hubbard_V>
67+
<Hubbard_V specie1="Fe1" index1="1" label1="3d" specie2="O1" index2="23" label2="2p">3.674932217565494E-002</Hubbard_V>
68+
<Hubbard_V specie1="Fe1" index1="1" label1="3d" specie2="O1" index2="44" label2="2p">3.674932217565494E-002</Hubbard_V>
69+
<Hubbard_V specie1="Fe1" index1="1" label1="3d" specie2="O1" index2="47" label2="2p">3.674932217565494E-002</Hubbard_V>
70+
<Hubbard_V specie1="Fe1" index1="1" label1="3d" specie2="O1" index2="55" label2="2p">3.674932217565494E-002</Hubbard_V>
71+
<Hubbard_V specie1="Fe2" index1="2" label1="3d" specie2="Fe2" index2="2" label2="3d">1.837466108782747E-001</Hubbard_V>
72+
<Hubbard_V specie1="Fe2" index1="2" label1="3d" specie2="O1" index2="24" label2="2p">3.674932217565494E-002</Hubbard_V>
73+
<Hubbard_V specie1="Fe2" index1="2" label1="3d" specie2="O1" index2="48" label2="2p">3.674932217565494E-002</Hubbard_V>
74+
<Hubbard_V specie1="Fe2" index1="2" label1="3d" specie2="O1" index2="56" label2="2p">3.674932217565494E-002</Hubbard_V>
75+
<Hubbard_V specie1="Fe2" index1="2" label1="3d" specie2="O1" index2="59" label2="2p">3.674932217565494E-002</Hubbard_V>
76+
<Hubbard_V specie1="Fe2" index1="2" label1="3d" specie2="O1" index2="67" label2="2p">3.674932217565494E-002</Hubbard_V>
77+
<Hubbard_V specie1="Fe2" index1="2" label1="3d" specie2="O1" index2="91" label2="2p">3.674932217565494E-002</Hubbard_V>
78+
<Hubbard_V specie1="O1" index1="3" label1="2p" specie2="Fe2" index2="22" label2="3d">3.674932217565494E-002</Hubbard_V>
79+
<Hubbard_V specie1="O1" index1="3" label1="2p" specie2="Fe2" index2="46" label2="3d">3.674932217565494E-002</Hubbard_V>
80+
<Hubbard_V specie1="O1" index1="3" label1="2p" specie2="Fe2" index2="54" label2="3d">3.674932217565494E-002</Hubbard_V>
81+
<Hubbard_V specie1="O1" index1="3" label1="2p" specie2="Fe1" index2="57" label2="3d">3.674932217565494E-002</Hubbard_V>
82+
<Hubbard_V specie1="O1" index1="3" label1="2p" specie2="Fe1" index2="65" label2="3d">3.674932217565494E-002</Hubbard_V>
83+
<Hubbard_V specie1="O1" index1="3" label1="2p" specie2="Fe1" index2="89" label2="3d">3.674932217565494E-002</Hubbard_V>
84+
<Hubbard_V specie1="O1" index1="4" label1="2p" specie2="Fe2" index2="58" label2="3d">3.674932217565494E-002</Hubbard_V>
85+
<Hubbard_V specie1="O1" index1="4" label1="2p" specie2="Fe2" index2="66" label2="3d">3.674932217565494E-002</Hubbard_V>
86+
<Hubbard_V specie1="O1" index1="4" label1="2p" specie2="Fe1" index2="69" label2="3d">3.674932217565494E-002</Hubbard_V>
87+
<Hubbard_V specie1="O1" index1="4" label1="2p" specie2="Fe2" index2="90" label2="3d">3.674932217565494E-002</Hubbard_V>
88+
<Hubbard_V specie1="O1" index1="4" label1="2p" specie2="Fe1" index2="93" label2="3d">3.674932217565494E-002</Hubbard_V>
89+
<Hubbard_V specie1="O1" index1="4" label1="2p" specie2="Fe1" index2="101" label2="3d">3.674932217565494E-002</Hubbard_V>
90+
<Hubbard_ns rank="2" dims=" 5 5" order="F" specie="Fe1" label="3d" spin="1" index="1">
91+
1.431570305581613E-001 -3.181832781422882E-003 -3.181832781422789E-003
92+
-9.860225835692585E-017 -6.363665562836351E-003
93+
-3.181832781422882E-003 9.897075558145182E-001 -2.248448016171220E-004
94+
-5.511096038605065E-003 2.248448016168599E-004
95+
-3.181832781422788E-003 -2.248448016171220E-004 9.897075558145136E-001
96+
5.511096038611794E-003 2.248448016205300E-004
97+
-1.000362851366274E-016 -5.511096038605065E-003 5.511096038611794E-003
98+
1.431570305581614E-001 -6.411637917098055E-015
99+
-6.363665562836351E-003 2.248448016168599E-004 2.248448016205300E-004
100+
-6.411721773359833E-015 9.897075558145104E-001
101+
</Hubbard_ns>
102+
<Hubbard_ns rank="2" dims=" 5 5" order="F" specie="Fe1" label="3d" spin="2" index="1">
103+
1.546224647641354E-001 2.170318829404453E-003 2.170318829404570E-003
104+
-1.328977753583052E-016 4.340637658818308E-003
105+
2.170318829404453E-003 9.904265886971974E-001 -1.315518044151248E-004
106+
3.759102481159509E-003 1.315518044148601E-004
107+
2.170318829404570E-003 -1.315518044151248E-004 9.904265886971928E-001
108+
-3.759102481152777E-003 1.315518044185255E-004
109+
-1.406633734187326E-016 3.759102481159509E-003 -3.759102481152777E-003
110+
1.546224647641355E-001 -6.417623898936301E-015
111+
4.340637658818309E-003 1.315518044148601E-004 1.315518044185255E-004
112+
-6.417317273009395E-015 9.904265886971896E-001
113+
</Hubbard_ns>
114+
<Hubbard_ns rank="2" dims=" 5 5" order="F" specie="Fe2" label="3d" spin="1" index="2">
115+
2.560526897758898E-001 -1.259454589280899E-001 -1.259454589280901E-001
116+
6.202059943090713E-016 -2.518909178561698E-001
117+
-1.259454589280899E-001 9.064164797303442E-001 4.256774686335565E-002
118+
-2.181439338460230E-001 -4.256774686335588E-002
119+
-1.259454589280901E-001 4.256774686335565E-002 9.064164797303399E-001
120+
2.181439338460301E-001 -4.256774686335209E-002
121+
6.236906878540754E-016 -2.181439338460230E-001 2.181439338460302E-001
122+
2.560526897758912E-001 -5.987479816013360E-015
123+
-2.518909178561697E-001 -4.256774686335588E-002 -4.256774686335209E-002
124+
-5.986583866913401E-015 9.064164797303373E-001
125+
</Hubbard_ns>
126+
<Hubbard_ns rank="2" dims=" 5 5" order="F" specie="Fe2" label="3d" spin="2" index="2">
127+
4.344218232317197E-001 1.754481478690268E-001 1.754481478690278E-001
128+
-1.466952644427200E-015 3.508962957380624E-001
129+
1.754481478690268E-001 7.685642619928010E-001 1.105972555071096E-001
130+
3.038851062030202E-001 -1.105972555071100E-001
131+
1.754481478690278E-001 1.105972555071096E-001 7.685642619927975E-001
132+
-3.038851062030133E-001 -1.105972555071068E-001
133+
-1.479765605941197E-015 3.038851062030202E-001 -3.038851062030133E-001
134+
4.344218232317189E-001 -5.580887060672350E-015
135+
3.508962957380623E-001 -1.105972555071100E-001 -1.105972555071068E-001
136+
-5.583709797969075E-015 7.685642619927947E-001
137+
</Hubbard_ns>
138+
<Hubbard_ns rank="2" dims=" 5 5" order="F" specie="O1" label="2p" spin="1" index="3">
139+
7.978457269308926E-001 4.252078822874832E-003 4.252078822875915E-003
140+
0.000000000000000E+000 0.000000000000000E+000
141+
4.252078822874832E-003 7.978457269308916E-001 -4.252078822875964E-003
142+
0.000000000000000E+000 0.000000000000000E+000
143+
4.252078822875915E-003 -4.252078822875964E-003 7.978457269308921E-001
144+
0.000000000000000E+000 0.000000000000000E+000
145+
0.000000000000000E+000 0.000000000000000E+000 0.000000000000000E+000
146+
0.000000000000000E+000 0.000000000000000E+000
147+
0.000000000000000E+000 0.000000000000000E+000 0.000000000000000E+000
148+
0.000000000000000E+000 0.000000000000000E+000
149+
</Hubbard_ns>
150+
<Hubbard_ns rank="2" dims=" 5 5" order="F" specie="O1" label="2p" spin="2" index="3">
151+
8.013153331087121E-001 -6.651877829553964E-003 -6.651877829552884E-003
152+
0.000000000000000E+000 0.000000000000000E+000
153+
-6.651877829553964E-003 8.013153331087111E-001 6.651877829552824E-003
154+
0.000000000000000E+000 0.000000000000000E+000
155+
-6.651877829552884E-003 6.651877829552824E-003 8.013153331087114E-001
156+
0.000000000000000E+000 0.000000000000000E+000
157+
0.000000000000000E+000 0.000000000000000E+000 0.000000000000000E+000
158+
0.000000000000000E+000 0.000000000000000E+000
159+
0.000000000000000E+000 0.000000000000000E+000 0.000000000000000E+000
160+
0.000000000000000E+000 0.000000000000000E+000
161+
</Hubbard_ns>
162+
<Hubbard_ns rank="2" dims=" 5 5" order="F" specie="O1" label="2p" spin="1" index="4">
163+
7.978457269308926E-001 4.252078822874831E-003 4.252078822875916E-003
164+
0.000000000000000E+000 0.000000000000000E+000
165+
4.252078822874831E-003 7.978457269308917E-001 -4.252078822875964E-003
166+
0.000000000000000E+000 0.000000000000000E+000
167+
4.252078822875916E-003 -4.252078822875964E-003 7.978457269308920E-001
168+
0.000000000000000E+000 0.000000000000000E+000
169+
0.000000000000000E+000 0.000000000000000E+000 0.000000000000000E+000
170+
0.000000000000000E+000 0.000000000000000E+000
171+
0.000000000000000E+000 0.000000000000000E+000 0.000000000000000E+000
172+
0.000000000000000E+000 0.000000000000000E+000
173+
</Hubbard_ns>
174+
<Hubbard_ns rank="2" dims=" 5 5" order="F" specie="O1" label="2p" spin="2" index="4">
175+
8.013153331087121E-001 -6.651877829553964E-003 -6.651877829552883E-003
176+
0.000000000000000E+000 0.000000000000000E+000
177+
-6.651877829553964E-003 8.013153331087109E-001 6.651877829552825E-003
178+
0.000000000000000E+000 0.000000000000000E+000
179+
-6.651877829552883E-003 6.651877829552825E-003 8.013153331087114E-001
180+
0.000000000000000E+000 0.000000000000000E+000
181+
0.000000000000000E+000 0.000000000000000E+000 0.000000000000000E+000
182+
0.000000000000000E+000 0.000000000000000E+000
183+
0.000000000000000E+000 0.000000000000000E+000 0.000000000000000E+000
184+
0.000000000000000E+000 0.000000000000000E+000
185+
</Hubbard_ns>
186+
<U_projection_type>ortho-atomic</U_projection_type>
187+
</dftU>
188+
</dft>
189+
</output>
190+
<exit_status>0</exit_status>
191+
<closed DATE="30 Jul 2025" TIME="20:58:46"></closed>
192+
</qes:espresso>

0 commit comments

Comments
 (0)