Extending models

This tutorial shows how to create extended (beyond-ΛCDM) models. As a simple example, we replace the cosmological constant with equation of state $w(τ) = -1$ by w₀wₐ dark energy with equation of state

\[w(τ) = \frac{P(τ)}{ρ(τ)} = w_0 + w_a (1 - a(τ)),\]

turning the ΛCDM model into the w₀wₐCDM model, as suggested by Chevallier, Polarski (2000) and Linder (2002).

SymBoltz represents each component of the Einstein-Boltzmann equations as a System. These are effectively incomplete "chunks" of logically related parameters, variables, equations and initial conditions that are composed together into a complete cosmological model. A cosmological model can be modified by adding or replacing components.

1. Create the reference model

First, we create the reference model that is to be extended. We will modify the standard ΛCDM model:

using SymBoltz
M1 = ΛCDM()
hierarchy(M1; describe = true)

ΛCDM: Standard cosmological constant and cold dark matter cosmological model
├─ g: Spacetime FLRW metric in Newtonian gauge
├─ G: General relativity gravity
├─ γ: Photon radiation
├─ ν: Massless neutrinos
├─ c: Cold dark matter
├─ b: Baryonic matter
│  ├─ rec: Baryon-photon recombination thermodynamics (RECFAST)
│  ├─ rei1: Reionization with tanh-like (activation function) contribution to the free electron fraction
│  └─ rei2: Reionization with tanh-like (activation function) contribution to the free electron fraction
├─ h: Massive neutrino
├─ Λ: Cosmological constant
└─ I: Harrison-Zel'dovich inflation

Everything belonging to the cosmological constant species is stored in the component M1.Λ:

equations(M1.Λ)

\[ \begin{align*} \delta\left( \tau, k \right) &= 0 \\ \theta\left( \tau, k \right) &= 0 \\ \sigma\left( \tau, k \right) &= 0 \\ w\left( \tau \right) &= -1 \\ P\left( \tau \right) &= w\left( \tau \right) \rho\left( \tau \right) \\ \rho\left( \tau \right) &= 0.11937 \left( a\left( \tau \right) \right)^{ - 3 \left( 1 + w\left( \tau \right) \right)} \mathtt{\Omega{_0}} \\ \Omega\left( \tau \right) &= 8.3776 \rho\left( \tau \right) \\ \mathtt{c_s^2}\left( \tau \right) &= w\left( \tau \right) \end{align*} \]

This is what we will replace.

2. Create the extended component

The background continuity equation $ρ̇(τ) = -3 H(τ) \, ρ(τ) \, (1 + w(τ))$ can be solved analytically with the w₀wₐ equation of state and the ansatz

\[ρ(τ) = ρ(τ₀) \, a(τ)ᵐ \exp(n (1 - a(τ))),\]

giving $m = -3 (1 + w₀ + wₐ)$ and $n = -3 wₐ$. The perturbation equations are

\[\begin{aligned} δ̇ &= -(1 + w) (θ - 3 Φ̇) - 3 \frac{ȧ}{a} (c_s^2 - w) δ, \\ θ̇ &= -\frac{ȧ}{a} (1 - 3w) θ - \frac{ẇ}{1+w} θ + \frac{c_s^2}{1+w} k^2 δ - k^2 σ + k^2 Ψ, \end{aligned}\]

with (adiabatic) initial conditions $δ = -\frac{3}{2} (1+w) Ψ$ and $θ = \frac{1}{2} k^2 τ Ψ$, following Bertschinger and Ma (equation 30).

Next, we must simply pack this into a symbolic component that represents the w₀wₐ dark energy species:

g, τ, k = M1.g, M1.τ, M1.k # reuse metric of original model
D = Differential(τ)

# 1. Create parameters (will resurface as tunable numbers in the full model)
pars = @parameters w₀ wₐ cₛ² Ω₀

# 2. Create variables
vars = @variables ρ(τ) P(τ) w(τ) δ(τ, k) θ(τ, k) σ(τ, k)

# 3. Specify equations (~ means equality in ModelingToolkit)
eqs = [
    # Background equations
    w ~ w₀ + wₐ * (1 - g.a) # equation of state
    ρ ~ 3/(8*Num(π))*Ω₀ * g.a^(-3*(1+w₀+wₐ)) # D(ρ) ~ -3 * g.ℰ * ρ * (1 + w) # energy density (ρ₀ => 3/(8*Num(π)) * exp(+3*wₐ) * Ω₀)
    P ~ w * ρ # pressure

    # Perturbation equations
    D(δ) ~ -(1 + w) * (θ - 3*g.Φ) - 3 * g.ℰ * (cₛ² - w) * δ # energy overdensity
    D(θ) ~ -g.ℰ * (1 - 3*w) * θ - D(w) / (1 + w) * θ + cₛ² / (1 + w) * k^2 * δ - k^2 * σ + k^2 * g.Ψ # momentum
    σ ~ 0 # shear stress
]

# 4. Specify initial conditions (for perturbations)
initialization_eqs = [
    δ ~ -3/2 * (1+w) * g.Ψ
    θ ~ 1/2 * (k^2/g.ℰ) * g.Ψ
]

# 5. Pack into an ODE system called "X"
description = "w₀wₐ (CPL) dynamical dark energy"
@named X = System(eqs, τ, vars, pars; initialization_eqs, description)

\[ \begin{align*} w\left( \tau \right) &= \mathtt{w{_0}} + \mathtt{w{_a}} \left( 1 - a\left( \tau \right) \right) \\ \rho\left( \tau \right) &= \frac{3 \left( a\left( \tau \right) \right)^{ - 3 \left( 1 + \mathtt{w_0} + \mathtt{w_a} \right)} \mathtt{\Omega_0}}{8 \pi} \\ P\left( \tau \right) &= w\left( \tau \right) \rho\left( \tau \right) \\ \frac{\mathrm{d}}{\mathrm{d}\tau} \delta\left( \tau, k \right) &= \left( - 3 \Phi\left( \tau, k \right) + \theta\left( \tau, k \right) \right) \left( -1 - w\left( \tau \right) \right) - 3 \left( \mathtt{c_s^2} - w\left( \tau \right) \right) \mathscr{E}\left( \tau \right) \delta\left( \tau, k \right) \\ \frac{\mathrm{d}}{\mathrm{d}\tau} \theta\left( \tau, k \right) &= \frac{k^{2} \mathtt{c_s^2} \delta\left( \tau, k \right)}{1 + w\left( \tau \right)} + \frac{ - \frac{\mathrm{d} w\left( \tau \right)}{\mathrm{d}\tau} \theta\left( \tau, k \right)}{1 + w\left( \tau \right)} + k^{2} \Psi\left( \tau, k \right) - k^{2} \sigma\left( \tau, k \right) - \left( 1 - 3 w\left( \tau \right) \right) \mathscr{E}\left( \tau \right) \theta\left( \tau, k \right) \\ \sigma\left( \tau, k \right) &= 0 \end{align*} \]

Note that the w₀wₐ component only knows about itself (and the metric), but is completely unaware of the theory of gravity, other species and other components. Its "job" is only to expose the variables like ρ, P, δ and σ that source the Einstein equations. This connection is made when the component is used to create a full cosmological model, as we will do next.

3. Create the extended model

Finally, we create a new ΛCDM model, but replace Λ by X, and call it the w0waCDM model:

M2 = ΛCDM(Λ = X, name = :w0waCDM)
hierarchy(M2; describe = true)

w0waCDM: Standard cosmological constant and cold dark matter cosmological model
├─ g: Spacetime FLRW metric in Newtonian gauge
├─ G: General relativity gravity
├─ γ: Photon radiation
├─ ν: Massless neutrinos
├─ c: Cold dark matter
├─ b: Baryonic matter
│  ├─ rec: Baryon-photon recombination thermodynamics (RECFAST)
│  ├─ rei1: Reionization with tanh-like (activation function) contribution to the free electron fraction
│  └─ rei2: Reionization with tanh-like (activation function) contribution to the free electron fraction
├─ h: Massive neutrino
├─ X: w₀wₐ (CPL) dynamical dark energy
└─ I: Harrison-Zel'dovich inflation

Now M2.Λ no longer exists, but M2.X contains our new dark energy species:

equations(M2.X)

Solve and compare the models

To test, let us set some parameters and solve both models with one perturbation mode. For the ΛCDM model:

θ1 = SymBoltz.parameters_Planck18(M1)
prob1 = CosmologyProblem(M1, θ1)
ks = 1.0
sol1 = solve(prob1, ks)

Cosmology solution for model ΛCDM
Stages:
  Background: return code Terminated; solved with Rodas4P; 1305 points
  Perturbations: return codes Success; solved with KenCarp4; 85-85 points; x1 k ∈ [1.0, 1.0] H₀/c (interpolation in-between)

And for the w₀wₐCDM model:

θ2 = merge(θ1, Dict(
    M2.X.w₀ => -0.9,
    M2.X.wₐ => 0.2,
    M2.X.cₛ² => 1.0
))
prob2 = CosmologyProblem(M2, θ2)
sol2 = solve(prob2, ks)

Cosmology solution for model w0waCDM
Stages:
  Background: return code Terminated; solved with Rodas4P; 1309 points
  Perturbations: return codes Success; solved with KenCarp4; 82-82 points; x1 k ∈ [1.0, 1.0] H₀/c (interpolation in-between)

Let us compare $H(τ)$ and $Ψ(k,τ)$ at equal scale factors $a(τ)$:

lgas = range(-3, 0, length=500) # log10(a)
H1s = sol1(M1.g.H, log10(M1.g.a) => lgas)
H2s = sol2(M2.g.H, log10(M2.g.a) => lgas)
Ψ1s = sol1(M1.g.Ψ, log10(M1.g.a) => lgas, ks)
Ψ2s = sol2(M2.g.Ψ, log10(M2.g.a) => lgas, ks)

using Plots
plot(lgas, H2s ./ H1s; xlabel = "lg(a)", label = "H₂ / H₁")
plot!(lgas, Ψ1s ./ Ψ2s; xlabel = "lg(a)", label = "Ψ₂ / Ψ₁")