#comoving volume calculator for massive neutrinos + w0-wa dark energy 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
from scipy.integrate import quad
import numpy as np
import matplotlib.pyplot as plt


 #where you want to compute the comoving distances
z1 = 1.
z2 = 2.
fsky  = 0.25

#dark energy equation of state = w0 + wa (1-a)
w0 = -1.
wa = 0.
#cold dark matter + baryon; this does not include massive neutrinos
Omega_m = 0.3
#spatial curvature
Omega_k = 0.
#effective number of species 
N_eff = 3.046 
#masses of species of neutrinos in eV
mass_nu = np.array([0.001, 0.009, 0.05])
#Hubble constant in km/s/Mpc
H0 = 70. 
#CMB temperature
T_CMB = 2.726


## integrals for massive neutrinos
def I_rho(lam):
    cutdown = 0.5
    cutup = 7.
    k = 1./6.9
    if lam < cutdown:
        return (7*np.pi**2/240.)*(1.+ (5./7./np.pi**2/k)*lam**2)**k
    elif lam > cutup:
        return 9.13453712e-02*lam + 5.90977955e-01/lam -4.52/(lam+0.86)**3
    else:
        lnlam = np.log(lam)
        return  0.30627596287 + lnlam*(0.033956 + lnlam*(0.02977417  + lnlam*(0.01602 + lnlam*(0.005584 + lnlam*1.15e-3))))

def I_p(lam):
    cutdown = 0.5
    cutup = 7.
    k = -1./14.
    if lam <= cutdown:
        return 7*np.pi**2/720.*(1.- (5./7./np.pi**2/k)*lam**2)**k
    elif lam > cutup:
        return 3.93985303e-01/lam -6.03424688/lam**3 + 2.54842801e2/(lam+1.15)**5
    else:
        lnlam = np.log(lam)
        return 0.09077334191+lnlam*(-0.00853080347+lnlam*(-0.0060953041+lnlam*( -0.00210631549374+lnlam*( -4.3851383e-05 + lnlam * 3.87e-4))))


    
#derived parameters
#photon density parameter
Omega_gamma = 2.4748e-5/(H0/100.)**2*(T_CMB/2.726)**4
#neutrino density parameter
eV_in_Kelvin = 11604.505
T_CNB = T_CMB*(4./11.)**(1./3.)
Omega_nu_reference = 0.56202e-5/(H0/100.)**2 * N_eff/len(mass_nu)/I_rho(0.) # this is, Omega_nu/I_rho(0) assuming massless
Omega_nu = np.array([I_rho(m*eV_in_Kelvin/T_CNB) * Omega_nu_reference for m in mass_nu ])
#dark energy density parameter
Omega_L = 1.0 - Omega_m - Omega_k - Omega_gamma - sum(Omega_nu)

## 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)))

##rho_nu/rho_critical as a function of z
def Omega_nu_scaled(z):
    s = 0.    
    for i in range(len(mass_nu)):
        s +=  I_rho(mass_nu[i]*eV_in_Kelvin/(1.+z)/T_CNB)
    return s*(1.+z)**4*Omega_nu_reference

#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_gamma*(1+z)**4 + Omega_nu_scaled(z))

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")