Halofit implementation

Project.toml

@@ -24,7 +24,10 @@ OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
+Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
+Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 ThreadPools = "b189fb0b-2eb5-4ed4-bc0c-d34c51242431"

scripts/halofit.jl

@@ -0,0 +1,35 @@
+using Interpolations
+using Bolt
+using Plots
+
+𝕡 = CosmoParams(Ω_c = 0.3) # set kwargs like so to change the default values
+# Compute expansion history quantities
+bg = Background(𝕡)
+# Compute ionization history (via RECFAST)
+𝕣 = Bolt.RECFAST(bg=bg, Yp=𝕡.Y_p, OmegaB=𝕡.Ω_b)
+ih = IonizationHistory(𝕣, 𝕡, bg)
+
+# Matter power spectrum
+kmin,kmax,nk = 10bg.H₀, 20000bg.H₀, 128
+ks = log10_k(kmin,kmax,nk) # k grid
+n_q,ℓᵧ,ℓ_ν,ℓ_mν,x=15,10,10,10,0
+@time pL = [plin(k,𝕡,bg,ih,n_q,ℓᵧ,ℓ_ν,ℓ_mν,x) for k in ks]
+p1 = plot(ks, vcat(pL...), xscale=:log10, yscale=:log10, label = "Linear")
+
+InterpPmm = Interpolations.interpolate( log10.(vcat(pL...)), BSpline(Cubic(Line(OnGrid()))))
+
+x = LinRange(log10(first(ks)), log10(last(ks)), length(ks))
+InterpPmm = scale(InterpPmm, x)
+InterpPmm = Interpolations.extrapolate(InterpPmm, Line())
+
+params_halofit = Bolt.halofit_params(InterpPmm)
+
+pNL_takahashi2012 =  [Bolt.takahashi2012(InterpPmm, params_halofit, 𝕡.Ω_b+𝕡.Ω_c, k) for k in ks]
+pNL_takabird =  [Bolt.takabird(InterpPmm, params_halofit, 𝕡.Ω_b+𝕡.Ω_c, 𝕡.Ω_r, 𝕡.Σm_ν, 𝕡.h, k) for k in ks]
+n_q,ℓᵧ,ℓ_ν,ℓ_mν,x=15,10,10,10,0
+pL_check, pNL_check = Bolt.plin_and_nonlin(ks,𝕡,bg,ih,n_q,ℓᵧ,ℓ_ν,ℓ_mν,x)
+
+plot!(p1, ks, vcat(pNL_takahashi2012...), label = "Takahashi2012")
+plot!(p1, ks, vcat(pNL_takabird...), label = "Takabird")
+
+println(pNL_check == vcat(pNL_takabird...))

src/Bolt.jl

@@ -25,6 +25,7 @@ using StaticArrays
 using DoubleFloats
 using MuladdMacro
 using LinearAlgebra
+using Roots
 
 
 using FFTW
@@ -76,5 +77,6 @@ include("ionization/ionization.jl")
 include("ionization/recfast.jl")
 include("perturbations.jl")
 include("spectra.jl")
+include("nonlinear.jl")
 
 end

src/nonlinear.jl

@@ -0,0 +1,173 @@
+function halofit_params(pk)
+    f(x) = √_S(pk, x) - 1
+    R = find_zero(f, (1e-2, 1e2), Bisection())
+    S_val = _S(pk, R)
+    #Although S should be one, I prefer to evaluate it to be consistent with the root
+    #finding algo result. S is a callable, while S_val is a fixed value
+    dlnSdlnR = R / S_val * _dSdR(pk, R)
+    d²lnSdlnR² = dlnSdlnR - dlnSdlnR^2 + (R^2 / S_val) * _d²SdR²(pk, R)
+    kσ = 1 / R
+    neff = -3 - dlnSdlnR
+    C = -d²lnSdlnR²
+    return [kσ, neff, C]
+end
+
+function _an(neff, C)
+    return 10^(1.5222 + 2.8553neff + 2.3706 * neff^2 + 0.9903 * neff^3 + 0.2250 * neff^4 -
+    0.6038C) #not added the w₀wₐ term, which is 0.1749ΩDEz(1 + wz)
+end
+
+function _bn(neff, C)
+    return 10^(-0.5642 + 0.5864neff + 0.5716 * neff^2 - 1.5474C)
+    # as above not added the w₀wₐ term, +0.2279*ΩDEz*(1.+w))
+end
+
+function _cn(neff, C)
+    return 10^(0.3698 + 2.0404neff + 0.8161 * neff^2 + 0.5869C)
+end
+
+function _γn(neff, C)
+    return 0.1971 - 0.0843neff + 0.8460C
+end
+
+function _αn(neff, C)
+    return abs(6.0835 + 1.3373neff - 0.1959 * neff^2 - 5.5274C)
+end
+
+function _βn(neff, C; fνz=0)
+    return 2.0379 - 0.7354neff + 0.3157 * neff^2 + 1.2490 * neff^3 +
+    0.3980 * neff^4 - 0.1682C + fνz * (1.081 + 0.395 * neff^2)
+end
+
+function _νn(neff)
+    return 10^(5.2105 + 3.6902neff)
+end
+
+function _n_coefficients(neff, C, fνz)
+    an = _an(neff, C)
+    bn = _bn(neff, C)
+    cn = _cn(neff, C)
+    γn = _γn(neff, C)
+    αn = _αn(neff, C)
+    βn = _βn(neff, C, fνz = fνz)
+    μn = 0
+    νn = _νn(neff)
+    return [an, bn, cn, γn, αn, βn, μn, νn]
+end
+
+function _f1a(Ωmz)
+    return Ωmz^-0.0732
+end
+
+function _f2a(Ωmz)
+    return Ωmz^-0.1423
+end
+
+function _f3a(Ωmz)
+    return Ωmz^0.0725
+end
+
+function _f1b(Ωmz)
+    return Ωmz^-0.0307
+end
+
+function _f2b(Ωmz)
+    return Ωmz^-0.0585
+end
+
+function _f3b(Ωmz)
+    return Ωmz^0.0743
+end
+
+function _f_coefficients(Ωmz, frac)
+    f1a = _f1a(Ωmz)
+    f2a = _f2a(Ωmz)
+    f3a = _f3a(Ωmz)
+    f1b = _f1b(Ωmz)
+    f2b = _f2b(Ωmz)
+    f3b = _f3b(Ωmz)
+    f1 = frac*f1b+(1-frac)*f1a
+    f2 = frac*f2b+(1-frac)*f2a
+    f3 = frac*f3b+(1-frac)*f3a
+    return [f1, f2, f3]
+end
+
+function _Δ²H(an, bn, cn, γn, μn, νn, f1, f2, f3, y)
+    return (an * y^(3 * f1)/(1 + bn * y^f2 + (cn * f3 * y)^(3-γn)))/(1 + μn / y + νn / y^2)
+    #TODO: μn is fixed to zero. Maybe we can remove it?
+end
+
+function _Δ²H(an, bn, cn, γn, μn, νn, f1, f2, f3, y, fνz)
+    return _Δ²H(an, bn, cn, γn, μn, νn, f1, f2, f3, y)*(1+fνz*0.977)
+end
+
+function _Δ²L(pk, k)
+    return k^3 / (2 * π^2) * 10^pk(log10(k))
+end
+
+function _Δ²̃L(pk, k, fνz, hz)
+    return _Δ²L(pk, k) * (1+fνz*47.48* (k/hz)^2/(1+1.5*(k/hz)^2))
+end
+
+function _Δ²Q(Δ²L, αn, βn, y)
+    return Δ²L * (1 + Δ²L)^βn / (1 + αn * Δ²L) * exp(-y / 4 - y^2 / 8)
+end
+
+function takahashi2012(pk, p::Array{T,1}, Ωmz::Real, k::Real) where {T<:Real}
+    #based on Takahashi 2012
+    kσ, neff, C = p
+    an = _an(neff, C)
+    bn = _bn(neff, C)
+    cn = _cn(neff, C)
+    γn = _γn(neff, C)
+    αn = _αn(neff, C)
+    βn = _βn(neff, C)
+    μn = 0
+    νn = _νn(neff)
+    f1 = _f1b(Ωmz)
+    f2 = _f2b(Ωmz)
+    f3 = _f3b(Ωmz)
+    y = k / kσ
+    Δ²H = _Δ²H(an, bn, cn, γn, μn, νn, f1, f2, f3, y)
+    Δ²L = _Δ²L(pk, k)
+    Δ²Q = _Δ²Q(Δ²L, αn, βn, y)
+    pk_nl = (Δ²Q + Δ²H) * 2 * π^2 / k^3
+    return pk_nl
+end
+
+function takabird(pk, p::Array{T,1}, Ωmz, Ωrz, Mν, hz, k) where {T<:Real}
+    #based on the revised Takahashi2012 present on the Class repository
+    Ωνz = Mν/(93.14 * hz^2)
+    fνz = Ωνz/Ωmz
+    ΩDEz = 1-Ωmz-Ωνz-Ωrz
+    #TODO: ask if Bolt Ω_r includes neutrinos or only radiation
+    frac = ΩDEz/(1-Ωmz)
+    # TODO: check if is required a z dependent fν, Ων etc
+    kσ, neff, C = p
+    an, bn, cn, γn, αn, βn, μn, νn = _n_coefficients(neff, C, fνz)
+    f1, f2, f3 = _f_coefficients(Ωmz, frac)
+    y = k / kσ
+    Δ²H = _Δ²H(an, bn, cn, γn, μn, νn, f1, f2, f3, y, fνz)
+    Δ²̃L = _Δ²̃L(pk, k, fνz, hz)
+    Δ²Q = _Δ²Q(Δ²̃L, αn, βn, y)
+    pk_nl = (Δ²Q + Δ²H) * 2 * π^2 / k^3
+    return pk_nl
+end
+
+function takabird(pL, ks, 𝕡)
+    InterpPmm = Interpolations.interpolate( log10.(vcat(pL...)),
+    BSpline(Cubic(Line(OnGrid()))))
+    x = LinRange(log10(first(ks)), log10(last(ks)), length(ks))
+    InterpPmm = scale(InterpPmm, x)
+    InterpPmm = Interpolations.extrapolate(InterpPmm, Line())
+    params_halofit = halofit_params(InterpPmm)
+    pNL_takabird = [takabird(InterpPmm, params_halofit, 𝕡.Ω_b+𝕡.Ω_c, 𝕡.Ω_r, 𝕡.Σm_ν, 𝕡.h, k)
+                    for k in ks]
+    return pNL_takabird
+end
+
+function plin_and_nonlin(ks,𝕡,bg,ih,n_q,ℓᵧ,ℓ_ν,ℓ_mν,x)
+    pL = [plin(k,𝕡,bg,ih,n_q,ℓᵧ,ℓ_ν,ℓ_mν,x) for k in ks]
+    pNL = takabird(pL, ks, 𝕡)
+    return vcat(pL...), vcat(pNL...)
+end

src/util.jl

@@ -106,3 +106,27 @@ function ldiv!(Y, pl::FFTLogPlan, A)
     pl.ifftplan! * Y
     Y .*= (pl.r).^(pl.q)
 end
+
+function _S(pk, R)
+    #assuming that pk has been interpolated in log-log space
+    dSdk(k) = exp(-(k * R)^2) * 10^pk(log10(k)) * k^2 / (2 * π^2)
+    res, _ = quadgk(dSdk, 0, 10 / R, rtol=1e-6)
+    return res
+end
+
+function _dSdRdk(k, R)
+    x = k * R
+    return -2R * exp(-x^2) * 10^pk(log10(k)) * k^4 / (2π^2)
+end
+
+function _dSdR(pk, R)
+    dSdRdk(k) = -2R * exp(-(k * R)^2) * 10^pk(log10(k)) * k^4 / (2π^2)
+    res, _ = quadgk(dSdRdk, 0, 10 / R)
+    return res
+ end
+
+ function _d²SdR²(pk, R)
+    d²σ²dR²dk(k)= (-2 + 4 * (k * R)^2) * exp(-(k * R)^2) * 10^pk(log10(k)) * k^4 / (2π^2)
+    res, _ = quadgk(d²σ²dR²dk, 0, 10 / R)
+    return res
+ end