Initialize ODEs with custom initialization problems

docs/src/automatic_differentiation.md

@@ -19,7 +19,7 @@ pars = [M.γ.T₀, M.c.Ω₀, M.b.Ω₀, M.ν.Neff, M.g.h, M.b.YHe, M.h.m_eV, M.
 prob0 = CosmologyProblem(M, Dict(pars .=> NaN); sparse = false) # dense Jacobian faster for AD
 
 probgen = parameter_updater(prob0, pars)
-P(k, θ) = spectrum_matter(probgen(θ), k; ptopts = (reltol = 1e-3, abstol = 1e-3))
+P(k, θ) = spectrum_matter(probgen(prob0, θ), k; ptopts = (reltol = 1e-3, abstol = 1e-3))
 ```
 It is now easy to evaluate the power spectrum:
 ```@example ad

docs/src/benchmarks.md

@@ -31,13 +31,14 @@ The points on each curve correspond to a sequence of tolerances.
 using DiffEqDevTools
 
 refalg = Rodas5P(linsolve = RFLUFactorization())
-bgsol = solve(prob.bg, refalg; abstol = 1e-12, reltol = 1e-12) # reference solution (results are similar compared to Rodas4/4P/5P/FBDF)
+bgprob = SymBoltz.setupbg(prob.bg, prob.bginit)
+bgsol = solve(bgprob, refalg; abstol = 1e-12, reltol = 1e-12) # reference solution (results are similar compared to Rodas4/4P/5P/FBDF)
 
 abstols = 1 ./ 10 .^ (7:11)
 reltols = 1 ./ 10 .^ (7:11)
 bgalgs = [Rodas4(), Rodas5(), Rodas4P(), Rodas5P(), Rodas6P(), FBDF(), QNDF()] # FBDF/QNDF unstable for some tolerances
 setups = [Dict(:alg => alg) for alg in bgalgs]
-wp = WorkPrecisionSet(prob.bg, abstols, reltols, setups; appxsol = bgsol, save_everystep = false, error_estimate = :l2)
+wp = WorkPrecisionSet(bgprob, abstols, reltols, setups; appxsol = bgsol, save_everystep = false, error_estimate = :l2)
 plot(wp; title = "Reference: $(SymBoltz.algname(refalg))", size = (800, 400), margin = 5*Plots.mm)
 ```
 
@@ -73,7 +74,7 @@ The points on each curve correspond to a sequence of tolerances.
 # TODO: test different nlsolve # hide
 # TODO: add AdaptiveRadau/RadauIIA5 when they support sparse J: https://github.com/SciML/OrdinaryDiffEq.jl/issues/2892 # hide
 ptalgs = [algtype(linsolve = KLUFactorization()) for algtype in [TRBDF2, KenCarp4, KenCarp47, KenCarp5, Kvaerno5, Rodas4P, Rodas5P, Rodas6P, QNDF, FBDF]]
-ptprobgen = SymBoltz.setuppt(prob.pt, bgsol)
+ptprobgen = SymBoltz.setuppt(prob.pt, prob.ptinit, bgsol)
 setups = [Dict(:alg => alg) for alg in ptalgs]
 refalg = Rodas5P(linsolve = KLUFactorization())
 abstols = 1 ./ 10 .^ (5:9)

docs/src/comparison.md

@@ -354,7 +354,7 @@ function P_class(k, pars)
     return P
 end
 function P_symboltz(k, pars)
-    prob′ = parameter_updater(prob, collect(keys(pars)))(collect(values(pars))) # TODO: move outside; common for Pk and Cl
+    prob′ = parameter_updater(prob, collect(keys(pars)))(prob, collect(values(pars))) # TODO: move outside; common for Pk and Cl
     P = spectrum_matter(prob′, k / u"Mpc") / u"Mpc^3"
     return P
 end
@@ -406,7 +406,7 @@ function Dl_class(modes, l, pars)
     return stack(Dl)
 end
 function Dl_symboltz(modes, jl, pars; kwargs...)
-    prob′ = parameter_updater(prob, collect(keys(pars)))(collect(values(pars)))
+    prob′ = parameter_updater(prob, collect(keys(pars)))(prob, collect(values(pars)))
     return spectrum_cmb(modes, prob′, jl; normalization = :Dl, kwargs...)
 end
 

docs/src/forecasting.md

@@ -33,7 +33,7 @@ Since $Cₗ$ is an expensive but smooth function of $l$, we make one function fo
 probgen = parameter_updater(prob0, pars_varying)
 jl = SphericalBesselCache(40:20:1000)
 ls = 40:1:1000
-Cl(θ) = spectrum_cmb(:TT, probgen(θ), jl, ls)
+Cl(θ) = spectrum_cmb(:TT, probgen(prob0, θ), jl, ls)
 ```
 We can now compute $Cₗ$ and the cosmic variance uncertainties
 ```math

docs/src/plot.md

@@ -98,8 +98,7 @@ function plot_interactive(prob::CosmologyProblem, xvar::SymBoltz.Num, yvar::SymB
     end
     probgen = parameter_updater(prob, [par for (par, _) in obspars])
     function xyfunc(θ)
-        prob = probgen(θ)
-        sol = solve(prob)
+        sol = solve(probgen(prob, θ))
         τ = τs(sol)
         xs = sol(xvar, τ)
         ys = sol(yvar, τ)
@@ -135,10 +134,9 @@ obspars = [
 ]
 probgen = parameter_updater(prob, [par for (par, _) in obspars])
 function xyfunc(θ)
-    prob = probgen(θ)
     lgks = unique([-4:0.5:-3; -3:0.2:-2; -2:0.05:0]) # as few points as possible
     ks = 10 .^ lgks / u"Mpc"
-    Ps = spectrum_matter(prob, ks; ptopts = (alg = SymBoltz.TRBDF2(), reltol = 1e-4, abstol = 1e-4))
+    Ps = spectrum_matter(probgen(prob, θ), ks; ptopts = (alg = SymBoltz.TRBDF2(), reltol = 1e-4, abstol = 1e-4))
     lgPs = log10.(Ps/u"Mpc^3")
 
     # smoothen with spline and sample more densely

docs/src/solve.md

@@ -29,7 +29,7 @@ To do so, use the function `parameter_updater` that returns a function that quic
 
 ```@example sol
 probmaker = parameter_updater(prob, [M.g.h, M.c.Ω₀]) # fast factory function
-prob = probmaker([0.70, 0.27]) # create updated problem
+prob = probmaker(prob, [0.70, 0.27]) # create updated problem
 ```
 
 ```@docs

src/observables/distances.jl

@@ -25,16 +25,15 @@ function distance_luminosity_function(M::System, pars_fixed, pars_varying, zs; b
     pars = merge(pars_fixed, Dict(pars_varying .=> NaN))
     as = @. 1 / (zs + 1)
     prob = CosmologyProblem(M, pars; pt = false, ivspan = (minimum(as), 1.0))
-    probgen = parameter_updater(prob, pars_varying; build_initializeprob = Val{false})
+    probgen = parameter_updater(prob, pars_varying)
 
     geta = getsym(prob, M.g.a)
     getτ = getsym(prob, M.τ)
     geth = getsym(prob, M.g.h)
     getΩk0 = getsym(prob, M.K.Ω₀)
 
-    return p -> begin
-        prob = probgen(p)
-        sol = solve(prob; bgopts, saveat = as, save_end = true)
+    return (p) -> begin
+        sol = solve(probgen(prob, p); bgopts, saveat = as, save_end = true)
         a = geta(sol)
         τ = getτ(sol)
         h = geth(sol)

src/observables/fourier.jl

@@ -246,7 +246,7 @@ Compute and evaluate source functions ``S(τ,k)`` with symbolic expressions `Ss`
 The options `bgopts` and `ptopts` are passed to the background and perturbation solves.
 """
 function source_grid(prob::CosmologyProblem, Ss::AbstractArray, τs, ks; bgopts = (), ptopts = (), thread = true, verbose = false)
-    bgsol = solvebg(prob.bg; bgopts..., verbose)
+    bgsol = solvebg(prob.bg, prob.bginit; bgopts..., verbose)
     getSs = map(S -> getsym(prob.pt, S), Ss)
     Ss = similar(bgsol, length(Ss), length(τs), length(ks))
     minimum(τs) ≥ bgsol.t[begin] && maximum(τs) ≤ bgsol.t[end] || error("input τs and computed background solution have different timespans")
@@ -256,7 +256,7 @@ function source_grid(prob::CosmologyProblem, Ss::AbstractArray, τs, ks; bgopts
         end
         return nothing
     end
-    solvept(prob.pt, bgsol, ks; output_func, saveat = τs, ptopts..., thread, verbose)
+    solvept(prob.pt, prob.ptinit, bgsol, ks; output_func, saveat = τs, ptopts..., thread, verbose)
     return Ss
 end
 
@@ -289,7 +289,7 @@ function source_grid_adaptive(prob::CosmologyProblem, Ss::AbstractVector, τs, k
         ptsaveopts = (saveat = τs,)
     end
 
-    ptprobgen = setuppt(prob.pt, bgsol)
+    ptprobgen = setuppt(prob.pt, prob.ptinit, bgsol)
 
     getSs = map(S -> getsym(prob.pt, S), Ss)
     function sourcek!(k, ik, Ss)
@@ -379,7 +379,7 @@ end
 
 # Dispatch without background solution
 function source_grid_adaptive(prob::CosmologyProblem, Ss::AbstractVector, τs, ks; bgopts = (), kwargs...)
-    bgsol = solvebg(prob.bg; bgopts...)
+    bgsol = solvebg(prob.bg, prob.bginit; bgopts...)
     return source_grid_adaptive(prob, Ss, τs, ks, bgsol; kwargs...)
 end