Coverage for ase / io / qbox.py: 96.67%
90 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1# fmt: off
3"""This module contains functions to read from QBox output files"""
5import re
6import xml.etree.ElementTree as ET
8from ase import Atom, Atoms
9from ase.calculators.singlepoint import SinglePointCalculator
10from ase.utils import reader
12# Compile regexs for fixing XML
13re_find_bad_xml = re.compile(r'<(/?)([A-z]+) expectation ([a-z]+)')
16@reader
17def read_qbox(f, index=-1):
18 """Read data from QBox output file
20 Inputs:
21 f - str or fileobj, path to file or file object to read from
22 index - int or slice, which frames to return
24 Returns
25 -------
26 list of Atoms or atoms, requested frame(s)
27 """
29 # Check whether this is a QB@all output
30 version = None
31 for line in f:
32 if '<release>' in line:
33 version = ET.fromstring(line)
34 break
35 if version is None:
36 raise Exception('Parse Error: Version not found')
37 is_qball = 'qb@LL' in version.text or 'qball' in version.text
39 # Load in atomic species
40 species = {}
41 if is_qball:
42 # Read all of the lines between release and the first call to `run`
43 species_data = []
44 for line in f:
45 if '<run' in line:
46 break
47 species_data.append(line)
48 species_data = '\n'.join(species_data)
50 # Read out the species information with regular expressions
51 symbols = re.findall('symbol_ = ([A-Z][a-z]?)', species_data)
52 masses = re.findall('mass_ = ([0-9.]+)', species_data)
53 names = re.findall('name_ = ([a-z]+)', species_data)
54 numbers = re.findall('atomic_number_ = ([0-9]+)', species_data)
56 # Compile them into a dictionary
57 for name, symbol, mass, number in zip(names, symbols, masses, numbers):
58 spec_data = dict(
59 symbol=symbol,
60 mass=float(mass),
61 number=float(number)
62 )
63 species[name] = spec_data
64 else:
65 # Find all species
66 species_blocks = _find_blocks(f, 'species', '<cmd>run')
68 for spec in species_blocks:
69 name = spec.get('name')
70 spec_data = dict(
71 symbol=spec.find('symbol').text,
72 mass=float(spec.find('mass').text),
73 number=int(spec.find('atomic_number').text))
74 species[name] = spec_data
76 # Find all of the frames
77 frames = _find_blocks(f, 'iteration', None)
79 # If index is an int, return one frame
80 if isinstance(index, int):
81 return _parse_frame(frames[index], species)
82 else:
83 return [_parse_frame(frame, species) for frame in frames[index]]
86def _find_blocks(fp, tag, stopwords='[qbox]'):
87 """Find and parse a certain block of the file.
89 Reads a file sequentially and stops when it either encounters the
90 end of the file, or until the it encounters a line that contains a
91 user-defined string *after it has already found at least one
92 desired block*. Use the stopwords ``[qbox]`` to read until the next
93 command is issued.
95 Groups the text between the first line that contains <tag> and the
96 next line that contains </tag>, inclusively. The function then
97 parses the XML and returns the Element object.
99 Inputs:
100 fp - file-like object, file to be read from
101 tag - str, tag to search for (e.g., 'iteration').
102 `None` if you want to read until the end of the file
103 stopwords - str, halt parsing if a line containing this string
104 is encountered
106 Returns
107 -------
108 list of xml.ElementTree, parsed XML blocks found by this class
109 """
111 start_tag = f'<{tag}'
112 end_tag = f'</{tag}>'
114 blocks = [] # Stores all blocks
115 cur_block = [] # Block being filled
116 in_block = False # Whether we are currently parsing
117 for line in fp:
119 # Check if the block has started
120 if start_tag in line:
121 if in_block:
122 raise Exception('Parsing failed: Encountered nested block')
123 else:
124 in_block = True
126 # Add data to block
127 if in_block:
128 cur_block.append(line)
130 # Check for stopping conditions
131 if stopwords is not None:
132 if stopwords in line and len(blocks) > 0:
133 break
135 if end_tag in line:
136 if in_block:
137 blocks.append(cur_block)
138 cur_block = []
139 in_block = False
140 else:
141 raise Exception('Parsing failed: End tag found before start '
142 'tag')
144 # Join strings in a block into a single string
145 blocks = [''.join(b) for b in blocks]
147 # Ensure XML compatibility. There are two specific tags in QBall that are
148 # not valid XML, so we need to run a
149 blocks = [re_find_bad_xml.sub(r'<\1\2_expectation_\3', b) for b in blocks]
151 # Parse the blocks
152 return [ET.fromstring(b) for b in blocks]
155def _parse_frame(tree, species):
156 """Parse a certain frame from QBOX output
158 Inputs:
159 tree - ElementTree, <iteration> block from output file
160 species - dict, data about species. Key is name of atom type,
161 value is data about that type
162 Return:
163 Atoms object describing this iteration"""
165 # Load in data about the system
166 energy = float(tree.find("etotal").text)
168 # Load in data about the cell
169 unitcell = tree.find('atomset').find('unit_cell')
170 cell = []
171 for d in ['a', 'b', 'c']:
172 cell.append([float(x) for x in unitcell.get(d).split()])
174 stress_tree = tree.find('stress_tensor')
175 if stress_tree is None:
176 stresses = None
177 else:
178 stresses = [float(stress_tree.find(f'sigma_{x}').text)
179 for x in ['xx', 'yy', 'zz', 'yz', 'xz', 'xy']]
181 # Create the Atoms object
182 atoms = Atoms(pbc=True, cell=cell)
184 # Load in the atom information
185 forces = []
186 for atom in tree.find('atomset').findall('atom'):
187 # Load data about the atom type
188 spec = atom.get('species')
189 symbol = species[spec]['symbol']
190 mass = species[spec]['mass']
192 # Get data about position / velocity / force
193 pos = [float(x) for x in atom.find('position').text.split()]
194 force = [float(x) for x in atom.find('force').text.split()]
195 momentum = [float(x) * mass
196 for x in atom.find('velocity').text.split()]
198 # Create the objects
199 atom = Atom(symbol=symbol, mass=mass, position=pos, momentum=momentum)
200 atoms += atom
201 forces.append(force)
203 # Create the calculator object that holds energy/forces
204 calc = SinglePointCalculator(atoms,
205 energy=energy, forces=forces, stress=stresses)
206 atoms.calc = calc
208 return atoms