def drawDOS(doscar="",atomNum=0,ax='',symbol='-',dos='tspd',legend=[],factor=1.):
if ax=='':ax=plt
#eSet,tDOSSet,sDOSSet,pDOSSet,dDOSSet = readPDOS(doscar,atomNum)
eSet,tDOSSet,sDOSSet,pDOSSet,dDOSSet = readDOSCAR(doscar,atomNum)
# plt.plot(eSet,sDOSSet)
# print eSet.shape, sDOSSet.shape
mozala = len(dos) - len(legend)
if mozala>0:
# legend += [''] *mozala
legend = [s for s in dos]
legend = legend[::-1]
if 't' in dos:
ax.plot(tDOSSet*factor,eSet,symbol,label=legend.pop())
if 's' in dos:
ax.plot(sDOSSet*factor,eSet,symbol,label=legend.pop()) #Si$_{}$
if 'p' in dos:
ax.plot(pDOSSet*factor,eSet,symbol,label=legend.pop()) #Si$_{}$
if 'd' in dos:
ax.plot(dDOSSet*factor,eSet,symbol,label=legend.pop()) #Si$_{}$
handles, labels = ax.get_legend_handles_labels()
max_dos = (np.average(tDOSSet) * 15 / 10.) * factor
ax.set_xlim((0, max_dos ))
leg=ax.legend(handles, labels, \
frameon=False ,loc='best',numpoints=1,handletextpad=0.5,borderpad=0.2, labelspacing=0.2, handlelength=1 )
if __name__=='__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f","--file", help="DOSCAR file path is needed. default is ./DOSCAR",default = "./DOSCAR")
parser.add_argument('-a', '--atoms', nargs='+', type=int, help="atom index for pdos. ex) 0: all atoms, averaged. [1, 2, 3] ")
parser.add_argument("-e", "--eRange", nargs=2, type=float, help = 'energy Range (eV)')
parser.add_argument("-w","--write", help="write data",default = None)
args = parser.parse_args()
doscar = args.file
atoms = args.atoms
l_write = args.write
eSet,tDOSSet,sDOSSet,pDOSSet,dDOSSet = readDOSCAR(doscar,atoms)
plt.plot(eSet,tDOSSet,'r-',label='DOS_total')
if atoms !=None:
plt.plot(eSet,sDOSSet,'g-',label=str(atoms)+'_s')
plt.plot(eSet,pDOSSet,'b-',label=str(atoms)+'_p')
plt.plot(eSet,dDOSSet,'y-',label=str(atoms)+'_d')
plt.legend()
if args.eRange != None:
eMin, eMax = args.eRange
plt.xlim(eMin, eMax)
if l_write:
