[交流] 处理态密度的脚本：vasp_pdos.py已有1人参与

#!/usr/bin/python

##### by jhshin #####
##
## 1. Extract DOS and from DOSCAR of VASP
## 2. Analysis DOS by distribution moment Eq.
##
## 2010.12.2
##########

import os
import sys

from getopt               import getopt
from numpy                import array
from ase.io               import *     
from ase.calculators.vasp import Vasp
from ase.calculators.vasp import VaspDos
from ase.dft              import get_distribution_moment

outcar_name_q = " OUTCAR name      > "
doscar_name_q = " DOSCAR name      > "
atom_list_q   = " Atom index       > "
orbit_list_q  = " Orbital          > "
en_range_q    = " Energy Min. Max. > "
screen_q      = " Screen on/off    > "

#####  Check and read options  #####
optlist, args = getopt(sys.argv[1:], 'hs')

help = 0 ; script = 0

for op,p in optlist:
  if op == '-h'   : help        = 1
  if op == '-s'   : script      = 1

if help==1:
  print ""
  print " Wellcome!"
  print ""
  print " Usage) ./vasp_pdos.py                : Input variables directly."
  print "        ./vasp_pdos.py < vasp_pdos.in : Use script file generated by user."
  print "        ./vasp_pdos.py -h             : Print explanation of this utility."
  print "        ./vasp_pdos.py -s             : Generate the sample script file named"
  print "                                        as 'vasp_pdos.in'."
  print "        ./vasp_pdos.py -hs (or) -sh   : Use 'h' and then 's' option, or inversly."
  print ""
  print " For all variables, see bellow. () is a default value"
  print ""
  print " - OUTCAR name : one string(OUTCAR), name of OUTCAR file"
  print "   ex) OUTCAR.xXX0"
  print ""
  print " - DOSCAR name : one string(DOSCAR), name of DOSCAR file"
  print "   ex) DOSCAR.xXX0"
  print ""
  print " - Atom index : list of integer, from 0 to # of total atoms"
  print "   ex) 0 1 2 3"
  print ""
  print " - Orbital : list of strings, which orbital to plot"
  print ""
  print "                     | Phase factor : X       | Phase factor : O          "
  print "   ==================.========================.==========================="
  print "                     | s                      | s                         "
  print "    Spin-unpolarized | p                      | px py pz                  "
  print "                     | d                      | dxy dyz dz2 dxz dx2       "
  print "   ------------------.------------------------.---------------------------"
  print "                     | s-up s-down (or) s+ s- | s-up s-down (or) s+ s-    "
  print "    Spin-polarized   | p-up p-down (or) p+ p- | px-up  ...  (or) px+ ...  "
  print "                     | d-up d-down (or) d+ d- | dxy-up ...  (or) dxy+ ... "
  print "   ex) s d"
  print ""
  print " - Energy Min. Max. : list of integer(* *), minimum and maximum of energy range"
  print "   ex) -10.0 5.0"
  print ""
  print " - Screen on/off : one integer(0), 0 is off and 1 is on"
  print ""
  print " Thanks!!\n"
  exit()

if script==1:
  print ""
  print " Edit 'vasp_pdos.in' file !!\n"
  in_script = open("vasp_pdos.in", 'wb')
  in_script.write( "%s\n"   % outcar_name_q)
  in_script.write( "%s\n"   % doscar_name_q)
  in_script.write( "%s\n"   % atom_list_q  )
  in_script.write( "%s\n"   % orbit_list_q )
  in_script.write( "%s\n"   % en_range_q   )
  in_script.write( "%s\n\n" % screen_q     )
  in_script.write( "# Enter values on the right side of '>'\n")
  in_script.write( "# Orbital list | (s) s  (p) px py pz  (d) dxy dyz dz2 dxz dx2")

  in_script.close()
  exit()
##########

#####  Input variables  #####
print ""
print " Wellcome! Enter bellow."
print " If you want to use a default value, do not enter anything."
print " ----------------------------------------------------------\n"
print " < Input variables >"
print ""
print " - Orbital list | (s) s  (p) px py pz  (d) dxy dyz dz2 dxz dx2"
print ""

print outcar_name_q, ; outcar_name = raw_input().replace(outcar_name_q,"")
if len(outcar_name) == outcar_name.count(" "): outcar_name="OUTCAR"
outcar_name = outcar_name.replace(" ","")
print "", outcar_name

print doscar_name_q, ; doscar_name = raw_input().replace(doscar_name_q,"")
if len(doscar_name) == doscar_name.count(" "): doscar_name="DOSCAR"
doscar_name = doscar_name.replace(" ","")
print "", doscar_name

print atom_list_q, ; atom_list   = raw_input().replace(atom_list_q,"").split()   
if len(atom_list) == 0:
  print "Enter list of integers!"
  exit()
print "", atom_list

print orbit_list_q, ; orbit_list  = raw_input().replace(orbit_list_q,"").split()   
if len(orbit_list) == 0:
  print "Enter list of strings!"
  exit()
print "", orbit_list

print en_range_q, ; en_range    = raw_input().replace(en_range_q,"")
if len(en_range) == en_range.count(" "): en_range="* *"
en_range = en_range.split()
print "", en_range

print screen_q, ; screen      = raw_input().replace(screen_q,"")
if len(screen) == screen.count(" "): screen="0"
screen = screen.replace(" ","")
print "", screen

print ""
print " ----------------------------------------------------------\n"
##########

#####  Read OUTCAR and DOSCAR  #####
fermiEn = (os.popen('grep E-fermi %s | cut -d":" -f 2' % outcar_name)).read().split()[0]
fermiEn = float(fermiEn)

doscar  = VaspDos(doscar=doscar_name)
doscar.read_doscar(doscar_name)

en_list  = array(doscar._get_energy()) -fermiEn

if en_range[0] == "*": en_min = min(en_list)
else: en_min = float(en_range[0])

if en_range[1] == "*": en_max = max(en_list)
else: en_max = float(en_range[1])
##########

#####  Prepare output  #####
outfile2_name  = doscar_name + ".info.out"
outfile2       = open(outfile2_name, 'wb')
gpInFile       = open("gnuplot.in", 'wb')

print          " < Output results >"
print          "\n# %10s %13s %16s %16s\n" % ("Atom index", "Elec. count", "Band center", "Band width")
outfile2.write("\n# %10s %13s %16s %16s\n" % ("Atom index", "Elec. count", "Band center", "Band width"))
##########

#####  Main loop : make DOS data and analysis  #####
sum=0
orbit0=""
dos_list0 = []

for atom in atom_list:
  for orbit in orbit_list:
    atom  = int(atom)

    if orbit == "+":
      sum = 1
      if len(orbit0)==0 : orbit0= orbit
      dos_list0 = dos_list

    else:
      dos_list = doscar.site_dos(atom, orbit)

      if sum==1:
        orbit = orbit0 + orbit
        orbit0 = orbit
        dos_list = dos_list + dos_list0
        sum = 0

      outfile1_name = str(doscar_name + ".atom" + str(atom) + orbit + "Orbit" + ".dat")
      outfile1      = open(outfile1_name, 'wb')

      en_list_mod  = []
      dos_list_mod = []
      i=0
      for en in en_list:
        dos = dos_list[i]
        if en_min < en < en_max:
          en_list_mod.append(en)
          dos_list_mod.append(dos)
          outfile1.write(' %16.6lf %9.6lf\n' % (en, dos))
        i += 1

      volume, center, width = get_distribution_moment(en_list_mod, dos_list_mod, (0,1,2))
      print          "  %10d %13.3lf %16.6lf %16.6lf     # %s"   % (atom, volume, center, width, orbit + " orbital")
      outfile2.write("  %10d %13.3lf %16.6lf %16.6lf     # %s\n" % (atom, volume, center, width, orbit + " orbital"))

      if str(atom) + str(orbit) == atom_list[0] + orbit_list[0]:
        gpInFile.write ("plot   '%s' w l\n" % outfile1_name)
      else:
        gpInFile.write ("replot '%s' w l\n" % outfile1_name)
##########

#####  Post-precessing  #####
outfile2.write('\n')

outfile1.close()
outfile2.close()
gpInFile.close()

if screen == '1': os.popen('gnuplot -persist gnuplot.in')

print ""
print " End\n"
##########