#Firts thing first we have to load pylab in order to be able to plot the results
%pylab inline
import numpy as np 

#We asssume that we have a BetaFunction.py file as well as a SolveRGEs.py file both produced by PyR@TE during a previous run

#Make the RGEclass available
import sys
sys.path.append('./Source/Output/')
#The RGEclass should be loaded
import RGEclass as rge
#The BetaFunction.py must be loaded (assumed in the current working directory, if not add the path to sys.path like above)
from BetaFunction import beta_function_SM as bSM

#Now the first thing is to create an RGE object of the RGEclass for or result
OurRGEs = rge.RGE(bSM,32,labels=['betaY_u[0,0]', 'betaY_u[0,1]', 'betaY_u[0,2]', 'betaY_u[1,0]', 'betaY_u[1,1]', 'betaY_u[1,2]', 'betaY_u[2,0]', 'betaY_u[2,1]', 'betaY_u[2,2]', 'betaY_e[0,0]', 'betaY_e[0,1]', 'betaY_e[0,2]', 'betaY_e[1,0]', 'betaY_e[1,1]', 'betaY_e[1,2]', 'betaY_e[2,0]', 'betaY_e[2,1]', 'betaY_e[2,2]', 'betaY_d[0,0]', 'betaY_d[0,1]', 'betaY_d[0,2]', 'betaY_d[1,0]', 'betaY_d[1,1]', 'betaY_d[1,2]', 'betaY_d[2,0]', 'betaY_d[2,1]', 'betaY_d[2,2]', 'bg_SU2L', 'bg1', 'bg_SU3c', 'bLambda_1', 'bmu_1'])
#note that the labels have to correspond to the ordering of the beta functions in BetaFunction.py

#list of all pre-defined values
OurRGEs

#Set the initial values, here we use the SM values predefine in the RGE class
lbd0 = 0.07
mu0 = -lbd0*OurRGEs.vev**2
#Initial all the values to zero and modify the one we want to be different
Y0 = [0]*32
#Up diagonal Yukawas
Y0[0],Y0[4],Y0[8]=OurRGEs.yU[0],OurRGEs.yU[1],OurRGEs.yU[2]
#Down diagonal Yukawas
Y0[9],Y0[13],Y0[17]=OurRGEs.yD[0],OurRGEs.yD[1],OurRGEs.yD[2]
#Gauge Couplings
Y0[-4],Y0[-5],Y0[-3]=OurRGEs.g10,OurRGEs.g20,OurRGEs.g30
#lambda and mu1
Y0[-2],Y0[-1] = lbd0,mu0
OurRGEs.Y0 = Y0

#Set initial and final values for the scale directly e.g. Mz to 10**16 (same as initial values)
tmin,tmax,tstep = np.log(OurRGEs.Mz),19,0.5
#By default the assumptions is set to 'two-loop': True,'diag':True which means that the two loop equations are used if provided and Yukawas are diagonal
#Note that the user is free to modify the betaFunction.py file to introduce new assumptions and then add them to the dictionnary.
OurRGEs.solve_rges(tmin,tmax,tstep,assumptions={'two-loop': False,'diag':True})

#Now we can acces all the results via rge.Sol
OurRGEs.Sol['bLambda_1'][:10]#ten first points of lambda(t)

#There is also an interpolated object accessible by,Can be called in any point
OurRGEs.Solint['bLambda_1'],OurRGEs.Solint['bLambda_1'](6)

#Let's first plot the quartic coupling. To do so we use the plot function of the pylab library
plot(OurRGEs.Sol['t'],OurRGEs.Solint['bLambda_1'](OurRGEs.Sol['t']),label=r'$\Lambda_1$')
#add a grid
grid(True)

#Let's now look at the gauge couplings 
plot(OurRGEs.Sol['t'],OurRGEs.Sol['bg1']**-2*np.sqrt(3/5.)**2,'k',label=r'$\alpha_1^{-1}$')
plot(OurRGEs.Sol['t'],OurRGEs.Sol['bg_SU2L']**-2,'b',label=r'$\alpha_2^{-1}$')
plot(OurRGEs.Sol['t'],OurRGEs.Sol['bg_SU3c']**-2,'r',label=r'$\alpha_3^{-1}$')
grid(True)
#titles
xlabel('$t=\log_{10}(Q\ \mathrm{GeV})$')
ylabel('Running in the SM')
title('Running of the gauge couplings in the SM obtained from PyR@TE')
#number of column of location of the legend, collects all the label entries above
legend(loc=4,ncol=3)

#To illustrate the difference between one loop and two loop we can create a new object with the same RGEs but soolve it for two loop.
#Note that we could use the same instance we have already created but we won't be able to compare the results then.
rge2L = rge.RGE(bSM,32,labels=['betaY_u[0,0]', 'betaY_u[0,1]', 'betaY_u[0,2]', 'betaY_u[1,0]', 'betaY_u[1,1]', 'betaY_u[1,2]', 'betaY_u[2,0]', 'betaY_u[2,1]', 'betaY_u[2,2]', 'betaY_e[0,0]', 'betaY_e[0,1]', 'betaY_e[0,2]', 'betaY_e[1,0]', 'betaY_e[1,1]', 'betaY_e[1,2]', 'betaY_e[2,0]', 'betaY_e[2,1]', 'betaY_e[2,2]', 'betaY_d[0,0]', 'betaY_d[0,1]', 'betaY_d[0,2]', 'betaY_d[1,0]', 'betaY_d[1,1]', 'betaY_d[1,2]', 'betaY_d[2,0]', 'betaY_d[2,1]', 'betaY_d[2,2]', 'bg_SU2L', 'bg1', 'bg_SU3c', 'bLambda_1', 'bmu_1'])

rge2L.Y0 = Y0

rge2L.solve_rges(tmin,tmax,tstep,assumptions={'two-loop': True,'diag':True})

#Now we can compare the top yukawa at one and two loop
plot(rge2L.Sol['t'],rge2L.Solint['betaY_u[2,2]'](rge2L.Sol['t']),label=r'Two loop $Y_t$')
plot(OurRGEs.Sol['t'],OurRGEs.Solint['betaY_u[2,2]'](OurRGEs.Sol['t']),label=r'One loop $Y_t$')
grid(True)
xlabel('$t=\log_{10}(Q\ \mathrm{GeV})$')
ylabel(r'$Y_t$')
title('Comparison 1L vs 2L of $Y_t$')
#number of column of location of the legend, collects all the label entries above
legend(loc=4,ncol=3)