Source code for storq.vasp.charge_density

import numpy as np


[docs]class VaspChargeDensity(object): """Class for representing VASP charge density""" def __init__(self, filename="CHG"): # Instance variables self.atoms = [] # List of Atoms objects self.chg = [] # Charge density self.chgdiff = [] # Charge density difference, if spin polarized self.aug = "" # Augmentation charges, not parsed just a big string self.augdiff = "" # Augmentation charge differece, is spin polarized # Note that the augmentation charge is not a list, since they # are needed only for CHGCAR files which store only a single image. if filename is not None: self.read(filename)
[docs] def is_spin_polarized(self): if len(self.chgdiff) > 0: return True
return False def _read_chg(self, fobj, chg, volume): """Read charge from file object Utility method for reading the actual charge density (or charge density difference) from a file object. On input, the file object must be at the beginning of the charge block, on output the file position will be left at the end of the block. The chg array must be of the correct dimensions. """ # VASP writes charge density as # WRITE(IU,FORM) (((C(NX,NY,NZ),NX=1,NGXC),NY=1,NGYZ),NZ=1,NGZC) # Fortran nested implied do loops; innermost index fastest # First, just read it in for zz in range(chg.shape[2]): for yy in range(chg.shape[1]): chg[:, yy, zz] = np.fromfile(fobj, count=chg.shape[0], sep=" ") chg /= volume
[docs] def read(self, filename="CHG"): """Read CHG or CHGCAR file. If CHG contains charge density from multiple steps all the steps are read and stored in the object. By default VASP writes out the charge density every 10 steps. chgdiff is the difference between the spin up charge density and the spin down charge density and is thus only read for a spin-polarized calculation. aug is the PAW augmentation charges found in CHGCAR. These are not parsed, they are just stored as a string so that they can be written again to a CHGCAR format file. """ import ase.io.vasp as aiv f = open(filename) self.atoms = [] self.chg = [] self.chgdiff = [] self.aug = "" self.augdiff = "" while True: try: atoms = aiv.read_vasp(f) except (IOError, ValueError, IndexError): # Probably an empty line, or we tried to read the # augmentation occupancies in CHGCAR break f.readline() ngr = f.readline().split() ng = (int(ngr[0]), int(ngr[1]), int(ngr[2])) chg = np.empty(ng) self._read_chg(f, chg, atoms.get_volume()) self.chg.append(chg) self.atoms.append(atoms) # Check if the file has a spin-polarized charge density part, and # if so, read it in. fl = f.tell() # First check if the file has an augmentation charge part (CHGCAR # file.) line1 = f.readline() if line1 == "": break elif line1.find("augmentation") != -1: augs = [line1] while True: line2 = f.readline() if line2.split() == ngr: self.aug = "".join(augs) augs = [] chgdiff = np.empty(ng) self._read_chg(f, chgdiff, atoms.get_volume()) self.chgdiff.append(chgdiff) elif line2 == "": break else: augs.append(line2) if len(self.aug) == 0: self.aug = "".join(augs) augs = [] else: self.augdiff = "".join(augs) augs = [] elif line1.split() == ngr: chgdiff = np.empty(ng) self._read_chg(f, chgdiff, atoms.get_volume()) self.chgdiff.append(chgdiff) else: f.seek(fl)
f.close() def _write_chg(self, fobj, chg, volume, format="chg"): """Write charge density Utility function similar to _read_chg but for writing. """ # Make a 1D copy of chg, must take transpose to get ordering right chgtmp = chg.T.ravel() # Multiply by volume chgtmp = chgtmp * volume # Must be a tuple to pass to string conversion chgtmp = tuple(chgtmp) # CHG format - 10 columns if format.lower() == "chg": # Write all but the last row for ii in range((len(chgtmp) - 1) // 10): fobj.write( " %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\ %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G\n" % chgtmp[ii * 10 : (ii + 1) * 10] ) # If the last row contains 10 values then write them without a # newline if len(chgtmp) % 10 == 0: fobj.write( " %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G" " %#11.5G %#11.5G %#11.5G %#11.5G %#11.5G" % chgtmp[len(chgtmp) - 10 : len(chgtmp)] ) # Otherwise write fewer columns without a newline else: for ii in range(len(chgtmp) % 10): fobj.write(" %#11.5G" % chgtmp[len(chgtmp) - len(chgtmp) % 10 + ii]) # Other formats - 5 columns else: # Write all but the last row for ii in range((len(chgtmp) - 1) // 5): fobj.write( " %17.10E %17.10E %17.10E %17.10E %17.10E\n" % chgtmp[ii * 5 : (ii + 1) * 5] ) # If the last row contains 5 values then write them without a # newline if len(chgtmp) % 5 == 0: fobj.write( " %17.10E %17.10E %17.10E %17.10E %17.10E" % chgtmp[len(chgtmp) - 5 : len(chgtmp)] ) # Otherwise write fewer columns without a newline else: for ii in range(len(chgtmp) % 5): fobj.write(" %17.10E" % chgtmp[len(chgtmp) - len(chgtmp) % 5 + ii]) # Write a newline whatever format it is fobj.write("\n") # Clean up del chgtmp
[docs] def write(self, filename="CHG", format=None): """Write VASP charge density in CHG format. filename: str Name of file to write to. format: str String specifying whether to write in CHGCAR or CHG format. """ import ase.io.vasp as aiv if format is None: if filename.lower().find("chgcar") != -1: format = "chgcar" elif filename.lower().find("chg") != -1: format = "chg" elif len(self.chg) == 1: format = "chgcar" else: format = "chg" f = open(filename, "w") for ii, chg in enumerate(self.chg): if format == "chgcar" and ii != len(self.chg) - 1: continue # Write only the last image for CHGCAR aiv.write_vasp(f, self.atoms[ii], direct=True, long_format=False) f.write("\n") for dim in chg.shape: f.write(" %4i" % dim) f.write("\n") vol = self.atoms[ii].get_volume() self._write_chg(f, chg, vol, format) if format == "chgcar": f.write(self.aug) if self.is_spin_polarized(): if format == "chg": f.write("\n") for dim in chg.shape: f.write(" %4i" % dim) self._write_chg(f, self.chgdiff[ii], vol, format) if format == "chgcar": f.write("\n") f.write(self.augdiff) if format == "chg" and len(self.chg) > 1: f.write("\n")
f.close()