Skip to content

Draft LOS integration with non-dual Bessel function argument #20

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Draft LOS integration with non-dual Bessel function argument #20

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 37 additions & 35 deletions src/observables/angular.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/observables/fourier.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading