Coverage for /builds/debichem-team/python-ase/ase/io/octopus/output.py: 97.67%

1import re 

2 

3import numpy as np 

4 

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

6 

7 

8class OctopusIOError(IOError): 

9 pass # Cannot find output files 

10 

11 

12def read_eigenvalues_file(fd): 

13 unit = None 

14 

15 for line in fd: 

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

17 if m is not None: 

18 unit = m.group(1) 

19 break 

20 line = next(fd) 

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

22 

23 # fermilevel = None 

24 kpts = [] 

25 eigs = [] 

26 occs = [] 

27 

28 for line in fd: 

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

30 if m: 

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

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

33 eigs.append({}) 

34 occs.append({}) 

35 else: 

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

37 if m is None: 

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

39 assert m is not None 

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

41 # it from the static/info instead. 

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

43 else: 

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

45 

46 if not eigs: 

47 # Only initialized if kpoint header was written 

48 eigs.append({}) 

49 occs.append({}) 

50 

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

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

53 

54 nkpts = len(kpts) 

55 nspins = len(eigs[0]) 

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

57 

58 kptsarr = np.array(kpts, float) 

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

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

61 

62 arrs = [eigsarr, occsarr] 

63 

64 for arr in arrs: 

65 arr.fill(np.nan) 

66 

67 for k in range(nkpts): 

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

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

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

71 

72 for arr in arrs: 

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

74 

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

76 return kptsarr, eigsarr, occsarr 

77 

78 

79def read_static_info_stress(fd): 

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

81 

82 headers = next(fd) 

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

84 for i in range(3): 

85 line = next(fd) 

86 tokens = line.split() 

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

88 stress_cv[i] = vec 

89 return stress_cv 

90 

91 

92def read_static_info_kpoints(fd): 

93 for line in fd: 

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

95 break 

96 

97 tokens = next(fd).split() 

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

99 bar = next(fd) 

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

101 

102 kpts = [] 

103 weights = [] 

104 

105 for line in fd: 

106 # Format: index kx ky kz weight 

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

108 if m is None: 

109 break 

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

111 weight = m.group(4) 

112 kpts.append(kxyz) 

113 weights.append(weight) 

114 

115 ibz_kpoints = np.array(kpts, float) 

116 kpoint_weights = np.array(weights, float) 

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

118 

119 

120def read_static_info_eigenvalues(fd, energy_unit): 

121 

122 values_sknx = {} 

123 

124 nbands = 0 

125 fermilevel = None 

126 for line in fd: 

127 line = line.strip() 

128 if line.startswith('#'): 

129 continue 

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

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

132 if m is not None: 

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

134 break 

135 

136 tokens = line.split() 

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

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

139 occupation = float(tokens[3]) 

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

141 

142 nspins = len(values_sknx) 

143 if nspins == 1: 

144 val = [values_sknx['--']] 

145 else: 

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

147 val = np.array(val, float) 

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

149 assert remainder == 0 

150 

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

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

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

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

155 assert eps_skn.flags.contiguous 

156 d = dict(nspins=nspins, 

157 nkpts=nkpts, 

158 nbands=nbands, 

159 eigenvalues=eps_skn, 

160 occupations=occ_skn) 

161 if fermilevel is not None: 

162 d.update(fermi_level=fermilevel) 

163 return d 

164 

165 

166def read_static_info_energy(fd, energy_unit): 

167 def get(name): 

168 for line in fd: 

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

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

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

172 

173 

174def read_static_info(fd): 

175 results = {} 

176 

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

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

179 

180 for line in fd: 

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

182 results.update(read_static_info_kpoints(fd)) 

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

184 unit = get_energy_unit(line) 

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

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

187 unit = get_energy_unit(line) 

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

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

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

191 stress = read_static_info_stress(fd) 

192 stress *= Hartree / Bohr**3 

193 results.update(stress=stress) 

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

195 line = next(fd) 

196 values = line.split() 

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

198 line = next(fd) 

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

200 line = next(fd) 

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

202 # Reading Local Magnetic Moments 

203 magmoms = [] 

204 for line in fd: 

205 if line == '\n': 

206 break # there is no more thing to search for 

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

208 values = line.split() 

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

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

211 elif line.startswith('Dipole'): 

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

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

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

215 elif line.startswith('Forces'): 

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

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

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

219 forces = [] 

220 line = next(fd) 

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

222 for line in fd: 

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

224 break 

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

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

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

228 

229 return results