diff --git a/docs/src/LCDM.md b/docs/src/LCDM.md index 01157f99d..41fd29ba7 100644 --- a/docs/src/LCDM.md +++ b/docs/src/LCDM.md @@ -38,6 +38,10 @@ x, W = SymBoltz.momentum_quadrature(f₀, nx) x² = x .^ 2 ∫dx_x²_f₀(f) = sum(collect(f .* W)) +# Curvature factors sₗ = √(1-(l²-1)K/k²) in perturbations +# max(…, 0) makes k support only multipoles l ≲ k/√K in closed universes +sl(l) = l ≤ 1 ? 1 : √(max(1 - (l^2-1)*K/k^2, 0)) + # 1) Independent variable for time evolution @independent_variables τ # conformal time D = Differential(τ) # derivative operator @@ -46,6 +50,7 @@ D = Differential(τ) # derivative operator pars = @parameters begin k, τ0, # wavenumber and conformal time today h, H0SI, # Hubble parameter in SI units (most equations have units where H0=1 and do not need these) + Ωk0, K, # curvature K = -Ωk0 (K>0 closed, K=0 flat, K<0 open) Ωc0, # cold dark matter Ωb0, YHe, fHe, κ0, # baryons and recombination Tγ0, Ωγ0, # photons @@ -87,8 +92,8 @@ eqs = [ χ ~ τ0 - τ # gravity equations - D(a) ~ √(8π/3 * ρ) * a^2 # 1st Friedmann equation - D(Φ) ~ -4π/3*a^2/ℋ*δρ - k^2/(3ℋ)*Φ - ℋ*Ψ + D(a) ~ √(8π/3 * ρ * a^2 - K) * a # 1st Friedmann equation + D(Φ) ~ -4π/3*a^2/ℋ*δρ - (k^2-3K)/(3ℋ)*Φ - ℋ*Ψ k^2 * (Φ - Ψ) ~ 12π * a^2 * Π ρ ~ ρc + ρb + ργ + ρν + ρh + ρΛ P ~ Pγ + Pν + Ph + PΛ @@ -165,17 +170,17 @@ eqs = [ wγ ~ 1/3 Pγ ~ wγ * ργ D(Fγ0) ~ -k*Fγ[1] + 4*D(Φ) - D(Fγ[1]) ~ k/3*(Fγ0-2Fγ[2]+4Ψ) - 4/3 * D(κ)/k * (θb - θγ) - [D(Fγ[l]) ~ k/(2l+1) * (l*Fγ[l-1] - (l+1)*Fγ[l+1]) + D(κ) * (Fγ[l] - δkron(l,2)/10*Πγ) for l in 2:lγmax-1]... - D(Fγ[lγmax]) ~ k*Fγ[lγmax-1] - (lγmax+1) / τ * Fγ[lγmax] + D(κ) * Fγ[lγmax] + D(Fγ[1]) ~ k/3*(Fγ0-2*sl(2)*Fγ[2]+4Ψ) - 4/3 * D(κ)/k * (θb - θγ) + [D(Fγ[l]) ~ k/(2l+1) * (l*sl(l)*Fγ[l-1] - (l+1)*sl(l+1)*Fγ[l+1]) + D(κ) * (Fγ[l] - δkron(l,2)/10*Πγ) for l in 2:lγmax-1]... + D(Fγ[lγmax]) ~ k*sl(lγmax)*Fγ[lγmax-1] - (lγmax+1) * cotK(K, τ) * Fγ[lγmax] + D(κ) * Fγ[lγmax] δγ ~ Fγ0 θγ ~ 3k*Fγ[1]/4 - σγ ~ Fγ[2]/2 + σγ ~ Fγ[2]/(2*sl(2)) Πγ ~ Fγ[2] + Gγ0 + Gγ[2] D(Gγ0) ~ k * (-Gγ[1]) + D(κ) * (Gγ0 - Πγ/2) - D(Gγ[1]) ~ k/(2*1+1) * (1*Gγ0 - 2*Gγ[2]) + D(κ) * Gγ[1] - [D(Gγ[l]) ~ k/(2l+1) * (l*Gγ[l-1] - (l+1)*Gγ[l+1]) + D(κ) * (Gγ[l] - δkron(l,2)/10*Πγ) for l in 2:lγmax-1]... - D(Gγ[lγmax]) ~ k*Gγ[lγmax-1] - (lγmax+1) / τ * Gγ[lγmax] + D(κ) * Gγ[lγmax] + D(Gγ[1]) ~ k/(2*1+1) * (1*Gγ0 - 2*sl(2)*Gγ[2]) + D(κ) * Gγ[1] + [D(Gγ[l]) ~ k/(2l+1) * (l*sl(l)*Gγ[l-1] - (l+1)*sl(l+1)*Gγ[l+1]) + D(κ) * (Gγ[l] - δkron(l,2)/10*Πγ) for l in 2:lγmax-1]... + D(Gγ[lγmax]) ~ k*sl(lγmax)*Gγ[lγmax-1] - (lγmax+1) * cotK(K, τ) * Gγ[lγmax] + D(κ) * Gγ[lγmax] # cold dark matter ρc ~ 3/8π * Ωc0 / a^3 @@ -189,12 +194,12 @@ eqs = [ Pν ~ wν * ρν Tν ~ Tν0 / a D(Fν0) ~ -k*Fν[1] + 4*D(Φ) - D(Fν[1]) ~ k/3*(Fν0-2Fν[2]+4Ψ) - [D(Fν[l]) ~ k/(2l+1) * (l*Fν[l-1] - (l+1)*Fν[l+1]) for l in 2:lνmax-1]... - D(Fν[lνmax]) ~ k*Fν[lνmax-1] - (lνmax+1) / τ * Fν[lνmax] + D(Fν[1]) ~ k/3*(Fν0-2*sl(2)*Fν[2]+4Ψ) + [D(Fν[l]) ~ k/(2l+1) * (l*sl(l)*Fν[l-1] - (l+1)*sl(l+1)*Fν[l+1]) for l in 2:lνmax-1]... + D(Fν[lνmax]) ~ k*sl(lνmax)*Fν[lνmax-1] - (lνmax+1) * cotK(K, τ) * Fν[lνmax] δν ~ Fν0 θν ~ 3k*Fν[1]/4 - σν ~ Fν[2]/2 + σν ~ Fν[2]/(2*sl(2)) # massive neutrinos Th ~ Th0 / a @@ -209,13 +214,13 @@ eqs = [ Δh ~ δh + 3ℋ*(1+wh)*θh/k^2 uh ~ ∫dx_x²_f₀(x .* ψh[:,1]) / (Iρh + IPh/3) θh ~ k * uh - σh ~ 2/3 * ∫dx_x²_f₀(x² ./ Eh .* ψh[:,2]) / (Iρh + IPh/3) + σh ~ 2/3 * ∫dx_x²_f₀(x² ./ Eh .* ψh[:,2]) / (Iρh + IPh/3) / sl(2) csh2 ~ ∫dx_x²_f₀(x² ./ Eh .* ψh0) / Iδρh [Eh[i] ~ √(x[i]^2 + yh^2) for i in 1:nx]... [D(ψh0[i]) ~ -k * x[i]/Eh[i] * ψh[i,1] - D(Φ) * dlnf₀_dlnx(x[i]) for i in 1:nx]... - [D(ψh[i,1]) ~ k/3 * x[i]/Eh[i] * (ψh0[i] - 2ψh[i,2]) - k/3 * Eh[i]/x[i] * Ψ * dlnf₀_dlnx(x[i]) for i in 1:nx]... - [D(ψh[i,l]) ~ k/(2l+1) * x[i]/Eh[i] * (l*ψh[i,l-1] - (l+1) * ψh[i,l+1]) for i in 1:nx, l in 2:lhmax-1]... - [D(ψh[i,lhmax]) ~ k/(2lhmax+1) * x[i]/Eh[i] * (lhmax*ψh[i,lhmax-1] - (lhmax+1) * ((2lhmax+1) * Eh[i]/x[i] * ψh[i,lhmax] / (k*τ) - ψh[i,lhmax-1])) for i in 1:nx]... + [D(ψh[i,1]) ~ k/3 * x[i]/Eh[i] * (ψh0[i] - 2*sl(2)*ψh[i,2]) - k/3 * Eh[i]/x[i] * Ψ * dlnf₀_dlnx(x[i]) for i in 1:nx]... + [D(ψh[i,l]) ~ k/(2l+1) * x[i]/Eh[i] * (l*sl(l)*ψh[i,l-1] - (l+1)*sl(l+1) * ψh[i,l+1]) for i in 1:nx, l in 2:lhmax-1]... + [D(ψh[i,lhmax]) ~ k * x[i]/Eh[i] * sl(lhmax)*ψh[i,lhmax-1] - (lhmax+1) * cotK(K, τ) * ψh[i,lhmax] for i in 1:nx]... # dark energy (cosmological constant or w0wa) wΛ ~ w0 + wa*(1-a) @@ -236,12 +241,12 @@ eqs = [ Δm ~ (ρb*Δb + ρc*Δc + ρh*Δh) / ρm # CMB source functions - ST_SW ~ v * (δγ/4 + Ψ + Πγ/16) + ST_SW ~ v * (δγ/4 + Ψ + Πγ/(16*sl(2))) # quadrupole source is the physical Πγ = Πγ_hierarchy/s₂ (see σγ) ST_ISW ~ exp(-κ) * D(Ψ + Φ) |> expand_derivatives ST_Doppler ~ D(v*θb) / k^2 |> expand_derivatives - ST_polarization ~ 3/(16k^2) * D(D(v*Πγ)) |> expand_derivatives + ST_polarization ~ 3/(16k^2*sl(2)) * D(D(v*Πγ)) |> expand_derivatives ST ~ ST_SW + ST_ISW + ST_Doppler + ST_polarization - SE ~ 3/16 * v*Πγ / (k*χ)^2 + SE ~ 3/16 * v*Πγ / (k*sinK(K,χ))^2 / sl(2) # angular size of a comoving scale is set by sinK(K,χ), not χ, in a curved universe Sψ ~ 0 # ifelse(τ ≥ τrec, -(g.Ψ+g.Φ) * (τ-τrec)/(τ0-τrec)/(τ0-τ), 0) # TODO # hide ] @@ -258,12 +263,12 @@ initialization_eqs = [ # photons Fγ0 ~ -2Ψ Fγ[1] ~ 2/3 * k*τ*Ψ - Fγ[2] ~ -8/15 * k/D(κ) * Fγ[1] - [Fγ[l] ~ -l/(2l+1) * k/D(κ) * Fγ[l-1] for l in 3:lγmax]... + Fγ[2] ~ -8/15 * k*sl(2)/D(κ) * Fγ[1] + [Fγ[l] ~ -l/(2l+1) * k*sl(l)/D(κ) * Fγ[l-1] for l in 3:lγmax]... Gγ0 ~ 5/16 * Fγ[2] - Gγ[1] ~ -1/16 * k/D(κ) * Fγ[2] + Gγ[1] ~ -(5-2*sl(2))/48 * k/D(κ) * Fγ[2] # → -1/16 as K → 0 Gγ[2] ~ 1/16 * Fγ[2] - [Gγ[l] ~ -l/(2l+1) * k/D(κ) * Gγ[l-1] for l in 3:lγmax]... + [Gγ[l] ~ -l/(2l+1) * k*sl(l)/D(κ) * Gγ[l-1] for l in 3:lγmax]... # cold dark matter δc ~ -3/2 * Ψ @@ -272,13 +277,13 @@ initialization_eqs = [ # massless neutrinos δν ~ -2 * Ψ θν ~ 1/2 * (k^2*τ) * Ψ - σν ~ 1/15 * (k*τ)^2 * Ψ - [Fν[l] ~ l/(2l+1) * k*τ * Fν[l-1] for l in 3:lνmax]... + Fν[2] ~ 2*sl(2)/15 * (k*τ)^2 * Ψ # i.e. σν = (kτ)²Ψ/15, curvature-independent at leading order; set on the hierarchy moment so it stays well-defined when s₂ = 0 + [Fν[l] ~ l/(2l+1) * k*τ*sl(l) * Fν[l-1] for l in 3:lνmax]... # massive neutrinos [ψh0[i] ~ -1/4 * (-2Ψ) * dlnf₀_dlnx(x[i]) for i in 1:nx]... [ψh[i,1] ~ -1/3 * Eh[i]/x[i] * (1/2*k*τ*Ψ) * dlnf₀_dlnx(x[i]) for i in 1:nx]... - [ψh[i,2] ~ -1/2 * (1/15*(k*τ)^2*Ψ) * dlnf₀_dlnx(x[i]) for i in 1:nx]... + [ψh[i,2] ~ -1/2 * (sl(2)/15*(k*τ)^2*Ψ) * dlnf₀_dlnx(x[i]) for i in 1:nx]... [ψh[i,l] ~ 0 for i in 1:nx, l in 3:lhmax]... # dark energy (w0wa) @@ -312,7 +317,8 @@ initial_conditions = [ Ων0 => Neff * 7/8 * (4/11)^(4/3) * Ωγ0 Nh => 3 Th0 => (4/11)^(1/3) * Tγ0 - ΩΛ0 => 1 - Ωγ0 - Ωc0 - Ωb0 - Ων0 - Ωh0 + K => -Ωk0 + ΩΛ0 => 1 - Ωγ0 - Ωc0 - Ωb0 - Ων0 - Ωh0 - Ωk0 Ωγ0 => π^2/15 * (kB*Tγ0)^4 / (ħ^3*c^5) * 8π*GN / (3*H0SI^2) mh => mh_eV * eV/c^2 yh0 => mh*c^2 / (kB*Th0) @@ -332,6 +338,7 @@ Now set remaining parameter values and compile the numerical problem: ```@example LCDM p = Dict( M.h => 0.7, + M.Ωk0 => 0.1, M.Ωc0 => 0.3, M.Ωb0 => 0.05, M.YHe => 0.25, @@ -369,10 +376,9 @@ plot(log10.(ks), log10.(transpose(Ps)), xlabel = "log10(k / (H₀/c))", ylabel = Now compute the CMB power spectrum: ```@example LCDM -jl = SphericalBesselCache(25:25:3000) -ls = 25:3000 modes = [:TT, :EE, :TE] -Dls = spectrum_cmb(modes, prob, jl, ls; normalization = :Dl) +ls = 25:3000 +Dls = spectrum_cmb(modes, prob, 25:25:3000, ls; normalization = :Dl) plot(ls, Dls[:,1]*1e12, ylabel = "10¹² D(ℓ)", label = "TT", subplot = 1, color = 1, layout = (3, 1), size = (600, 1000), left_margin=5*Plots.mm) plot!(ls, Dls[:,2]*1e12, ylabel = "10¹² D(ℓ)", label = "EE", subplot = 2, color = 2) plot!(ls, Dls[:,3]*1e12, ylabel = "10¹² D(ℓ)", label = "TE", subplot = 3, color = 3, xlabel = "ℓ") diff --git a/src/SymBoltz.jl b/src/SymBoltz.jl index a917943c1..8880bc4b5 100644 --- a/src/SymBoltz.jl +++ b/src/SymBoltz.jl @@ -45,6 +45,7 @@ include("models/cosmologies.jl") include("solve.jl") include("observables/distances.jl") include("observables/fourier.jl") +include("bessel.jl") include("observables/angular.jl") include("parameters.jl") include("plot.jl") @@ -59,6 +60,7 @@ export AbstractInterpolator, EquispacedInterpolator, CubicSplineInterpolator, Ch export express_derivatives export lingrid, loggrid, cosgrid, chebgrid, joingrids!, kτ0grid_default export plot_interactive +export sinK, cotK, asinK using PrecompileTools: @compile_workload @compile_workload begin diff --git a/src/bessel.jl b/src/bessel.jl new file mode 100644 index 000000000..0cc5c8f5e --- /dev/null +++ b/src/bessel.jl @@ -0,0 +1,403 @@ +using Bessels: besselj!, sphericalbesselj +using ForwardDiff +using ForwardDiffChainRules +import ChainRulesCore + +struct SphericalBesselCache{Tl, Tdy <: Union{Matrix{Float64}, Nothing}} + l::Tl + y::Matrix{Float64} + dy::Tdy + dx::Float64 + invdx::Float64 + x::Vector{Float64} +end + +""" + SphericalBesselCache(ls; xmax = 20*maximum(ls), dx = 2π/15, hermite = true, thread = true) + +Create interpolation cache for the spherical Bessel function ``jₗ(x)`` for orders `ls` for `0 ≤ x ≤ xmax` with uniform spacing `dx`. +If `hermite`, cubic Hermite interpolation is used with the analytical derivative ``jₗ′(x)`` instead of linear interpolation. +The computation uses fast recurrence relations when `ls` contains `Integer` orders only, +and otherwise falls back to explicit ``jₗ(x)`` evaluation for every ``l`` and ``x``. +If `thread`, the tabulation is parallellized over independent ``x``. +""" +function SphericalBesselCache(ls; xmax = 20*maximum(ls), dx = 2π/15, hermite = true, thread = true) + xmin = 0.0 + xs = range(xmin, xmax, length = trunc(Int, (xmax - xmin) / dx)) # fixed length (so endpoints are exact) that gives step as close to dx as possible + invdx = 1.0 / step(xs) # using the resulting step, which need not be exactly dx + xs = collect([xs; xs[end]]) # pad with 1 extra duplicate point to avoid bounds check during interpolation + ys = Matrix{Float64}(undef, length(ls), length(xs)) # contiguous in l + dys = hermite ? similar(ys) : nothing + jl_table!(ys, dys, ls, xs; thread) + return SphericalBesselCache{typeof(ls), typeof(dys)}(ls, ys, dys, dx, invdx, xs) +end + +# First argument is the cache index il, not the multipole l +@inline Base.@propagate_inbounds @fastmath function (jl::SphericalBesselCache{Tl, Nothing})(il::Int, x) where {Tl} + w = x * jl.invdx # 0-based float index (assume x0 = 0) + i = trunc(Int, w) # 0-based integer index of left interval point; faster than searchsortedfirst(jl.x, x) + w = w - i # remainder ∈ [0, 1] + y₋ = jl.y[il, i+1] # +1 for 1-based indexing + y₊ = jl.y[il, i+2] + return muladd(w, y₊ - y₋, y₋) # i.e. y₋ + (y₊ - y₋) * (x - x₋) * jl.invdx +end + +@inline Base.@propagate_inbounds @fastmath function (jl::SphericalBesselCache{Tl, Matrix{Float64}})(il::Int, x) where {Tl} + w = x * jl.invdx + i = trunc(Int, w) + w = w - i + wm1 = w - 1.0 + y₋ = jl.y[il, i+1] + y₊ = jl.y[il, i+2] + dy₋ = jl.dy[il, i+1] + dy₊ = jl.dy[il, i+2] + return (1+2w)*wm1*wm1 * y₋ + w*w*(3-2w) * y₊ + w*wm1 * (wm1 * dy₋ + w * dy₊) * jl.dx # https://en.wikipedia.org/wiki/Cubic_Hermite_spline +end + +function Base.show(io::IO, jl::SphericalBesselCache{Tl, Tdy}) where {Tl, Tdy} + method = Tdy == Nothing ? "linear" : "Hermite" + print(io, "jₗ(x) $method interpolation cache ") + print(io, "for $(minimum(jl.l)) ≤ l ≤ $(maximum(jl.l)) and ") + print(io, "$(jl.x[begin]) ≤ x ≤ $(jl.x[end]) ") + print(io, "($(Base.format_bytes(Base.summarysize(jl))))\n") +end + +# Out-of-place spherical Bessel function variants +jl(l, x) = sphericalbesselj(l, x) # for l ≥ 0, from Bessels.jl +jl′(l, x) = iszero(l) ? -jl(one(l), x) : l/(2l+1)*jl(l-1,x) - (l+1)/(2l+1)*jl(l+1,x) # analytical relation with special case for j₀′(x) = -j₁(x), where the general relation would evaluate j₋₁, which diverges at x = 0, with zero weight) + +# In-place spherical Bessel function variants +# TODO: contribute back to Bessels.jl +function jl!(out, l::AbstractRange, x::Number) + besselj!(out, l .+ 0.5, x) + if x == 0.0 && l[begin] == 0 + out[begin] = 1.0 + elseif x != 0.0 + out .*= √(π/(2*x)) + end + return out +end +function jlsafe!(out, l::AbstractRange, x::Number) + out .= jl.(l, x) + return out +end +function jl′(l, ls::AbstractRange, Jls) + i = 1 + l - ls[begin] # ls[i] == l (assuming step of ls is 1) + return l/(2l+1)*Jls[i-1] - (l+1)/(2l+1)*Jls[i+1] # analytical result (see e.g. https://arxiv.org/pdf/astro-ph/9702170 eq. (13)-(15)) +end + +# Overload chain rule for spherical Bessel function +ChainRulesCore.frule((_, _, Δx), ::typeof(jl), l, x) = jl(l, x), jl′(l, x) * Δx # (value, derivative) +@ForwardDiff_frule jl(l::Integer, x::ForwardDiff.Dual) # define dispatch + +#= +Recurrence for spherical Bessel functions jₗ(k*χ) and hyperspherical Bessel functions Φₗ(χ, k) +inspired by arXiv:1312.2697 and arXiv:1311.0839 and CLASS' hyperspherical.c. + +Forward/upward recursion is stable for χ above the turning point where k sinK(K,χ) = √(l(l+1)): + 1. seed Φ₀ = sin(kχ) / (k sinK(K,χ)) + 2. seed Φ₁ = (cotK(K,χ) Φ₀ - cos(kχ)/sinK(K,χ)) / √K₁ + 3. iterate √Kₗ Φₗ = (2l-1) cotK(K,χ) Φₗ₋₁ - √Kₗ₋₁ Φₗ₋₂ (where √Kₗ = √(k² - K l²)) + +Backward/downward recursion is stable for χ below the turning point: + 1. seed Φₗₘₐₓ₊₁/Φₗₘₐₓ from a continued fraction expansion + 2. iterate √Kₗ Φₗ₋₁ = (2l+1) cotK(K,χ) Φₗ - √Kₗ₊₁ Φₗ₊₁ + 3. rescale all Φₗ depending on Φ₀ +=# + +# Tabulate spherical Bessel functions jₗ(x) and its derivative jₗ′(x) (if passed) using integer-l recurrence (faster) +function jl_table!(ys::AbstractMatrix, dys::Union{AbstractMatrix, Nothing}, ls::AbstractArray{<:Integer}, xs::AbstractVector; thread = true) + minimum(ls) ≥ 0 || throw(ArgumentError("multipoles must be non-negative, but got minimum(ls) = $(minimum(ls))")) + lmax = maximum(ls) + @inbounds @tasks for ix in eachindex(xs) # parallelized over independent x + @set scheduler = thread ? :dynamic : :serial + @local jls = zeros(Float64, lmax+2) # one multipole higher than requested, because jₗ′(x) depends on jₗₘₐₓ₊₁(x) + jl_recurrence!(jls, lmax+1, xs[ix]) + for (il, l) in enumerate(ls) + ys[il, ix] = jls[l+1] + if !isnothing(dys) + dys[il, ix] = if l == 0 + -jls[2] # j₀′(x) = -j₁(x) (general relation would index jls[0] = j₋₁ with zero weight) + else + (l*jls[l] - (l+1)*jls[l+2]) / (2l+1) # jₗ′(x) = l/(2l+1) * jₗ₋₁(x) - (l+1)/(2l+1) * jₗ₊₁(x) + end + end + end + end + return ys, dys +end + +# When l is not integer, fall back to evaluating jₗ(x) and jₗ′(x) with Bessels.jl for arbitrary l (slower) +function jl_table!(ys::AbstractMatrix, dys::Union{AbstractMatrix, Nothing}, ls, xs::AbstractVector; thread = true) + @inbounds @tasks for ix in eachindex(xs) # parallelized over independent x + @set scheduler = thread ? :dynamic : :serial + x = xs[ix] + for (il, l) in enumerate(ls) + ys[il, ix] = jl(l, x) + if !isnothing(dys) + dys[il, ix] = jl′(l, x) + end + end + end + return ys, dys +end + +""" + jl_recurrence!(out, lmax, x) + +Compute spherical Bessel functions ``jₗ(x)`` for every integer ``l = 0, \\dots, l_\\mathrm{max}`` and save them +in `out[l+1]`. Uses forward/backward recursion above/below the turning point ``√(lmax*(lmax+1))``. +""" +function jl_recurrence!(out::AbstractVector, lmax::Integer, x) + if x == 0.0 + fill!(out, zero(eltype(x))) # jₗ(0) = 0 for l > 0 (the recursions would divide by 0) + out[1] = one(eltype(x)) # j₀(0) = 1 + elseif x ≥ √(lmax*(lmax+1)) # "turning point": above it jₗ(x) is oscillatory, below it exponentially suppressed + jl_forward!(out, lmax, x) + else + jl_backward!(out, lmax, x) + end + return out +end + +function jl_forward!(out::AbstractVector, lmax::Integer, x) + invx = 1 / x + out[1] = sin(x) * invx # j₀(x) = sin(x) / x + out[2] = (out[1] - cos(x)) * invx # j₁(x) = sin(x)/x² - cos(x)/x + @inbounds for l in 1:lmax-1 + out[l+2] = (2l+1) * invx * out[l+1] - out[l] # jₗ₊₁(x) = (2l+1)/x jₗ(x) - jₗ₋₁(x) + end + return out +end + +function jl_backward!(out::AbstractVector, lmax::Integer, x) + invx = 1 / x # precompute; cheaper to multiply by 1/x instead of dividing by x + jₗ = 2.0^-900 # arbitrary nonzero jₗₘₐₓ seed; renormalized at the end + jₗ′_jₗ = jl_logderiv(lmax, x) # continued fraction of jₗₘₐₓ′/jₗₘₐₓ + jₗ₊₁ = jₗ * (lmax*invx - jₗ′_jₗ) # jₗₘₐₓ₊₁ from derivative relation jₗ′ = l/x jₗ - jₗ₊₁ + out[lmax+1] = jₗ + @inbounds for l in lmax:-1:1 + jₗ₋₁ = (2l+1) * invx * jₗ - jₗ₊₁ + if abs(jₗ₋₁) > 2.0^900 # renormalize previously computed jₗ before overflow + s = 2.0^-900 + jₗ₋₁ *= s + jₗ *= s + @views out[l+1:lmax+1] .*= s + end + jₗ₊₁ = jₗ + jₗ = jₗ₋₁ + out[l] = jₗ + end + s = (sin(x) * invx) / out[1] + out .*= s # renormalize all jₗ to match analytical j₀ + return out +end + +# Evaluate jₗ′(x)/jₗ(x) from its continued fraction expansion with the modified Lentz algorithm +function jl_logderiv(l::Integer, x) + invx = 1 / x + f = l * invx + C = f + D = zero(typeof(x)) + @inbounds for j in 1:1000 # increase max iterations if not converging + b = (2*(l+j)+1) * invx + D = 1 / (b - D) + C = b - 1/C + CD = C * D + f *= CD # fₗ = jₗ′/jₗ + abs(CD - 1) < eps(typeof(x)) && return f # converged + end + error("Continued fraction for jₗ′(x)/jₗ(x) did not converge at (l, x) = ($l, $x)") +end + +""" + Φl_recurrence!(out, lmax::Integer, χ, k, K, sqK::AbstractVector, invsqK::AbstractVector) + +Compute hyperspherical Bessel functions ``Φₗ(χ, k)`` of a universe with curvature `K` for every integer +``l = 0, \\dots, l_\\mathrm{max}`` and store them in `out[l+1]`. Uses forward/backward recurrence +above/below the turning point, where ``k * sinK(K, χ) ≥ √(lmax*(lmax+1))``. +""" +function Φl_recurrence!(out::AbstractVector, lmax::Integer, χ, k, K, sqK::AbstractVector, invsqK::AbstractVector) + if χ == 0.0 + fill!(out, zero(eltype(χ))) # Φₗ(0) = 0 for l > 0 (the recursions would divide by 0) + out[1] = one(eltype(χ)) # Φ₀(0) = 1 + elseif k * sinK(K, χ) ≥ √(lmax*(lmax+1)) # above or below turning point? + Φl_forward!(out, lmax, χ, k, K, sqK, invsqK) + else + Φl_backward!(out, lmax, χ, k, K, sqK, invsqK) + end + return out +end + +""" + Φl_recurrence!(out::AbstractMatrix, ls::AbstractArray{<:Integer}, χs::AbstractVector, k, K, sqK::AbstractVector, invsqK::AbstractVector; Φmin = 1e-20) + +Same as above for every radial coordinate in `χs` at once, storing ``Φ_l(χ_i)`` in `out[i, l+1]` (so `out` must +be `length(χs)` × ``l_\\mathrm{max}+1``). `χs` must be sorted in descending order, as it is in the line-of-sight +integral, so that the parts above and below the turning point are contiguous views that can be handed to the +(vectorized) forward and backward recursion. Tables for `sqK` and `invsqK` are precomputed unless given explictly. +""" +function Φl_recurrence!(out::AbstractMatrix, ls::AbstractArray{<:Integer}, χs::AbstractVector, k, K, sqK::AbstractVector, invsqK::AbstractVector; Φmin = 1e-20) + K ≤ 0 || throw(ArgumentError("K must be non-positive (flat or open universe), but got $K")) + size(out) == (length(χs), length(ls)) || throw(ArgumentError("out size is $(size(out)), but should be $((length(χs), length(ls)))")) + issorted(χs; rev = true) || throw(ArgumentError("χs is not sorted in descending order")) + issorted(ls) || throw(ArgumentError("ls is not sorted in descending order")) + lmax = ls[end] + turningpoint = √(lmax*(lmax+1)) + ifwd = 1 : searchsortedlast(k .* sinK.(K, χs), turningpoint; rev = true) # LHS monotonically increasing for K < 0 + ibwd = (last(ifwd) + 1) : (iszero(χs[end]) ? length(χs) - 1 : length(χs)) + izero = (last(ibwd) + 1) : length(χs) + if !isempty(ifwd) + @views Φl_forward!(out[ifwd, :], ls, χs[ifwd], k, K, sqK, invsqK) + end + if !isempty(ibwd) + @views Φl_backward!(out[ibwd, :], ls, χs[ibwd], k, K, sqK, invsqK; Φmin) + end + if !isempty(izero) + @views out[izero, :] .= 0.0 # Φₗ(0) = 0 for l > 0 (the recursions would divide by 0) + if ls[1] == 0 + @views out[izero, 1] .= 1.0 # Φ₀(0) = 1, only if l = 0 was requested + end + end + return out +end + +# Same, but without sqK and invsqK arrays +function Φl_recurrence!(out::AbstractMatrix, ls::AbstractArray{<:Integer}, χs::AbstractVector, k, K; kwargs...) + K ≤ 0 || throw(ArgumentError("K must be non-positive (flat or open universe), but got $K")) + lmax = ls[end] + sqK = zeros(Float64, lmax+2) # TODO: don't require Float64 + invsqK = zeros(Float64, lmax+2) + sqrtK_table!(sqK, invsqK, K, k) + return Φl_recurrence!(out, ls, χs, k, K, sqK, invsqK; kwargs...) +end + +function Φl_forward!(out::AbstractVector, lmax::Integer, χ, k, K, sqK::AbstractVector, invsqK::AbstractVector) + ck = cotK(K, χ) + sk = sinK(K, χ) + out[1] = sin(k*χ) / (k * sk) # Φ₀ + out[2] = (ck * out[1] - cos(k*χ) / sk) * invsqK[2] # Φ₁ = (cotK Φ₀ - k cot(kχ) Φ₀) / √K₁ + @inbounds for l in 1:lmax-1 + out[l+2] = ((2l+1) * ck * out[l+1] - sqK[l+1] * out[l]) * invsqK[l+2] # sqK[l+1] = √Kₗ, invsqK[l+2] = 1/√Kₗ₊₁ + end + return out +end + +@inline function Φₗ_save_if_requested!(out, ls, l, Φ, lstoreidx) + if l == ls[lstoreidx] + @inbounds for i in 1:size(out, 1) + out[i, lstoreidx] = Φ[i] + end + return lstoreidx + 1 + end + return lstoreidx +end + +# Same for many radial coordinates at once, with out[i, l+1] = Φₗ(χs[i]). The recursion in l is serial, but +# different χ are independent, so the χ-loop is innermost and vectorizes (hence χ is the contiguous dimension). +function Φl_forward!(out::AbstractMatrix, ls::AbstractArray{<:Integer}, χs::AbstractVector, k, K, sqK::AbstractVector, invsqK::AbstractVector) + ck = cotK.(K, χs) # depends on χ but not on l, so tabulate it once instead of recomputing it every step + Φₗ₋₁ = similar(χs) + Φₗ = similar(χs) + for i in eachindex(χs) # TODO: dot syntax + sk = sinK(K, χs[i]) + Φₗ₋₁[i] = sin(k*χs[i]) / (k * sk) # Φ₀ + Φₗ[i] = (ck[i] * Φₗ₋₁[i] - cos(k*χs[i]) / sk) * invsqK[2] # Φ₁ + end + + lstoreidx = 1 # ls[lstoreidx] is the next multipole to output + lstoreidx = Φₗ_save_if_requested!(out, ls, 0, Φₗ₋₁, lstoreidx) # save Φ₀? + lstoreidx = Φₗ_save_if_requested!(out, ls, 1, Φₗ, lstoreidx) # save Φ₁? + + @inbounds for l in 1:ls[end]-1 + Φₗ₊₁ = Φₗ₋₁ # reuse array name (Φₗ₋₁[i] is not needed anymore writing Φₗ₊₁[i]) + @simd for i in eachindex(χs) # since it does not depend on χ + Φₗ₊₁[i] = ((2l+1) * ck[i] * Φₗ[i] - sqK[l+1] * Φₗ₋₁[i]) * invsqK[l+2] + end + lstoreidx = Φₗ_save_if_requested!(out, ls, l+1, Φₗ₊₁, lstoreidx) # save Φₗ₊₁? + Φₗ, Φₗ₋₁ = Φₗ₋₁, Φₗ # swap array pointers for next iteration + end + return out +end + +function Φl_backward!(out::AbstractVector, lmax::Integer, χ, k, K, sqK::AbstractVector, invsqK::AbstractVector) + ck = cotK(K, χ) + Φₗ = 2.0^-900 # Φₗₘₐₓ, arbitrary nonzero seed (fixed by the final normalization) + Φₗ′_Φₗ = Φl_logderiv(lmax, χ, k, K) # continued fraction of Φₗₘₐₓ′/Φₗₘₐₓ + Φₗ₊₁ = Φₗ * (lmax*ck - Φₗ′_Φₗ) * invsqK[lmax+2] # Φₗₘₐₓ₊₁ from the derivative recursion Φₗ′ = l cotK Φₗ - √Kₗ₊₁ Φₗ₊₁ + out[lmax+1] = Φₗ + @inbounds for l in lmax:-1:1 + Φₗ₋₁ = ((2l+1) * ck * Φₗ - sqK[l+2] * Φₗ₊₁) * invsqK[l+1] # sqK[l+2] = √Kₗ₊₁, invsqK[l+1] = 1/√Kₗ + if abs(Φₗ₋₁) > 2.0^900 # renormalize previously computed Φₗ before overflow + s = 2.0^-900 + Φₗ₋₁ *= s + Φₗ *= s + @views out[l+1:lmax+1] .*= s + end + Φₗ₊₁ = Φₗ + Φₗ = Φₗ₋₁ + out[l] = Φₗ + end + s = sin(k*χ) / (k * sinK(K, χ)) / out[1] + out .*= s + return out +end + +# Same for many radial coordinates at once, with out[i, l+1] = Φₗ(χs[i]); see Φl_forward! above. The recurrence +# and the overflow check are split into separate loops so that the recurrence can be @simd-vectorized over χ. +function Φl_backward!(out::AbstractMatrix, ls::AbstractArray{<:Integer}, χs::AbstractVector, k, K, sqK::AbstractVector, invsqK::AbstractVector; Φmin = 0.0) + lmax = ls[end] + _out = zeros(lmax + 1) # TODO: take as input argument? + @inbounds for i in eachindex(χs) + #println("χ = $(χs[i]), lmax = $lmax") + Φl_backward!(_out, lmax, χs[i], k, K, sqK, invsqK) + + # Lower lmax to skip the recurrence Φₗ where it is below a tiny threshold + # Immediately set Φₗ = 0 there + # Works because Φₗ decreases monotonically as l decreases when χ is below the turning point + while lmax ≥ 1 && abs(_out[1+lmax]) < Φmin + _out[1+lmax] = 0.0 + lmax -= 1 + end + + for j in eachindex(ls) + out[i, j] = _out[1 + ls[j]] + end + end + return out +end + +# Evaluate Φₗ'(χ; k)/Φₗ(χ; k) from its continued fraction expansion with the modified Lentz algorithm +function Φl_logderiv(l::Integer, χ, k, K) + ck = cotK(K, χ) + σ(m) = √(k^2 - K*m^2) # √Kₘ, computed on the fly since the fraction runs to arbitrarily high m + nudge(x) = copysign(max(abs(x), 1e-100), x) # keep x not too close to zero + f = nudge(l * ck) # b₀ + C = f + D = zero(typeof(χ)) + σj = σ(l+1) # √K_{l+j} at j = 1 + @inbounds for j in 1:1000000 # increase max iterations if not converging + σj1 = σ(l+j+1) # √K_{l+j+1} + a = -σj / σj1 + j == 1 && (a *= σj) # the leading √K_{l+1} of the fraction, folded into the first numerator + b = (2*(l+j)+1) * ck / σj1 + D = 1 / nudge(b + a*D) + C = nudge(b + a/C) + CD = C * D + f *= CD # fₗ = Φₗ′/Φₗ + abs(CD - 1) < eps(typeof(χ)) && return f # converged + σj = σj1 + end + error("Continued fraction for Φₗ′(χ)/Φₗ(χ) did not converge at (l, χ, k, K) = ($l, $χ, $k, $K)") +end + +# Tabulate l-dependent coefficients ``√(Kₗ) = √(k² - K l²)`` of the hyperspherical Bessel recurrence and their inverses. +function sqrtK_table!(sqK::AbstractVector, invsqK::AbstractVector, K, k) + @inbounds for i in eachindex(sqK, invsqK) + l = i - 1 + Kₗ = k^2 - K*l^2 + sqK[i] = √(Kₗ) + invsqK[i] = 1 / sqK[i] + end + return sqK, invsqK +end diff --git a/src/interpolation.jl b/src/interpolation.jl index b3f2a62d4..aeb37ab45 100644 --- a/src/interpolation.jl +++ b/src/interpolation.jl @@ -1,4 +1,4 @@ -abstract type AbstractInterpolator{T} end +abstract type AbstractInterpolator{T} <: AbstractVector{T} end # interpolation x-grid is indexable/iterable/broadcastable like a vector struct CubicSplineInterpolator{T, F <: Function} <: AbstractInterpolator{T} xs::Vector{T} # points in input domain: x = f⁻¹(y) (e.g. wavenumbers k) @@ -136,14 +136,15 @@ function ChebyshevIntegerInterpolator(xmin, xmax, order::Integer) return ChebyshevIntegerInterpolator(xs, ys, ws, identity) end -Base.eltype(::Type{<:AbstractInterpolator{T}}) where {T} = T # type of x-points -Base.extrema(interp::AbstractInterpolator) = (minimum(interp), maximum(interp)) -Base.firstindex(interp::AbstractInterpolator) = firstindex(interp.xs) -Base.lastindex(interp::AbstractInterpolator) = lastindex(interp.xs) +# Array interface over the x-points, so an interpolator can be indexed/iterated/broadcast like a vector +Base.size(interp::AbstractInterpolator) = size(interp.xs) +Base.getindex(interp::AbstractInterpolator, i::Int) = interp.xs[i] +Base.IndexStyle(::Type{<:AbstractInterpolator}) = IndexLinear() + +# The x-points are sorted in ascending order, so their extrema are the endpoints (instead of a linear scan) Base.minimum(interp::AbstractInterpolator) = interp[begin] Base.maximum(interp::AbstractInterpolator) = interp[end] -Base.getindex(interp::AbstractInterpolator, i::Int) = interp.xs[i] -Base.iterate(interp::AbstractInterpolator, args...; kwargs...) = iterate(interp.xs, args...; kwargs...) +Base.extrema(interp::AbstractInterpolator) = (minimum(interp), maximum(interp)) # Compute Barycentric interpolation weights wᵢ = 1 / ∏_{j≠i}(xᵢ - xⱼ) for arbitrary points # See https://people.maths.ox.ac.uk/trefethen/barycentric.pdf (section 7) @@ -155,7 +156,6 @@ function baryweights(x::AbstractVector) return w end -Base.length(interp::AbstractInterpolator) = length(interp.xs) order(interp::AbstractInterpolator) = length(interp) - 1 # Barycentric interpolation formula https://epubs.siam.org/doi/10.1137/S0036144502417715 @@ -207,6 +207,7 @@ interpolate(x::AbstractInterpolator, y, x′) = x(y, x.f.(x′)) interpolate(x::PiecewiseChebyshevInterpolator, y, x′) = x(y, x′) # no f field interpolate(x::AbstractVector, y, x′) = interpolate(CubicSplineInterpolator(x), y, x′) +Base.show(io::IO, ::MIME"text/plain", interp::AbstractInterpolator) = show(io, interp) # summarize instead of listing all x-points Base.show(io::IO, interp::CubicSplineInterpolator) = print(io, "Cubic spline interpolator: type = $(eltype(interp)), domain = $(extrema(interp)), order = $(order(interp))") Base.show(io::IO, interp::EquispacedInterpolator) = print(io, "Equispaced polynomial interpolator: type = $(eltype(interp)), domain = $(extrema(interp)), order = $(order(interp))") Base.show(io::IO, interp::ChebyshevInterpolator) = print(io, "Chebyshev polynomial interpolator: type = $(eltype(interp)), domain = $(extrema(interp)), order = $(order(interp))") diff --git a/src/observables/angular.jl b/src/observables/angular.jl index af2acb1fd..5aaa0226a 100644 --- a/src/observables/angular.jl +++ b/src/observables/angular.jl @@ -1,86 +1,6 @@ using Bessels: besselj!, sphericalbesselj using DataInterpolations using MatterPower -using ForwardDiff -using ForwardDiffChainRules -import ChainRulesCore - -struct SphericalBesselCache{Tl, Tdy <: Union{Matrix{Float64}, Nothing}} - l::Tl - y::Matrix{Float64} - dy::Tdy - dx::Float64 - invdx::Float64 - x::Vector{Float64} -end - -function SphericalBesselCache(ls; xmax = 20*maximum(ls), dx = 2π/15, hermite = true) - xmin = 0.0 - xs = range(xmin, xmax, length = trunc(Int, (xmax - xmin) / dx)) # fixed length (so endpoints are exact) that gives step as close to dx as possible - invdx = 1.0 / step(xs) # using the resulting step, which need not be exactly dx - xs = collect([xs; xs[end]]) # pad with 1 extra duplicate point to avoid bounds check during interpolation - ys = jl.(ls, xs') # contiguous in l - dys = hermite ? jl′.(ls, xs') : nothing - return SphericalBesselCache{typeof(ls), typeof(dys)}(ls, ys, dys, dx, invdx, xs) -end - -# First argument is the cache index il, not the multipole l -@inline Base.@propagate_inbounds @fastmath function (jl::SphericalBesselCache{Tl, Nothing})(il::Int, x) where {Tl} - w = x * jl.invdx # 0-based float index (assume x0 = 0) - i = trunc(Int, w) # 0-based integer index of left interval point; faster than searchsortedfirst(jl.x, x) - w = w - i # remainder ∈ [0, 1] - y₋ = jl.y[il, i+1] # +1 for 1-based indexing - y₊ = jl.y[il, i+2] - return muladd(w, y₊ - y₋, y₋) # i.e. y₋ + (y₊ - y₋) * (x - x₋) * jl.invdx -end - -@inline Base.@propagate_inbounds @fastmath function (jl::SphericalBesselCache{Tl, Matrix{Float64}})(il::Int, x) where {Tl} - w = x * jl.invdx - i = trunc(Int, w) - w = w - i - wm1 = w - 1.0 - y₋ = jl.y[il, i+1] - y₊ = jl.y[il, i+2] - dy₋ = jl.dy[il, i+1] - dy₊ = jl.dy[il, i+2] - return (1+2w)*wm1*wm1 * y₋ + w*w*(3-2w) * y₊ + w*wm1 * (wm1 * dy₋ + w * dy₊) * jl.dx # https://en.wikipedia.org/wiki/Cubic_Hermite_spline -end - -function Base.show(io::IO, jl::SphericalBesselCache{Tl, Tdy}) where {Tl, Tdy} - method = Tdy == Nothing ? "linear" : "Hermite" - print(io, "jₗ(x) $method interpolation cache ") - print(io, "for $(minimum(jl.l)) ≤ l ≤ $(maximum(jl.l)) and ") - print(io, "$(jl.x[begin]) ≤ x ≤ $(jl.x[end]) ") - print(io, "($(Base.format_bytes(Base.summarysize(jl))))\n") -end - -# Out-of-place spherical Bessel function variants -jl(l, x) = sphericalbesselj(l, x) # for l ≥ 0, from Bessels.jl -jl′(l, x) = l/(2l+1)*jl(l-1,x) - (l+1)/(2l+1)*jl(l+1,x) # for l ≥ 1, analytical relation - -# In-place spherical Bessel function variants -# TODO: contribute back to Bessels.jl -function jl!(out, l::AbstractRange, x::Number) - besselj!(out, l .+ 0.5, x) - if x == 0.0 && l[begin] == 0 - out[begin] = 1.0 - elseif x != 0.0 - out .*= √(π/(2*x)) - end - return out -end -function jlsafe!(out, l::AbstractRange, x::Number) - out .= jl.(l, x) - return out -end -function jl′(l, ls::AbstractRange, Jls) - i = 1 + l - ls[begin] # ls[i] == l (assuming step of ls is 1) - return l/(2l+1)*Jls[i-1] - (l+1)/(2l+1)*Jls[i+1] # analytical result (see e.g. https://arxiv.org/pdf/astro-ph/9702170 eq. (13)-(15)) -end - -# Overload chain rule for spherical Bessel function -ChainRulesCore.frule((_, _, Δx), ::typeof(jl), l, x) = jl(l, x), jl′(l, x) * Δx # (value, derivative) -@ForwardDiff_frule jl(l::Integer, x::ForwardDiff.Dual) # define dispatch # TODO: line-of-sight integrate Θl using ODE for evolution of Jl? # TODO: spline sphericalbesselj for each l, from x=0 to x=kmax*(τ0-τini) @@ -91,7 +11,7 @@ ChainRulesCore.frule((_, _, Δx), ::typeof(jl), l, x) = jl(l, x), jl′(l, x) * # TODO: use u = k*χ as integration variable, so oscillations of Bessel functions are the same for every k? # TODO: define and document symbolic dispatch! """ - los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, jl::SphericalBesselCache; l_limber = typemax(Int), thread = true, verbose = false) where {T} + los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, jl::SphericalBesselCache, K::Real = 0.0; l_limber = typemax(Int), thread = true, verbose = false) where {T} For the given `ls` and `ks`, compute the line-of-sight integrals ```math @@ -105,7 +25,8 @@ Iₗ ≈ √(π/(2l+1)) S(τ₀-(l+1/2)/k, k) ``` is used for `l ≥ l_limber`. """ -function los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, jl::SphericalBesselCache; l_limber = typemax(Int), thread = true, verbose = false) where {T} +function los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, jl::SphericalBesselCache, K::Real = 0.0; l_limber = typemax(Int), thread = true, verbose = false) where {T} + @assert K == 0 "the spherical Bessel function cache only covers a flat universe (K = 0), but got K = $K; pass the multipoles `ls` instead of a `SphericalBesselCache` to use the hyperspherical recursion" @assert size(Ss, 1) == length(τs) "size(Ss, 1) = $(size(Ss, 1)) and length(τs) = $(length(τs)) differ" @assert size(Ss, 2) == length(ks) "size(Ss, 2) = $(size(Ss, 2)) and length(ks) = $(length(ks)) differ" @assert collect(ls) == collect(jl.l) "ls must match the l-values stored in the Bessel cache" @@ -184,6 +105,118 @@ function los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractV return Is end +""" + los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, ::Nothing, K::Real = 0.0; kwargs...) where {T} + +Same as `los_integrate(Ss, ls, τs, ks, K; ...)`. This lets callers that hold a `jl` which is either a +`SphericalBesselCache` or `nothing` dispatch on it directly, instead of branching on `isnothing(jl)`. +""" +function los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, ::Nothing, K::Real = 0.0; kwargs...) where {T} + return los_integrate(Ss, ls, τs, ks, K; kwargs...) +end + +""" + los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, K::Real = 0.0; l_limber = typemax(Int), thread = true, verbose = false) where {T} + +Same as `los_integrate(Ss, ls, τs, ks, jl::SphericalBesselCache; ...)`, but compute the hyperspherical Bessel function +``Φₗ^k(χ)`` on the fly with recurrences for integer `ls` (without a precomputed cache) and with curvature ``K``. +""" +function los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, K::Real = 0.0; l_limber = typemax(Int), thread = true, verbose = false) where {T} + K ≤ 0 || throw(ArgumentError("line-of-sight integration is only implemented for flat (K = 0) and open (K < 0) universes, but got K = $K (closed)")) + @assert size(Ss, 1) == length(τs) "size(Ss, 1) = $(size(Ss, 1)) and length(τs) = $(length(τs)) differ" + @assert size(Ss, 2) == length(ks) "size(Ss, 2) = $(size(Ss, 2)) and length(ks) = $(length(ks)) differ" + @assert issorted(τs) "τs must be sorted in ascending order" + @assert issorted(ks) "ks must be sorted in ascending order" + @assert issorted(ls) "ls must be sorted in ascending order" # necessary for Limber indexing logic + all(l -> l == round(l), ls) || throw(ArgumentError( + "ls must be exact integer multipoles for the recursion-based los_integrate (got e.g. $(first(Iterators.filter(l -> l != round(l), ls)))). " * + "A coarse/fractional l-grid for interpolation is only supported by the SphericalBesselCache-based method.")) + ls = round.(Int, ls) + error_if_nonfinite(Ss) + + τs = collect(τs) # force array to avoid floating point errors with ranges in following χs due to (e.g. tiny negative χ) + τ0 = τs[end] + χs = τ0 .- τs + nτ = length(τs) + + ws = similar(τs) # precompute trapezoidal rule weights + ws[1] = 0.5 * (τs[2] - τs[1]) + @inbounds for iτ in 2:nτ-1 + ws[iτ] = 0.5 * (τs[iτ+1] - τs[iτ-1]) + end + ws[nτ] = 0.5 * (τs[nτ] - τs[nτ-1]) + + nl = length(ls) + Is = similar(Ss, length(ks), nl) + il_limber = searchsortedfirst(ls, l_limber) # First il index with l ≥ l_limber (=nl+1 when l_limber = typemax, i.e. no Limber modes) + ls_nonlimber = @view ls[1:il_limber-1] # only these need the (non-Limber) recursion; Φl is sized/indexed to match + lmax = il_limber > 1 ? ls[il_limber-1] : 0 # largest l that needs the full (non-Limber) recursion + + verbose && l_limber < typemax(Int) && println("Using Limber approximation for l ≥ $l_limber") + + il_limber == 1 || lmax ≥ 2 || throw(ArgumentError("the recursion needs multipoles l ≥ 2, but got lmax = $lmax below l_limber = $l_limber")) + + @fastmath @inbounds @tasks for ik in eachindex(ks) + @set scheduler = thread ? :dynamic : :serial + @local begin + tmp = zeros(T, nl) + Φl = zeros(Float64, nτ, length(ls_nonlimber)) # Φₗ(χᵢ) for every τᵢ and l ≤ lmax at once (see Φl_recurrence!) + sqK = zeros(Float64, lmax+2) # recursion coefficients √Kₗ = √(k²-Kl²) and 1/√Kₗ for l = 0, …, lmax+1, + invsqK = zeros(Float64, lmax+2) # which depend on k but not on χ (so are tabulated once per k) + end + k = ks[ik] + sqrtK_table!(sqK, invsqK, K, k) + verbose && print("\rLOS integrating k-mode $ik / $(length(ks))") + + # Full line-of-sight integrals for l < l_limber (skipped entirely when every requested l uses Limber) + fill!(tmp, zero(T)) + if il_limber > 1 + Φl_recurrence!(Φl, ls_nonlimber, χs, k, K, sqK, invsqK) # one recursion sweep gives Φₗ(χᵢ) for every τᵢ and l ≤ lmax at once + @inbounds for iτ in eachindex(τs) + Sw = ws[iτ] * Ss[iτ, ik] + @inbounds @simd for il in 1:il_limber-1 + tmp[il] += Sw * Φl[iτ, il] + end + end + end + + # Limber approximation for l ≥ l_limber (does not use Φₗ at all). Φₗ acts like a delta function at its + # turning point, which now sits at k⋅sinK(K,χ) = l+½ instead of at kχ = l+½. Matching turning points is + # exactly CLASS' flat approximation Φₗ(χ) ≈ A⋅jₗ(γkχ) with γk = (l+½)/χₗ (transfer.c, `rescale_argument` + # and `rescale_amplitude`), so the flat Limber result carries over with 1/k → A⋅χₗ/(l+½). + @inbounds for il in il_limber:nl + l = ls[il] + χ = asinK(K, (l + 1/2) / k) # turning point χₗ, where k⋅sinK(K,χ) = l+½ + if χ ≤ χs[1] # otherwise source is zero before recombination + i₋ = searchsortedfirst(τs, τ0 - χ) + χ₋ = χs[i₋] + S₋ = Ss[i₋, ik] + if i₋ == 1 + S = S₋ + else + i₊ = i₋ - 1 # χs is descending, so χ₋ < χ < χ₊ + χ₊ = χs[i₊] + S₊ = Ss[i₊, ik] + Δχ = χ₊ - χ₋ + S′₋ = i₋ ≤ nτ-1 ? (Ss[i₋+1, ik] - S₊) / (χs[i₋+1] - χ₊) : (S₊ - S₋) / Δχ + S′₊ = i₊ ≥ 2 ? (S₋ - Ss[i₋-2, ik]) / (χ₋ - χs[i₋-2]) : (S₊ - S₋) / Δχ + t = (χ - χ₋) / Δχ + t² = t*t + t³ = t²*t + S = (2t³-3t²+1)*S₋ + (t³-2t²+t)*Δχ*S′₋ + (-2t³+3t²)*S₊ + (t³-t²)*Δχ*S′₊ + A = (1 - K*l*(l+1)/k^2)^(-1/12) # amplitude of the same rescaling (CLASS' `rescale_amplitude`); → 1 when K = 0 + tmp[il] = √(π/(2l+1)) * S * A * χ / (l + 1/2) # → the flat √(π/(2l+1))⋅S/k, since χ = (l+½)/k there + end + end + end + + Is[ik, :] .= tmp + end + verbose && println() + + return Is +end + # TODO: integrate splines instead of trapz! https://discourse.julialang.org/t/how-to-speed-up-the-numerical-integration-with-interpolation/96223/5 @doc raw""" spectrum_cmb(ΘlAs::AbstractMatrix, ΘlBs::AbstractMatrix, P0s::AbstractVector, ls::AbstractVector, ks::AbstractVector; normalization = :Cl, thread = true) @@ -250,25 +283,54 @@ jl = SphericalBesselCache(ls) modes = [:TT, :TE, :ψψ, :ψT] Dls = spectrum_cmb(modes, prob, jl; normalization = :Dl, unit = u"μK") ``` + +Alternatively, pass a plain vector of integer multipoles instead of a `jl::SphericalBesselCache`: +```julia +Dls = spectrum_cmb(modes, prob, ls; normalization = :Dl, unit = u"μK") +``` +to compute ``Φ_l(χ)`` on the fly with [`Φl_recurrence!`](@ref) instead of from a precomputed cache (see that method's docstring for the tradeoffs). """ -function spectrum_cmb(modes::AbstractVector{<:Symbol}, prob::CosmologyProblem, jl::SphericalBesselCache; normalization = :Cl, unit = nothing, kinterp = nothing, Δkτ0 = 2π/4, xs = cosgrid(0.0, 1.0; length=300), τcut = 1e-2, l_limber = 10, bgopts = (alg = bgalg(prob), reltol = 1e-7, abstol = 1e-7), ptopts = (alg = ptalg(prob), reltol = 1e-5, abstol = 1e-5), thread = true, verbose = false, kwargs...) +function spectrum_cmb(modes::AbstractVector{<:Symbol}, prob::CosmologyProblem, jl::SphericalBesselCache; kwargs...) + return spectrum_cmb(modes, prob, collect(jl.l); jl, kwargs...) # the cache already knows which l's it holds +end + +""" + spectrum_cmb(modes::AbstractVector{<:Symbol}, prob::CosmologyProblem, ls::AbstractVector; jl = nothing, kwargs...) + +Same as `spectrum_cmb(modes, prob, jl::SphericalBesselCache; kwargs...)`, but without a precomputed spherical +Bessel function cache: ``Φ_l(χ)`` is instead computed on the fly with [`Φl_recurrence!`](@ref) for every +wavenumber and time, which also covers curved universes. `ls` must consist of actual (sorted, ascending) integer +multipoles, since a recursion in `l` can only ever produce values at integer `l` (unlike the cache-based method, +which also supports a coarse, fractional `ls` grid for later interpolation). + +This method carries the implementation shared with the cache-based one, which calls it with `ls = jl.l` and the +cache passed as `jl`. +""" +function spectrum_cmb(modes::AbstractVector{<:Symbol}, prob::CosmologyProblem, ls::AbstractArray; jl::Union{SphericalBesselCache, Nothing} = nothing, normalization = :Cl, unit = nothing, kinterp = nothing, Δkτ0 = 2π/4, xs = cosgrid(0.0, 1.0; length=300), τcut = 1e-2, l_limber = 10, bgopts = (alg = bgalg(prob), reltol = 1e-7, abstol = 1e-7), ptopts = (alg = ptalg(prob), reltol = 1e-5, abstol = 1e-5), thread = true, verbose = false, kwargs...) + ls = collect(ls) # Define 1-2-3 indices corresponding for present modes iT = 'T' in join(modes) ? 1 : 0 iE = 'E' in join(modes) ? iT + 1 : 0 iψ = 'ψ' in join(modes) ? max(iE, iT) + 1 : 0 + sol = solve(prob; bgopts, verbose) + K = curvature(sol) # spatial curvature constant of the model, passed on to the line-of-sight integration + τ0 = getsym(sol, prob.M.τ0)(sol) + + # An open universe has no super-curvature modes: its scalar spectrum is bounded below by k² ≥ -K. + # Including k < √(-K) anyway inflates the lowest multipoles by an order of magnitude. + kmin = max(1e-2, K < 0 ? √(-K) : 0.0) + # Automatically determine grid if not provided manually if isnothing(kinterp) if iψ > 0 - kinterp = ChebyshevInterpolator(1e-2, 1e4, 130; f = fk_tanh, f⁻¹ = fk⁻¹_tanh) # higher kmax for lensing; f that stretches acoustic oscillations for k ≲ 2000 with higher sampling density + kinterp = ChebyshevInterpolator(kmin, 1e4, 130; f = fk_tanh, f⁻¹ = fk⁻¹_tanh) # higher kmax for lensing; f that stretches acoustic oscillations for k ≲ 2000 with higher sampling density else - kinterp = ChebyshevInterpolator(1e-2, 2e3, 60) # lower kmax for T/E-only; sample uniform acoustic oscillations in linear k + kinterp = ChebyshevInterpolator(kmin, 2e3, 60) # lower kmax for T/E-only; sample uniform acoustic oscillations in linear k end + elseif minimum(kinterp) < kmin + throw(ArgumentError("k-grid starts at k = $(minimum(kinterp)), below the smallest wavenumber √(-K) = $kmin supported by this open universe")) end - - ls = collect(jl.l) - sol = solve(prob; bgopts, verbose) - τ0 = getsym(sol, prob.M.τ0)(sol) ks_fine = lingrid(minimum(kinterp), maximum(kinterp); step=Δkτ0/τ0) # for k-quadrature after LOS integration τs = sol.bg.t # by default, use background (thermodynamics) time points for line of sight integration @@ -290,7 +352,7 @@ function spectrum_cmb(modes::AbstractVector{<:Symbol}, prob::CosmologyProblem, j Ss[end, :] .= Ref(zero(eltype(Ss))) # remove any Inf/NaN at last time χ=0; weighted by jₗ(0)=0 anyway # Integrate all sources simultaneously without Limber approximation - Θls = los_integrate(Ss, ls, τs, ks_fine, jl; verbose, thread, kwargs...) + Θls = los_integrate(Ss, ls, τs, ks_fine, jl, K; verbose, thread, kwargs...) Θls = stack(Θls) # to 3D array if iT > 0 Θls[iT, :, :] ./= ks_fine @@ -299,7 +361,7 @@ function spectrum_cmb(modes::AbstractVector{<:Symbol}, prob::CosmologyProblem, j Θls[iE, :, :] .*= transpose(@. √((ls+2)*(ls+1)*(ls+0)*(ls-1))) ./ (ks_fine .^ 2) end if iψ > 0 && l_limber ≤ ls[end] - Θls[iψ, :, :] .= los_integrate(getindex.(Ss, iψ), ls, τs, ks_fine, jl; l_limber, verbose, thread, kwargs...) # overwrite with Limber result + Θls[iψ, :, :] .= los_integrate(getindex.(Ss, iψ), ls, τs, ks_fine, jl, K; l_limber, verbose, thread, kwargs...) # overwrite with Limber result end P0s = spectrum_primordial(ks_fine, sol) # more accurate @@ -335,17 +397,29 @@ function spectrum_cmb(modes::AbstractVector{<:Symbol}, prob::CosmologyProblem, j end """ - spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, jl::SphericalBesselCache, ls::AbstractVector; kwargs...) + spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, jl::Union{SphericalBesselCache, AbstractVector}, ls::AbstractVector; kwargs...) -Same, but compute the spectrum properly only for `jl.l` and then interpolate the results to all `ls`. +Same, but compute the spectrum properly only for `jl` (a `SphericalBesselCache`, `AbstractInterpolator` over ``l`` or an `AbstractArray` of multipoles) +and then interpolate the results to all `ls`. """ +function spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, ls_coarse::AbstractArray, ls::AbstractVector; normalization = :Cl, linterp_normalization = l -> l^5, kwargs...) + minimum(ls) ≥ minimum(ls_coarse) && maximum(ls) ≤ maximum(ls_coarse) || throw(ArgumentError("l-range $(extrema(ls)) is outside the l-range $(extrema(ls_coarse)) of the coarse l-grid")) + spectra_coarse = spectrum_cmb(modes, prob, ls_coarse; kwargs...) + spectra_fine = similar(spectra_coarse, (length(ls), size(spectra_coarse)[2])) + for imode in eachindex(modes) + spectra_fine[:, imode] = interpolate(ls_coarse, spectra_coarse[:, imode] .* linterp_normalization.(ls_coarse), ls) ./ linterp_normalization.(ls) + spectra_fine[:, imode] = normalize_spectrum_cmb(normalization, ls, spectra_fine[:, imode]) + end + return spectra_fine +end function spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, jl::SphericalBesselCache, ls::AbstractVector; normalization = :Cl, linterp_normalization = l -> l^5, kwargs...) - minimum(ls) ≥ minimum(jl.l) && maximum(ls) ≤ maximum(jl.l) || throw(ArgumentError("l-range $(extrema(ls)) is outside the l-range $(extrema(jl.l)) of the spherical Bessel function")) + ls_coarse = jl.l + minimum(ls) ≥ minimum(ls_coarse) && maximum(ls) ≤ maximum(ls_coarse) || throw(ArgumentError("l-range $(extrema(ls)) is outside the l-range $(extrema(ls_coarse)) of the coarse l-grid")) spectra_coarse = spectrum_cmb(modes, prob, jl; kwargs...) spectra_fine = similar(spectra_coarse, (length(ls), size(spectra_coarse)[2])) for imode in eachindex(modes) - spectra_fine[:, imode] = interpolate(jl.l, spectra_coarse[:, imode] .* linterp_normalization.(jl.l), ls) ./ linterp_normalization.(ls) # interpolate l⁵*Cₗ (by default) for smoothness - spectra_fine[:, imode] = normalize_spectrum_cmb(normalization, ls, spectra_fine[:, imode]) # normalize AFTER interpolation + spectra_fine[:, imode] = interpolate(ls_coarse, spectra_coarse[:, imode] .* linterp_normalization.(ls_coarse), ls) ./ linterp_normalization.(ls) + spectra_fine[:, imode] = normalize_spectrum_cmb(normalization, ls, spectra_fine[:, imode]) end return spectra_fine end diff --git a/src/solve.jl b/src/solve.jl index 8d4b5f75b..4664fd224 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -617,6 +617,20 @@ function issuccess(sol::CosmologySolution) return successful_retcode(sol.bg) && (isnothing(sol.pts) || all(successful_retcode(pt) for pt in sol.pts)) end +""" + curvature(sol::CosmologySolution) + +Spatial curvature constant ``K = -Ω_{k0}`` of the solved model, in units where ``H₀ = c = 1`` +(``K > 0`` closed, ``K = 0`` flat, ``K < 0`` open). Returns `0.0` for a model without curvature. +""" +function curvature(sol::CosmologySolution) + M = sol.prob.M + have(M, :K) && return -sol[M.K.Ω₀] # curvature as an effective species (see `curvature(g; ...)`) + hasproperty(M, :K) && return sol[M.K] # curvature as a bare parameter (as in the minimal ΛCDM model) + hasproperty(M, :Ωk0) && return -sol[M.Ωk0] + return 0.0 +end + function integrate(xs, ys; integrator = Trapezoidal()) return NumericalIntegration.integrate(xs, ys, integrator) end diff --git a/src/utils.jl b/src/utils.jl index c2c0acd54..8155a0a15 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -97,6 +97,53 @@ function filter_system(f::Function, sys::System) return System(eqs, iv, vars, pars; initialization_eqs=ieqs, initial_conditions=ics, bindings, guesses=guesses, name=nameof(sys), description=get_description(sys)) end +# Generalized trigonometric functions for a universe with curvature `K`: +# sinK(K, χ) = sin(√K χ) / √K (→ χ as K → 0), +# cotK(K, χ) = √K cot(√K χ) (→ 1/χ as K → 0), +# asinK(K, y) = inverse of sinK in χ for fixed K. +# All are real and analytic in K (√K is imaginary for an open universe, but trig functions then turn into their hyperbolic counterparts). +# Maclaurin series in u = Kχ² cover the removable singularities at K = 0 and the small-|u| region where the closed forms would cancel. +function sinK(K::Real, χ::Real) + u = K * χ^2 + if abs(u) < 1e-2 + return χ * evalpoly(u, (1.0, -1/6, 1/120, -1/5040, 1/362880)) # χ⋅sin(√u)/√u + elseif K > 0 + β = √K + return sin(β*χ) / β + else + β = √(-K) + return sinh(β*χ) / β + end +end +function cotK(K::Real, χ::Real) + u = K * χ^2 + if abs(u) < 1e-2 + return evalpoly(u, (1.0, -1/3, -1/45, -2/945, -1/4725, -2/93555)) / χ # √u⋅cot(√u)/χ + elseif K > 0 + β = √K + return β * cot(β*χ) + else + β = √(-K) + return β * coth(β*χ) + end +end +function asinK(K::Real, y::Real) + u = K * y^2 + if abs(u) < 1e-3 + return y * evalpoly(u, (1.0, 1/6, 3/40, 15/336, 105/3456)) # y⋅asin(√u)/√u + elseif K > 0 + β = √K + return asin(β*y) / β + else + β = √(-K) + return asinh(β*y) / β + end +end + +# Define for symbolic use +@register_symbolic sinK(K, χ) +@register_symbolic cotK(K, χ) + have(sys, s::Symbol) = s in nameof.(ModelingToolkit.get_systems(sys)) have(s) = !isnothing(s) # shorthand for checking if we have a given species diff --git a/test/runtests.jl b/test/runtests.jl index 66106b0f4..3d08de9f2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -580,13 +580,21 @@ end @test isapprox(Dls_cheb, Dls; rtol = 1e-4) @test isapprox(Dls_chebint, Dls; rtol = 1e-4) + # Without precomputed Bessel (i.e. using Bessel recurrence on the fly) + @test_throws "must be exact integer" spectrum_cmb(:TT, prob, jl_cubic.l, ls; normalization = :Dl) # TODO: error + @test_throws "must be exact integer" spectrum_cmb(:TT, prob, jl_cheb.l, ls; normalization = :Dl) # TODO: error + Dls_chebint_recurrence = spectrum_cmb(:TT, prob, jl_chebint.l, ls; normalization = :Dl) + @test isapprox(Dls_chebint, Dls_chebint_recurrence; rtol = 1e-4) + # Error with bad input @test_throws "outside the l-range" spectrum_cmb(:TT, prob, jl, 1:3000; normalization = :Dl) end @testset "Toggle threading" begin - @test length(unique(fetch.(map(i -> SymBoltz.@spawnif(threadid(), true), 1:10)))) > 1 - @test only(unique(fetch.(map(i -> SymBoltz.@spawnif(threadid(), false), 1:10)))) == 1 + if nthreads() > 1 + @test length(unique(fetch.(map(i -> SymBoltz.@spawnif(threadid(), true), 1:10)))) > 1 + @test only(unique(fetch.(map(i -> SymBoltz.@spawnif(threadid(), false), 1:10)))) == 1 + end end @testset "Sparse Jacobian" begin @@ -963,6 +971,13 @@ end interp = CubicSplineInterpolator(x) @test eltype(interp) == eltype(x) @test issorted(interp) + + # interpolators are vectors of their x-points satisfying the AbstractArray interface + @test interp isa AbstractVector{eltype(x)} + @test size(interp) == size(x) && length(interp) == length(x) + @test collect(interp) == interp.xs == x + @test interp[begin] == minimum(interp) == x[begin] && interp[end] == maximum(interp) == x[end] + @test sprint(show, MIME"text/plain"(), interp) == sprint(show, interp) # summary, not all x-points x′ = range(x[begin], x[end]; length = 1000) y′ = interpolate(interp, sin.(x), x′) @test all(interpolate(x, sin.(x), x′) .== y′) # should fall exactly back to cubic spline interpolation @@ -1004,3 +1019,19 @@ end @test isapprox(y′, sin.(π/30 .* x′); atol = 1e-10) @test_throws "collide" ChebyshevIntegerInterpolator(0, 100, 23) end + +@testset "hyperspherical Bessel function recursion" begin + K = -0.1 + ls = unique(Int.(round.(exp.(range(log(2), log(2500); length = 100))))) + + χs = collect(reverse(range(0.0, 3.0; length = 500))) + Φl = zeros(length(χs), length(ls)) + + k = 1e1 + Φmin = 1e-20 + @time SymBoltz.Φl_recurrence!(Φl, ls, χs, k, K; Φmin) + + #using Plots + #heatmap(reverse(χs), ls, max.(log10.(abs.(reverse(transpose(Φl); dims=2))), log10(Φmin)); xlabel = "χ", ylabel = "ℓ", title = "Φ(χ, ℓ, k=$k)") + #plot(k .* χs, Φl[:, 1]; xlims = (0, 10)) +end