Merge branch 'main' of https://github.com/hersle/SymBoltz.jl

docs/src/benchmarks.md

@@ -13,13 +13,18 @@ ks = 10 .^ range(-5, 1, length=N) / u"Mpc"
 # TODO: make proper; sort; test accuracy; work-precision diagram?
 # TODO: test different linsolve and nlsolve
 # TODO: use BenchmarkTools.BenchmarkGroup
-solvers = [TRBDF2(), AutoTsit5(TRBDF2()), KenCarp4(), AutoTsit5(KenCarp4()), KenCarp47(), Kvaerno5(), Rodas5P(), QNDF()] # TODO: Rodas5?
+solvers = [TRBDF2(), AutoTsit5(TRBDF2()), KenCarp4(), AutoTsit5(KenCarp4()), KenCarp47(), Kvaerno5(), Rodas5P(), Rodas5(), Rodas4(), Rodas4P(), QNDF()]
 timings = []
 solve_with(solver; reltol = 1e-5) = solve(M, pars, ks; solver, reltol, thread=false)
 for solver in solvers
     timing = @benchmark solve_with($solver)
     push!(timings, timing)
 end
+
+# Sort by efficiency
+idxs = sortperm(map(t -> mean(t.times), timings))
+solvers, timings = solvers[idxs], timings[idxs]
+
 bar(
     SymBoltz.solvername.(solvers),
     map(t -> mean(t.times/N)/1e6, timings),

docs/src/getting_started.md

@@ -57,7 +57,7 @@ To solve only the background, you can simply omit the `ks` argument: `solve(prob
 
 ## 3. Access the solution
 
-You are now free to [do whatever you want with the solution object](@ref "Solution interface").
+You are now free to [do whatever you want with the solution object](@ref "Solution handling").
 For example, to get the reduced Hubble function $E(t) = H(t) / H_0$ for 300 log-spaced conformal times:
 ```@example getting_started
 ts = exp.(range(log.(extrema(sol[M.t]))..., length=300))
@@ -95,7 +95,7 @@ plot!(p[2], sol, log10(M.g.a), [M.b.w, M.c.w, M.γ.w, M.ν.w, M.h.w, M.Λ.w])
 plot!(p[3], sol, log10(M.g.a), log10(M.g.H / M.g.H₀))
 plot!(p[4], sol, log10(M.g.a), [M.b.rec.XHe⁺⁺, M.b.rec.XHe⁺, M.b.rec.XH⁺, M.b.rec.Xe])
 plot!(p[5], sol, log10(M.g.a), log10.([M.b.rec.Tγ, M.b.rec.Tb] ./ M.γ.T₀))
-plot!(p[6], sol, ks[1], log10(M.g.a), log10(abs(M.b.rec.τ)), klabel = false)
+plot!(p[6], sol, log10(M.g.a), log10(abs(M.b.rec.τ)))
 plot!(p[7], sol, ks_plot, log10(M.g.a), [M.g.Φ, M.g.Ψ])
 plot!(p[8], sol, ks_plot, log10(M.g.a), log10.(abs.([M.b.δ, M.c.δ, M.γ.δ, M.ν.δ, M.h.δ])), klabel = false)
 plot!(p[9], sol, ks_plot, log10(M.g.a), log10.(abs.([M.b.θ, M.c.θ, M.γ.θ, M.ν.θ])), klabel = false)

docs/src/solution.md

@@ -41,13 +41,26 @@ as = sol(ts, M.g.a) # scale factors
 nothing # hide
 ```
 
+## Shooting method parametrization
+
+It is common to parametrize some models not by initial conditions or constant parameters, but by values of variables at some (non-initial) time, like today.
+This is exactly like the boundary conditions of a boundary value problem.
+SymBoltz.jl supports such parametrizations with the shooting method:
+
+```@docs
+shoot
+```
+
+Example usage is shown on the [models page](@ref "Models").
+
 ## Choice of solver
 
 In principle, models can be solved with any [DifferentialEquations.jl ODE solver](https://docs.sciml.ai/DiffEqDocs/stable/solvers/ode_solve/).
 But most cosmological models have very stiff Einstein-Boltzmann equations that can only be solved by implicit solvers, while explicit solvers usually fail.
 For the [standard ΛCDM model](@ref "Standard ΛCDM"), some good solvers are:
 
-- `Rodas5P`: Slow for large systems. Very accurate. Handles severe stiffness. Default background solver.
+- `Rodas4P`: Slow for large systems. Very accurate. Handles extreme stiffness. Default background solver.
+- `Rodas5P`: Slow for large systems. Very accurate. Handles severe stiffness.
 - `KenCarp4` (and `KenCarp47`): Fast. Handles medium stiffness. Default perturbation solver.
 - `Kvaerno5`: Behaves similar to `KenCarp4`. Slightly more accurate. Slightly slower.
 - `TRBDF2`: Very fast. Decent accuracy. Handles severe stiffness.

src/components/thermodynamics.jl

@@ -70,7 +70,7 @@ function thermodynamics_recombination_recfast(g; kwargs...)
         βHe ~ 4 * αHe / λe^3 * exp(-βb*E_He_∞_2s)
         KHe0⁻¹ ~ (8π*g.H) / λ_He_2p_1s^3 # RECFAST He flag 0
         KHe1⁻¹ ~ -exp(-3*A2Ps*nHe*(1-XHe⁺+1e-10)/KHe0⁻¹) * KHe0⁻¹ # RECFAST He flag 1 (close to zero modification?) # TODO: not that good, reliability depends on ϵ to avoid division by 0; try to use proper Saha ICs with XHe⁺ ≠ 1.0 and remove it
-        γ2Ps ~ 3*A2Ps*fHe*(1-XHe⁺+1e-10)*c^2 / (1.436289e-22*8π*√(2π/(βb*mHe*c^2))*(1-XH⁺+1e-10)*(f_He_2p_1s)^3) # TODO: introduce ν_He_2p_1s?
+        γ2Ps ~ abs(3*A2Ps*fHe*(1-XHe⁺+1e-10)*c^2 / (1.436289e-22*8π*√(2π/(βb*mHe*c^2))*(1-XH⁺+1e-10)*(f_He_2p_1s)^3)) # abs to reduce chance for early-time crash # TODO: address properly # TODO: introduce ν_He_2p_1s?
         KHe2⁻¹ ~ A2Ps/(1+0.36*γ2Ps^0.86)*3*nHe*(1-XHe⁺) # RECFAST He flag 2 (Doppler correction) # TODO: increase reliability, particularly at initial time
         KHe ~ 1 / (KHe0⁻¹ + KHe1⁻¹ + KHe2⁻¹) # corrections to inverse KHe are additive
         CHe ~ (exp(-βb*E_He_2p_2s) + KHe*ΛHe*nHe*(1-XHe⁺)) /
@@ -82,7 +82,7 @@ function thermodynamics_recombination_recfast(g; kwargs...)
         βHe3 ~ 4/3 * αHe3 / λe^3 * exp(-βb*Ewtf)
         τHe3 ~ A2Pt*nHe*(1-XHe⁺+1e-10)*3 * λ_He_2p_1s_tri^3/(8π*g.H)
         pHe3 ~ (1 - exp(-τHe3)) / τHe3
-        γ2Pt ~ 3*A2Pt*fHe*(1-XHe⁺+1e-10)*c^2 / (8π*1.484872e-22*f_He_2p_1s_tri*√(2π/(βb*mHe*c^2))*(1-XH⁺+1e-10)) / (f_He_2p_1s_tri)^2 # TODO: fails at extremely early times
+        γ2Pt ~ abs(3*A2Pt*fHe*(1-XHe⁺+1e-10)*c^2 / (8π*1.484872e-22*f_He_2p_1s_tri*√(2π/(βb*mHe*c^2))*(1-XH⁺+1e-10)) / (f_He_2p_1s_tri)^2) # abs to reduce chance for early-time crash # TODO: address properly
         CHe3 ~ (1e-10 + A2Pt*(pHe3+1/(1+0.66*γ2Pt^0.9)/3)*exp(-βb*E_He_2p_2s_tri)) /
                (1e-10 + A2Pt*(pHe3+1/(1+0.66*γ2Pt^0.9)/3)*exp(-βb*E_He_2p_2s_tri) + βHe3) # added 1e-10 to avoid NaN at late times (does not change early behavior) # TODO: is sign in p-s exponentials wrong/different to what it is in just CHe?
         DXHe⁺_triplet ~ -g.a/g.H₀ * CHe3 * (αHe3*XHe⁺*ne - βHe3*(1-XHe⁺)*3*exp(-βb*E_He_2s_1s_tri))

src/solve.jl

@@ -94,11 +94,11 @@ end
 # TODO: solve thermodynamics only if parameters contain thermodynamics parameters?
 # TODO: shoot to reach E = 1 today when integrating forwards
 """
-    solve(M::CosmologyModel, pars; aini = 1e-8, solver = Rodas5P(), reltol = 1e-10, kwargs...)
+    solve(M::CosmologyModel, pars; aini = 1e-8, solver = Rodas4P(), reltol = 1e-10, kwargs...)
 
 Solve `CosmologyModel` with parameters `pars` at the background level.
 """
-function solve(M::CosmologyModel, pars; aini = 1e-8, solver = Rodas5P(), reltol = 1e-15, backwards = true, thermo = true, debug_initialization = false, guesses = Dict(), jac = false, sparse = false, kwargs...)
+function solve(M::CosmologyModel, pars; aini = 1e-8, solver = Rodas4P(), reltol = 1e-10, backwards = true, thermo = true, debug_initialization = false, guesses = Dict(), jac = false, sparse = false, kwargs...)
     # Split parameters into DifferentialEquations' "u0" and "p" convention # TODO: same in perturbations
     params = merge(pars, Dict(M.k => 0.0)) # k is unused, but must be set https://github.com/SciML/ModelingToolkit.jl/issues/3013 # TODO: remove
     pars = intersect(keys(params), parameters(M)) # separate parameters from initial conditions
@@ -163,11 +163,11 @@ end
 
 # TODO: pass background solution to avoid recomputing it
 """
-    solve(M::CosmologyModel, pars, ks; aini = 1e-8, solver = KenCarp4(), reltol = 1e-9, backwards = true, verbose = false, thread = true, jac = false, sparse = false, kwargs...)
+    solve(M::CosmologyModel, pars, ks; aini = 1e-8, solver = KenCarp4(), reltol = 1e-8, backwards = true, verbose = false, thread = true, jac = false, sparse = false, kwargs...)
 
 Solve `CosmologyModel` with parameters `pars` up to the perturbative level for wavenumbers `ks`.
 """
-function solve(M::CosmologyModel, pars, ks::AbstractArray; aini = 1e-8, solver = KenCarp4(), reltol = 1e-9, backwards = true, verbose = false, thread = true, jac = false, sparse = false, kwargs...)
+function solve(M::CosmologyModel, pars, ks::AbstractArray; aini = 1e-8, solver = KenCarp4(), reltol = 1e-8, backwards = true, verbose = false, thread = true, jac = false, sparse = false, kwargs...)
     ks = k_dimensionless(ks, pars[M.g.h])
 
     !issorted(ks) && throw(error("ks = $ks are not sorted in ascending order"))
@@ -176,6 +176,7 @@ function solve(M::CosmologyModel, pars, ks::AbstractArray; aini = 1e-8, solver =
     tini, tend = extrema(th_sol.th[t])
     if M.spline_thermo
         th_sol_spline = isempty(kwargs) ? th_sol : solve(M, pars; aini, backwards, jac, sparse) # should solve again if given keyword arguments, like saveat
+        # TODO: when solving thermo with low reltol: even though the solution is correct, just taking its points for splining can be insufficient. should increase number of points, so it won't mess up the perturbations
         pars = merge(pars, Dict(
             M.pt.b.rec.τspline => spline(th_sol_spline[M.b.rec.τ], th_sol_spline[M.t]), # TODO: more time points, spline log(t)?
             M.pt.b.rec.τ̇spline => spline(th_sol_spline[M.b.rec.τ̇], th_sol_spline[M.t]),
@@ -307,6 +308,11 @@ function (sol::CosmologySolution)(tvar::Num, k, t, idxs)
 end
 
 
+"""
+    shoot(M::CosmologyModel, pars_fixed, pars_varying, conditions; solver = TrustRegion(), verbose = false, kwargs...)
+
+Solve a cosmological model with fixed parameters, while varying some parameters (from their initial guesses) until the given `conditions` at the end time are satisfied.
+"""
 function shoot(M::CosmologyModel, pars_fixed, pars_varying, conditions; solver = TrustRegion(), verbose = false, kwargs...)
     funcs = [eq.lhs - eq.rhs for eq in conditions] .|> ModelingToolkit.wrap # expressions that should be 0 # TODO: shouldn't have to wrap