from QuantumATK import *
import pylab
import numpy as np
from pylab import *

# List of c(tb09) parameters used in the MGGA bandstructure calculations
tb09=[0.9,1.0,1.1,1.2]

# Read the bandstructures objects calculated at different tb09 parameters and calculate the
# indirect band gap.
ind_gap=[]
for n in range(0,4):
    bandstructure=nlread("MGGA_bulk.hdf5", Bandstructure)[n]
    ind_gap.append(bandstructure.indirectBandGap().inUnitsOf(eV))

# Calculate the Slater-Koster bindirect band gap
bandstructure_SK=nlread("bulk_SK.hdf5", Bandstructure)[-1]
SK = bandstructure_SK.indirectBandGap().inUnitsOf(eV)

# Perform a linear fit of the calculated MGGA band gap vs tb09 parameter and
# calculate the c parameters which will give the SK band gap.
(m,b)=polyfit(tb09,ind_gap,1)
fitted_c = (SK - b) / m

print "Fitted c parameter: ",  fitted_c

#  Make the plot
pylab.plot([min(tb09),max(tb09)],[SK,SK] , linewidth=1.0, linestyle='-', color='black', label="SK band gap" )
pylab.plot([fitted_c, fitted_c],[min(ind_gap),max(ind_gap)] , linewidth=1.0, linestyle='-', color='b', label="fitted c" )

pylab.plot(tb09,ind_gap,'*b',label='MGGA gap')

yp=polyval([m,b],tb09)
pylab.plot(tb09,yp, '--')

pylab.xlabel("TB09 - c")
pylab.ylabel("Band Gap (eV)")
pylab.grid(True)
pylab.legend(loc="lower right")
pylab.show()