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

[hersle]

Dirty draft testing script:

using SymBoltz, Plots

M = ΛCDM()
p = parameters_Planck18(M)
prob = CosmologyProblem(M, p)
sol = solve(prob)

# Compute τ integration points and weights for integrating ∫ dτ S(τ) jₗ(k(τ₀-τ)) from τi to τ0
function weights_trapezoidal(l, k, τi, τ0, N = 1024)
    τ = range(τi, τ0, length = N)
    Δτ = step(τ)
    τ = collect(τ)
    w = [i == 1 || i == N ? Δτ/2 : Δτ for i in eachindex(τ)]
    w .*= SymBoltz.jl.(l, k .* (τ0 .- τ))
    return τ, w
end

# Same, but for Clenshaw-Curtis quadrature instead of trapezoidal
function weights_clenshaw_curtis(l, k, τi, τ0, N = 1024)
    τ = zeros(N)
    w = zeros(N)
    
    for n in 0:N-1
        x = cos(π*n/(N-1)) # Chebyshev node in [-1, 1]
        τ[n+1] = (τ0+τi)/2 + (τ0-τi)/2 * x # transform to [τi, τ0]
        s = sum(cos(2π*n*m/(N-1)) / (1 - 4m^2) for m in 1:div(N-1, 2)) # sum over Fourier modes
        cₙ = (n == 0 || n == N-1) ? 1 : 2 # cₙ coefficient: 1 for endpoints, 2 for interior points
        w[n+1] = cₙ / (N-1) * (1 + 2s)
    end

    w .*= (τ0 - τi) / 2 # rescale weights for the interval [τi, τ0]
    w .*= SymBoltz.jl.(l, k.*(τ0.-τ)) # multiply by spherical Bessel function weights

    return reverse(τ), reverse(w) # make τ ascending
end

# Same, but for Filon-Clenshaw-Curtis quadrature
# (decomposing the function to be integrated by into chebyshev polynomials)
function weights_filon_clenshaw_curtis(l, k, τi, τ0, N = 1024; N_moment = 1024)
    # Calculate rescaled Chebyshev nodes into the interval [τi, τ0]
    x = cos.((0:N-1) .* π / (N-1)) # Chebyshev nodes in [-1, 1]
    τ = (τ0 + τi) / 2 .+ (τ0 - τi) / 2 .* x
    
    # To compute weights for Filon-type method, we need to compute moments:
    # M_n = ∫_{-1}^{1} T_n(x) j_l(k(τ0 - τ(x))) dx
    # where τ(x) = (τ0 + τi)/2 + (τ0 - τi)/2 * x
    
    # Compute the moments M_n for n = 0, 1, ..., N-1
    moments = zeros(N)
    
    # Use numerical integration to compute moments
    # For better accuracy, use a finer grid for moment computation
    div_N_moment_m_1 = 1.0 / (N_moment - 1)
    x_fine = cos.((0:N_moment-1) .* π / (N_moment-1))
    τ_fine = (τ0 + τi) / 2 .+ (τ0 - τi) / 2 .* x_fine
    
    Tₙ = zeros(N_moment)
    Tₙ₋₁ = zeros(N_moment)
    Tₙ₋₂ = zeros(N_moment)
    integrand = zeros(N_moment)
    w_fine = zeros(N_moment)
    @fastmath for n in 0:N-1
        # Evaluate Chebyshev polynomial T_n at fine grid points
        # Using recurrence: T_0(x) = 1, T_1(x) = x, T_{n+1}(x) = 2x*T_n(x) - T_{n-1}(x)
        if n == 0
            Tₙ .= 1.0
        elseif n == 1
            Tₙ .= x_fine
        else
            Tₙ .= 2 .* x_fine .* Tₙ₋₁ .- Tₙ₋₂
        end
        
        # Integrate T_n(x) * j_l(k(τ0 - τ(x))) using Clenshaw-Curtis on fine grid
        integrand .= Tₙ .* SymBoltz.jl.(l, k .* (τ0 .- τ_fine))
        
        # Apply Clenshaw-Curtis weights
        for j in 0:N_moment-1
            c_j = (j == 0 || j == N_moment-1) ? 1.0 : 2.0
            sum_term = 0.0
            for m in 1:div(N_moment-1, 2)
                sum_term += cos(2π*j*m * div_N_moment_m_1) / (1 - 4m^2)
            end
            w_fine[j+1] = c_j * div_N_moment_m_1 * (1 + 2 * sum_term)
        end
        
        moments[n+1] = sum(integrand .* w_fine) * (τ0 - τi) / 2

        Tₙ₋₂ .= Tₙ₋₁
        Tₙ₋₁ .= Tₙ
    end
    
    # Now compute the weights for evaluating the integral
    # The integral is: sum_n a_n * M_n, where a_n are Chebyshev coefficients
    # The Chebyshev coefficients are: a_n = (2/(N-1)) * c_n * sum_j c_j * f_j * cos(π*n*j/(N-1))
    # where c_0 = c_{N-1} = 1/2, c_j = 1 otherwise
    # So: integral = sum_n M_n * a_n = sum_j f_j * w_j
    # where w_j = sum_n M_n * (2/(N-1)) * c_n * c_j * cos(π*n*j/(N-1))
    
    w = zeros(N)
    for j in 0:N-1
        c_j = (j == 0 || j == N-1) ? 0.5 : 1.0
        for n in 0:N-1
            c_n = (n == 0 || n == N-1) ? 0.5 : 1.0
            w[j+1] += (2.0 / (N-1)) * c_n * c_j * cos(π * n * j / (N-1)) * moments[n+1]
        end
    end

    return reverse(τ), reverse(w) # make τ ascending
end

function integrate(weightsfunc, l, k, N)
    sol = solve(prob)
    τi = sol.bg.t[begin]
    τ0 = sol.bg.t[end]
    τ, w = weightsfunc(l, k, τi, τ0, N)
    sol = solve(prob, k; ptextraopts = (saveat = τ,))
    S = only(sol.pts)[M.ST]
    I = sum(S .* w)
    return I
end

τ0 = sol.bg.t[end]
l = 1000
k = l / τ0
Ns_trapezoidal = [2^N for N in 4:13]
Is_trapezoidal = map(N -> integrate(weights_trapezoidal, l, k, N), Ns_trapezoidal)
Ns_clenshaw_curtis = [2^N for N in 4:11]
Is_clenshaw_curtis = map(N -> integrate(weights_clenshaw_curtis, l, k, N), Ns_clenshaw_curtis)
Is_filon_clenshaw_curtis = map(N -> integrate(weights_filon_clenshaw_curtis, l, k, N), Ns_clenshaw_curtis)

Plots.default(marker = :circle)
plot(xlabel = "log₂(N)", ylabel = "I", xticks = 0:20)
plot!(log2.(Ns_trapezoidal), Is_trapezoidal; label = "Trapezoidal")
plot!(log2.(Ns_clenshaw_curtis), Is_clenshaw_curtis; label = "Clenshaw-Curtis")
plot!(log2.(Ns_clenshaw_curtis), Is_filon_clenshaw_curtis; label = "Filon-Clenshaw-Curtis")

CC seems to converge faster than trapezoidal.