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 def __init__(self, filename): 

29 # Instance variables 

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

31 self.chg = [] # Charge density 

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

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

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

35 

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

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

38 # image. 

39 if filename is not None: 

40 self.read(filename) 

41 

42 def is_spin_polarized(self): 

43 if len(self.chgdiff) > 0: 

44 return True 

45 return False 

46 

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

48 """Read charge from file object 

49 

50 Utility method for reading the actual charge density (or 

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

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

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

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

55 

56 """ 

57 # VASP writes charge density as 

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

59 # Fortran nested implied do loops; innermost index fastest 

60 # First, just read it in 

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

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

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

64 chg /= volume 

65 

66 def read(self, filename): 

67 """Read CHG or CHGCAR file. 

68 

69 If CHG contains charge density from multiple steps all the 

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

71 writes out the charge density every 10 steps. 

72 

73 chgdiff is the difference between the spin up charge density 

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

75 spin-polarized calculation. 

76 

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

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

79 be written again to a CHGCAR format file. 

80 

81 """ 

82 import ase.io.vasp as aiv 

83 fd = open(filename) 

84 self.atoms = [] 

85 self.chg = [] 

86 self.chgdiff = [] 

87 self.aug = '' 

88 self.augdiff = '' 

89 while True: 

90 try: 

91 atoms = aiv.read_vasp(fd) 

92 except (IOError, ValueError, IndexError): 

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

94 # augmentation occupancies in CHGCAR 

95 break 

96 fd.readline() 

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

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

99 chg = np.empty(ng) 

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

101 self.chg.append(chg) 

102 self.atoms.append(atoms) 

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

104 # if so, read it in. 

105 fl = fd.tell() 

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

107 # file.) 

108 line1 = fd.readline() 

109 if line1 == '': 

110 break 

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

112 augs = [line1] 

113 while True: 

114 line2 = fd.readline() 

115 if line2.split() == ngr: 

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

117 augs = [] 

118 chgdiff = np.empty(ng) 

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

120 self.chgdiff.append(chgdiff) 

121 elif line2 == '': 

122 break 

123 else: 

124 augs.append(line2) 

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

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

127 augs = [] 

128 else: 

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

130 augs = [] 

131 elif line1.split() == ngr: 

132 chgdiff = np.empty(ng) 

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

134 self.chgdiff.append(chgdiff) 

135 else: 

136 fd.seek(fl) 

137 fd.close() 

138 

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

140 """Write charge density 

141 

142 Utility function similar to _read_chg but for writing. 

143 

144 """ 

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

146 chgtmp = chg.T.ravel() 

147 # Multiply by volume 

148 chgtmp = chgtmp * volume 

149 # Must be a tuple to pass to string conversion 

150 chgtmp = tuple(chgtmp) 

151 # CHG format - 10 columns 

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

153 # Write all but the last row 

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

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

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

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

158 # newline 

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

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

161 ' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G' % 

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

163 # Otherwise write fewer columns without a newline 

164 else: 

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

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

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

168 # Other formats - 5 columns 

169 else: 

170 # Write all but the last row 

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

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

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

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

175 # newline 

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

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

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

179 # Otherwise write fewer columns without a newline 

180 else: 

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

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

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

184 # Write a newline whatever format it is 

185 fobj.write('\n') 

186 

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

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

189 

190 filename: str 

191 Name of file to write to. 

192 format: str 

193 String specifying whether to write in CHGCAR or CHG 

194 format. 

195 

196 """ 

197 import ase.io.vasp as aiv 

198 if format is None: 

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

200 format = 'chgcar' 

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

202 format = 'chg' 

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

204 format = 'chgcar' 

205 else: 

206 format = 'chg' 

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

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

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

210 continue # Write only the last image for CHGCAR 

211 aiv.write_vasp(fd, 

212 self.atoms[ii], 

213 direct=True, 

214 long_format=False) 

215 fd.write('\n') 

216 for dim in chg.shape: 

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

218 fd.write('\n') 

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

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

221 if format == 'chgcar': 

222 fd.write(self.aug) 

223 if self.is_spin_polarized(): 

224 if format == 'chg': 

225 fd.write('\n') 

226 for dim in chg.shape: 

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

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

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

230 if format == 'chgcar': 

231 # a new line is always provided self._write_chg 

232 fd.write(self.augdiff) 

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

234 fd.write('\n') 

235 

236 

237class VaspDos: 

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

239 

240 The energies are in property self.energy 

241 

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

243 

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

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

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

247 NDOS). 

248 

249 The self.efermi property contains the currently set Fermi 

250 level. Changing this value shifts the energies. 

251 

252 """ 

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

254 """Initialize""" 

255 self._efermi = 0.0 

256 self.read_doscar(doscar) 

257 self.efermi = efermi 

258 

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

260 # POSCAR. 

261 self.sort = [] 

262 self.resort = [] 

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

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

265 lines = file.readlines() 

266 file.close() 

267 for line in lines: 

268 data = line.split() 

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

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

271 

272 def _set_efermi(self, efermi): 

273 """Set the Fermi level.""" 

274 ef = efermi - self._efermi 

275 self._efermi = efermi 

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

277 try: 

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

279 except IndexError: 

280 pass 

281 

282 def _get_efermi(self): 

283 return self._efermi 

284 

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

286 

287 def _get_energy(self): 

288 """Return the array with the energies.""" 

289 return self._total_dos[0, :] 

290 

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

292 

293 def site_dos(self, atom, orbital): 

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

295 

296 atom: int 

297 Atom index 

298 orbital: int or str 

299 Which orbital to plot 

300 

301 If the orbital is given as an integer: 

302 If spin-unpolarized calculation, no phase factors: 

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

304 Spin-polarized, no phase factors: 

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

306 If phase factors have been calculated, orbitals are 

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

308 double in the above fashion if spin polarized. 

309 

310 """ 

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

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

313 if self.resort: 

314 atom = self.resort[atom] 

315 

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

317 # since the 0th column contains the energies 

318 if isinstance(orbital, int): 

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

320 n = self._site_dos.shape[1] 

321 

322 from .vasp_data import PDOS_orbital_names_and_DOSCAR_column 

323 norb = PDOS_orbital_names_and_DOSCAR_column[n] 

324 

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

326 

327 def _get_dos(self): 

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

329 return self._total_dos[1, :] 

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

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

332 

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

334 

335 def _get_integrated_dos(self): 

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

337 return self._total_dos[2, :] 

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

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

340 

341 integrated_dos = property(_get_integrated_dos, None, None, 

342 'Integrated average DOS in cell') 

343 

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

345 """Read a VASP DOSCAR file""" 

346 fd = open(fname) 

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

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

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

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

351 dos = [] 

352 for nd in range(ndos): 

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

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

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

356 # necessary for generating site-projected DOS 

357 dos = [] 

358 for na in range(natoms): 

359 line = fd.readline() 

360 if line == '': 

361 # No site-projected DOS 

362 break 

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

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

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

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

367 for nd in range(1, ndos): 

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

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

370 dos.append(cdos.T) 

371 self._site_dos = np.array(dos) 

372 fd.close() 

373 

374 

375class xdat2traj: 

376 def __init__(self, 

377 trajectory=None, 

378 atoms=None, 

379 poscar=None, 

380 xdatcar=None, 

381 sort=None, 

382 calc=None): 

383 """ 

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

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

386 """ 

387 if not poscar: 

388 self.poscar = 'POSCAR' 

389 else: 

390 self.poscar = poscar 

391 

392 if not atoms: 

393 # This reads the atoms sorted the way VASP wants 

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

395 resort_reqd = True 

396 else: 

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

398 self.atoms = atoms 

399 resort_reqd = False 

400 

401 if not calc: 

402 self.calc = Vasp() 

403 else: 

404 self.calc = calc 

405 if not sort: 

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

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

408 else: 

409 self.calc.sort = sort 

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

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

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

413 

414 if not xdatcar: 

415 self.xdatcar = 'XDATCAR' 

416 else: 

417 self.xdatcar = xdatcar 

418 

419 if not trajectory: 

420 self.trajectory = 'out.traj' 

421 else: 

422 self.trajectory = trajectory 

423 

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

425 

426 if resort_reqd: 

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

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

429 # Forces are read with the atoms sorted using resort 

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

431 

432 def convert(self): 

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

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

435 del (lines[0:8]) 

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

437 del (lines[0:6]) 

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

439 del (lines[0:5]) 

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

441 del (lines[0:8]) 

442 step = 0 

443 iatom = 0 

444 scaled_pos = [] 

445 for line in lines: 

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

447 if step == 0: 

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

449 scaled_pos = np.array(scaled_pos) 

450 # Now resort the positions to match self.atoms 

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

452 

453 calc = SinglePointCalculator(self.atoms, 

454 energy=self.energies[step], 

455 forces=self.forces[step]) 

456 self.atoms.calc = calc 

457 self.out.write(self.atoms) 

458 scaled_pos = [] 

459 iatom = 0 

460 step += 1 

461 else: 

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

463 iatom += 1 

464 scaled_pos.append( 

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

466 

467 # Write also the last image 

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

469 if step == 0: 

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

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

472 self.atoms.set_scaled_positions(scaled_pos) 

473 calc = SinglePointCalculator(self.atoms, 

474 energy=self.energies[step], 

475 forces=self.forces[step]) 

476 self.atoms.calc = calc 

477 self.out.write(self.atoms) 

478 

479 self.out.close()