Skip to content

Commit

Permalink
cp2k_pdos: implement multi-file (UKS) PDOS parsing on same grid
Browse files Browse the repository at this point in the history
  • Loading branch information
@dev-zero
dev-zero committed Nov 20, 2018
1 parent 6c36246 commit bc930c7
Showing 1 changed file with 55 additions and 37 deletions.
92 changes: 55 additions & 37 deletions scripts/cp2k_pdos.py
View file
Expand Up @@ -38,6 +38,10 @@
HEADER_MATCH = re.compile(
r'\# Projected DOS for atomic kind (?P<element>\w+) at iteration step i = \d+, E\(Fermi\) = [ \t]* (?P<Efermi>[^\t ]+) a\.u\.')

# Column indexes, starting from 0
EIGENVALUE_COLUMN = 1
DENSITY_COLUMN = 3


# from https://stackoverflow.com/a/17603000/1400465
@contextlib.contextmanager
Expand All @@ -56,8 +60,8 @@ def smart_open(filename=None):

def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('pdosfilename', metavar='<pdos-file>', type=str,
help="the PDOS file generated by CP2K")
parser.add_argument('pdosfilenames', metavar='<PDOS-file>', type=str, nargs='+',
help="the PDOS file generated by CP2K, specify two (alpha/beta) files for a common grid")
parser.add_argument('--sigma', '-s', type=float, default=0.02,
help="sigma for the gaussian distribution (default: 0.02)")
parser.add_argument('--de', '-d', type=float, default=0.001,
Expand All @@ -72,43 +76,57 @@ def main():
help="write output to specified file (default: write to standard output)")
args = parser.parse_args()

with open(args.pdosfilename, 'r') as fhandle:
header = HEADER_MATCH.match(fhandle.readline().rstrip())
if not header:
print(("The file '{}' does not look like a CP2K PDOS output.\n"
"If it is indeed a correct output file, please report an issue at\n"
" https://github.com/dev-zero/cp2k-tools/issues").format(args.pdosfilename))
sys.exit(1)
alldata = []
orb_headers = []

for pdosfilename in args.pdosfilenames:
with open(pdosfilename, 'r') as fhandle:
match = HEADER_MATCH.match(fhandle.readline().rstrip())
if not match:
print(("The file '{}' does not look like a CP2K PDOS output.\n"
"If it is indeed a correct output file, please report an issue at\n"
" https://github.com/dev-zero/cp2k-tools/issues").format(pdosfilename))
sys.exit(1)

efermi = float(match.group('Efermi'))
# header is originally: ['#', 'MO', 'Eigenvalue', '[a.u.]', 'Occupation', 's', 'py', ...]
header = fhandle.readline().rstrip().split()[1:] # remove the comment directly
header[1:3] = [' '.join(header[1:3])] # rejoin "Eigenvalue" and its unit
data = np.loadtxt(fhandle) # load the rest directly with numpy

alldata.append(data)

orb_headers += header[DENSITY_COLUMN:]

# take the boundaries over all energy eigenvalues (not guaranteed to be the same)
# add a margin to not cut-off Gaussians at the borders
margin = 10 * args.sigma
emin = min(np.min(data[:, EIGENVALUE_COLUMN]) for data in alldata) - margin
emax = max(np.max(data[:, EIGENVALUE_COLUMN]) for data in alldata) + margin
ncols = sum(data.shape[1] - DENSITY_COLUMN for data in alldata)
nmesh = int((emax-emin)/args.de)+1 # calculate manually instead of using np.arange to ensure emax inside the mesh
xmesh = np.linspace(emin, emax, nmesh)
ymesh = np.zeros((nmesh, ncols))

efermi = float(header.group('Efermi'))
# header is originally: ['#', 'MO', 'Eigenvalue', '[a.u.]', 'Occupation', 's', 'py', ...]
header = fhandle.readline().rstrip().split()[1:] # remove the comment directly
header[1:3] = [' '.join(header[1:3])] # rejoin "Eigenvalue" and its unit
data = np.loadtxt(fhandle) # load the rest directly with numpy
# printing to stderr makes it possible to simply redirect the stdout to a file
print("Integration step: {:14.5f}".format(args.de), file=sys.stderr)
print("Sigma: {:14.5f}".format(args.sigma), file=sys.stderr)
print("Minimum energy: {:14.5f} - {:.5f}".format(emin+margin, margin), file=sys.stderr)
print("Maximum energy: {:14.5f} + {:.5f}".format(emax-margin, margin), file=sys.stderr)
print("Nr of mesh points: {:14d}".format(nmesh), file=sys.stderr)

npnts, ncols = data.shape
ncols -= 3 # ignore energy, occupation and #MO for mesh
fact = args.de/(args.sigma*np.sqrt(2.0*np.pi))

emin = min(data[:, 1]) - 0.2 # energies are in the second column
emax = max(data[:, 1]) + 0.2
nmesh = int((emax-emin)/args.de)+1
coloffset = 0
for fname, data in zip(args.pdosfilenames, alldata):
print("Nr of lines: {:14d} in {}".format(data.shape[0], fname), file=sys.stderr)
ncol = data.shape[1] - DENSITY_COLUMN

# printing to stderr makes it possible to simply redirect the stdout to a file
print("# of lines: {:14d}".format(npnts), file=sys.stderr)
print("Integration step: {:14.3f}".format(args.de), file=sys.stderr)
print("Sigma: {:14.3f}".format(args.sigma), file=sys.stderr)
print("Minimum energy: {:14.3f}".format(emin), file=sys.stderr)
print("Maximum energy: {:14.3f}".format(emax), file=sys.stderr)
print("# of mesh points: {:14d}".format(nmesh), file=sys.stderr)

# reproduce exactly the previous Fortran-based code
xmesh = np.array([emin+i*args.de for i in range(1, nmesh+1)])
ymesh = np.zeros((nmesh, ncols))
for idx in range(nmesh):
func = np.exp(-(xmesh[idx]-data[:, EIGENVALUE_COLUMN])**2/(2.0*args.sigma**2))*fact
ymesh[idx, coloffset:(coloffset+ncol)] = func.dot(data[:, DENSITY_COLUMN:])

fact = args.de/(args.sigma*np.sqrt(2.0*np.pi))
for mpnt in range(nmesh):
func = np.exp(-(xmesh[mpnt]-data[:, 1])**2/(2.0*args.sigma**2))*fact
ymesh[mpnt, :] = func.dot(data[:, 3:])
coloffset += ncol

if args.total_sum:
finalsum = np.sum(ymesh, 0)*args.de
Expand All @@ -121,9 +139,9 @@ def main():

with smart_open(args.output) as fhandle:
if not args.no_header:
print(("{:>16}" + " {:>16}"*ncols).format("Energy_[eV]", *header[3:]), file=fhandle)
for mpnt in range(nmesh):
print(("{:16.8f}" + " {:16.8f}"*ncols).format(xmesh[mpnt], *ymesh[mpnt, :]), file=fhandle)
print(("{:>16}" + " {:>16}"*ncols).format("Energy_[eV]", *orb_headers), file=fhandle)
for idx in range(nmesh):
print(("{:16.8f}" + " {:16.8f}"*ncols).format(xmesh[idx], *ymesh[idx, :]), file=fhandle)


if __name__ == '__main__':
Expand Down

0 comments on commit bc930c7

Please sign in to comment.