Merge pull request #41 from hersle/ksampling

Improve sampling of perturbation modes for angular spectra

Author: hersle

docs/src/comparison.md

@@ -409,10 +409,10 @@ l = 20:20:2000 # CLASS default is lmax = 2500
 Dl1 = Dl_class([:TT, :TE, :EE], l, pars)
 jl = SphericalBesselCache(l)
 Dl2 = Dl_symboltz([:TT, :TE, :EE], jl, pars)
-plot_compare(l, l, Dl1[:, 1], Dl2[:, 1], "l", "Dₗ(TT)"; tol = 2e-12)
+plot_compare(l, l, Dl1[:, 1], Dl2[:, 1], "l", "Dₗ(TT)"; tol = 1e-12)
 ```
 ```@example class
-plot_compare(l, l, Dl1[:, 2], Dl2[:, 2], "l", "Dₗ(TE)"; tol = 3e-14)
+plot_compare(l, l, Dl1[:, 2], Dl2[:, 2], "l", "Dₗ(TE)"; tol = 2e-14)
 ```
 ```@example class
 plot_compare(l, l, Dl1[:, 3], Dl2[:, 3], "l", "Dₗ(EE)"; tol = 6e-15)

docs/src/observables.md

@@ -157,6 +157,36 @@ rs = sound_horizon(sol)
 plot(τs, rs; xlabel = "τ / H₀⁻¹", ylabel = "rₛ / (c/H₀)")
 ```
 
+## Source functions
+
+```@docs
+source_grid
+source_grid_adaptive
+```
+
+```@example
+using SymBoltz, Plots, DataInterpolations
+M = ΛCDM(h = nothing, ν = nothing)
+pars = parameters_Planck18(M)
+prob = CosmologyProblem(M, pars)
+sol = solve(prob)
+
+τs = sol[M.τ] # conformal times in background solution
+ks = [1.0, 2000.0] # initial coarse grid
+ks, Ss = source_grid_adaptive(prob, [M.ST0], τs, ks; atol = 5.0)
+iτ = argmax(sol[M.b.v]) # index of decoupling time
+iτs = iτ-75:iτ+75 # indices around decoupling
+p1 = surface(ks, τs[iτs], Ss[1, iτs, :]; camera = (45, 25), xlabel = "k", ylabel = "τ", zlabel = "S", colorbar = false)
+
+lgas = -6.0:0.2:0.0
+τs = LinearInterpolation(sol[M.τ], sol[log10(M.g.a)])(lgas) # τ at given lg(a)
+ks = 5.0:5.0:100.0
+Ss = source_grid(prob, [M.g.Ψ], τs, ks)
+p2 = wireframe(ks, lgas, Ss[1, :, :]; camera = (75, 20), xlabel = "k", ylabel = "lg(a)", zlabel = "Φ")
+
+plot(p1, p2)
+```
+
 ## Line-of-sight integration
 
 ```@docs

src/SymBoltz.jl

@@ -45,7 +45,7 @@ export RMΛ, ΛCDM, w0waCDM, QCDM, GRΛCDM, BDΛCDM
 export CosmologyProblem, CosmologySolution
 export solve, solvebg, solvept, remake, issuccess, parameter_updater
 export parameters_Planck18
-export spectrum_primordial, spectrum_matter, spectrum_matter_nonlinear, spectrum_cmb, correlation_function, variance_matter, stddev_matter, los_integrate, source_grid, sound_horizon, SphericalBesselCache
+export spectrum_primordial, spectrum_matter, spectrum_matter_nonlinear, spectrum_cmb, correlation_function, variance_matter, stddev_matter, los_integrate, source_grid, source_grid_adaptive, sound_horizon, SphericalBesselCache
 export express_derivatives
 
 end

src/observables/angular.jl

@@ -181,17 +181,16 @@ function spectrum_cmb(ΘlAs::AbstractMatrix, ΘlBs::AbstractMatrix, P0s::Abstrac
 end
 
 """
-    spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, jl::SphericalBesselCache; normalization = :Cl, unit = nothing, Δkτ0 = 2π/2, Δkτ0_S = 8.0, kτ0min = 0.1*ls[begin], kτ0max = 3*ls[end], u = (τ->tanh(τ)), u⁻¹ = (u->atanh(u)), Nlos = 768, integrator = TrapezoidalRule(), bgopts = (alg = Rodas4P(), reltol = 1e-9, abstol = 1e-9), ptopts = (alg = KenCarp4(), reltol = 1e-8, abstol = 1e-8), thread = true, verbose = false, kwargs...)
+    spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, jl::SphericalBesselCache; normalization = :Cl, unit = nothing, kτ0s = 0.1*jl.l[begin]:2π/2:2*jl.l[end], u = (τ->tanh(τ)), u⁻¹ = (u->atanh(u)), Nlos = 768, integrator = TrapezoidalRule(), bgopts = (alg = Rodas4P(), reltol = 1e-9, abstol = 1e-9), ptopts = (alg = KenCarp4(), reltol = 1e-8, abstol = 1e-8), sourceopts = (atol = 1e0,), verbose = false, kwargs...)
 
 Compute the CMB power spectra `modes` (`:TT`, `:EE`, `:TE` or an array thereof) ``C_l^{AB}``'s at angular wavenumbers `ls` from the cosmological solution `sol`.
 If `unit` is `nothing` the spectra are of dimensionless temperature fluctuations relative to the present photon temperature; while if `unit` is a temperature unit the spectra are of dimensionful temperature fluctuations.
 """
-function spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, jl::SphericalBesselCache; normalization = :Cl, unit = nothing, Δkτ0 = 2π/2, Δkτ0_S = 8.0, kτ0min = 0.1*jl.l[begin], kτ0max = 3*jl.l[end], u = (τ->tanh(τ)), u⁻¹ = (u->atanh(u)), Nlos = 768, integrator = TrapezoidalRule(), bgopts = (alg = Rodas4P(), reltol = 1e-9, abstol = 1e-9), ptopts = (alg = KenCarp4(), reltol = 1e-8, abstol = 1e-8), thread = true, verbose = false, kwargs...)
+function spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, jl::SphericalBesselCache; normalization = :Cl, unit = nothing, kτ0s = 0.1*jl.l[begin]:2π/2:2*jl.l[end], u = (τ->tanh(τ)), u⁻¹ = (u->atanh(u)), Nlos = 768, integrator = TrapezoidalRule(), bgopts = (alg = Rodas4P(), reltol = 1e-9, abstol = 1e-9), ptopts = (alg = KenCarp4(), reltol = 1e-8, abstol = 1e-8), sourceopts = (atol = 1e0,), verbose = false, kwargs...)
     ls = jl.l
-    kτ0s_coarse, kτ0s_fine = cmb_kτ0s(ls[begin], ls[end]; Δkτ0, Δkτ0_S, kτ0min, kτ0max)
     sol = solve(prob; bgopts, verbose)
     τ0 = getsym(sol, prob.M.τ0)(sol)
-    ks_coarse = kτ0s_coarse ./ τ0
+    ks_fine = collect(kτ0s ./ τ0)
     τs = sol.bg.t # by default, use background (thermodynamics) time points for line of sight integration
     τi = τs[begin]
     if Nlos != 0 # instead choose Nlos time points τ = τ(u) corresponding to uniformly spaced u
@@ -209,11 +208,10 @@ function spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, jl::Spheric
     Ss = Num[]
     iT > 0 && push!(Ss, prob.M.ST0)
     iE > 0 && push!(Ss, prob.M.ST2_polarization)
-    Ss_coarse = source_grid(prob, Ss, τs, ks_coarse; bgopts, ptopts, thread, verbose) # TODO: pass kτ0 and x # TODO: pass bgsol
+    ks_coarse = range(ks_fine[begin], ks_fine[end]; length = 2)
+    ks_coarse, Ss_coarse = source_grid_adaptive(prob, Ss, τs, ks_coarse; bgopts, ptopts, verbose, sourceopts...) # TODO: thread flag # TODO: pass kτ0 and x # TODO: pass bgsol
 
     # Interpolate source function to finer k-grid
-    ks_fine = collect(kτ0s_fine ./ τ0)
-    ks_fine = clamp.(ks_fine, ks_coarse[begin], ks_coarse[end]) # TODO: ideally avoid
     Ss_fine = source_grid(Ss_coarse, ks_coarse, ks_fine)
     Ss_fine[:, end, :] .= 0.0
 
@@ -267,7 +265,7 @@ function spectrum_cmb(mode::Symbol, args...; kwargs...)
     return spectrum_cmb([mode], args...; kwargs...)[:, begin]
 end
 
-function cmb_kτ0s(lmin, lmax; Δkτ0 = 2π/2, Δkτ0_S = 8.0, kτ0min = 0.1*lmin, kτ0max = 3*lmax)
+function cmb_kτ0s(lmin, lmax; Δkτ0 = 2π/2, Δkτ0_S = 8.0, kτ0min = 0.1*lmin, kτ0max = 2*lmax)
     kτ0s_fine = range(kτ0min, kτ0max, step=Δkτ0) # use integer multiple so endpoints are the same
     kτ0s_coarse = range(kτ0s_fine[begin], kτ0s_fine[end], length = Int(floor((kτ0max-kτ0min)/Δkτ0_S+1)))
     kτ0s_coarse[begin] == kτ0s_fine[begin] && kτ0s_coarse[end] == kτ0s_fine[end] || error("different wavenumber endpoints")

src/observables/fourier.jl

@@ -152,6 +152,12 @@ end
 # TODO: take in ks and N_skip_interp = 1, 2, 3, ... for more efficient? same for τ?
 # TODO: source_grid(prob::CosmologyProblem seemed type-stable?
 # TODO: return getter for (ks, τs)
+"""
+    source_grid(sol::CosmologySolution, Ss::AbstractVector, τs)
+
+Evaluate and return source functions ``S(τ,k)`` from the solution `sol`.
+The source functions are given by symbolic expressions `Ss`, and evaluated on a grid with conformal times `τs` and the same wavenumbers as `sol` is solved for.
+"""
 function source_grid(sol::CosmologySolution, Ss::AbstractVector, τs)
     # Evaluate integrated perturbations on coarse grid
     ks = sol.ks
@@ -165,6 +171,13 @@ function source_grid(sol::CosmologySolution, Ss::AbstractVector, τs)
     return Ss
 end
 
+"""
+    source_grid(Ss_coarse::AbstractArray, ks_coarse, ks_fine; ktransform = identity)
+
+Interpolate values `Ss_coarse` of source functions ``S(τ,k)`` from a coarse wavenumber grid `ks_coarse` to a fine grid `ks_fine`.
+The interpolation is linear in `ktransform(k)` (e.g. `identity` for interpolation in ``k`` or `log` for interpolation in ``\\ln k``.
+Conformal times are unchanged.
+"""
 function source_grid(Ss_coarse::AbstractArray, ks_coarse, ks_fine; ktransform = identity)
     size_coarse = size(Ss_coarse)
     size_fine = (size_coarse[1], size_coarse[2], length(ks_fine))
@@ -182,12 +195,18 @@ function source_grid(Ss_coarse::AbstractArray, ks_coarse, ks_fine; ktransform =
     return Ss_fine
 end
 
-# TODO: take in kτ0s and xs
-function source_grid(prob::CosmologyProblem, S::AbstractArray, τs, ks; bgopts = (), ptopts = (), thread = true, verbose = false)
+"""
+    source_grid(prob::CosmologyProblem, Ss::AbstractArray, τs, ks; bgopts = (), ptopts = (), thread = true, verbose = false)
+
+Compute and evaluate source functions ``S(τ,k)`` with symbolic expressions `Ss` on a grid with conformal times `τs` and wavenumbers `ks` from the problem `prob`.
+
+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)
-    getSs = map(s -> getsym(prob.pt, s), S)
-    Ss = similar(bgsol, length(S), length(τs), length(ks))
-    extrema(τs) == extrema(bgsol.t) || error("input τs and computed background solution have different timespans") # TODO: don't rely on
+    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")
     function output_func(sol, ik)
         for iS in eachindex(getSs)
             Ss[iS, :, ik] .= getSs[iS](sol)
@@ -197,3 +216,87 @@ function source_grid(prob::CosmologyProblem, S::AbstractArray, τs, ks; bgopts =
     solvept(prob.pt, bgsol, ks, prob.bgspline; output_func, saveat = τs, ptopts..., thread, verbose)
     return Ss
 end
+
+# Called internally by source_grid_adaptive
+function source_grid_refine(i1, i2, ks, Ss, Sk, savelock; verbose = false, kwargs...)
+    k1 = ks[i1]
+    k2 = ks[i2]
+    S1 = Ss[:, :, i1]
+    S2 = Ss[:, :, i2]
+
+    k2 ≥ k1 || error("shit")
+
+    k = (k1 + k2) / 2
+    S = Sk(k)
+
+    i = 0 # define in outer scope, then set inside lock block
+    @lock savelock begin
+        push!(ks, k)
+        i = length(ks) # copy, since first refine below can modify i before call to second refine
+        Ss[:, :, i] .= S
+        verbose && println("Refined k-grid between [$k1, $k2] on thread $(threadid()) to $i total points")
+    end
+
+    Sint = (S1 .+ S2) ./ 2 # linear interpolation
+
+    # check if interpolation is close enough for all sources
+    # (equivalent to finding the source grid of each source separately)
+    @sync if !all(isapprox(S[iS, :], Sint[iS, :]; kwargs...) for iS in 1:size(Ss, 1))
+        @spawn source_grid_refine(i1, i, ks, Ss, Sk, savelock; verbose, kwargs...) # refine left subinterval
+        @spawn source_grid_refine(i, i2, ks, Ss, Sk, savelock; verbose, kwargs...) # refine right subinterval
+    end
+end
+
+# TODO: Hermite interpolation
+# TODO: create SourceFunction type that does k and τ interpolation?
+"""
+    source_grid_adaptive(prob::CosmologyProblem, Ss::AbstractVector, τs, ks; bgopts = (), ptopts = (), sort = true, kwargs...)
+
+Adaptively compute and evaluate source functions ``S(τ,k)`` with symbolic expressions `Ss` on a grid with fixed conformal times `τs`, but adaptively refined grid of wavenumbers from the problem `prob`.
+The source functions are first evaluated on the (coarse) initial grid `ks`.
+Each subinterval ``(k₁, k₂)`` of `ks` is then adaptively refined until the linear interpolation ``Sᵢ = (S(k₁)+S(k₂))/2`` to the midpoint ``k=(k₁+k₂)/2`` approximates the actual value ``S(k)`` there within some tolerance.
+The comparison ``Sᵢ ≈ S`` is done with `isapprox(Sᵢ, S; kwargs...)`, where `S` and `Sᵢ` are vectors with the (conformal) timeseries of the source function for that wavenumber.
+It receives the keyword arguments `kwargs` passed to this function, so `atol`, `rtol` and/or `norm` can be specified to tune the tolerance.
+
+Returns the refined wavenumbers sorted in ascending order and a grid with the corresponding source function values.
+If not `sort`, the wavenumbers and source function values are instead left in the order in which they were inserted in the refinement process.
+
+The options `bgopts` and `ptopts` are passed to the background and perturbation solves.
+"""
+function source_grid_adaptive(prob::CosmologyProblem, Ss::AbstractVector, τs, ks; bgopts = (), ptopts = (), sort = true, kwargs...)
+    bgsol = solvebg(prob.bg; bgopts...)
+    ptprob0, ptprobgen = setuppt(prob.pt, bgsol, prob.bgspline)
+
+    getSs = map(S -> getsym(prob.pt, S), Ss)
+    function Sk(k)
+        ptprob = ptprobgen(ptprob0, k)
+        ptsol = solvept(ptprob; saveat = τs, ptopts...)
+        S = transpose(stack(map(getS -> getS(ptsol), getSs)))
+        S[:, end] .= 0.0 # avoid NaN/Infs from 1/χ today; fine in LOS integration, since always weighted by 0-valued Bessel function # TODO: avoid
+        return S
+    end
+
+    length(ks) ≥ 2 || error("Initial k-grid must have at least 2 values")
+    ks = Vector(ks) # common to input a range
+    Ss = similar(bgsol, length(Ss), length(τs), 1024)
+
+    savelock = ReentrantLock()
+
+    @threads for i in 1:length(ks)
+        Ss[:, :, i] .= Sk(ks[i])
+    end
+    @threads for i in 1:length(ks)-1
+        source_grid_refine(i, i+1, ks, Ss, Sk, savelock; kwargs...)
+    end
+
+    # sort according to k
+    if sort
+        is = sortperm(ks)
+        ks = ks[is]
+        Ss = Ss[:, :, is]
+    else
+        Ss = Ss[:, :, 1:length(ks)]
+    end
+
+    return ks, Ss
+end

src/solve.jl

@@ -427,6 +427,27 @@ function solvept(ptprob::ODEProblem, bgsol::ODESolution, ks::AbstractArray, bgsp
     verbose && println()
     return ptsols
 end
+"""
+    solvept(ptprob::ODEProblem; alg = DEFAULT_PTALG, reltol = 1e-8, abstol = 1e-8, kwargs...)
+
+Solve the perturbation problem `ptprob` and return the solution.
+Its wavenumber and background spline must already be initialized, for example with `setuppt`.
+
+# Examples
+
+```julia
+# ...
+prob = CosmologyProblem(M, pars)
+bgsol = solvebg(prob.bg)
+ptprob0, ptprobgen = SymBoltz.setuppt(prob.pt, bgsol, prob.bgspline)
+k = 1.0
+ptprob = ptprobgen(ptprob0, k)
+ptsol = solvept(ptprob)
+```
+"""
+function solvept(ptprob::ODEProblem; alg = DEFAULT_PTALG, reltol = 1e-8, abstol = 1e-8, kwargs...)
+    return solve(ptprob, alg; reltol, abstol, kwargs...)
+end
 
 function time_today(prob::CosmologyProblem)
     getτ0 = SymBoltz.getsym(prob.bg, :τ0)