python绘图小程序_【分享】一个绘制band的python小程序 - 第一原理 - 小木虫 - 学术 科研 互动社区...

CODE:

!/usr/bin/env python

#make a simple plot of band structure from vasp calculation

# -Brandon Cook April 7, 2010

from numpy import *

import pylab

class OUTCAR():

def __init__(self):

f = open('OUTCAR', 'r')

s = f.read()

f.close()

i = s.find('E-fermi')

self.ef = float(s[i:i+19].split()[-1])

i = s.find('NBANDS')

self.nbands = int(s[i:i+15].split()[-1])

i = s.find('NKPTS')

self.nkpts = int(s[i:i+15].split()[-1])

class KPOINTS():

def __init__(self):

# kpts file should be in line mode

# only take 1 line for now

f = open('KPOINTS', 'r')

s = f.readlines()

f.close()

l1 = s[-2].split()

l2 = s[-1].split()

self.p1 = array([ l1[0], l1[1], l1[2] ], float)

self.p2 = array([ l2[0], l2[1], l2[2] ], float)

class BandStructure():

# Just take all the points and make a scatter plot

# todo: connect the lines

def __init__(self):

f = open('EIGENVAL', 'r')

s = f.readlines()

f.close()

oc = OUTCAR()

kp = KPOINTS()

axis = kp.p2 - kp.p1

axis /= linalg.norm(axis)

nbands = oc.nbands

nkpts  = oc.nkpts

kpts = zeros( nkpts*nbands )

en   = zeros( nkpts*nbands )

istart = 7

for k in range(0, nkpts):

ik = istart + k*(nbands+2)

kpt = s[ik].split()

kpt = array([kpt[0], kpt[1], kpt[2]], float)

x = dot(kpt, axis)

i1 = k*nbands

i2 = (k+1)*nbands

kpts[i1:i2] = x

for ib in range(0, nbands):

j1 = ik + 1 + ib

j2 = i1 + ib

en[j2] = float(s[j1].split()[-1])

pylab.scatter(kpts, en, s=1)

# add a line for the fermi energy

xmin = 0.0

xmax = linalg.norm(kp.p2 - kp.p1)

x = linspace(xmin, xmax, num=100, endpoint=True)

y = oc.ef * ones(100)

pylab.plot(x,y)

pylab.show()

o = OUTCAR()

print 'fermi energy, nbands, nkpts'

print o.ef,o.nbands,o.nkpts

b = BandStructure()