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

2 

3from ase import Atom, Atoms 

4from ase.calculators.singlepoint import SinglePointCalculator 

5from ase.utils import reader 

6 

7import re 

8import xml.etree.ElementTree as ET 

9 

10 

11# Compile regexs for fixing XML 

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

13 

14 

15@reader 

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

17 """Read data from QBox output file 

18 

19 Inputs: 

20 f - str or fileobj, path to file or file object to read from 

21 index - int or slice, which frames to return 

22 

23 Returns: 

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

25 """ 

26 

27 # Check whether this is a QB@all output 

28 version = None 

29 for line in f: 

30 if '<release>' in line: 

31 version = ET.fromstring(line) 

32 break 

33 if version is None: 

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

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

36 

37 # Load in atomic species 

38 species = dict() 

39 if is_qball: 

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

41 species_data = [] 

42 for line in f: 

43 if '<run' in line: 

44 break 

45 species_data.append(line) 

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

47 

48 # Read out the species information with regular expressions 

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

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

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

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

53 

54 # Compile them into a dictionary 

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

56 spec_data = dict( 

57 symbol=symbol, 

58 mass=float(mass), 

59 number=float(number) 

60 ) 

61 species[name] = spec_data 

62 else: 

63 # Find all species 

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

65 

66 for spec in species_blocks: 

67 name = spec.get('name') 

68 spec_data = dict( 

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

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

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

72 species[name] = spec_data 

73 

74 # Find all of the frames 

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

76 

77 # If index is an int, return one frame 

78 if isinstance(index, int): 

79 return _parse_frame(frames[index], species) 

80 else: 

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

82 

83 

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

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

86 

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

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

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

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

91 command is issued. 

92 

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

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

95 parses the XML and returns the Element object. 

96 

97 Inputs: 

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

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

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

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

102 is encountered 

103 

104 Returns: 

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

106 """ 

107 

108 start_tag = '<%s' % tag 

109 end_tag = '</%s>' % tag 

110 

111 blocks = [] # Stores all blocks 

112 cur_block = [] # Block being filled 

113 in_block = False # Whether we are currently parsing 

114 for line in fp: 

115 

116 # Check if the block has started 

117 if start_tag in line: 

118 if in_block: 

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

120 else: 

121 in_block = True 

122 

123 # Add data to block 

124 if in_block: 

125 cur_block.append(line) 

126 

127 # Check for stopping conditions 

128 if stopwords is not None: 

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

130 break 

131 

132 if end_tag in line: 

133 if in_block: 

134 blocks.append(cur_block) 

135 cur_block = [] 

136 in_block = False 

137 else: 

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

139 'tag') 

140 

141 # Join strings in a block into a single string 

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

143 

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

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

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

147 

148 # Parse the blocks 

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

150 

151 

152def _parse_frame(tree, species): 

153 """Parse a certain frame from QBOX output 

154 

155 Inputs: 

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

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

158 value is data about that type 

159 Return: 

160 Atoms object describing this iteration""" 

161 

162 # Load in data about the system 

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

164 

165 # Load in data about the cell 

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

167 cell = [] 

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

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

170 

171 stress_tree = tree.find('stress_tensor') 

172 if stress_tree is None: 

173 stresses = None 

174 else: 

175 stresses = [float(stress_tree.find('sigma_%s' % x).text) 

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

177 

178 # Create the Atoms object 

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

180 

181 # Load in the atom information 

182 forces = [] 

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

184 # Load data about the atom type 

185 spec = atom.get('species') 

186 symbol = species[spec]['symbol'] 

187 mass = species[spec]['mass'] 

188 

189 # Get data about position / velocity / force 

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

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

192 momentum = [float(x) * mass 

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

194 

195 # Create the objects 

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

197 atoms += atom 

198 forces.append(force) 

199 

200 # Create the calculator object that holds energy/forces 

201 calc = SinglePointCalculator(atoms, 

202 energy=energy, forces=forces, stress=stresses) 

203 atoms.calc = calc 

204 

205 return atoms