[交流] 【分享】一个绘制band的python小程序 已有4人参与

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