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()