from QuantumATK import *
import pylab
import sys

filename = "au_vacuumgap"

# Calculate capacitance from charges
# Read the data
data_m = numpy.load('%s_mulliken_data.npz' % filename)
data_e = numpy.load('%s_energy_data.npz' % filename)
biases_m = data_m['v']
net_charges = data_m['m']

# Symmetrize energy data for better fit around V=0
biases_e = numpy.concatenate((-data_e['v'],data_e['v']))
energies = numpy.concatenate((data_e['e'],data_e['e']))

plot_v = numpy.arange(0.01,1.0,0.01)

# Calculate capacitance from total energies simply as Q/V
q = [-numpy.cumsum(c-net_charges[0])[10] for c in net_charges]
function_m = SplineInterpolation1D(biases_m,q)
cq = [function_m(v)/v for v in plot_v]*elementary_charge/Volt

# Calculate capacitance from total energies simply as 2E/V^2
function_e = SplineInterpolation1D(biases_e, energies)
ce = [2.0*function_e(v)/(v*v) for v in plot_v]*elementary_charge/Volt

# Plotting the line
pylab.figure()
pylab.plot(plot_v,cq.inUnitsOf(Coulomb/Volt), label="From Mulliken")
pylab.plot(plot_v,ce.inUnitsOf(Coulomb/Volt), label="From Energy")
pylab.xlabel('Bias (Volt)')
pylab.ylabel('Capacitance (Farad)')
pylab.legend()

pylab.savefig("%s_cv.png" % filename, dpi=120)
pylab.show()