Coverage for /builds/ase/ase/ase/calculators/vasp/vasp_auxiliary.py: 29.74%
195 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3import os
4import re
6import numpy as np
9def get_vasp_version(string):
10 """Extract version number from header of stdout.
12 Example::
14 >>> get_vasp_version('potato vasp.6.1.2 bumblebee')
15 '6.1.2'
17 """
18 match = re.search(r'vasp\.(\S+)', string, re.M)
19 return match.group(1)
22class VaspChargeDensity:
23 """Class for representing VASP charge density.
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
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
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)
43 def is_spin_polarized(self):
44 if len(self.chgdiff) > 0:
45 return True
46 return False
48 def _read_chg(self, fobj, chg, volume):
49 """Read charge from file object
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.
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
67 def read(self, filename):
68 """Read CHG or CHGCAR file.
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.
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.
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.
82 """
83 import ase.io.vasp as aiv
84 with open(filename) as fd:
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_configuration(fd)
93 except (KeyError, RuntimeError, ValueError):
94 # Probably an empty line, or we tried to read the
95 # augmentation occupancies in CHGCAR
96 break
98 # Note: We continue reading from the same file, and
99 # this relies on read_vasp() to read no more lines
100 # than it currently does.
101 fd.readline()
103 ngr = fd.readline().split()
104 ng = (int(ngr[0]), int(ngr[1]), int(ngr[2]))
105 chg = np.empty(ng)
106 self._read_chg(fd, chg, atoms.get_volume())
107 self.chg.append(chg)
108 self.atoms.append(atoms)
109 # Check if the file has a spin-polarized charge density part,
110 # and if so, read it in.
111 fl = fd.tell()
112 # First check if the file has an augmentation charge part
113 # (CHGCAR file.)
114 line1 = fd.readline()
115 if line1 == '':
116 break
117 elif line1.find('augmentation') != -1:
118 augs = [line1]
119 while True:
120 line2 = fd.readline()
121 if line2.split() == ngr:
122 self.aug = ''.join(augs)
123 augs = []
124 chgdiff = np.empty(ng)
125 self._read_chg(fd, chgdiff, atoms.get_volume())
126 self.chgdiff.append(chgdiff)
127 elif line2 == '':
128 break
129 else:
130 augs.append(line2)
131 if len(self.aug) == 0:
132 self.aug = ''.join(augs)
133 augs = []
134 else:
135 self.augdiff = ''.join(augs)
136 augs = []
137 elif line1.split() == ngr:
138 chgdiff = np.empty(ng)
139 self._read_chg(fd, chgdiff, atoms.get_volume())
140 self.chgdiff.append(chgdiff)
141 else:
142 fd.seek(fl)
144 def _write_chg(self, fobj, chg, volume, format='chg'):
145 """Write charge density
147 Utility function similar to _read_chg but for writing.
149 """
150 # Make a 1D copy of chg, must take transpose to get ordering right
151 chgtmp = chg.T.ravel()
152 # Multiply by volume
153 chgtmp = chgtmp * volume
154 # Must be a tuple to pass to string conversion
155 chgtmp = tuple(chgtmp)
156 # CHG format - 10 columns
157 if format.lower() == 'chg':
158 # Write all but the last row
159 for ii in range((len(chgtmp) - 1) // 10):
160 fobj.write(' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\
161 %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\n' % chgtmp[ii * 10:(ii + 1) * 10])
162 # If the last row contains 10 values then write them without a
163 # newline
164 if len(chgtmp) % 10 == 0:
165 fobj.write(' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G'
166 ' %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G' %
167 chgtmp[len(chgtmp) - 10:len(chgtmp)])
168 # Otherwise write fewer columns without a newline
169 else:
170 for ii in range(len(chgtmp) % 10):
171 fobj.write((' %#11.5G') %
172 chgtmp[len(chgtmp) - len(chgtmp) % 10 + ii])
173 # Other formats - 5 columns
174 else:
175 # Write all but the last row
176 for ii in range((len(chgtmp) - 1) // 5):
177 fobj.write(' %17.10E %17.10E %17.10E %17.10E %17.10E\n' %
178 chgtmp[ii * 5:(ii + 1) * 5])
179 # If the last row contains 5 values then write them without a
180 # newline
181 if len(chgtmp) % 5 == 0:
182 fobj.write(' %17.10E %17.10E %17.10E %17.10E %17.10E' %
183 chgtmp[len(chgtmp) - 5:len(chgtmp)])
184 # Otherwise write fewer columns without a newline
185 else:
186 for ii in range(len(chgtmp) % 5):
187 fobj.write((' %17.10E') %
188 chgtmp[len(chgtmp) - len(chgtmp) % 5 + ii])
189 # Write a newline whatever format it is
190 fobj.write('\n')
192 def write(self, filename, format=None):
193 """Write VASP charge density in CHG format.
195 filename: str
196 Name of file to write to.
197 format: str
198 String specifying whether to write in CHGCAR or CHG
199 format.
201 """
202 import ase.io.vasp as aiv
203 if format is None:
204 if filename.lower().find('chgcar') != -1:
205 format = 'chgcar'
206 elif filename.lower().find('chg') != -1:
207 format = 'chg'
208 elif len(self.chg) == 1:
209 format = 'chgcar'
210 else:
211 format = 'chg'
212 with open(filename, 'w') as fd:
213 for ii, chg in enumerate(self.chg):
214 if format == 'chgcar' and ii != len(self.chg) - 1:
215 continue # Write only the last image for CHGCAR
216 aiv.write_vasp(fd,
217 self.atoms[ii],
218 direct=True)
219 fd.write('\n')
220 for dim in chg.shape:
221 fd.write(' %4i' % dim)
222 fd.write('\n')
223 vol = self.atoms[ii].get_volume()
224 self._write_chg(fd, chg, vol, format)
225 if format == 'chgcar':
226 fd.write(self.aug)
227 if self.is_spin_polarized():
228 if format == 'chg':
229 fd.write('\n')
230 for dim in chg.shape:
231 fd.write(' %4i' % dim)
232 fd.write('\n') # a new line after dim is required
233 self._write_chg(fd, self.chgdiff[ii], vol, format)
234 if format == 'chgcar':
235 # a new line is always provided self._write_chg
236 fd.write(self.augdiff)
237 if format == 'chg' and len(self.chg) > 1:
238 fd.write('\n')
241class VaspDos:
242 """Class for representing density-of-states produced by VASP
244 The energies are in property self.energy
246 Site-projected DOS is accesible via the self.site_dos method.
248 Total and integrated DOS is accessible as numpy.ndarray's in the
249 properties self.dos and self.integrated_dos. If the calculation is
250 spin polarized, the arrays will be of shape (2, NDOS), else (1,
251 NDOS).
253 The self.efermi property contains the currently set Fermi
254 level. Changing this value shifts the energies.
256 """
258 def __init__(self, doscar='DOSCAR', efermi=0.0):
259 """Initialize"""
260 self._efermi = 0.0
261 self.read_doscar(doscar)
262 self.efermi = efermi
264 # we have determine the resort to correctly map ase atom index to the
265 # POSCAR.
266 self.sort = []
267 self.resort = []
268 if os.path.isfile('ase-sort.dat'):
269 with open('ase-sort.dat') as file:
270 lines = file.readlines()
271 for line in lines:
272 data = line.split()
273 self.sort.append(int(data[0]))
274 self.resort.append(int(data[1]))
276 def _set_efermi(self, efermi):
277 """Set the Fermi level."""
278 ef = efermi - self._efermi
279 self._efermi = efermi
280 self._total_dos[0, :] = self._total_dos[0, :] - ef
281 try:
282 self._site_dos[:, 0, :] = self._site_dos[:, 0, :] - ef
283 except IndexError:
284 pass
286 def _get_efermi(self):
287 return self._efermi
289 efermi = property(_get_efermi, _set_efermi, None, "Fermi energy.")
291 def _get_energy(self):
292 """Return the array with the energies."""
293 return self._total_dos[0, :]
295 energy = property(_get_energy, None, None, "Array of energies")
297 def site_dos(self, atom, orbital):
298 """Return an NDOSx1 array with dos for the chosen atom and orbital.
300 atom: int
301 Atom index
302 orbital: int or str
303 Which orbital to plot
305 If the orbital is given as an integer:
306 If spin-unpolarized calculation, no phase factors:
307 s = 0, p = 1, d = 2
308 Spin-polarized, no phase factors:
309 s-up = 0, s-down = 1, p-up = 2, p-down = 3, d-up = 4, d-down = 5
310 If phase factors have been calculated, orbitals are
311 s, py, pz, px, dxy, dyz, dz2, dxz, dx2
312 double in the above fashion if spin polarized.
314 """
315 # Correct atom index for resorting if we need to. This happens when the
316 # ase-sort.dat file exists, and self.resort is not empty.
317 if self.resort:
318 atom = self.resort[atom]
320 # Integer indexing for orbitals starts from 1 in the _site_dos array
321 # since the 0th column contains the energies
322 if isinstance(orbital, int):
323 return self._site_dos[atom, orbital + 1, :]
324 n = self._site_dos.shape[1]
326 from .vasp_data import PDOS_orbital_names_and_DOSCAR_column
327 norb = PDOS_orbital_names_and_DOSCAR_column[n]
329 return self._site_dos[atom, norb[orbital.lower()], :]
331 def _get_dos(self):
332 if self._total_dos.shape[0] == 3:
333 return self._total_dos[1, :]
334 elif self._total_dos.shape[0] == 5:
335 return self._total_dos[1:3, :]
337 dos = property(_get_dos, None, None, 'Average DOS in cell')
339 def _get_integrated_dos(self):
340 if self._total_dos.shape[0] == 3:
341 return self._total_dos[2, :]
342 elif self._total_dos.shape[0] == 5:
343 return self._total_dos[3:5, :]
345 integrated_dos = property(_get_integrated_dos, None, None,
346 'Integrated average DOS in cell')
348 def read_doscar(self, fname="DOSCAR"):
349 """Read a VASP DOSCAR file"""
350 with open(fname) as fd:
351 natoms = int(fd.readline().split()[0])
352 [fd.readline() for _ in range(4)]
353 # First we have a block with total and total integrated DOS
354 ndos = int(fd.readline().split()[2])
355 dos = []
356 for _ in range(ndos):
357 dos.append(np.array([float(x) for x in fd.readline().split()]))
358 self._total_dos = np.array(dos).T
359 # Next we have one block per atom, if INCAR contains the stuff
360 # necessary for generating site-projected DOS
361 dos = []
362 for _ in range(natoms):
363 line = fd.readline()
364 if line == '':
365 # No site-projected DOS
366 break
367 ndos = int(line.split()[2])
368 line = fd.readline().split()
369 cdos = np.empty((ndos, len(line)))
370 cdos[0] = np.array(line)
371 for nd in range(1, ndos):
372 line = fd.readline().split()
373 cdos[nd] = np.array([float(x) for x in line])
374 dos.append(cdos.T)
375 self._site_dos = np.array(dos)