# sij.py

# Calculate the overlap matrix elements of the transformation waves, i.e., Sij
# Sij = < phi_j | \tilde{phi}_i >^AE - < phi_j | \tilde{phi}_i > ^ PS
# i, j loop over different angular momentum l

# usage:
# sij.py
# or 
# sij.py valence-gs.dat valence-x.dat
# valence-gs(x).dat are generated by the vanderbilt uspp code

# output: Sij.dat

# Yufeng Liang, Jan 2017, LBNL

import sys

try:
	import scipy as sp
except ImportError:
	import numpy as sp

import matplotlib.pyplot as plt

# defaulted
fil_valgs = "valence-gs.dat"
fil_valx  = "valence-x.dat"

if len(sys.argv) == 3:
	fil_valgs = sys.argv[1]
	fil_valx  = sys.argv[2]
elif len(sys.argv) != 2:
	print("usage: python sij.py valence-gs.dat valence-x.dat")
	sys.exit(0)

# import the radial | phi_i > ^ AE and | phi_i > ^ PS from the file fname
def read_val(fname):
	f = open(fname, 'r')
	np = int(f.readline()) # number of projectors
	l = [0] * np # angular momentum
	for ip in range(np):
		nl = f.readline()
		ngrid, l[ip] = int(nl.split()[0]), int(nl.split()[1]) # ngrid, l
		if ip == 0:
			# initialize arrays
			phi_ae = sp.zeros([ngrid, np + 1]) # grid phi_1 phi_2 ... phi_np
			phi_ps = sp.array(phi_ae)
		for ig in range(ngrid):
			nl = f.readline().split()
			phi_ps[ig, 0], phi_ps[ig, ip + 1] = float(nl[0]), float(nl[1])
			# r phi(r) --> phi(r)
			#if ig > 0:
			#	phi_ps[ig, ip + 1] /= phi_ps[ig, 0]
		nl = f.readline()
		ngrid, l[ip] = int(nl.split()[0]), int(nl.split()[1]) # ngrid, l
		for ig in range(ngrid):
			nl = f.readline().split()
			phi_ae[ig, 0], phi_ae[ig, ip + 1] = float(nl[0]), float(nl[1])
			# r phi(r) --> phi(r)
			#if ig > 0:
			#	phi_ae[ig, ip + 1] /= phi_ae[ig, 0]
	f.close()
	return l, phi_ps, phi_ae

l_gs, phi_gs_ps, phi_gs_ae = read_val(fil_valgs)
l_x, phi_x_ps, phi_x_ae = read_val(fil_valx)

# plot the radial wavefunctions

def draw_wave(l, phi_ps, phi_ae):

	# color listd
	clist = ['r', 'g', 'b', 'm', 'c', 'y', 'k', '#eeefff']

	np = len(l)
	for ip in range(np):
		plt.plot(phi_ps[:, 0], phi_ps[:, ip + 1], clist[ip] + '--', linewidth = 3)
		plt.plot(phi_ae[:, 0], phi_ae[:, ip + 1], clist[ip],  label = 'l = ' + str(l[ip]))

	plt.legend(loc = 'upper center')


draw_wave(l_gs, phi_gs_ps, phi_gs_ae)
#plt.savefig('wave_gs.eps', format = 'eps', dpi = 1000)
plt.savefig('wave_gs.png', format = 'png')
plt.close()

draw_wave(l_x, phi_x_ps, phi_x_ae)
#plt.savefig('wave_x.eps', format = 'eps', dpi = 1000)
plt.savefig('wave_x.png', format = 'png')
plt.close()

# Calculate the overlap S_ij
# i - final state, j - initial state

np_gs = len(l_gs)
np_x = len(l_x)

S = sp.zeros([np_x, np_gs])

def radial_int(rgrid, phi_1, phi_2):
	
	nr = len(rgrid)
	phi_prod = sp.multiply(phi_1.conjugate(), phi_2)
	# trapezoidal rule
	return sp.sum ( sp.multiply( ( phi_prod[1 : nr] + phi_prod[0 : nr - 1] ) / 2.0, rgrid[1 : nr] - rgrid[0 : nr - 1]) )

if ( len(phi_gs_ps) != len(phi_x_ps) ) or ( len(phi_gs_ae) != len(phi_x_ae) ):
	print("error: radial grid size of pseudo and all-electron wave mismatch. ")
	sys.exit(0)
elif ( any(phi_gs_ps[:, 0] - phi_x_ps[:, 0]) > 1e-12) \
or   ( any(phi_gs_ae[:, 0] - phi_x_ae[:, 0]) > 1e-12):
	print("error: radial grids of pseudo and all-electron wave mismatch. ")
	sys.exit(0)

for i in range(np_x):
	for j in range(np_gs):
		if l_x[i] == l_gs[j]:
			S[i, j] = radial_int(phi_gs_ae[:, 0], phi_gs_ae[:, j + 1], phi_x_ae[:, i + 1]) \
			        - radial_int(phi_gs_ps[:, 0], phi_gs_ps[:, j + 1], phi_x_ps[:, i + 1])

sp.savetxt('Sij.dat', S)


