Use saving callback in source_grid to avoid saving unused ODE unknowns

Project.toml

@@ -8,6 +8,7 @@ Bessels = "0e736298-9ec6-45e8-9647-e4fc86a2fe38"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
+DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 FastChebInterp = "cf66c380-9a80-432c-aff8-4f9c79c0bdde"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ForwardDiffChainRules = "c9556dd2-1aed-4cfe-8560-1557cf593001"
@@ -52,6 +53,7 @@ Bessels = "0.2.8"
 ChainRulesCore = "1.25.1"
 CommonSolve = "0.2.4"
 DataInterpolations = "8.0.0"
+DiffEqCallbacks = "4.18.3"
 FastChebInterp = "1"
 ForwardDiff = "0.10, 1.0.0"
 ForwardDiffChainRules = "0.3.0"

src/SymBoltz.jl

@@ -7,6 +7,7 @@ using Symbolics, SymbolicUtils
 using OrdinaryDiffEqRosenbrock, OrdinaryDiffEqSDIRK, OrdinaryDiffEqTsit5
 using OrdinaryDiffEqNonlinearSolve: NLNewton
 using NonlinearSolve
+using DiffEqCallbacks: SavedValues, SavingCallback
 using OhMyThreads
 using Base.Threads
 using Setfield

src/observables/fourier.jl

@@ -270,17 +270,20 @@ The options `bgopts` and `ptopts` are passed to the background and perturbation
 function source_grid(prob::CosmologyProblem, Ss, τs, ks, bgsol::ODESolution; ptopts = (), thread = true, verbose = false)
     getSs = getsym(prob.pt, Ss)
     T = source_eltype(Ss, eltype(bgsol))
-    Ss = Matrix{T}(undef, length(τs), length(ks))
     minimum(τs) ≥ bgsol.t[begin] && maximum(τs) ≤ bgsol.t[end] || error("input τs and computed background solution have different timespans")
-    function output_func(sol, ik)
-        vals = getSs(sol)
-        @inbounds for iτ in eachindex(vals)
-            Ss[iτ, ik] = T(vals[iτ])
-        end
-        return nothing
+
+    # Save only the requested source values instead of the full ODE solution with all unknowns
+    # Save callback similar to https://github.com/SciML/SciMLBase.jl/blob/97f6d4aff88ab5f2dedc90ef503edabe72f00e93/src/solutions/ode_solutions.jl#L369-L373
+    save_func(u, t, integrator) = getSs(SciMLBase.ProblemState(; u, p = SciMLBase.parameter_values(integrator), t))
+    savedvalues = [SavedValues(eltype(τs), T) for _ in ks] # one save container per k
+    callback(ik) = SavingCallback(save_func, savedvalues[ik]; saveat = τs)
+    solvept(prob.pt, bgsol, ks; callback, save_everystep = false, save_start = false, dense = false, ptopts..., thread, verbose)
+
+    out = Matrix{T}(undef, length(τs), length(ks))
+    @inbounds for ik in eachindex(ks), iτ in eachindex(τs)
+        out[iτ, ik] = savedvalues[ik].saveval[iτ]
     end
-    solvept(prob.pt, bgsol, ks; output_func, saveat = τs, ptopts..., thread, verbose)
-    return Ss
+    return out
 end
 function source_grid(prob::CosmologyProblem, Ss, τs, ks; bgopts = (), verbose = false, kwargs...)
     bgsol = solvebg(prob.bg; bgopts..., verbose)

src/solve.jl

@@ -540,7 +540,7 @@ If `thread` and Julia is running with multiple threads, the solution of independ
 `ptivini` is a number or a function of ``k`` that sets the initial time of integration for each perturbation mode, but is always clamped to the background timespan.
 The return value is a vector with one `ODESolution` per wavenumber, or its mapping through `output_func` if a custom transformation is passed.
 """
-function solvept(ptprob::ODEProblem, bgsol::ODESolution, ks::AbstractArray, ptivini = -Inf; alg = ptalg(ptprob), reltol = 1e-5, abstol = 1e-5, output_func = (sol, i) -> sol, thread = true, verbose = false, kwargs...)
+function solvept(ptprob::ODEProblem, bgsol::ODESolution, ks::AbstractArray, ptivini = -Inf; alg = ptalg(ptprob), reltol = 1e-5, abstol = 1e-5, output_func = (sol, i) -> sol, callback = (i -> nothing), thread = true, verbose = false, kwargs...)
     check_solve_args(ptprob, alg)
 
     if thread && Threads.nthreads() == 1
@@ -563,7 +563,7 @@ function solvept(ptprob::ODEProblem, bgsol::ODESolution, ks::AbstractArray, ptiv
         return output_func(sol, i)
     end
 
-    ptsols = fetch.(@spawnif output_func_warn(solve(ptprobgen(ks[i]), alg; verbose = verbosity(verbose), reltol, abstol, kwargs...), i) thread for i in eachindex(ks)) # wait for all tasks to finish and get the returned solutions
+    ptsols = fetch.(@spawnif output_func_warn(solve(ptprobgen(ks[i]), alg; verbose = verbosity(verbose), reltol, abstol, callback = callback(i), kwargs...), i) thread for i in eachindex(ks)) # wait for all tasks to finish and get the returned solutions
     verbose && println()
     return ptsols
 end

test/runtests.jl

@@ -745,8 +745,8 @@ end
     vars = @variables a(τ) ℋ(τ)
     pars = @parameters Ωr0 Ωm0 ΩΛ0 [shoot=true]
     eqs = [ℋ ~ √(Ωr0/a^4 + Ωm0/a^3 + ΩΛ0) * a, D(a) ~ a*ℋ]
-    initialization_eqs = [ℋ ~ 1/τ]
-    guesses = Dict(ΩΛ0 => 1 - Ωr0 - Ωm0, a => τ)
+    initialization_eqs = [a^2*ℋ^2 ~ a^2/τ^2] # this form avoids 1/a and is more stable
+    guesses = Dict(ΩΛ0 => 1 - Ωr0 - Ωm0, a => √(Ωr0) * τ)
     constraints = [ℋ ~ 1]
     @named M = System(eqs, τ, vars, pars; initialization_eqs, guesses, constraints)
     @test Set(keys(SymBoltz.shootvars(M))) == Set(ΩΛ0)