Improve parameter fit performance

Author: hersle

docs/src/parameter_fitting.md

@@ -81,7 +81,7 @@ function dL(z, prob::CosmologyProblem, Ωm0, h; Ωr0 = 5e-5)
         M.m.Ω₀ => Ωm0,
         M.Λ.Ω₀ => 1 - Ωr0 - Ωm0,
         M.g.h => h, M.r.T₀ => 0.0 # TODO: don't set
-    ))
+    ), build_initializeprob = false) # bypass unnecessary initialization for performance
     sol = solve(prob)
     return dL(z, sol)
 end
@@ -112,7 +112,7 @@ end
 
 # TODO: speed up: https://discourse.julialang.org/t/modelingtoolkit-odesystem-in-turing/115700/
 sn = supernova(data, prob) # condition model on data
-chain = sample(sn, NUTS(), MCMCSerial(), 200, 1)
+chain = sample(sn, NUTS(; init_ϵ = 0.0125), MCMCSerial(), 2000, 1)
 ```
 As we see above, the MCMC `chain` displays a summary with information about the fitted parameters, including their posterior means and standard deviations.
 We can also plot the chains:

src/components/metric.jl

@@ -9,9 +9,6 @@ function metric(; name = :g, kwargs...)
     defs = [
         H₀ => H100 * h
     ]
-    initialization_eqs = [
-        O(ϵ^0)(a ~ D(a) * t)
-    ]
     guesses = [
         a => 1e-5
     ]
@@ -24,5 +21,5 @@ function metric(; name = :g, kwargs...)
         H ~ E * H₀
         g₁₁ ~ a^2 * (-1 - 2*ϵ*Ψ)
         g₂₂ ~ a^2 * (+1 - 2*ϵ*Φ)
-    ], t, vars, pars; defaults = defs, initialization_eqs, guesses, name, description, kwargs...)
+    ], t, vars, pars; defaults = defs, guesses, name, description, kwargs...)
 end

src/models.jl

@@ -74,9 +74,12 @@ function ΛCDM(;
     if all.(have(species, :Ω₀)) && startswith(ModelingToolkit.description(G), "General relativity")
         push!(defs, species[end].Ω₀ => 1 - sum(s.Ω₀ for s in species[begin:end-1]))
     end
+    initialization_eqs = [
+        O(ϵ^0)(g.a ~ D(g.a) * t)
+    ]
 
     description = "Standard cosmological constant and cold dark matter cosmological model"
-    connections = ODESystem([eqs0; eqs1], t, vars, [pars; k]; defaults = defs, name, description)
+    connections = ODESystem([eqs0; eqs1], t, vars, [pars; k]; defaults = defs, initialization_eqs, name, description)
     components = filter(!isnothing, [g; G; species; I])
     M = compose(connections, components...)
     return complete(M; flatten = false)
@@ -118,6 +121,7 @@ function RMΛ(;
     ] .|> O(ϵ^1)
     defs = [
         g.Ψ => 20 // 15,
+        g.a => √(r.Ω₀) * t
     ]
     connections = ODESystem([eqs0; eqs1], t, [], [k]; defaults = defs, name)
     M = compose(connections, g, G, species...)

src/solve.jl

@@ -161,12 +161,12 @@ end
 Return an updated `CosmologyProblem` where parameters in `prob` are updated to values specified in `pars`.
 Parameters that are not specified in `pars` keep their values from `prob`.
 """
-function remake(prob::CosmologyProblem, pars::Dict; bg = true, th = true, pt = true, shoot = true)
+function remake(prob::CosmologyProblem, pars::Dict; bg = true, th = true, pt = true, shoot = true, kwargs...)
     pars_full = merge(prob.pars, pars) # save full dictionary for constructor
     vars, pars = split_vars_pars(prob.M, pars)
-    bg = bg && !isnothing(prob.bg) ? remake(prob.bg; u0 = vars, p = pars) : nothing
-    th = th && !isnothing(prob.th) ? remake(prob.th; u0 = vars, p = pars) : nothing
-    pt = pt && !isnothing(prob.pt) ? remake(prob.pt; u0 = vars, p = pars) : nothing
+    bg = bg && !isnothing(prob.bg) ? remake(prob.bg; u0 = vars, p = pars, kwargs...) : nothing
+    th = th && !isnothing(prob.th) ? remake(prob.th; u0 = vars, p = pars, kwargs...) : nothing
+    pt = pt && !isnothing(prob.pt) ? remake(prob.pt; u0 = vars, p = pars, kwargs...) : nothing
     shoot_pars = shoot ? prob.shoot : keys(Dict())
     shoot_conditions = shoot ? prob.conditions : []
     return CosmologyProblem(prob.M, bg, th, pt, pars_full, shoot_pars, shoot_conditions, prob.spline_thermo)