#comoving volume calculator for w0wa model
#by Zhiqi Huang, huangzhq25@mail.sysu.edu.cn 
#use c/H_0 as the length unit, and 1/H_0 the time unit; but finally convert to Gpc^3 for output
#if you want the result to be in unit of h^{-3}Gpc^3, set H0 = 100 (i.e. h=1)
import numpy as np
from scipy.integrate import quad

z1 = 1.
z2 = 2.
fsky = 0.25


#density parameters
Omega_m = 0.3
Omega_k = 0.
Omega_r = 0. #for z < 100 ignore the radiation
Omega_L = 1.0 - Omega_m - Omega_k - Omega_r

w0 = -0.8
wa = 0.2

H0 = 70  #km/s/Mpc


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


#Hubble parameter H = dot a / a as a function of z
def Hofz(z):
    return np.sqrt(Omega_L*rhoDE_ratio(z) + Omega_m*(1+z)**3 + Omega_k * (1+z)**2 + Omega_r*(1+z)**4)

def dchibydz(z): 
    return 1.0/Hofz(z)

def rofchi(chi):
    if(abs(Omega_k)<1.e-20):
        return chi
    if(Omega_k < 0):
        return np.sin(np.sqrt(-Omega_k)*chi)/np.sqrt(-Omega_k)
    return np.sinh(np.sqrt(Omega_k)*chi)/np.sqrt(Omega_k)

def ComovingAngularDiameterDistance(z): #r(z)
    chi, err = quad(dchibydz, 0., z)
    r = rofchi(chi)
    return r

def dV_by_dz_dSolidAngle(z): #comoving volume per solid angle per redshift
    return ComovingAngularDiameterDistance(z)**2*dchibydz(z)

def ComovingVolume(z1, z2, fsky):
    V_per_SolidAngle, err = quad(dV_by_dz_dSolidAngle, z1, z2)
    return V_per_SolidAngle*(4.*np.pi*fsky)

#comoving volume betwee z1, z2 with sky fraction fsky 
V= ComovingVolume(z1, z2, fsky)

#finally, convert to Gpc^3
speed_of_light = 2.998e5 #in km/s
H0Gpc = (H0 / speed_of_light) * 1.e3
print("The comoving volume between z=", z1, " and z=", z2, ", with sky fraction ", fsky, " is ", V/H0Gpc**3, " Gpc^3")