#sample script to compute growth factor and growth rate for (flat or nonflat) w0wa model 
# by Zhiqi Huang,  huangzhq25@mail.sysu.edu.cn

import numpy as np
from scipy.integrate import quad, odeint
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt

#density parameters
Omega_k = 0.
Omega_r = 8.55e-5
Omega_m = 0.3
Omega_L = 1-Omega_k - Omega_r - Omega_m
w0 = -1.
wa = 0.

## rho_DE/rho_{DE0} as a function of a
def rhoDE_ratio(a):
    return np.exp(-3.*((1.+w0+wa)*np.log(a)+wa*(1.-a)))

#H^2 as a function of a, in unit of H_0^2
def Hsqofa(a):
    return Omega_r/a**4 + Omega_m/a**3 + Omega_k/a**2 + Omega_L*rhoDE_ratio(a)

#dH/dt as a function of a, in unit of H_0^2
def dotHofa(a):
    return -2*Omega_r/a**4 - 1.5*Omega_m/a**3 - Omega_k/a**2 - 1.5*(1.+w0+wa*(1-a))*Omega_L * rhoDE_ratio(a)


def growth_eq(y, lna):  #y = [G , dG/dN]
    a = np.exp(lna)
    Hsq = Hsqofa(a)
    dotH = dotHofa(a)
    return np.array( [ y[1], -(4+dotH/Hsq)*y[1]-(3+dotH/Hsq-1.5/Hsq/a**3*Omega_m)*y[0] ] )


#set initial conditions
a_ini = 1.e-2 #it can be arbitrary, as long as 3.e-4 << a_ini << 1
Hsq_ini = Hsqofa(a_ini)
dotH_ini = dotHofa(a_ini)
coef_1 = 4+dotH_ini/Hsq_ini
coef_0 = 3+dotH_ini/Hsq_ini-1.5/Hsq_ini/a_ini**3*Omega_m
eps = coef_0/(coef_1 - coef_0 - 1.)
G_ini = 1.+ eps #this is arbitrary normalization; will renormalize in the end
dGdN_ini = -eps

nsteps = 500
lna = np.linspace(np.log(a_ini), 0., nsteps)
a = np.exp(lna)
Dbya = odeint(growth_eq, [ G_ini, dGdN_ini ], lna)
#now renormalize to D(z=0) = 1
D = Dbya[:, 0] / Dbya[nsteps-1, 0] * a
#compute the growth rate f = d ln D/d ln a
f = 1. + Dbya[:,1]/Dbya[:,0]

#interpolate to get functions D(a), f(a)
Dofa = interp1d(a, D, kind="cubic")
fofa = interp1d(a, f, kind="cubic")

a_test = 0.5
print("D(", a_test, ") = ", Dofa(a_test))
print("f(", a_test, ") = ", fofa(a_test))

#make plot
#plt.plot(lna, Dbya[:,0], '-', lna, f, '--')
#plt.xlabel(r'$\ln a$')
#plt.legend([r'$D/a$ (not normalized)', r'$f$'], loc='best')
#plt.show()