From e3dc14033130c57a0125f271150e56ede4412904 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Thu, 11 Sep 2025 15:00:26 +0200 Subject: [PATCH] Draft LOS integration with non-dual Bessel function argument --- src/observables/angular.jl | 72 ++++++++++++++++++++------------------ src/observables/fourier.jl | 2 +- 2 files changed, 38 insertions(+), 36 deletions(-) diff --git a/src/observables/angular.jl b/src/observables/angular.jl index fcc95d04..55e21325 100644 --- a/src/observables/angular.jl +++ b/src/observables/angular.jl @@ -55,7 +55,7 @@ ChainRulesCore.frule((_, _, Δx), ::typeof(jl_x2), l, x) = jl_x2(l, x), jl_x2′ # 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, Rl::Function; integrator = TrapezoidalRule(), verbose = false) where {T <: Real} + los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, xs::AbstractVector, ys::AbstractVector, Rl::Function; integrator = TrapezoidalRule(), verbose = false) where {T <: Real} For the given `ls` and `ks`, compute the line-of-sight-integrals ```math @@ -64,59 +64,61 @@ Iₗ(k) = ∫dτ S(k,τ) Rₗ(k(τ₀-τ)) over the source function values `Ss` against the radial functions `Rl` (e.g. the spherical Bessel functions ``jₗ(x)``). The element `Ss[i,j]` holds the source function value ``S(kᵢ, τⱼ)``. """ -function los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, Rl::Function = jl; integrator = TrapezoidalRule(), verbose = false) where {T <: Real} +function los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, xs::AbstractVector, ys::AbstractVector, Rl::Function = jl; integrator = TrapezoidalRule(), verbose = false) where {T <: Real} # Julia is column-major; make sure innermost loop indices appear first in slice expressions (https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-column-major) - @assert size(Ss) == (length(τs), length(ks)) "size(Ss) = $(size(Ss)) ≠ $((length(τs), length(ks)))" - τs = collect(τs) # force array to avoid floating point errors with ranges in following χs due to (e.g. tiny negative χ) - χs = τs[end] .- τs - Is = similar(Ss, length(ks), length(ls)) + @assert size(Ss) == (length(xs), length(ys)) "size(Ss) = $(size(Ss)) ≠ $((length(xs), length(ys)))" + xs = collect(xs) # force array to avoid floating point errors with ranges in following χs due to (e.g. tiny negative χ) + Is = similar(Ss, length(ys), length(ls)) # TODO: skip and set Rl to zero if l ≳ kτ0 or another cutoff? @tasks for il in eachindex(ls) # parallellize independent loop iterations @local begin # define task-local values (declared once for all loop iterations) - ∂I_∂τ = similar(Ss, length(τs)) + ∂I_∂τ = similar(Ss, length(xs)) end l = ls[il] verbose && print("\rLOS integrating with l = $l") - for ik in eachindex(ks) - k = ks[ik] - for iτ in eachindex(τs) - S = Ss[iτ,ik] - χ = χs[iτ] - kχ = k * χ - ∂I_∂τ[iτ] = S * Rl(l, kχ) # TODO: rewrite LOS integral to avoid evaluating Rl with dual numbers? # TODO: rewrite LOS integral with y = kτ0 and x=τ/τ0 to cache jls independent of cosmology + for iy in eachindex(ys) + y = ys[iy] + for ix in eachindex(xs) + S = Ss[ix,iy] + x = xs[ix] + kχ = y * (1 - x) + ∂I_∂τ[ix] = S * Rl(l, kχ) # TODO: rewrite LOS integral to avoid evaluating Rl with dual numbers? # TODO: rewrite LOS integral with y = kτ0 and x=τ/τ0 to cache jls independent of cosmology end - Is[ik,il] = integrate(τs, ∂I_∂τ; integrator) # integrate over τ # TODO: add starting I(τini) to fix small l? + Is[iy,il] = integrate(xs, ∂I_∂τ; integrator) # integrate over τ # TODO: add starting I(τini) to fix small l? end end verbose && println() return Is end -function los_integrate(sol::CosmologySolution, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, S, Rl::Function = jl; ktransform = identity, kwargs...) # TODO: Ss +function los_integrate(sol::CosmologySolution, ls::AbstractVector, xs::AbstractVector, ys::AbstractVector, S, Rl::Function = jl; ktransform = identity, kwargs...) # TODO: Ss + τ0 = getsym(sol, prob.M.τ0)(sol) + τs = xs .* τ0 + ks = ys / τ0 Ss = [S] Ss = source_grid(sol, Ss, τs) Ss = source_grid(Ss, sol.ks, ks; ktransform) Ss = @view Ss[1, :, :] - return los_integrate(Ss, ls, τs, ks, Rl; kwargs...) + return los_integrate(Ss, ls, xs, ys, Rl; kwargs...) end """ - los_temperature(sol::CosmologySolution, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector; ktransform = identity, kwargs...) + los_temperature(sol::CosmologySolution, ls::AbstractVector, xs::AbstractVector, ys::AbstractVector; ktransform = identity, kwargs...) Calculate photon temperature multipoles today by line-of-sight integration. """ -function los_temperature(sol::CosmologySolution, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector; ktransform = identity, kwargs...) - return los_integrate(sol, ls, τs, ks, sol.prob.M.ST0, jl; ktransform, kwargs...) +function los_temperature(sol::CosmologySolution, ls::AbstractVector, xs::AbstractVector, ys::AbstractVector; ktransform = identity, kwargs...) + return los_integrate(sol, ls, xs, ys, sol.prob.M.ST0, jl; ktransform, kwargs...) end """ - los_polarization(sol::CosmologySolution, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector; ktransform = identity, kwargs...) + los_polarization(sol::CosmologySolution, ls::AbstractVector, xs::AbstractVector, ys::AbstractVector; ktransform = identity, kwargs...) Calculate photon E-mode polarization multipoles today by line-of-sight integration. """ -function los_polarization(sol::CosmologySolution, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector; ktransform = identity, kwargs...) - return los_integrate(sol, ls, τs, ks, sol.prob.M.ST2_polarization, jl_x2; ktransform, kwargs...) .* transpose(@. √((ls+2)*(ls+1)*(ls+0)*(ls-1))) +function los_polarization(sol::CosmologySolution, ls::AbstractVector, xs::AbstractVector, ys::AbstractVector; ktransform = identity, kwargs...) + return los_integrate(sol, ls, xs, ys, sol.prob.M.ST2_polarization, jl_x2; ktransform, kwargs...) .* transpose(@. √((ls+2)*(ls+1)*(ls+0)*(ls-1))) end # TODO: integrate splines instead of trapz! https://discourse.julialang.org/t/how-to-speed-up-the-numerical-integration-with-interpolation/96223/5 @@ -157,26 +159,26 @@ function spectrum_cmb(ΘlAs::AbstractMatrix, ΘlBs::AbstractMatrix, P0s::Abstrac end """ - spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, ls::AbstractVector; 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, ls::AbstractVector; 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 = (x->tanh(3.19x)), u⁻¹ = (u->atanh(u)/3.19), 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...) 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, ls::AbstractVector; 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...) +function spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, ls::AbstractVector; 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 = (x->tanh(3.19x)), u⁻¹ = (u->atanh(u)/3.19), 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...) 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 - τs = sol.bg.t # by default, use background (thermodynamics) time points for line of sight integration - τi = τs[begin] + xs = sol.bg.t ./ τ0 # by default, use background (thermodynamics) time points for line of sight integration if Nlos != 0 # instead choose Nlos time points τ = τ(u) corresponding to uniformly spaced u - τmin, τmax = extrema(τs) - umin, umax = u(τmin), u(τmax) + xmin, xmax = xs[begin], xs[end] # TODO: non-dual e.g. 4e-7, 1.0 + umin, umax = u(xmin), u(xmax) us = range(umin, umax, length = Nlos) - τs = u⁻¹.(us) - τs[begin] = τi - τs[end] = τ0 + xs = u⁻¹.(us) + xs[begin] = xmin + xs[end] = xmax end + τs = τ0 .* xs # Integrate perturbations to calculate source function on coarse k-grid iT = 'T' in join(modes) ? 1 : 0 @@ -189,10 +191,10 @@ function spectrum_cmb(modes::AbstractVector, prob::CosmologyProblem, ls::Abstrac # 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 = source_grid(Ss_coarse, ks_coarse, ks_fine) .* τ0 - ΘlTs = iT > 0 ? los_integrate(@view(Ss_fine[iT, :, :]), ls, τs, ks_fine, jl; integrator, verbose, kwargs...) : nothing - ΘlEs = iE > 0 ? los_integrate(@view(Ss_fine[iE, :, :]), ls, τs, ks_fine, jl_x2; integrator, verbose, kwargs...) .* transpose(@. √((ls+2)*(ls+1)*(ls+0)*(ls-1))) : nothing + ΘlTs = iT > 0 ? los_integrate(@view(Ss_fine[iT, :, :]), ls, xs, kτ0s_fine, jl; integrator, verbose, kwargs...) : nothing + ΘlEs = iE > 0 ? los_integrate(@view(Ss_fine[iE, :, :]), ls, xs, kτ0s_fine, jl_x2; integrator, verbose, kwargs...) .* transpose(@. √((ls+2)*(ls+1)*(ls+0)*(ls-1))) : nothing P0s = spectrum_primordial(ks_fine, sol) # more accurate diff --git a/src/observables/fourier.jl b/src/observables/fourier.jl index 12345fd5..945c2c20 100644 --- a/src/observables/fourier.jl +++ b/src/observables/fourier.jl @@ -187,7 +187,7 @@ function source_grid(prob::CosmologyProblem, S::AbstractArray, τs, ks; bgopts = 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 + τs[begin] ≥ bgsol.t[begin] && τs[end] ≤ bgsol.t[end] || error("input τs is not a subset of computed background solution timespans ($(extrema(τs)) vs $(extrema(bgsol.t)))") # TODO: don't rely on function output_func(sol, ik) for iS in eachindex(getSs) Ss[iS, :, ik] .= getSs[iS](sol)