Coverage for /builds/ase/ase/ase/io/octopus/output.py: 97.67%

1# fmt: off 

2 

3import re 

4 

5import numpy as np 

6 

7from ase.units import Angstrom, Bohr, Debye, Hartree, eV 

8 

9 

10class OctopusIOError(IOError): 

11 pass # Cannot find output files 

12 

13 

14def read_eigenvalues_file(fd): 

15 unit = None 

16 

17 for line in fd: 

18 m = re.match(r'Eigenvalues\s*\[(.+?)\]', line) 

19 if m is not None: 

20 unit = m.group(1) 

21 break 

22 line = next(fd) 

23 assert line.strip().startswith('#st'), line 

24 

25 # fermilevel = None 

26 kpts = [] 

27 eigs = [] 

28 occs = [] 

29 

30 for line in fd: 

31 m = re.match(r'#k.*?\(\s*(.+?),\s*(.+?),\s*(.+?)\)', line) 

32 if m: 

33 k = m.group(1, 2, 3) 

34 kpts.append(np.array(k, float)) 

35 eigs.append({}) 

36 occs.append({}) 

37 else: 

38 m = re.match(r'\s*\d+\s*(\S+)\s*(\S+)\s*(\S+)', line) 

39 if m is None: 

40 m = re.match(r'Fermi energy\s*=\s*(\S+)\s*', line) 

41 assert m is not None 

42 # We can also return the fermilevel but so far we just read 

43 # it from the static/info instead. 

44 # fermilevel = float(m.group(1)) 

45 else: 

46 spin, eig, occ = m.group(1, 2, 3) 

47 

48 if not eigs: 

49 # Only initialized if kpoint header was written 

50 eigs.append({}) 

51 occs.append({}) 

52 

53 eigs[-1].setdefault(spin, []).append(float(eig)) 

54 occs[-1].setdefault(spin, []).append(float(occ)) 

55 

56 nkpts = len(kpts) 

57 nspins = len(eigs[0]) 

58 nbands = len(eigs[0][spin]) 

59 

60 kptsarr = np.array(kpts, float) 

61 eigsarr = np.empty((nkpts, nspins, nbands)) 

62 occsarr = np.empty((nkpts, nspins, nbands)) 

63 

64 arrs = [eigsarr, occsarr] 

65 

66 for arr in arrs: 

67 arr.fill(np.nan) 

68 

69 for k in range(nkpts): 

70 for arr, lst in [(eigsarr, eigs), (occsarr, occs)]: 

71 arr[k, :, :] = [lst[k][sp] for sp 

72 in (['--'] if nspins == 1 else ['up', 'dn'])] 

73 

74 for arr in arrs: 

75 assert not np.isnan(arr).any() 

76 

77 eigsarr *= {'H': Hartree, 'eV': eV}[unit] 

78 return kptsarr, eigsarr, occsarr 

79 

80 

81def read_static_info_stress(fd): 

82 stress_cv = np.empty((3, 3)) 

83 

84 headers = next(fd) 

85 assert headers.strip().startswith('T_{ij}') 

86 for i in range(3): 

87 line = next(fd) 

88 tokens = line.split() 

89 vec = np.array(tokens[1:4]).astype(float) 

90 stress_cv[i] = vec 

91 return stress_cv 

92 

93 

94def read_static_info_kpoints(fd): 

95 for line in fd: 

96 if line.startswith('List of k-points'): 

97 break 

98 

99 tokens = next(fd).split() 

100 assert tokens == ['ik', 'k_x', 'k_y', 'k_z', 'Weight'] 

101 bar = next(fd) 

102 assert bar.startswith('---') 

103 

104 kpts = [] 

105 weights = [] 

106 

107 for line in fd: 

108 # Format: index kx ky kz weight 

109 m = re.match(r'\s*\d+\s*(\S+)\s*(\S+)\s*(\S+)\s*(\S+)', line) 

110 if m is None: 

111 break 

112 kxyz = m.group(1, 2, 3) 

113 weight = m.group(4) 

114 kpts.append(kxyz) 

115 weights.append(weight) 

116 

117 ibz_kpoints = np.array(kpts, float) 

118 kpoint_weights = np.array(weights, float) 

119 return dict(ibz_kpoints=ibz_kpoints, kpoint_weights=kpoint_weights) 

120 

121 

122def read_static_info_eigenvalues(fd, energy_unit): 

123 

124 values_sknx = {} 

125 

126 nbands = 0 

127 fermilevel = None 

128 for line in fd: 

129 line = line.strip() 

130 if line.startswith('#'): 

131 continue 

132 if not line[:1].isdigit(): 

133 m = re.match(r'Fermi energy\s*=\s*(\S+)', line) 

134 if m is not None: 

135 fermilevel = float(m.group(1)) * energy_unit 

136 break 

137 

138 tokens = line.split() 

139 nbands = max(nbands, int(tokens[0])) 

140 energy = float(tokens[2]) * energy_unit 

141 occupation = float(tokens[3]) 

142 values_sknx.setdefault(tokens[1], []).append((energy, occupation)) 

143 

144 nspins = len(values_sknx) 

145 if nspins == 1: 

146 val = [values_sknx['--']] 

147 else: 

148 val = [values_sknx['up'], values_sknx['dn']] 

149 val = np.array(val, float) 

150 nkpts, remainder = divmod(len(val[0]), nbands) 

151 assert remainder == 0 

152 

153 eps_skn = val[:, :, 0].reshape(nspins, nkpts, nbands) 

154 occ_skn = val[:, :, 1].reshape(nspins, nkpts, nbands) 

155 eps_skn = eps_skn.transpose(1, 0, 2).copy() 

156 occ_skn = occ_skn.transpose(1, 0, 2).copy() 

157 assert eps_skn.flags.contiguous 

158 d = dict(nspins=nspins, 

159 nkpts=nkpts, 

160 nbands=nbands, 

161 eigenvalues=eps_skn, 

162 occupations=occ_skn) 

163 if fermilevel is not None: 

164 d.update(fermi_level=fermilevel) 

165 return d 

166 

167 

168def read_static_info_energy(fd, energy_unit): 

169 def get(name): 

170 for line in fd: 

171 if line.strip().startswith(name): 

172 return float(line.split('=')[-1].strip()) * energy_unit 

173 return dict(energy=get('Total'), free_energy=get('Free')) 

174 

175 

176def read_static_info(fd): 

177 results = {} 

178 

179 def get_energy_unit(line): # Convert "title [unit]": ---> unit 

180 return {'[eV]': eV, '[H]': Hartree}[line.split()[1].rstrip(':')] 

181 

182 for line in fd: 

183 if line.strip('*').strip().startswith('Brillouin zone'): 

184 results.update(read_static_info_kpoints(fd)) 

185 elif line.startswith('Eigenvalues ['): 

186 unit = get_energy_unit(line) 

187 results.update(read_static_info_eigenvalues(fd, unit)) 

188 elif line.startswith('Energy ['): 

189 unit = get_energy_unit(line) 

190 results.update(read_static_info_energy(fd, unit)) 

191 elif line.startswith('Total stress tensor ['): 

192 if '[H/b^3]' in line: 

193 stress = read_static_info_stress(fd) 

194 stress *= Hartree / Bohr**3 

195 results.update(stress=stress) 

196 elif line.startswith('Total Magnetic Moment'): 

197 line = next(fd) 

198 values = line.split() 

199 results['magmom'] = float(values[-1]) 

200 line = next(fd) 

201 assert line.startswith('Local Magnetic Moments') 

202 line = next(fd) 

203 assert line.split() == ['Ion', 'mz'] 

204 # Reading Local Magnetic Moments 

205 magmoms = [] 

206 for line in fd: 

207 if line == '\n': 

208 break # there is no more thing to search for 

209 line = line.replace('\n', ' ') 

210 values = line.split() 

211 magmoms.append(float(values[-1])) 

212 results['magmoms'] = np.array(magmoms) 

213 elif line.startswith('Dipole'): 

214 assert line.split()[-1] == '[Debye]' 

215 dipole = [float(next(fd).split()[-1]) for i in range(3)] 

216 results['dipole'] = np.array(dipole) * Debye 

217 elif line.startswith('Forces'): 

218 forceunitspec = line.split()[-1] 

219 forceunit = {'[eV/A]': eV / Angstrom, 

220 '[H/b]': Hartree / Bohr}[forceunitspec] 

221 forces = [] 

222 line = next(fd) 

223 assert line.strip().startswith('Ion') 

224 for line in fd: 

225 if line.strip().startswith('---'): 

226 break 

227 tokens = line.split()[-3:] 

228 forces.append([float(f) for f in tokens]) 

229 results['forces'] = np.array(forces) * forceunit 

230 

231 return results