QE xml parser

parse_qe_xml.py

import xml.etree.ElementTree as ET
from io import StringIO
import numpy as np
from ase import units, Atoms
from ase.calculators.singlepoint import SinglePointCalculator

def parse_qe_xml(filename):
    tree = ET.parse(filename)
    root = tree.getroot()

    cell = np.stack([
        np.loadtxt(StringIO(root.find("input").find("atomic_structure").find("cell").find(ax).text))
        for ax in ["a1", "a2", "a3"]
    ], axis=0) * units.Bohr

    elements = []
    positions = []
    for child in root.find("output").find("atomic_structure").find("atomic_positions"):
        elements.append(child.attrib["name"])
        positions.append(np.loadtxt(StringIO(child.text)))
    positions = np.stack(positions, axis=0) * units.Bohr
    structure = Atoms(elements, positions, cell=cell, pbc=True)

    forces = np.loadtxt(StringIO(root.find("output").find("forces").text)) * units.Ha / units.Bohr
    energy = float(root.find("output").find("total_energy").find("etot").text) * units.Ha
    structure.calc = SinglePointCalculator(atoms=structure, energy=energy, forces=forces)

    return structure

if __name__ == "__main__":
    import sys

    structure = parse_qe_xml(sys.argv[1])
    print(structure.symbols)
    print(structure.cell)
    print(structure.positions)
    print(structure.get_forces())
    print(structure.get_potential_energy())