Use arbitrary quadrature nodes and weights in line-of-sight integral

Project.toml

@@ -9,10 +9,10 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 FastChebInterp = "cf66c380-9a80-432c-aff8-4f9c79c0bdde"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ForwardDiffChainRules = "c9556dd2-1aed-4cfe-8560-1557cf593001"
 GraphRecipes = "bd48cda9-67a9-57be-86fa-5b3c104eda73"
-Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 MatterPower = "743831d1-0e6a-4e86-a0c4-1a098a68aca3"
@@ -54,10 +54,10 @@ ChainRulesCore = "1.25.1"
 CommonSolve = "0.2.4"
 DataInterpolations = "8.0.0"
 FastChebInterp = "1"
+FFTW = "1.10.0"
 ForwardDiff = "0.10, 1.0.0"
 ForwardDiffChainRules = "0.3.0"
 GraphRecipes = "0.5.13"
-Integrals = "4.5.0, 5.0.0"
 LinearAlgebra = "1.10.0"
 LinearSolve = "3.19.0"
 MatterPower = "0.5.0"

src/SymBoltz.jl

@@ -41,6 +41,7 @@ include("models/curvature.jl")
 include("models/inflation.jl")
 include("models/cosmologies.jl")
 include("solve.jl")
+include("quadrature.jl")
 include("observables/distances.jl")
 include("observables/fourier.jl")
 include("observables/angular.jl")
@@ -56,6 +57,7 @@ export spectrum_primordial, spectrum_matter, spectrum_matter_nonlinear, spectrum
 export AbstractInterpolator, EquispacedInterpolator, CubicSplineInterpolator, ChebyshevInterpolator, PiecewiseChebyshevInterpolator, order, interpolate
 export express_derivatives
 export lingrid, loggrid, cosgrid, chebgrid, joingrids!, kτ0grid_default
+export Quadrature, TrapezoidalQuadrature, SimpsonQuadrature, ClenshawCurtisQuadrature, GaussQuadrature, GaussKronrodQuadrature, nodes, weights, integrate, transform
 
 using PrecompileTools: @compile_workload
 @compile_workload begin

src/observables/angular.jl

@@ -1,4 +1,3 @@
-using Integrals
 using Bessels: besselj!, sphericalbesselj
 using DataInterpolations
 using MatterPower
@@ -15,7 +14,7 @@ struct SphericalBesselCache{Tl, Tdy <: Union{Matrix{Float64}, Nothing}}
     x::Vector{Float64}
 end
 
-function SphericalBesselCache(ls; xmax = 20*maximum(ls), dx = 2π/15, hermite = true)
+function SphericalBesselCache(ls; dx = 2π/15, xmax = 20*maximum(ls)+dx, hermite = true) # add one dx to xmax to prevent floating-point bounds checking issues
     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
@@ -92,7 +91,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), integrator = TrapezoidalRule(), thread = true, verbose = false) where {T}
+    los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, jl::SphericalBesselCache[, quad::Quadrature]; l_limber = typemax(Int), thread = true, verbose = false) where {T}
 
 For the given `ls` and `ks`, compute the line-of-sight integrals
 ```math
@@ -106,7 +105,7 @@ 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), integrator = TrapezoidalRule(), thread = true, verbose = false) where {T}
+function los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractVector, ks::AbstractVector, jl::SphericalBesselCache, quad::Quadrature; l_limber = typemax(Int), thread = true, verbose = false) where {T}
     @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"
@@ -122,19 +121,14 @@ function los_integrate(Ss::AbstractMatrix{T}, ls::AbstractVector, τs::AbstractV
     χ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)
 
     verbose && l_limber < typemax(Int) && println("Using Limber approximation for l ≥ $l_limber")
 
+    ws = quad.w .* (quad.b-quad.a)/2
+
     # Loop order k → τ → l to get SIMD on the innermost l-loop
     @fastmath @inbounds @tasks for ik in eachindex(ks)
         @set scheduler = thread ? :dynamic : :serial
@@ -187,7 +181,7 @@ 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; integrator = TrapezoidalRule(), normalization = :Cl, thread = true)
+    spectrum_cmb(ΘlAs::AbstractMatrix, ΘlBs::AbstractMatrix, P0s::AbstractVector, ls::AbstractVector, ks::AbstractVector[, quad::Quadrature]; normalization = :Cl, thread = true)
 
 Compute the angular power spectrum
 ```math
@@ -196,22 +190,20 @@ Cₗᴬᴮ = (2/π) ∫\mathrm{d}k \, k² P₀(k) Θₗᴬ(k,τ₀) Θₗᴮ(k,
 for the given `ls`.
 If `normalization == :Dl`, compute ``Dₗ = Cₗ l (l+1) / 2π`` instead.
 """
-function spectrum_cmb(ΘlAs::AbstractMatrix, ΘlBs::AbstractMatrix, P0s::AbstractVector, ls::AbstractVector, ks::AbstractVector; integrator = TrapezoidalRule(), normalization = :Cl, thread = true)
+function spectrum_cmb(ΘlAs::AbstractMatrix, ΘlBs::AbstractMatrix, P0s::AbstractVector, ls::AbstractVector, ks::AbstractVector, quad::Quadrature; normalization = :Cl, thread = true)
     size(ΘlAs) == size(ΘlBs) || error("ΘlAs and ΘlBs have different sizes")
     eltype(ΘlAs) == eltype(ΘlBs) || error("ΘlAs and ΘlBs have different types")
 
     Cls = similar(ΘlAs, length(ls))
-    ks_with0 = [0.0; ks] # add dummy value with k=0 for integration
 
     @tasks for il in eachindex(ls)
         # TODO: skip kτ0 ≲ l?
         @set scheduler = thread ? :dynamic : :static
-        @local dCl_dks_with0 = zeros(eltype(ΘlAs), length(ks_with0)) # local task workspace (must zero first element)
+        @local dCl_dks = zeros(eltype(ΘlAs), length(ks)) # local task workspace (must zero first element)
         ΘlA = @view ΘlAs[:, il]
         ΘlB = @view ΘlBs[:, il]
-        @. dCl_dks_with0[2:end] = 2/π * ks^2 * P0s * ΘlA * ΘlB
-        spline = CubicSpline(dCl_dks_with0, ks_with0)
-        Cls[il] = DataInterpolations.integral(spline, ks_with0[begin], ks_with0[end]) # integrate over k (_with0 adds one additional point at (0,0))
+        dCl_dks .= 2/π .* ks .^ 2 .* P0s .* ΘlA .* ΘlB
+        Cls[il] = quad(dCl_dks) # integrate over k (_with0 adds one additional point at (0,0))
     end
 
     return normalize_spectrum_cmb(normalization, ls, Cls)
@@ -221,7 +213,7 @@ fk_tanh(k, k0=2000.0) = tanh(k/k0)
 fk⁻¹_tanh(k, k0=2000.0) = k0*atanh(k)
 
 """
-    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, integrator = TrapezoidalRule(), 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...)
+    spectrum_cmb(modes::AbstractVector{<:Symbol}, prob::CosmologyProblem, jl::SphericalBesselCache; normalization = :Cl, unit = nothing, kinterp = nothing, τquad = TrapezoidalQuadrature(cosgrid(0.0, 1.0; length=300)), kquad = ClenshawCurtisQuadrature(8192), τ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...)
 
 Compute angular CMB power spectra ``Cₗᴬᴮ`` at angular wavenumbers `ls` from the cosmological problem `prob`.
 The requested `modes` are specified as a vector of symbols in the form `:AB`, where `A` and `B` are `T` (temperature), `E` (E-mode polarization) or `ψ` (lensing).
@@ -252,7 +244,7 @@ modes = [:TT, :TE, :ψψ, :ψT]
 Dls = spectrum_cmb(modes, prob, jl; normalization = :Dl, unit = u"μK")
 ```
 """
-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, integrator = TrapezoidalRule(), 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; normalization = :Cl, unit = nothing, kinterp = nothing, τquad = TrapezoidalQuadrature(cosgrid(0.0, 1.0; length=300)), kquad = ClenshawCurtisQuadrature(8192), τ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...)
     # Define 1-2-3 indices corresponding for present modes
     iT = 'T' in join(modes) ? 1 : 0
     iE = 'E' in join(modes) ? iT + 1 : 0
@@ -269,20 +261,11 @@ function spectrum_cmb(modes::AbstractVector{<:Symbol}, prob::CosmologyProblem, j
 
     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
-    τs = τs[τs .≥ τcut]
-    if xs isa AbstractArray
-        # explicit fractional grid x = (τ-τi)/(τ0-τi) ∈ [0,1]
-        xs[begin] == 0 || error("xs begins with $(xs[begin]), but should begin with 0")
-        xs[end] == 1 || error("xs ends with $(xs[end]), but should end with 1")
-        τs = τs[begin] .+ (τs[end] .- τs[begin]) .* xs
-    elseif xs isa Int
-        # interpolate xs points from background time grid, preserving its density structure
-        τs = LinearInterpolation(τs, 1.0:length(τs)).(range(1.0, length(τs), length = xs))
-    end
+    τquad = transform(τquad, (τs[begin], τs[end])) # map from [-1, 1] to [τi, τ0]
+    τs = nodes(τquad)
+    kquad = transform(kquad, extrema(kinterp))
+    ks_fine = nodes(kquad) # for k-quadrature after LOS integration
 
     # Integrate perturbations to calculate source function on coarse k-grid
     Ss = [S for (S, i) in [(prob.M.k*prob.M.ST, iT), (prob.M.k^2*prob.M.SE, iE), (prob.M.Sψ, iψ)] if i > 0]
@@ -291,7 +274,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; integrator, verbose, thread, kwargs...)
+    Θls = los_integrate(Ss, ls, τs, ks_fine, jl, τquad; verbose, thread, kwargs...)
     Θls = stack(Θls) # to 3D array
     if iT > 0
         Θls[iT, :, :] ./= ks_fine
@@ -300,7 +283,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, integrator, verbose, thread, kwargs...) # overwrite with Limber result
+        Θls[iψ, :, :] .= los_integrate(getindex.(Ss, iψ), ls, τs, ks_fine, jl, τquad; l_limber, verbose, thread, kwargs...) # overwrite with Limber result
     end
 
     P0s = spectrum_primordial(ks_fine, sol) # more accurate
@@ -327,7 +310,7 @@ function spectrum_cmb(modes::AbstractVector{<:Symbol}, prob::CosmologyProblem, j
         iB = geti(Symbol(mode[lastindex(mode)]))
         ΘlAs = @view(Θls[iA, :, :])
         ΘlBs = @view(Θls[iB, :, :])
-        spectrum = spectrum_cmb(ΘlAs, ΘlBs, P0s, ls, ks_fine; integrator, normalization, thread)
+        spectrum = spectrum_cmb(ΘlAs, ΘlBs, P0s, ls, ks_fine, kquad; normalization, thread)
         spectrum *= factor^2 # possibly make dimensionful
         spectra[:, i] .= spectrum
     end

src/quadrature.jl

@@ -0,0 +1,110 @@
+using FFTW
+using QuadGK
+
+struct Quadrature{T, U}
+    x::Vector{T} # integration points on [-1, +1]
+    w::Vector{T} # integration weights on [-1, +1]
+    a::U
+    b::U
+    name::Symbol
+
+    function Quadrature(x::AbstractVector{T}, w::AbstractVector{T}, interval = (-1.0, 1.0); name = Symbol()) where {T}
+        all(-1 .≤ x .≤ 1) || throw(ArgumentError("Quadrature points must be on the canonical interval [-1, 1]"))
+        issorted(x) || throw(ArgumentError("Quadrature points must be sorted in ascending order"))
+        sum(w) ≈ 2 || throw(ArgumentError("Quadrature weights must sum to 2, but sums to $(sum(w))"))
+        length(x) == length(w) || throw(ArgumentError("Quadrature nodes and weights must have same length"))
+        new{T, eltype(interval)}(x, w, interval[begin], interval[end], name)
+    end
+end
+
+function TrapezoidalQuadrature(N::Integer, args...)
+    N ≥ 2 || throw(ArgumentError("Trapezoidal rule needs at least 2 points"))
+    x = collect(range(-1, 1, length = N))
+    w = fill(2/(N-1), N)
+    w[begin] = 1/(N-1)
+    w[end] = 1/(N-1)
+    return Quadrature(x, w, args...; name = Symbol("Trapezoidal"))
+end
+
+function TrapezoidalQuadrature(x::AbstractArray, args...)
+    N = length(x)
+    N ≥ 2 || throw(ArgumentError("Trapezoidal quadrature needs at least 2 points"))
+    xmin, xmax = extrema(x)
+    x = collect(x)
+    x .= -1 .+ 2 .* (x .- xmin) / (xmax - xmin) # normalize to canonical [-1, 1]
+    w = zeros(N)
+    # each interval contributes (y1+y2)*(x2-x1)/2
+    for i in 1:N-1
+        dx = x[i+1] - x[i]
+        w[i] += dx / 2
+        w[i+1] += dx / 2
+    end
+    return Quadrature(x, w, (xmin, xmax), args...; name = Symbol("Trapezoidal"))
+end
+
+function SimpsonQuadrature(x::AbstractArray, args...)
+    N = length(x)
+    isodd(N) && N ≥ 3 || throw(ArgumentError("Simpson's rule needs an odd number of points ≥ 3"))
+    xmin, xmax = extrema(x)
+    x = collect(x)
+    x .= -1 .+ 2 .* (x .- xmin) / (xmax - xmin) # normalize to canonical [-1, 1]
+    w = zeros(N)
+    @inbounds for i in 1:2:N-2
+        # fit a quadratic through each triple of points (possibly unequally spaced) and integrate it exactly
+        x0, x1, x2 = x[i], x[i+1], x[i+2]
+        h1 = x1 - x0
+        h2 = x2 - x1
+        H = h1 + h2
+        w[i] += H/6 * (2h1 - h2) / h1
+        w[i+1] += H/6 * H^2 / (h1*h2)
+        w[i+2] += H/6 * (2h2 - h1) / h2
+    end
+    return Quadrature(x, w, (xmin, xmax), args...; name = Symbol("Simpson"))
+end
+
+SimpsonQuadrature(N::Integer, args...) = SimpsonQuadrature(range(-1, 1, length = N), args...) # reduces to the standard uniform-grid rule (1, 4, 2, 4, 2, ..., 4, 1)
+
+function ClenshawCurtisQuadrature(N::Integer, args...)
+    N ≥ 2 || throw(ArgumentError("Clenshaw-Curtis quadrature needs at least 2 points"))
+    n = N - 1 # number of FFT points
+    x = [-cos(π*m/n) for m in 0:n] # Chebyshev nodes in ascending order
+    v = [1 / (1 - 4*min(m, n-m)^2) for m in 0:n-1] # 1/(1-4k^2) and it's mirror image
+    w = similar(x)
+    w[begin:end-1] .= 2/n .* real(fft(v)) # Discrete Cosine Transform (DCT) for O(N log N) time instead of O(N²)
+    w[begin] = w[end] = w[begin]/2 # modify endpoint factors
+    return Quadrature(x, w, args...; name = Symbol("Clenshaw-Curtis"))
+end
+
+function GaussQuadrature(N::Integer, args...)
+    N ≥ 1 || throw(ArgumentError("The number of Gauss quadrature points must be positive"))
+    x, w = QuadGK.gauss(N)
+    return Quadrature(x, w, args...; name = Symbol("Gauss"))
+end
+
+function GaussKronrodQuadrature(N::Integer, args...)
+    N ≥ 3 && isodd(N) || throw(ArgumentError("The number of Gauss-Kronrod quadrature points must be at least 3 and odd"))
+    x, w = QuadGK.kronrod(div(N-1, 2)) # only returns left half
+    x = [x; -reverse(x[begin:end-1])] # add right half (antisymmetric)
+    w = [w; reverse(w[begin:end-1])] # add right half (symmetric)
+    return Quadrature(x, w, args...; name = Symbol("Gauss-Kronrod"))
+end
+
+transform(q::Quadrature, interval) = Quadrature(q.x, q.w, interval; name = q.name) # arrays are reuse and shared!
+Base.eltype(::Quadrature{T, U}) where {T, U} = Base.promote_type(T, U)
+Base.nameof(q::Quadrature) = q.name
+Base.show(io::IO, q::Quadrature) = print(io, q.name == Symbol() ? "Q" : "$(q.name) q", "uadrature rule: on [$(q.a), $(q.b)], $(length(q)) points, eltype = $(eltype(q))")
+Base.length(q::Quadrature) = length(q.x)
+Base.eachindex(q::Quadrature) = eachindex(q.x)
+Base.:(==)(q1::Quadrature, q2::Quadrature) = q1.x == q2.x && q1.w == q2.w && q1.a == q2.a && q1.b == q2.b
+Base.:(≈)(q1::Quadrature, q2::Quadrature) = q1.x ≈ q2.x && q1.w ≈ q2.w && q1.a ≈ q2.a && q1.b ≈ q2.b
+node(q::Quadrature, i) = (q.b+q.a)/2 + (q.b-q.a)/2 * q.x[i]
+nodes(q::Quadrature) = (q.b+q.a)/2 .+ (q.b-q.a)/2 .* q.x
+weights(q::Quadrature) = (q.b-q.a)/2 .* q.w # weights for integration on [a, b]
+@inbounds @fastmath function integrate(q::Quadrature, f::AbstractVector)
+    length(f) == length(q) || throw(ArgumentError("$(length(f))-point f is incompatible with $(length(q))-point quadrature rule"))
+    return (q.b-q.a)/2 * sum(q.w[i] * f[i] for i in eachindex(q)) # integrate samples of f on [a, b]
+end
+@inbounds @fastmath function integrate(q::Quadrature, f::Function)
+    return (q.b-q.a)/2 * sum(q.w[i] * f(node(q, i)) for i in eachindex(q)) # integrate f(x) on [a, b]
+end
+(q::Quadrature)(args...) = integrate(q, args...) # calling is equivalent to integrating

src/solve.jl

@@ -605,11 +605,6 @@ function issuccess(sol::CosmologySolution)
     return successful_retcode(sol.bg) && (isnothing(sol.pts) || all(successful_retcode(pt) for pt in sol.pts))
 end
 
-function integrate(xs, ys; integrator = TrapezoidalRule())
-    prob = SampledIntegralProblem(ys, xs)
-    sol = solve(prob, integrator)
-    return sol.u
-end
 integrate_cumulative(sol::CosmologySolution, x, y) = cumul_integrate(sol[x], sol[y])
 integrate_cumulative(sol::CosmologySolution, y) = integrate_cumulative(sol, sol.prob.M.τ, y)