Adds ODEs for structure growth. Makes WCDM and LCDM initialization more uniform.

Project.toml

@@ -3,6 +3,7 @@ uuid = "76746363-e552-5dba-9a5a-cef6fa9cc5ab"
 version = "1.0.2"
 
 [deps]
+DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAstro = "6112ee07-acf9-5e0f-b108-d242c714bf9f"

src/Cosmology.jl

@@ -3,6 +3,8 @@ module Cosmology
 using QuadGK, Unitful
 import Unitful: km, s, Gyr
 using UnitfulAstro: Mpc, Gpc
+using DifferentialEquations
+
 
 export cosmology,
        age,
@@ -17,88 +19,90 @@ export cosmology,
        hubble_time,
        luminosity_dist,
        lookback_time,
-       scale_factor
+       scale_factor,
+       f_DE,
+       growth_factor,
+       σ_T, m_H, m_He, m_p
+
+
 
 abstract type AbstractCosmology end
 abstract type AbstractClosedCosmology <: AbstractCosmology end
 abstract type AbstractFlatCosmology <: AbstractCosmology end
 abstract type AbstractOpenCosmology <: AbstractCosmology end
 
-struct FlatLCDM{T <: Real} <: AbstractFlatCosmology
-    h::T
-    Ω_Λ::T
-    Ω_m::T
-    Ω_r::T
-end
-FlatLCDM(h::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real) =
-    FlatLCDM(promote(float(h), float(Ω_Λ), float(Ω_m), float(Ω_r))...)
-
-
-a2E(c::FlatLCDM, a) = sqrt(c.Ω_r + c.Ω_m * a + c.Ω_Λ * a^4)
-
-struct ClosedLCDM{T <: Real} <: AbstractClosedCosmology
-    h::T
-    Ω_k::T
-    Ω_Λ::T
-    Ω_m::T
-    Ω_r::T
-end
-ClosedLCDM(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real) =
-    ClosedLCDM(promote(float(h), float(Ω_k), float(Ω_Λ), float(Ω_m),
-                       float(Ω_r))...)
-
-
-struct OpenLCDM{T <: Real} <: AbstractOpenCosmology
-    h::T
-    Ω_k::T
-    Ω_Λ::T
-    Ω_m::T
-    Ω_r::T
-end
-OpenLCDM(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real) =
-    OpenLCDM(promote(float(h), float(Ω_k), float(Ω_Λ), float(Ω_m),
-                     float(Ω_r))...)
-
 
-function a2E(c::Union{ClosedLCDM,OpenLCDM}, a)
-    a2 = a * a
-    sqrt(c.Ω_r + c.Ω_m * a + (c.Ω_k + c.Ω_Λ * a2) * a2)
-end
-
-for c in ("Flat", "Open", "Closed")
-    name = Symbol("$(c)WCDM")
+for model in ("LCDM", "WCDM")
+    for curv in ("Flat", "Open", "Closed")
+        name = Symbol("$(curv)$(model)")
+        @eval begin
+            Base.@kwdef struct $(name){N <: Real} <: $(Symbol("Abstract$(curv)Cosmology"))
+                h::N = 0.67
+                Ω_k::N = 0
+                Ω_c::N = 0.3
+                Ω_b::N = 0
+                Neff::N = 3.04
+                T_cmb::N = 2.755
+                w0::N = -1
+                wa::N = 0
+
+                # Derived densities
+                Ω_γ::N = 4.48131e-7 * T_cmb^4 / h^2
+                Ω_ν::N = Neff * Ω_γ * (7 / 8) * (4 / 11)^(4 / 3)
+                Ω_r::N = Ω_γ + Ω_ν
+                Ω_m::N = Ω_c + Ω_b
+                Ω_Λ::N = 1. - Ω_m - Ω_r - Ω_k  
+
+
+            end
+            function $(name)(args...)
+                $(name)(promote(map(float, args)...)...)
+            end
+            function $(name)(;kwargs...)
+                $(name)(;promote(map(float, kwargs)...)...)
+            end
+
+        end
+    end
     @eval begin
-        struct $(name){T <: Real} <: $(Symbol("Abstract$(c)Cosmology"))
-            h::T
-            Ω_k::T
-            Ω_Λ::T
-            Ω_m::T
-            Ω_r::T
-            w0::T
-            wa::T
+        function $(Symbol("$model"))(;kwargs...)
+            kwargs = Dict(kwargs)
+            Ω_k = haskey(kwargs, :Ω_k) ? kwargs[:Ω_k] : 0.
+            if Ω_k < 0
+                $(Symbol("Closed$(model)"))(;kwargs...)
+            elseif Ω_k > 0
+                $(Symbol("Open$(model)"))(;kwargs...)
+            else
+                $(Symbol("Flat$(model)"))(;kwargs...)
+            end
         end
-        function $(name)(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real,
-                         w0::Real, wa::Real)
-            $(name)(promote(float(h), float(Ω_k), float(Ω_Λ), float(Ω_m),
-                            float(Ω_r), float(w0), float(wa))...)
+    end
+    @eval begin
+        function $(Symbol("$model"))(args...)
+            Ω_k = args[2]
+            if Ω_k < 0
+                $(Symbol("Closed$(model)"))(args...)
+            elseif Ω_k > 0
+                $(Symbol("Open$(model)"))(args...)
+            else
+                $(Symbol("Flat$(model)"))(args...)
+            end
         end
     end
+
 end
 
-function WCDM(h::Real, Ω_k::Real, Ω_Λ::Real, Ω_m::Real, Ω_r::Real, w0::Real, wa::Real)
-    if Ω_k < 0
-        ClosedWCDM(h, Ω_k, Ω_Λ, Ω_m, Ω_r, w0, wa)
-    elseif Ω_k > 0
-        OpenWCDM(h, Ω_k, Ω_Λ, Ω_m, Ω_r, w0, wa)
-    else
-        FlatWCDM(h, Ω_k, Ω_Λ, Ω_m, Ω_r, w0, wa)
-    end
+function a2E(c::Union{FlatLCDM,ClosedLCDM,OpenLCDM}, a)
+    a2 = a * a
+    sqrt(c.Ω_r + c.Ω_m * a + (c.Ω_k + c.Ω_Λ * a2) * a2)
 end
+f_DE(c::Union{FlatLCDM,ClosedLCDM,OpenLCDM}, a) = 0
 
 function a2E(c::Union{FlatWCDM,ClosedWCDM,OpenWCDM}, a)
     ade = exp((1 - 3 * (c.w0 + c.wa)) * log(a) + 3 * c.wa * (a - 1))
     sqrt(c.Ω_r + (c.Ω_m + c.Ω_k * a) * a + c.Ω_Λ * ade)
 end
+f_DE(c::Union{FlatWCDM,ClosedWCDM,OpenWCDM}, a) = (-3 * (1 + c.w0) + 3 * c.wa * ((a - 1) / log(a - 1e-5) - 1))
 
 """
     cosmology(;h = 0.69,
@@ -136,31 +140,32 @@ Cosmology.ClosedWCDM{Float64}(0.69, -0.1, 0.8099122024007929, 0.29, 8.7797599207
 function cosmology(;h = 0.69,
                    Neff = 3.04,
                    OmegaK = 0,
-                   OmegaM = 0.29,
+                   OmegaC = 0.25,
+                   OmegaB = 0.05,
                    OmegaR = nothing,
+                   OmegaM = nothing,
                    Tcmb = 2.7255,
                    w0 = -1,
-    wa = 0)
+                   wa = 0)
+
+    if !(OmegaR === nothing)
+        Tcmb = 0.
+        Neff = 0.
+    end
 
-    if OmegaR === nothing
-        OmegaG = 4.48131e-7 * Tcmb^4 / h^2
-        OmegaN = Neff * OmegaG * (7 / 8) * (4 / 11)^(4 / 3)
-        OmegaR = OmegaG + OmegaN
+    if !(OmegaM === nothing)
+        OmegaC = OmegaM
+        OmegaB = 0.
     end
 
-    OmegaL = 1 - OmegaK - OmegaM - OmegaR
 
     if !(w0 == -1 && wa == 0)
-        return WCDM(h, OmegaK, OmegaL, OmegaM, OmegaR, w0, wa)
+        return WCDM(;h = h, Ω_k = OmegaK, Ω_c = OmegaC, Ω_b = OmegaB, Neff = Neff, T_cmb = Tcmb, w0 = w0, wa = wa)
+    else
+        return LCDM(;h = h, Ω_k = OmegaK, Ω_c = OmegaC, Ω_b = OmegaB, Neff = Neff, T_cmb = Tcmb, w0 = w0, wa = wa)
     end
 
-    if OmegaK < 0
-        return ClosedLCDM(h, OmegaK, OmegaL, OmegaM, OmegaR)
-        elseif OmegaK > 0
-        return OpenLCDM(h, OmegaK, OmegaL, OmegaM, OmegaR)
-        else
-        return FlatLCDM(h, OmegaL, OmegaM, OmegaR)
-    end
+
 end
 
 # hubble rate
@@ -287,4 +292,71 @@ for f in (:hubble_dist0, :hubble_dist, :hubble_time0, :hubble_time, :comoving_ra
     @eval $f(u::Unitful.Unitlike, args...; kws...) = uconvert(u, $f(args...; kws...))
 end
 
+# Density evolution
+
+OmegaM(c::AbstractCosmology, z) = c.Ω_m * (1 + z)^3 / E(c, z)^2
+OmegaDE(c::AbstractCosmology, z) = c.Ω_Λ * scale_factor(z)^f_DE(c, scale_factor(z)) / E(c, z)^2
+
+
+
+
+
+
+# Growth of structure
+
+function growth_derivatives!(du, u, c, a)
+
+    # Ex 1.118 in Leclerq's thesis shows the equation
+    # satisfied by the second order growth factor
+
+    D_1, D_1′, D_2, D_2′ = u
+    z = 1 / a - 1
+    Omega_m = OmegaM(c, z)
+    Omega_de = OmegaDE(c, z)
+
+    D_1′′ = 1.5 * Omega_m * D_1 / a^2 - (Omega_de - 0.5 * Omega_m + 2) * D_1′ / a
+    D_2′′ = 1.5 * Omega_m * (D_2 - D_1^2) / a^2 - (Omega_de - 0.5 * Omega_m + 2) * D_2′ / a
+
+    du[1] = D_1′
+    du[2] = D_1′′
+    du[3] = D_2′
+    du[4] = D_2′′
+    return du
+
+end
+
+function growth_factor(c::AbstractCosmology, a::Vector{<:Real})
+    """
+    Returns growth factors and rates at scale factor a
+    Returns: D1(a), D1′(a), D1''(a), D2(a), D2′(a), D2''(a)
+    """
+    a0 = 1e-2
+    aspan = (a0,1.)
+    a_save = a
+    push!(a_save, 1.)
+    D1_in = a0 # EdS conditions
+    D1′_in = 1.
+    D2_in = -3. * D1_in^2 / 7 # Approx eq. 1.119 at tau = 0, Eds implies Om = 1
+    D2′_in = -6* D1_in / 7
+    u0 = [D1_in, D1′_in, D2_in, D2′_in]
+    prob = ODEProblem(growth_derivatives!,u0,aspan,c)
+    sol = solve(prob, saveat = a_save)
+    D1 = sol[1,:] ./ sol[1,end]
+    D1′ = sol[2,:] ./ sol[1,end]
+    D2 = sol[3,:] ./ sol[3,end]
+    D2′ = sol[4,:] ./ sol[1,end]
+    du = similar(u0)
+    D1′′ = similar(a_save)
+    D2′′ = similar(a_save)
+    for i in eachindex(a_save)
+        growth_derivatives!(du, sol[:,i], c, a_save[i])
+        D1′′[i] = du[2] / sol[1,end]
+        D2′′[i] = du[4] / sol[3,end]
+    end
+    return D1, D1′, D1′′, D2, D2′, D2′′
+end
+
+
+
 end # module
+

test/runtests.jl

@@ -9,6 +9,7 @@ age_rtol = 2e-4
 # strip the units away from the integrand function
 integrand(c, z) = 4pi*ustrip(comoving_volume_element(c, z))
 
+
 @testset "FlatLCDM" begin
     c = cosmology(h=0.7, OmegaM=0.3, OmegaR=0)
     @test angular_diameter_dist(c,1,rtol=dist_rtol) ≈ 1651.9145u"Mpc" rtol = dist_rtol
@@ -23,6 +24,15 @@ integrand(c, z) = 4pi*ustrip(comoving_volume_element(c, z))
     @test age(c,1,rtol=age_rtol) ≈ 5.7527u"Gyr" rtol = age_rtol
     @test lookback_time(c,1,rtol=age_rtol) ≈ (13.4694-5.7527)u"Gyr" rtol = age_rtol
     @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
+    z = 0.5
+    a = Cosmology.scale_factor(z)
+    @test Cosmology.E(c, z) ≈ sqrt(c.Ω_r * a^-4 + c.Ω_m * a^-3 + c.Ω_Λ * a^f_DE(c, a))
+    growth_res = growth_factor(c,[a])
+    @test growth_res[1][1] ≈ 0.7729844 rtol = 1e-3
+    @test growth_res[2][1] ≈ 0.86867464 rtol = 1e-3
+    @test growth_res[3][1] ≈ -1.2082956 rtol = 1e-1
+
+
 end
 
 @testset "OpenLCDM" begin
@@ -39,6 +49,9 @@ end
     @test age(c,1,rtol=age_rtol) ≈ 5.5466u"Gyr" rtol = age_rtol
     @test lookback_time(c,1,rtol=age_rtol) ≈ (13.064-5.5466)u"Gyr" rtol = age_rtol
     @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
+    z = 0.5
+    a = Cosmology.scale_factor(z)
+    @test Cosmology.E(c, z) ≈ sqrt(c.Ω_k * a^-2 + c.Ω_r * a^-4 + c.Ω_m * a^-3 + c.Ω_Λ * a^f_DE(c, a))
 end
 
 @testset "ClosedLCDM" begin
@@ -55,6 +68,9 @@ end
     @test age(c,1,rtol=age_rtol) ≈ 5.9868u"Gyr" rtol = age_rtol
     @test lookback_time(c,1,rtol=age_rtol) ≈ (13.925-5.9868)u"Gyr" rtol = age_rtol
     @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
+    z = 0.5
+    a = Cosmology.scale_factor(z)
+    @test Cosmology.E(c, z) ≈ sqrt(c.Ω_k * a^-2 + c.Ω_r * a^-4 + c.Ω_m * a^-3 + c.Ω_Λ * a^f_DE(c, a))
 end
 
 @testset "FlatWCDM" begin
@@ -71,6 +87,9 @@ end
     @test age(c,1,rtol=age_rtol) ≈ 5.6464u"Gyr" rtol = age_rtol
     @test lookback_time(c,1,rtol=age_rtol) ≈ (13.1915-5.6464)u"Gyr" rtol = age_rtol
     @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
+    z = 0.5
+    a = Cosmology.scale_factor(z)
+    @test Cosmology.E(c, z) ≈ sqrt(c.Ω_k * a^-2 + c.Ω_r * a^-4 + c.Ω_m * a^-3 + c.Ω_Λ * a^f_DE(c, a)) rtol = 1e-4
 end
 
 @testset "OpenWCDM" begin
@@ -87,6 +106,9 @@ end
     @test age(c,1,rtol=age_rtol) ≈ 5.4659u"Gyr" rtol = age_rtol
     @test lookback_time(c,1,rtol=age_rtol) ≈ (12.8488-5.4659)u"Gyr" rtol = age_rtol
     @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
+    z = 0.5
+    a = Cosmology.scale_factor(z)
+    @test Cosmology.E(c, z) ≈ sqrt(c.Ω_k * a^-2 + c.Ω_r * a^-4 + c.Ω_m * a^-3 + c.Ω_Λ * a^f_DE(c, a)) rtol = 1e-4
 end
 
 @testset "ClosedWCDM" begin
@@ -103,6 +125,9 @@ end
     @test age(c,1,rtol=age_rtol) ≈ 5.8482u"Gyr" rtol = age_rtol
     @test lookback_time(c,1,rtol=age_rtol) ≈ (13.5702-5.8482)u"Gyr" rtol = age_rtol
     @test age(c, 1) + lookback_time(c, 1) ≈ age(c, 0)
+    z = 0.5
+    a = Cosmology.scale_factor(z)
+    @test Cosmology.E(c, z) ≈ sqrt(c.Ω_k * a^-2 + c.Ω_r * a^-4 + c.Ω_m * a^-3 + c.Ω_Λ * a^f_DE(c, a)) rtol = 1e-4
 end
 
 @testset "Non-Float64" begin
@@ -134,7 +159,12 @@ end
 
 @testset "Utilities" begin
     c = cosmology(h = 0.7)
+    cw = cosmology(h = 0.7, w0 = -1, wa = 0.1)
     @test hubble_time(c, 0) ≈ Cosmology.hubble_time0(c)
     @test hubble_dist(c, 0) ≈ Cosmology.hubble_dist0(c)
     @test H(c, 0) ≈ 70u"km/s/Mpc"
+
 end
+
+
+