[交流] 处理能带的脚本：vasp_band.py

#!/usr/bin/python
################################################################################
##                                                                            ##
##  vasp_band.py:                                                             ##
##                                                                            ##
##    Plot energy bands from VASP in/out files                                ##
##    You need files as following,                                            ##
##                                                                            ##
##        - EIGENVAL (VASP output file)                                       ##
##        - OUTCAR (VASP output file)                                         ##
##                                                                            ##
##    cf. './vasp_band.py -h'                                             ##
##                                                                            ##
##                                                    Designed by Jungho Shin ##
################################################################################

import os
import sys
import getopt

from numpy import *


##### default variables to be controled  #####
inEigenvalName = 'EIGENVAL'  ## the file name of "EIGENVAL" format
inOutcarName = 'OUTCAR'  ## the file name of "OUTCAR" format
outDataName = 'band.dat'  ## the name of a data(2D coordinate) file

optPersist = 1  ## (1 or 0) if 1, the plot remains on the screen
gpTitle = '< Energy Band Structure >'  ## graph's title
gpSize = '0.7, 1.0'  ## size of the plot
gpXrange = '[*:*]'  ## X axis range
gpXlabel = 'K-points'  ## X axis label
gpXticLabel = 0  ## X tic's label, if 0, expresstion by x value
gpYrange = '[*:*]'  ## Y axis range
gpYlabel = 'Energy(eV)'  ## Y axis label
gpTerm = 'x11'  ## a type of terminal and command
gpPlotMethod = 'with line'  ## plot method
gpPicForm = 'pos'  ## image format
gpPosDefault = 'landscape enhanced color solid "Times-Roman" 15'  ## postscript's defaults
gpPicName = 'band.eps'  ## image file's name
gpKey = 'nokey'  ## 'key' or 'nokey'
                 ## if 'nokey', don't display the label of lines
##########


##### default variables to be fixed #####
blnEigenval = 7  ## the number of comment lines in "EIGENVAL" file
##########


##### options #####
optlist, args = getopt.getopt( sys.argv[1:],
                               'o:p:k:h', ['kpoints='    , 'eigenval='   ,
                                           'outcar='     , 'blnKpoint='  ,
                                           'title='      , 'size='       ,
                                           'xRange='     , 'xLabel='     ,
                                           'xTicLabel='  , 'yRange='     ,
                                           'yLabel='     , 'term='       ,
                                           'plotMethod=' , 'form='       ,
                                           'epsDefault=' , 'picName='] )

for op,p in optlist:
    if op == '-o' : outDataName = p
    if op == '-p' : optPersist = int(p)
    if op == '-k' : gpKey = p
    if op == '-h' : help = 1

    if op == '--eigenval'   : inEigenvalName = p
    if op == '--outcar'     : inOutcarName   = p
    if op == '--title'      : gpTitle        = p
    if op == '--size'       : gpSize         = p
    if op == '--xRange'     : gpXrange       = p
    if op == '--xLabel'     : gpXlabel       = p
    if op == '--xTicLabel'  : gpXticLabel    = p.split()
    if op == '--yRange'     : gpYrange       = p
    if op == '--yLabel'     : gpYlabel       = p
    if op == '--term'       : gpTerm         = p
    if op == '--plotMethod' : gpPlotMethod   = p
    if op == '--form'       : gpPicForm      = p
    if op == '--epsDefault' : gpPosDefault   = p
    if op == '--picName'    : gpPicName      = p
##########


##### print default variables #####
print ""
print ""
print "     ** Printing Variables **"
print ""

print " -o            : the file name of output data"
print " -p            : if 1, the plot remains on the screen (0 or 1)"
print " -k            : if 'nokey', don't display the label of lines",
print "('key' or 'nokey')"
print ""
print " --eigenval=   : the file name of EIGENVAL format"
print " --outcar=     : the file name of OUTCAR format"
print " --title=      : graph's title"
print " --size=       : size of the plot, ratio of x and y axis"
print " --xRange=     : X axis range"
print " --xLabel=     : X axis label"
print " --xTicLabel=  : X tic label, about greek letter, add 'gk' at the front of alphabet."
print "               : usage) --xTicLabel=' X W gkG L '"
print " --yRange=     : Y axis range"
print " --yLabel=     : Y axis label"
print " --term=       : a type of terminal"
print " --plotMethod= : plotting method by Gnuplot command(plot)"
print " --form=       : image format"
print " --epsDefault= : postscript's defaults"
print " --picName=    : image file's name"
print ""

print " -o '%-s'"          % outDataName
print " -p %-d"            % optPersist
print " -k '%-s'"          % gpKey
print " --eigenval='%-s'"  % inEigenvalName
print " --outcar='%-s'"    % inOutcarName
print " --title='%-s'"     % gpTitle
print " --size='%-s'"      % gpSize
print " --xRange='%-s'"    % gpXrange
print " --xLabel='%-s'"    % gpXlabel

if gpXticLabel==0: print " --xTicLabel=%d" % gpXticLabel
else:
    print " --xTicLabel='",
    for string in gpXticLabel: print "%s" % string,
    print "'"

print " --yRange='%-s'"    % gpYrange
print " --yLabel='%-s'"    % gpYlabel
print " --term='%-s'"      % gpTerm
print " --plotMethod='%s'" % gpPlotMethod
print " --form='%-s'"      % gpPicForm
print " --epsDefault='%-s'"% gpPosDefault
print " --picName='%-s'"   % gpPicName
print ""
print ""

if help==1: sys.exit()
##########


##### open file  #####
inEigenval = open(inEigenvalName, 'r')
outData = open(outDataName, 'wb')
outGpInput = open('gnuinput.in', 'wb')
##########


#####  extract 3D coordinates from the input line  #####
def extract(inLine):
    list = inLine.split()
    x = float(list[0])
    y = float(list[1])
    z = float(list[2])
    return x, y, z
##########


##### calculate distance between two points #####
def distCrd(x1, y1, z1, x2, y2, z2):
    x21 = abs(x2 - x1)
    y21 = abs(y2 - y1)
    z21 = abs(z2 - z1)
    dist21 = sqrt(x21*x21 + y21*y21 + z21*z21)
    return dist21
##########


##### obtain y values(energy) from the "EIGENVAL" file #####
def obtainy():
    xMinorCrd = inEigenval.readline()  ## xyz
    x, y, z = extract(xMinorCrd)
    i=0 ; yValueList=[]
    while i < nplotLine:
        list = inEigenval.readline().split()
        yValue = float(list[1])
        yValueList.append(yValue)
        i = i+1
    next = inEigenval.readline()  ## blank line
    return x, y, z, yValueList, next
##########


##### append x tic label #####
def appendXticLabel(order):
    if gpXticLabel==0: xMajorList.append('%s, ' % xMinor)
    else:
        if len(gpXticLabel) < order+1: print 'ERROR => Needed more tic labels of x axis!!'; sys.exit()
        if gpXticLabel[order][0:2]=='gk': gpXticLabel[order] = "{/Symbol "+gpXticLabel[order][2]+"}"
        xMajorList.append('"%s" %s, ' % (gpXticLabel[order], xMinor))
        order = order+1
    return order
##########


##### read comment(blank) lines #####
i=0
while i < blnEigenval:
    if i==5:
        list = inEigenval.readline().split()
        nplotLine = int(list[2])
    else:
        line = inEigenval.readline()
    i = i+1
##########


##### transfer format and write into the output file #####
fermiE = (os.popen('grep E-fermi %s | cut -d":" -f 2' % inOutcarName)).read().split()[0]
print 'Fermi energy is %.4f(eV)\n' % float(fermiE)

j=1; l=0
xMajorList=[]
xMinor, xMinorXb, xMinorYb, xMinorZb = 0, 0, 0, 0

l = appendXticLabel(l)

while 1:
    xMinorXa, xMinorYa, xMinorZa, yValueList, next = obtainy()
    dist = distCrd(xMinorXb, xMinorYb, xMinorZb, xMinorXa, xMinorYa, xMinorZa)

    if (xMinorXa == xMinorXb and xMinorYa == xMinorYb and xMinorZa == xMinorZb): k=0
    else: k=1

    xMinorXb, xMinorYb, xMinorZb = xMinorXa, xMinorYa, xMinorZa

    if j==1: dist=0
    xMinor = xMinor + dist

    if not (j>1 and k==0):
        outData.write('%9.4f' % xMinor)
    else:
        l = appendXticLabel(l)

    for yValue in yValueList:
        yValue = yValue-float(fermiE)
        if not (j>1 and k==0): outData.write('%9.4f' % yValue)

    if not (j>1 and k==0): outData.write('\n')
    if not next: break
    j=2

l = appendXticLabel(l)
##########


##### use Gnuplot module #####
xticsStr = ''.join(xMajorList)

if optPersist==1: optPersist='persist'
else: optPersist='nopersist'
outGpInput.write("set term %s %s\n" % (gpTerm, optPersist))

outGpInput.write("set title '%s'\n" % gpTitle)
outGpInput.write("set size %s\n" % gpSize)

if gpKey=='key': outGpInput.write("set key\n")
else: outGpInput.write("unset key\n")

outGpInput.write("set xlabel '%s'\n" % gpXlabel)
if gpXrange=='[*:*]': gpXrange = '[*:'+str(xMinor+0.001*xMinor)+']'
outGpInput.write("set xrange %s\n" % gpXrange)
outGpInput.write("set xtics ( %s )\n" % xticsStr[0:-2])
outGpInput.write("set ylabel '%s'\n" % gpYlabel)
outGpInput.write("set yrange %s\n" % gpYrange)

outGpInput.write("plot '%s' using 1:2 %s\n" % (outDataName, gpPlotMethod))

i=3
while i < nplotLine+2:
    outGpInput.write("replot '%s' using 1:%d %s\n" % (outDataName, i, gpPlotMethod))
    i = i+1

if gpPicForm == 'pos': outGpInput.write("set term %s %s\n" % (gpPicForm, gpPosDefault))
else: outGpInput.write('set term %s\n' % gpPicForm)

outGpInput.write("set output '%s'\n" % gpPicName)
outGpInput.write("replot '%s' using 1:%d %s\n" % (outDataName, nplotLine+1, gpPlotMethod))
outGpInput.write("set term %s %s\n" % (gpTerm, optPersist))

inEigenval.close()
outData.close()
outGpInput.close()

os.popen('gnuplot gnuinput.in')
##########