1import re 

2import os 

3import numpy as np 

4import ase 

5from .vasp import Vasp 

6from ase.calculators.singlepoint import SinglePointCalculator 

7 

8 

9def get_vasp_version(string): 

10 """Extract version number from header of stdout. 

11 

12 Example:: 

13 

14 >>> get_vasp_version('potato vasp.6.1.2 bumblebee') 

15 '6.1.2' 

16 

17 """ 

18 match = re.search(r'vasp\.(\S+)', string, re.M) 

19 return match.group(1) 

20 

21 

22class VaspChargeDensity: 

23 """Class for representing VASP charge density. 

24 

25 Filename is normally CHG.""" 

26 # Can the filename be CHGCAR? There's a povray tutorial 

27 # in doc/tutorials where it's CHGCAR as of January 2021. --askhl 

28 

29 def __init__(self, filename): 

30 # Instance variables 

31 self.atoms = [] # List of Atoms objects 

32 self.chg = [] # Charge density 

33 self.chgdiff = [] # Charge density difference, if spin polarized 

34 self.aug = '' # Augmentation charges, not parsed just a big string 

35 self.augdiff = '' # Augmentation charge differece, is spin polarized 

36 

37 # Note that the augmentation charge is not a list, since they 

38 # are needed only for CHGCAR files which store only a single 

39 # image. 

40 if filename is not None: 

41 self.read(filename) 

42 

43 def is_spin_polarized(self): 

44 if len(self.chgdiff) > 0: 

45 return True 

46 return False 

47 

48 def _read_chg(self, fobj, chg, volume): 

49 """Read charge from file object 

50 

51 Utility method for reading the actual charge density (or 

52 charge density difference) from a file object. On input, the 

53 file object must be at the beginning of the charge block, on 

54 output the file position will be left at the end of the 

55 block. The chg array must be of the correct dimensions. 

56 

57 """ 

58 # VASP writes charge density as 

59 # WRITE(IU,FORM) (((C(NX,NY,NZ),NX=1,NGXC),NY=1,NGYZ),NZ=1,NGZC) 

60 # Fortran nested implied do loops; innermost index fastest 

61 # First, just read it in 

62 for zz in range(chg.shape[2]): 

63 for yy in range(chg.shape[1]): 

64 chg[:, yy, zz] = np.fromfile(fobj, count=chg.shape[0], sep=' ') 

65 chg /= volume 

66 

67 def read(self, filename): 

68 """Read CHG or CHGCAR file. 

69 

70 If CHG contains charge density from multiple steps all the 

71 steps are read and stored in the object. By default VASP 

72 writes out the charge density every 10 steps. 

73 

74 chgdiff is the difference between the spin up charge density 

75 and the spin down charge density and is thus only read for a 

76 spin-polarized calculation. 

77 

78 aug is the PAW augmentation charges found in CHGCAR. These are 

79 not parsed, they are just stored as a string so that they can 

80 be written again to a CHGCAR format file. 

81 

82 """ 

83 import ase.io.vasp as aiv 

84 fd = open(filename) 

85 self.atoms = [] 

86 self.chg = [] 

87 self.chgdiff = [] 

88 self.aug = '' 

89 self.augdiff = '' 

90 while True: 

91 try: 

92 atoms = aiv.read_vasp(fd) 

93 except (IOError, ValueError, IndexError): 

94 # Probably an empty line, or we tried to read the 

95 # augmentation occupancies in CHGCAR 

96 break 

97 fd.readline() 

98 ngr = fd.readline().split() 

99 ng = (int(ngr[0]), int(ngr[1]), int(ngr[2])) 

100 chg = np.empty(ng) 

101 self._read_chg(fd, chg, atoms.get_volume()) 

102 self.chg.append(chg) 

103 self.atoms.append(atoms) 

104 # Check if the file has a spin-polarized charge density part, and 

105 # if so, read it in. 

106 fl = fd.tell() 

107 # First check if the file has an augmentation charge part (CHGCAR 

108 # file.) 

109 line1 = fd.readline() 

110 if line1 == '': 

111 break 

112 elif line1.find('augmentation') != -1: 

113 augs = [line1] 

114 while True: 

115 line2 = fd.readline() 

116 if line2.split() == ngr: 

117 self.aug = ''.join(augs) 

118 augs = [] 

119 chgdiff = np.empty(ng) 

120 self._read_chg(fd, chgdiff, atoms.get_volume()) 

121 self.chgdiff.append(chgdiff) 

122 elif line2 == '': 

123 break 

124 else: 

125 augs.append(line2) 

126 if len(self.aug) == 0: 

127 self.aug = ''.join(augs) 

128 augs = [] 

129 else: 

130 self.augdiff = ''.join(augs) 

131 augs = [] 

132 elif line1.split() == ngr: 

133 chgdiff = np.empty(ng) 

134 self._read_chg(fd, chgdiff, atoms.get_volume()) 

135 self.chgdiff.append(chgdiff) 

136 else: 

137 fd.seek(fl) 

138 fd.close() 

139 

140 def _write_chg(self, fobj, chg, volume, format='chg'): 

141 """Write charge density 

142 

143 Utility function similar to _read_chg but for writing. 

144 

145 """ 

146 # Make a 1D copy of chg, must take transpose to get ordering right 

147 chgtmp = chg.T.ravel() 

148 # Multiply by volume 

149 chgtmp = chgtmp * volume 

150 # Must be a tuple to pass to string conversion 

151 chgtmp = tuple(chgtmp) 

152 # CHG format - 10 columns 

153 if format.lower() == 'chg': 

154 # Write all but the last row 

155 for ii in range((len(chgtmp) - 1) // 10): 

156 fobj.write(' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\ 

157 %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\n' % chgtmp[ii * 10:(ii + 1) * 10]) 

158 # If the last row contains 10 values then write them without a 

159 # newline 

160 if len(chgtmp) % 10 == 0: 

161 fobj.write(' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G' 

162 ' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G' % 

163 chgtmp[len(chgtmp) - 10:len(chgtmp)]) 

164 # Otherwise write fewer columns without a newline 

165 else: 

166 for ii in range(len(chgtmp) % 10): 

167 fobj.write((' %#11.5G') % 

168 chgtmp[len(chgtmp) - len(chgtmp) % 10 + ii]) 

169 # Other formats - 5 columns 

170 else: 

171 # Write all but the last row 

172 for ii in range((len(chgtmp) - 1) // 5): 

173 fobj.write(' %17.10E %17.10E %17.10E %17.10E %17.10E\n' % 

174 chgtmp[ii * 5:(ii + 1) * 5]) 

175 # If the last row contains 5 values then write them without a 

176 # newline 

177 if len(chgtmp) % 5 == 0: 

178 fobj.write(' %17.10E %17.10E %17.10E %17.10E %17.10E' % 

179 chgtmp[len(chgtmp) - 5:len(chgtmp)]) 

180 # Otherwise write fewer columns without a newline 

181 else: 

182 for ii in range(len(chgtmp) % 5): 

183 fobj.write((' %17.10E') % 

184 chgtmp[len(chgtmp) - len(chgtmp) % 5 + ii]) 

185 # Write a newline whatever format it is 

186 fobj.write('\n') 

187 

188 def write(self, filename, format=None): 

189 """Write VASP charge density in CHG format. 

190 

191 filename: str 

192 Name of file to write to. 

193 format: str 

194 String specifying whether to write in CHGCAR or CHG 

195 format. 

196 

197 """ 

198 import ase.io.vasp as aiv 

199 if format is None: 

200 if filename.lower().find('chgcar') != -1: 

201 format = 'chgcar' 

202 elif filename.lower().find('chg') != -1: 

203 format = 'chg' 

204 elif len(self.chg) == 1: 

205 format = 'chgcar' 

206 else: 

207 format = 'chg' 

208 with open(filename, 'w') as fd: 

209 for ii, chg in enumerate(self.chg): 

210 if format == 'chgcar' and ii != len(self.chg) - 1: 

211 continue # Write only the last image for CHGCAR 

212 aiv.write_vasp(fd, 

213 self.atoms[ii], 

214 direct=True, 

215 long_format=False) 

216 fd.write('\n') 

217 for dim in chg.shape: 

218 fd.write(' %4i' % dim) 

219 fd.write('\n') 

220 vol = self.atoms[ii].get_volume() 

221 self._write_chg(fd, chg, vol, format) 

222 if format == 'chgcar': 

223 fd.write(self.aug) 

224 if self.is_spin_polarized(): 

225 if format == 'chg': 

226 fd.write('\n') 

227 for dim in chg.shape: 

228 fd.write(' %4i' % dim) 

229 fd.write('\n') # a new line after dim is required 

230 self._write_chg(fd, self.chgdiff[ii], vol, format) 

231 if format == 'chgcar': 

232 # a new line is always provided self._write_chg 

233 fd.write(self.augdiff) 

234 if format == 'chg' and len(self.chg) > 1: 

235 fd.write('\n') 

236 

237 

238class VaspDos: 

239 """Class for representing density-of-states produced by VASP 

240 

241 The energies are in property self.energy 

242 

243 Site-projected DOS is accesible via the self.site_dos method. 

244 

245 Total and integrated DOS is accessible as numpy.ndarray's in the 

246 properties self.dos and self.integrated_dos. If the calculation is 

247 spin polarized, the arrays will be of shape (2, NDOS), else (1, 

248 NDOS). 

249 

250 The self.efermi property contains the currently set Fermi 

251 level. Changing this value shifts the energies. 

252 

253 """ 

254 

255 def __init__(self, doscar='DOSCAR', efermi=0.0): 

256 """Initialize""" 

257 self._efermi = 0.0 

258 self.read_doscar(doscar) 

259 self.efermi = efermi 

260 

261 # we have determine the resort to correctly map ase atom index to the 

262 # POSCAR. 

263 self.sort = [] 

264 self.resort = [] 

265 if os.path.isfile('ase-sort.dat'): 

266 file = open('ase-sort.dat', 'r') 

267 lines = file.readlines() 

268 file.close() 

269 for line in lines: 

270 data = line.split() 

271 self.sort.append(int(data[0])) 

272 self.resort.append(int(data[1])) 

273 

274 def _set_efermi(self, efermi): 

275 """Set the Fermi level.""" 

276 ef = efermi - self._efermi 

277 self._efermi = efermi 

278 self._total_dos[0, :] = self._total_dos[0, :] - ef 

279 try: 

280 self._site_dos[:, 0, :] = self._site_dos[:, 0, :] - ef 

281 except IndexError: 

282 pass 

283 

284 def _get_efermi(self): 

285 return self._efermi 

286 

287 efermi = property(_get_efermi, _set_efermi, None, "Fermi energy.") 

288 

289 def _get_energy(self): 

290 """Return the array with the energies.""" 

291 return self._total_dos[0, :] 

292 

293 energy = property(_get_energy, None, None, "Array of energies") 

294 

295 def site_dos(self, atom, orbital): 

296 """Return an NDOSx1 array with dos for the chosen atom and orbital. 

297 

298 atom: int 

299 Atom index 

300 orbital: int or str 

301 Which orbital to plot 

302 

303 If the orbital is given as an integer: 

304 If spin-unpolarized calculation, no phase factors: 

305 s = 0, p = 1, d = 2 

306 Spin-polarized, no phase factors: 

307 s-up = 0, s-down = 1, p-up = 2, p-down = 3, d-up = 4, d-down = 5 

308 If phase factors have been calculated, orbitals are 

309 s, py, pz, px, dxy, dyz, dz2, dxz, dx2 

310 double in the above fashion if spin polarized. 

311 

312 """ 

313 # Correct atom index for resorting if we need to. This happens when the 

314 # ase-sort.dat file exists, and self.resort is not empty. 

315 if self.resort: 

316 atom = self.resort[atom] 

317 

318 # Integer indexing for orbitals starts from 1 in the _site_dos array 

319 # since the 0th column contains the energies 

320 if isinstance(orbital, int): 

321 return self._site_dos[atom, orbital + 1, :] 

322 n = self._site_dos.shape[1] 

323 

324 from .vasp_data import PDOS_orbital_names_and_DOSCAR_column 

325 norb = PDOS_orbital_names_and_DOSCAR_column[n] 

326 

327 return self._site_dos[atom, norb[orbital.lower()], :] 

328 

329 def _get_dos(self): 

330 if self._total_dos.shape[0] == 3: 

331 return self._total_dos[1, :] 

332 elif self._total_dos.shape[0] == 5: 

333 return self._total_dos[1:3, :] 

334 

335 dos = property(_get_dos, None, None, 'Average DOS in cell') 

336 

337 def _get_integrated_dos(self): 

338 if self._total_dos.shape[0] == 3: 

339 return self._total_dos[2, :] 

340 elif self._total_dos.shape[0] == 5: 

341 return self._total_dos[3:5, :] 

342 

343 integrated_dos = property(_get_integrated_dos, None, None, 

344 'Integrated average DOS in cell') 

345 

346 def read_doscar(self, fname="DOSCAR"): 

347 """Read a VASP DOSCAR file""" 

348 fd = open(fname) 

349 natoms = int(fd.readline().split()[0]) 

350 [fd.readline() for nn in range(4)] # Skip next 4 lines. 

351 # First we have a block with total and total integrated DOS 

352 ndos = int(fd.readline().split()[2]) 

353 dos = [] 

354 for nd in range(ndos): 

355 dos.append(np.array([float(x) for x in fd.readline().split()])) 

356 self._total_dos = np.array(dos).T 

357 # Next we have one block per atom, if INCAR contains the stuff 

358 # necessary for generating site-projected DOS 

359 dos = [] 

360 for na in range(natoms): 

361 line = fd.readline() 

362 if line == '': 

363 # No site-projected DOS 

364 break 

365 ndos = int(line.split()[2]) 

366 line = fd.readline().split() 

367 cdos = np.empty((ndos, len(line))) 

368 cdos[0] = np.array(line) 

369 for nd in range(1, ndos): 

370 line = fd.readline().split() 

371 cdos[nd] = np.array([float(x) for x in line]) 

372 dos.append(cdos.T) 

373 self._site_dos = np.array(dos) 

374 fd.close() 

375 

376 

377class xdat2traj: 

378 def __init__(self, 

379 trajectory=None, 

380 atoms=None, 

381 poscar=None, 

382 xdatcar=None, 

383 sort=None, 

384 calc=None): 

385 """ 

386 trajectory is the name of the file to write the trajectory to 

387 poscar is the name of the poscar file to read. Default: POSCAR 

388 """ 

389 if not poscar: 

390 self.poscar = 'POSCAR' 

391 else: 

392 self.poscar = poscar 

393 

394 if not atoms: 

395 # This reads the atoms sorted the way VASP wants 

396 self.atoms = ase.io.read(self.poscar, format='vasp') 

397 resort_reqd = True 

398 else: 

399 # Assume if we pass atoms that it is sorted the way we want 

400 self.atoms = atoms 

401 resort_reqd = False 

402 

403 if not calc: 

404 self.calc = Vasp() 

405 else: 

406 self.calc = calc 

407 if not sort: 

408 if not hasattr(self.calc, 'sort'): 

409 self.calc.sort = list(range(len(self.atoms))) 

410 else: 

411 self.calc.sort = sort 

412 self.calc.resort = list(range(len(self.calc.sort))) 

413 for n in range(len(self.calc.resort)): 

414 self.calc.resort[self.calc.sort[n]] = n 

415 

416 if not xdatcar: 

417 self.xdatcar = 'XDATCAR' 

418 else: 

419 self.xdatcar = xdatcar 

420 

421 if not trajectory: 

422 self.trajectory = 'out.traj' 

423 else: 

424 self.trajectory = trajectory 

425 

426 self.out = ase.io.trajectory.Trajectory(self.trajectory, mode='w') 

427 

428 if resort_reqd: 

429 self.atoms = self.atoms[self.calc.resort] 

430 self.energies = self.calc.read_energy(all=True)[1] 

431 # Forces are read with the atoms sorted using resort 

432 self.forces = self.calc.read_forces(self.atoms, all=True) 

433 

434 def convert(self): 

435 lines = open(self.xdatcar).readlines() 

436 if len(lines[7].split()) == 0: 

437 del (lines[0:8]) 

438 elif len(lines[5].split()) == 0: 

439 del (lines[0:6]) 

440 elif len(lines[4].split()) == 0: 

441 del (lines[0:5]) 

442 elif lines[7].split()[0] == 'Direct': 

443 del (lines[0:8]) 

444 step = 0 

445 iatom = 0 

446 scaled_pos = [] 

447 for line in lines: 

448 if iatom == len(self.atoms): 

449 if step == 0: 

450 self.out.write_header(self.atoms[self.calc.resort]) 

451 scaled_pos = np.array(scaled_pos) 

452 # Now resort the positions to match self.atoms 

453 self.atoms.set_scaled_positions(scaled_pos[self.calc.resort]) 

454 

455 calc = SinglePointCalculator(self.atoms, 

456 energy=self.energies[step], 

457 forces=self.forces[step]) 

458 self.atoms.calc = calc 

459 self.out.write(self.atoms) 

460 scaled_pos = [] 

461 iatom = 0 

462 step += 1 

463 else: 

464 if not line.split()[0] == 'Direct': 

465 iatom += 1 

466 scaled_pos.append( 

467 [float(line.split()[n]) for n in range(3)]) 

468 

469 # Write also the last image 

470 # I'm sure there is also more clever fix... 

471 if step == 0: 

472 self.out.write_header(self.atoms[self.calc.resort]) 

473 scaled_pos = np.array(scaled_pos)[self.calc.resort] 

474 self.atoms.set_scaled_positions(scaled_pos) 

475 calc = SinglePointCalculator(self.atoms, 

476 energy=self.energies[step], 

477 forces=self.forces[step]) 

478 self.atoms.calc = calc 

479 self.out.write(self.atoms) 

480 

481 self.out.close()