Coverage for /builds/ericyuan00000/ase/ase/io/elk.py: 86.33%

256 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-06-18 01:20 +0000

1# fmt: off 

2 

3import collections 

4from pathlib import Path 

5 

6import numpy as np 

7 

8from ase import Atoms 

9from ase.units import Bohr, Hartree 

10from ase.utils import reader, writer 

11 

12elk_parameters = {'swidth': Hartree} 

13 

14 

15@reader 

16def read_elk(fd): 

17 """Import ELK atoms definition. 

18 

19 Reads unitcell, atom positions, magmoms from elk.in/GEOMETRY.OUT file. 

20 """ 

21 

22 lines = fd.readlines() 

23 

24 scale = np.ones(4) # unit cell scale 

25 positions = [] 

26 cell = [] 

27 symbols = [] 

28 magmoms = [] 

29 

30 # find cell scale 

31 for n, line in enumerate(lines): 

32 if line.split() == []: 

33 continue 

34 if line.strip() == 'scale': 

35 scale[0] = float(lines[n + 1]) 

36 elif line.startswith('scale'): 

37 scale[int(line.strip()[-1])] = float(lines[n + 1]) 

38 for n, line in enumerate(lines): 

39 if line.split() == []: 

40 continue 

41 if line.startswith('avec'): 

42 cell = np.array( 

43 [[float(v) * scale[1] for v in lines[n + 1].split()], 

44 [float(v) * scale[2] for v in lines[n + 2].split()], 

45 [float(v) * scale[3] for v in lines[n + 3].split()]]) 

46 if line.startswith('atoms'): 

47 lines1 = lines[n + 1:] # start subsearch 

48 spfname = [] 

49 natoms = [] 

50 atpos = [] 

51 bfcmt = [] 

52 for n1, line1 in enumerate(lines1): 

53 if line1.split() == []: 

54 continue 

55 if 'spfname' in line1: 

56 spfnamenow = lines1[n1].split()[0] 

57 spfname.append(spfnamenow) 

58 natomsnow = int(lines1[n1 + 1].split()[0]) 

59 natoms.append(natomsnow) 

60 atposnow = [] 

61 bfcmtnow = [] 

62 for line in lines1[n1 + 2:n1 + 2 + natomsnow]: 

63 atposnow.append([float(v) for v in line.split()[0:3]]) 

64 if len(line.split()) == 6: # bfcmt present 

65 bfcmtnow.append( 

66 [float(v) for v in line.split()[3:]]) 

67 atpos.append(atposnow) 

68 bfcmt.append(bfcmtnow) 

69 # symbols, positions, magmoms based on ELK spfname, atpos, and bfcmt 

70 symbols = '' 

71 positions = [] 

72 magmoms = [] 

73 for n, s in enumerate(spfname): 

74 symbols += str(s[1:].split('.')[0]) * natoms[n] 

75 positions += atpos[n] # assumes fractional coordinates 

76 if len(bfcmt[n]) > 0: 

77 # how to handle cases of magmoms being one or three dim array? 

78 magmoms += [m[-1] for m in bfcmt[n]] 

79 atoms = Atoms(symbols, scaled_positions=positions, cell=[1, 1, 1]) 

80 if len(magmoms) > 0: 

81 atoms.set_initial_magnetic_moments(magmoms) 

82 # final cell scale 

83 cell = cell * scale[0] * Bohr 

84 

85 atoms.set_cell(cell, scale_atoms=True) 

86 atoms.pbc = True 

87 return atoms 

88 

89 

90@writer 

91def write_elk_in(fd, atoms, parameters=None): 

92 if parameters is None: 

93 parameters = {} 

94 

95 parameters = dict(parameters) 

96 species_path = parameters.pop('species_dir', None) 

97 

98 if parameters.get('spinpol') is None: 

99 if atoms.get_initial_magnetic_moments().any(): 

100 parameters['spinpol'] = True 

101 

102 if 'xctype' in parameters: 

103 if 'xc' in parameters: 

104 raise RuntimeError("You can't use both 'xctype' and 'xc'!") 

105 

106 if parameters.get('autokpt'): 

107 if 'kpts' in parameters: 

108 raise RuntimeError("You can't use both 'autokpt' and 'kpts'!") 

109 if 'ngridk' in parameters: 

110 raise RuntimeError( 

111 "You can't use both 'autokpt' and 'ngridk'!") 

112 if 'ngridk' in parameters: 

113 if 'kpts' in parameters: 

114 raise RuntimeError("You can't use both 'ngridk' and 'kpts'!") 

115 

116 if parameters.get('autoswidth'): 

117 if 'smearing' in parameters: 

118 raise RuntimeError( 

119 "You can't use both 'autoswidth' and 'smearing'!") 

120 if 'swidth' in parameters: 

121 raise RuntimeError( 

122 "You can't use both 'autoswidth' and 'swidth'!") 

123 

124 inp = {} 

125 inp.update(parameters) 

126 

127 if 'xc' in parameters: 

128 xctype = {'LDA': 3, # PW92 

129 'PBE': 20, 

130 'REVPBE': 21, 

131 'PBESOL': 22, 

132 'WC06': 26, 

133 'AM05': 30, 

134 'mBJLDA': (100, 208, 12)}[parameters['xc']] 

135 inp['xctype'] = xctype 

136 del inp['xc'] 

137 

138 if 'kpts' in parameters: 

139 # XXX should generalize kpts handling. 

140 from ase.calculators.calculator import kpts2mp 

141 mp = kpts2mp(atoms, parameters['kpts']) 

142 inp['ngridk'] = tuple(mp) 

143 vkloff = [] # is this below correct? 

144 for nk in mp: 

145 if nk % 2 == 0: # shift kpoint away from gamma point 

146 vkloff.append(0.5) 

147 else: 

148 vkloff.append(0) 

149 inp['vkloff'] = vkloff 

150 del inp['kpts'] 

151 

152 if 'smearing' in parameters: 

153 name = parameters.smearing[0].lower() 

154 if name == 'methfessel-paxton': 

155 stype = parameters.smearing[2] 

156 else: 

157 stype = {'gaussian': 0, 

158 'fermi-dirac': 3, 

159 }[name] 

160 inp['stype'] = stype 

161 inp['swidth'] = parameters.smearing[1] 

162 del inp['smearing'] 

163 

164 # convert keys to ELK units 

165 for key, value in inp.items(): 

166 if key in elk_parameters: 

167 inp[key] /= elk_parameters[key] 

168 

169 # write all keys 

170 for key, value in inp.items(): 

171 fd.write(f'{key}\n') 

172 if isinstance(value, bool): 

173 fd.write(f'.{("false", "true")[value]}.\n\n') 

174 elif isinstance(value, (int, float)): 

175 fd.write(f'{value}\n\n') 

176 else: 

177 fd.write('%s\n\n' % ' '.join([str(x) for x in value])) 

178 

179 # cell 

180 fd.write('avec\n') 

181 for vec in atoms.cell: 

182 fd.write('%.14f %.14f %.14f\n' % tuple(vec / Bohr)) 

183 fd.write('\n') 

184 

185 # atoms 

186 species = {} 

187 symbols = [] 

188 for a, (symbol, m) in enumerate( 

189 zip(atoms.get_chemical_symbols(), 

190 atoms.get_initial_magnetic_moments())): 

191 if symbol in species: 

192 species[symbol].append((a, m)) 

193 else: 

194 species[symbol] = [(a, m)] 

195 symbols.append(symbol) 

196 fd.write('atoms\n%d\n' % len(species)) 

197 # scaled = atoms.get_scaled_positions(wrap=False) 

198 scaled = np.linalg.solve(atoms.cell.T, atoms.positions.T).T 

199 for symbol in symbols: 

200 fd.write(f"'{symbol}.in' : spfname\n") 

201 fd.write('%d\n' % len(species[symbol])) 

202 for a, m in species[symbol]: 

203 fd.write('%.14f %.14f %.14f 0.0 0.0 %.14f\n' % 

204 (tuple(scaled[a]) + (m,))) 

205 

206 # if sppath is present in elk.in it overwrites species blocks! 

207 

208 # Elk seems to concatenate path and filename in such a way 

209 # that we must put a / at the end: 

210 if species_path is not None: 

211 fd.write(f"sppath\n'{species_path}/'\n\n") 

212 

213 

214class ElkReader: 

215 def __init__(self, path): 

216 self.path = Path(path) 

217 

218 def _read_everything(self): 

219 yield from self._read_energy() 

220 

221 with (self.path / 'INFO.OUT').open() as fd: 

222 yield from parse_elk_info(fd) 

223 

224 with (self.path / 'EIGVAL.OUT').open() as fd: 

225 yield from parse_elk_eigval(fd) 

226 

227 with (self.path / 'KPOINTS.OUT').open() as fd: 

228 yield from parse_elk_kpoints(fd) 

229 

230 def read_everything(self): 

231 dct = dict(self._read_everything()) 

232 

233 # The eigenvalue/occupation tables do not say whether there are 

234 # two spins, so we have to reshape them from 1 x K x SB to S x K x B: 

235 spinpol = dct.pop('spinpol') 

236 if spinpol: 

237 for name in 'eigenvalues', 'occupations': 

238 array = dct[name] 

239 _, nkpts, nbands_double = array.shape 

240 assert _ == 1 

241 assert nbands_double % 2 == 0 

242 nbands = nbands_double // 2 

243 newarray = np.empty((2, nkpts, nbands)) 

244 newarray[0, :, :] = array[0, :, :nbands] 

245 newarray[1, :, :] = array[0, :, nbands:] 

246 if name == 'eigenvalues': 

247 # Verify that eigenvalues are still sorted: 

248 diffs = np.diff(newarray, axis=2) 

249 assert all(diffs.flat[:] > 0) 

250 dct[name] = newarray 

251 return dct 

252 

253 def _read_energy(self): 

254 txt = (self.path / 'TOTENERGY.OUT').read_text() 

255 tokens = txt.split() 

256 energy = float(tokens[-1]) * Hartree 

257 yield 'free_energy', energy 

258 yield 'energy', energy 

259 

260 

261def parse_elk_kpoints(fd): 

262 header = next(fd) 

263 lhs, rhs = header.split(':', 1) 

264 assert 'k-point, vkl, wkpt' in rhs, header 

265 nkpts = int(lhs.strip()) 

266 

267 kpts = np.empty((nkpts, 3)) 

268 weights = np.empty(nkpts) 

269 

270 for ikpt in range(nkpts): 

271 line = next(fd) 

272 tokens = line.split() 

273 kpts[ikpt] = np.array(tokens[1:4]).astype(float) 

274 weights[ikpt] = float(tokens[4]) 

275 yield 'ibz_kpoints', kpts 

276 yield 'kpoint_weights', weights 

277 

278 

279def parse_elk_info(fd): 

280 dct = collections.defaultdict(list) 

281 fd = iter(fd) 

282 

283 spinpol = None 

284 converged = False 

285 actually_did_not_converge = False 

286 # Legacy code kept track of both these things, which is strange. 

287 # Why could a file both claim to converge *and* not converge? 

288 

289 # We loop over all lines and extract also data that occurs 

290 # multiple times (e.g. in multiple SCF steps) 

291 for line in fd: 

292 # "name of quantity : 1 2 3" 

293 tokens = line.split(':', 1) 

294 if len(tokens) == 2: 

295 lhs, rhs = tokens 

296 dct[lhs.strip()].append(rhs.strip()) 

297 

298 elif line.startswith('Convergence targets achieved'): 

299 converged = True 

300 elif 'reached self-consistent loops maximum' in line.lower(): 

301 actually_did_not_converge = True 

302 

303 if 'Spin treatment' in line: 

304 # (Somewhat brittle doing multi-line stuff here.) 

305 line = next(fd) 

306 spinpol = line.strip() == 'spin-polarised' 

307 

308 yield 'converged', converged and not actually_did_not_converge 

309 if spinpol is None: 

310 raise RuntimeError('Could not determine spin treatment') 

311 yield 'spinpol', spinpol 

312 

313 if 'Fermi' in dct: 

314 yield 'fermi_level', float(dct['Fermi'][-1]) * Hartree 

315 

316 if 'total force' in dct: 

317 forces = [] 

318 for line in dct['total force']: 

319 forces.append(line.split()) 

320 yield 'forces', np.array(forces, float) * (Hartree / Bohr) 

321 

322 

323def parse_elk_eigval(fd): 

324 

325 def match_int(line, word): 

326 number, colon, word1 = line.split() 

327 assert word1 == word 

328 assert colon == ':' 

329 return int(number) 

330 

331 def skip_spaces(line=''): 

332 while not line.strip(): 

333 line = next(fd) 

334 return line 

335 

336 line = skip_spaces() 

337 nkpts = match_int(line, 'nkpt') # 10 : nkpts 

338 line = next(fd) 

339 nbands = match_int(line, 'nstsv') # 15 : nstsv 

340 

341 eigenvalues = np.empty((nkpts, nbands)) 

342 occupations = np.empty((nkpts, nbands)) 

343 kpts = np.empty((nkpts, 3)) 

344 

345 for ikpt in range(nkpts): 

346 line = skip_spaces() 

347 tokens = line.split() 

348 assert tokens[-1] == 'vkl', tokens 

349 assert ikpt + 1 == int(tokens[0]) 

350 kpts[ikpt] = np.array(tokens[1:4]).astype(float) 

351 

352 line = next(fd) # "(state, eigenvalue and occupancy below)" 

353 assert line.strip().startswith('(state,'), line 

354 for iband in range(nbands): 

355 line = next(fd) 

356 tokens = line.split() # (band number, eigenval, occ) 

357 assert iband + 1 == int(tokens[0]) 

358 eigenvalues[ikpt, iband] = float(tokens[1]) 

359 occupations[ikpt, iband] = float(tokens[2]) 

360 

361 yield 'ibz_kpoints', kpts 

362 yield 'eigenvalues', eigenvalues[None] * Hartree 

363 yield 'occupations', occupations[None]