Coverage for /builds/ase/ase/ase/io/qbox.py: 96.67%

1# fmt: off 

2 

3"""This module contains functions to read from QBox output files""" 

4 

5import re 

6import xml.etree.ElementTree as ET 

7 

8from ase import Atom, Atoms 

9from ase.calculators.singlepoint import SinglePointCalculator 

10from ase.utils import reader 

11 

12# Compile regexs for fixing XML 

13re_find_bad_xml = re.compile(r'<(/?)([A-z]+) expectation ([a-z]+)') 

14 

15 

16@reader 

17def read_qbox(f, index=-1): 

18 """Read data from QBox output file 

19 

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 

23 

24 Returns: 

25 list of Atoms or atoms, requested frame(s) 

26 """ 

27 

28 # Check whether this is a QB@all output 

29 version = None 

30 for line in f: 

31 if '<release>' in line: 

32 version = ET.fromstring(line) 

33 break 

34 if version is None: 

35 raise Exception('Parse Error: Version not found') 

36 is_qball = 'qb@LL' in version.text or 'qball' in version.text 

37 

38 # Load in atomic species 

39 species = {} 

40 if is_qball: 

41 # Read all of the lines between release and the first call to `run` 

42 species_data = [] 

43 for line in f: 

44 if '<run' in line: 

45 break 

46 species_data.append(line) 

47 species_data = '\n'.join(species_data) 

48 

49 # Read out the species information with regular expressions 

50 symbols = re.findall('symbol_ = ([A-Z][a-z]?)', species_data) 

51 masses = re.findall('mass_ = ([0-9.]+)', species_data) 

52 names = re.findall('name_ = ([a-z]+)', species_data) 

53 numbers = re.findall('atomic_number_ = ([0-9]+)', species_data) 

54 

55 # Compile them into a dictionary 

56 for name, symbol, mass, number in zip(names, symbols, masses, numbers): 

57 spec_data = dict( 

58 symbol=symbol, 

59 mass=float(mass), 

60 number=float(number) 

61 ) 

62 species[name] = spec_data 

63 else: 

64 # Find all species 

65 species_blocks = _find_blocks(f, 'species', '<cmd>run') 

66 

67 for spec in species_blocks: 

68 name = spec.get('name') 

69 spec_data = dict( 

70 symbol=spec.find('symbol').text, 

71 mass=float(spec.find('mass').text), 

72 number=int(spec.find('atomic_number').text)) 

73 species[name] = spec_data 

74 

75 # Find all of the frames 

76 frames = _find_blocks(f, 'iteration', None) 

77 

78 # If index is an int, return one frame 

79 if isinstance(index, int): 

80 return _parse_frame(frames[index], species) 

81 else: 

82 return [_parse_frame(frame, species) for frame in frames[index]] 

83 

84 

85def _find_blocks(fp, tag, stopwords='[qbox]'): 

86 """Find and parse a certain block of the file. 

87 

88 Reads a file sequentially and stops when it either encounters the 

89 end of the file, or until the it encounters a line that contains a 

90 user-defined string *after it has already found at least one 

91 desired block*. Use the stopwords ``[qbox]`` to read until the next 

92 command is issued. 

93 

94 Groups the text between the first line that contains <tag> and the 

95 next line that contains </tag>, inclusively. The function then 

96 parses the XML and returns the Element object. 

97 

98 Inputs: 

99 fp - file-like object, file to be read from 

100 tag - str, tag to search for (e.g., 'iteration'). 

101 `None` if you want to read until the end of the file 

102 stopwords - str, halt parsing if a line containing this string 

103 is encountered 

104 

105 Returns: 

106 list of xml.ElementTree, parsed XML blocks found by this class 

107 """ 

108 

109 start_tag = f'<{tag}' 

110 end_tag = f'</{tag}>' 

111 

112 blocks = [] # Stores all blocks 

113 cur_block = [] # Block being filled 

114 in_block = False # Whether we are currently parsing 

115 for line in fp: 

116 

117 # Check if the block has started 

118 if start_tag in line: 

119 if in_block: 

120 raise Exception('Parsing failed: Encountered nested block') 

121 else: 

122 in_block = True 

123 

124 # Add data to block 

125 if in_block: 

126 cur_block.append(line) 

127 

128 # Check for stopping conditions 

129 if stopwords is not None: 

130 if stopwords in line and len(blocks) > 0: 

131 break 

132 

133 if end_tag in line: 

134 if in_block: 

135 blocks.append(cur_block) 

136 cur_block = [] 

137 in_block = False 

138 else: 

139 raise Exception('Parsing failed: End tag found before start ' 

140 'tag') 

141 

142 # Join strings in a block into a single string 

143 blocks = [''.join(b) for b in blocks] 

144 

145 # Ensure XML compatibility. There are two specific tags in QBall that are 

146 # not valid XML, so we need to run a 

147 blocks = [re_find_bad_xml.sub(r'<\1\2_expectation_\3', b) for b in blocks] 

148 

149 # Parse the blocks 

150 return [ET.fromstring(b) for b in blocks] 

151 

152 

153def _parse_frame(tree, species): 

154 """Parse a certain frame from QBOX output 

155 

156 Inputs: 

157 tree - ElementTree, <iteration> block from output file 

158 species - dict, data about species. Key is name of atom type, 

159 value is data about that type 

160 Return: 

161 Atoms object describing this iteration""" 

162 

163 # Load in data about the system 

164 energy = float(tree.find("etotal").text) 

165 

166 # Load in data about the cell 

167 unitcell = tree.find('atomset').find('unit_cell') 

168 cell = [] 

169 for d in ['a', 'b', 'c']: 

170 cell.append([float(x) for x in unitcell.get(d).split()]) 

171 

172 stress_tree = tree.find('stress_tensor') 

173 if stress_tree is None: 

174 stresses = None 

175 else: 

176 stresses = [float(stress_tree.find(f'sigma_{x}').text) 

177 for x in ['xx', 'yy', 'zz', 'yz', 'xz', 'xy']] 

178 

179 # Create the Atoms object 

180 atoms = Atoms(pbc=True, cell=cell) 

181 

182 # Load in the atom information 

183 forces = [] 

184 for atom in tree.find('atomset').findall('atom'): 

185 # Load data about the atom type 

186 spec = atom.get('species') 

187 symbol = species[spec]['symbol'] 

188 mass = species[spec]['mass'] 

189 

190 # Get data about position / velocity / force 

191 pos = [float(x) for x in atom.find('position').text.split()] 

192 force = [float(x) for x in atom.find('force').text.split()] 

193 momentum = [float(x) * mass 

194 for x in atom.find('velocity').text.split()] 

195 

196 # Create the objects 

197 atom = Atom(symbol=symbol, mass=mass, position=pos, momentum=momentum) 

198 atoms += atom 

199 forces.append(force) 

200 

201 # Create the calculator object that holds energy/forces 

202 calc = SinglePointCalculator(atoms, 

203 energy=energy, forces=forces, stress=stresses) 

204 atoms.calc = calc 

205 

206 return atoms