From cd77039fe99bd2c3a380b7094413a0cbd885f031 Mon Sep 17 00:00:00 2001 From: marcobonici Date: Tue, 3 Mar 2026 16:38:05 -0500 Subject: [PATCH 1/5] More tests; added chebyshev functionality --- Project.toml | 16 +++++-- src/Capse.jl | 110 +++++++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 94 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 216 insertions(+), 4 deletions(-) diff --git a/Project.toml b/Project.toml index c0e126f..d556ee2 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Capse" uuid = "994f66c3-d2c7-4ba6-88fb-4a10f50800ba" -authors = ["marcobonici "] version = "0.3.7" +authors = ["marcobonici "] [deps] AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6" @@ -10,17 +10,25 @@ JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605" [compat] -julia = "1.10" -AbstractCosmologicalEmulators = "0.8" +AbstractCosmologicalEmulators = "0.8, 0.9" Adapt = "3, 4" +DifferentiationInterface = "0.7.16" +ForwardDiff = "1.3.2" JSON = "1" +Mooncake = "0.5.12" NPZ = "0.4" +Zygote = "0.7.10" +julia = "1.10" [extras] +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605" SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5" Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["NPZ", "SimpleChains", "Static", "Test"] +test = ["DifferentiationInterface", "ForwardDiff", "Mooncake", "NPZ", "SimpleChains", "Static", "Test", "Zygote"] diff --git a/src/Capse.jl b/src/Capse.jl index a1a1841..0c4480a 100644 --- a/src/Capse.jl +++ b/src/Capse.jl @@ -8,6 +8,7 @@ import JSON: parsefile import NPZ: npzread export get_Cℓ +export ChebyshevInterpolPlan, prepare_Cℓ_interpolation, interp_Cℓ abstract type AbstractCℓEmulators end @@ -105,6 +106,115 @@ function get_Cℓ(input_params, Cℓemu::AbstractCℓEmulators) return Cℓemu.Postprocessing(input_params, norm_output, Cℓemu) end +""" + ChebyshevInterpolPlan{P, T} + +Pre-computed plan for interpolating Cℓ spectra from a Chebyshev ℓ-grid onto a user-defined ℓ-grid. +Construct with `prepare_Cℓ_interpolation(emulator, ℓgrid_new)`. +""" +struct ChebyshevInterpolPlan{P, T} + cheb_plan::ChebyshevPlan{1, P, T} # FFT plan for decomposition + T_mat::Matrix{T} # Chebyshev basis at ℓgrid_new, shape (n_new, K+1) + ℓ_min::T + ℓ_max::T + K::Int + ascending::Bool # true if the emulator ℓgrid was stored ascending +end + +""" + prepare_Cℓ_interpolation(Cℓemu, ℓgrid_new; tol=1e-6) -> ChebyshevInterpolPlan + +Prepare a reusable interpolation plan for `interp_Cℓ`. + +# Arguments +- `Cℓemu::AbstractCℓEmulators`: The emulator whose `ℓgrid` defines the Chebyshev nodes. +- `ℓgrid_new::AbstractVector`: Target ℓ-values for evaluation. +- `tol::Real=1e-6`: Tolerance for Chebyshev grid validation. + +# Grid-orientation logic +1. Retrieve `ℓgrid` from the emulator. +2. If ascending, set `ascending=true` and reverse internally (Chebyshev nodes are descending). +3. Compute `ℓ_min`, `ℓ_max`, `K = length(ℓgrid) - 1`. +4. Generate the *expected* Chebyshev grid via `chebpoints(K, ℓ_min, ℓ_max)` and compare + element-wise to the (possibly reversed) emulator grid. + - If `norm(expected - observed) / norm(expected) > tol`, emit a `@warn` + informing the user that the ℓ-grid may not be a Chebyshev grid and accuracy + could be degraded. +5. Prepare the FFT plan and precompute `T_mat = chebyshev_polynomials(ℓgrid_new, ℓ_min, ℓ_max, K)`. +""" +function prepare_Cℓ_interpolation(Cℓemu::AbstractCℓEmulators, + ℓgrid_new::AbstractVector; + tol::Real = 1e-6) + ℓgrid = get_ℓgrid(Cℓemu) + ascending = issorted(ℓgrid) # ascending ↔ needs reversal for FFTW + ℓgrid_desc = ascending ? reverse(ℓgrid) : ℓgrid + + ℓ_min = Float64(last(ℓgrid_desc)) # smallest ℓ (tail of descending vector) + ℓ_max = Float64(first(ℓgrid_desc)) # largest ℓ (head of descending vector) + K = length(ℓgrid) - 1 + + # Validate against theoretical Chebyshev nodes + expected = chebpoints(K, ℓ_min, ℓ_max) # descending, K+1 points + + # We use LinearAlgebra.norm so we need to make sure LinearAlgebra is available + # but to avoid adding dependencies we can just use maximum absolute difference + rel_err = maximum(abs.(expected .- ℓgrid_desc)) / maximum(abs.(expected)) + + if rel_err > tol + @warn "The emulator ℓ-grid does not appear to be a Chebyshev grid " * + "(relative max deviation = $(round(rel_err; sigdigits=3))). " * + "Interpolation accuracy may be degraded." + end + + cheb_plan = prepare_chebyshev_plan(ℓ_min, ℓ_max, K) + T_mat = chebyshev_polynomials(Float64.(ℓgrid_new), ℓ_min, ℓ_max, K) + + return ChebyshevInterpolPlan(cheb_plan, T_mat, ℓ_min, ℓ_max, K, ascending) +end + +""" + interp_Cℓ(Cℓ_vals, plan) -> Vector + +Interpolate a single Cℓ spectrum from its Chebyshev ℓ-grid onto the target ℓ-grid +baked into `plan`. + +# Arguments +- `Cℓ_vals::AbstractVector`: Spectrum values on the emulator's Chebyshev ℓ-grid (length K+1). + Must be in the *same orientation* as the emulator's stored ℓ-grid (ascending or descending). +- `plan::ChebyshevInterpolPlan`: Prepared interpolation plan. + +# Returns +- `Vector`: Spectrum evaluated at the target ℓ-grid. +""" +function interp_Cℓ(Cℓ_vals::AbstractVector, plan::ChebyshevInterpolPlan) + # If emulator stored ascending, reverse so that values match descending Chebyshev nodes + vals_desc = plan.ascending ? reverse(Cℓ_vals) : Cℓ_vals + coeffs = chebyshev_decomposition(plan.cheb_plan, vals_desc) + return plan.T_mat * coeffs +end + +""" + interp_Cℓ(Cℓ_mat, plan) -> Matrix + +Interpolate multiple Cℓ spectra (columns of a matrix) onto the target ℓ-grid +baked into `plan`. + +# Arguments +- `Cℓ_mat::AbstractMatrix`: Shape `(n_ℓ, n_spectra)`. Each column is a spectrum on the + emulator's Chebyshev ℓ-grid, in the same orientation as the stored ℓ-grid. +- `plan::ChebyshevInterpolPlan` + +# Returns +- `Matrix`: Shape `(length(ℓgrid_new), n_spectra)`. +""" +function interp_Cℓ(Cℓ_mat::AbstractMatrix, plan::ChebyshevInterpolPlan) + # Reverse rows if the emulator's ℓ-grid was ascending + mat_desc = plan.ascending ? reverse(Cℓ_mat; dims=1) : Cℓ_mat + # chebyshev_decomposition already handles batched (matrix) input column-wise + coeffs = chebyshev_decomposition(plan.cheb_plan, mat_desc) # (K+1, n_spectra) + return plan.T_mat * coeffs # (n_new, n_spectra) +end + """ get_ℓgrid(CℓEmulator::AbstractCℓEmulators) -> AbstractVector diff --git a/test/runtests.jl b/test/runtests.jl index d4f692f..99ddfcf 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,6 +3,9 @@ using NPZ using SimpleChains using Static using Capse +using AbstractCosmologicalEmulators: chebpoints +using DifferentiationInterface +import ForwardDiff, Zygote, Mooncake mlpd = SimpleChain( static(6), @@ -40,3 +43,94 @@ capse_loaded_emu = Capse.load_emulator("emu/") @test_logs (:warn, "No emulator description found!") Capse.get_emulator_description(capse_emu) @test Capse.get_Cℓ(cosmo_vec, capse_emu) == Capse.get_Cℓ(cosmo_vec, capse_loaded_emu) end + +@testset "Chebyshev interpolation" begin + # 1. Correctness on smooth function + K = 39 + ℓ_min, ℓ_max = 2.0, 2500.0 + ℓ_cheb_desc = chebpoints(K, ℓ_min, ℓ_max) # Descending + + # Test function: smooth spectrum-like shape + f_test(l) = 1e4 * exp(-l / 500) + + Cℓ_true_desc = f_test.(ℓ_cheb_desc) + + # Create fake emulator with descending grid + emu_desc = Capse.CℓEmulator(TrainedEmulator = emu, ℓgrid=ℓ_cheb_desc, InMinMax = inminmax, + OutMinMax = outminmax, Postprocessing = postprocessing) + + ℓ_new = collect(LinRange(ℓ_min, ℓ_max, 100)) + Cℓ_exact = f_test.(ℓ_new) + + # Prepare plan + plan_desc = prepare_Cℓ_interpolation(emu_desc, ℓ_new) + @test plan_desc.ascending == false + + # Vector interp + Cℓ_interp_desc = interp_Cℓ(Cℓ_true_desc, plan_desc) + @test length(Cℓ_interp_desc) == 100 + @test isapprox(Cℓ_interp_desc, Cℓ_exact; rtol=1e-10) + + # 2. Ascending grid + ℓ_cheb_asc = reverse(ℓ_cheb_desc) + Cℓ_true_asc = reverse(Cℓ_true_desc) + + emu_asc = Capse.CℓEmulator(TrainedEmulator = emu, ℓgrid=ℓ_cheb_asc, InMinMax = inminmax, + OutMinMax = outminmax, Postprocessing = postprocessing) + + plan_asc = prepare_Cℓ_interpolation(emu_asc, ℓ_new) + @test plan_asc.ascending == true + + Cℓ_interp_asc = interp_Cℓ(Cℓ_true_asc, plan_asc) + @test isapprox(Cℓ_interp_asc, Cℓ_interp_desc; rtol=1e-14) + + # 3. Matrix layout + # Columns are spectra + Cℓ_mat_desc = hcat(Cℓ_true_desc, Cℓ_true_desc .* 1.1) + Cℓ_mat_interp = interp_Cℓ(Cℓ_mat_desc, plan_desc) + @test size(Cℓ_mat_interp) == (100, 2) + @test isapprox(Cℓ_mat_interp[:, 1], Cℓ_interp_desc; rtol=1e-14) + @test isapprox(Cℓ_mat_interp[:, 2], Cℓ_interp_desc .* 1.1; rtol=1e-14) + + # 4. Warnings on non-Chebyshev grid + ℓ_uniform = collect(LinRange(ℓ_min, ℓ_max, K+1)) + emu_uniform = Capse.CℓEmulator(TrainedEmulator = emu, ℓgrid=ℓ_uniform, InMinMax = inminmax, + OutMinMax = outminmax, Postprocessing = postprocessing) + @test_logs (:warn, r"The emulator ℓ-grid does not appear to be a Chebyshev grid") prepare_Cℓ_interpolation(emu_uniform, ℓ_new) +end + +@testset "Chebyshev AD tests" begin + K = 39 + ℓ_min, ℓ_max = 2.0, 2500.0 + ℓ_cheb_desc = chebpoints(K, ℓ_min, ℓ_max) + ℓ_new = collect(LinRange(ℓ_min, ℓ_max, 100)) + + emu = Capse.CℓEmulator(TrainedEmulator = capse_emu.TrainedEmulator, ℓgrid=ℓ_cheb_desc, + InMinMax = inminmax, OutMinMax = outminmax, Postprocessing = postprocessing) + plan = prepare_Cℓ_interpolation(emu, ℓ_new) + + v = rand(K+1) + M = rand(K+1, 3) + + backends = [ + AutoForwardDiff(), + AutoZygote(), + AutoMooncake(config=nothing) + ] + + for b in backends + @testset "AD backend: $(typeof(b))" begin + # Vector case + f_vec(x) = sum(interp_Cℓ(x, plan)) + g_vec = DifferentiationInterface.gradient(f_vec, b, v) + @test size(g_vec) == size(v) + @test all(isfinite, g_vec) + + # Matrix case + f_mat(X) = sum(interp_Cℓ(X, plan)) + g_mat = DifferentiationInterface.gradient(f_mat, b, M) + @test size(g_mat) == size(M) + @test all(isfinite, g_mat) + end + end +end From 85561d8c147be96823c5e2664484c509a08dacd8 Mon Sep 17 00:00:00 2001 From: marcobonici Date: Tue, 3 Mar 2026 16:40:49 -0500 Subject: [PATCH 2/5] Updated project --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index d556ee2..c22a42c 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Capse" uuid = "994f66c3-d2c7-4ba6-88fb-4a10f50800ba" version = "0.3.7" -authors = ["marcobonici "] +authors = ["Marco Bonici "] [deps] AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6" From 16d2aa030a3a782af96103b0b41748b2b58f1bf3 Mon Sep 17 00:00:00 2001 From: marcobonici Date: Tue, 3 Mar 2026 17:05:17 -0500 Subject: [PATCH 3/5] Updating version --- Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index c22a42c..0d59b2e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Capse" uuid = "994f66c3-d2c7-4ba6-88fb-4a10f50800ba" -version = "0.3.7" +version = "0.4.0" authors = ["Marco Bonici "] [deps] @@ -10,7 +10,7 @@ JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605" [compat] -AbstractCosmologicalEmulators = "0.8, 0.9" +AbstractCosmologicalEmulators = "0.9" Adapt = "3, 4" DifferentiationInterface = "0.7.16" ForwardDiff = "1.3.2" From cccf563bfb0a63f3fa10a9cbd126217cdfb3782b Mon Sep 17 00:00:00 2001 From: marcobonici Date: Tue, 3 Mar 2026 17:11:20 -0500 Subject: [PATCH 4/5] Updating test --- .github/workflows/test.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 7366b33..1476c44 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -43,8 +43,8 @@ jobs: ${{ runner.os }}- - uses: julia-actions/julia-runtest@v1 - uses: julia-actions/julia-processcoverage@v1 - - uses: codecov/codecov-action@v4 - env: - CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} + - uses: codecov/codecov-action@v5 with: - file: lcov.info + files: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + fail_ci_if_error: false From ad34df2598b6ed40f473d54c052ae6dbba788113 Mon Sep 17 00:00:00 2001 From: marcobonici Date: Tue, 3 Mar 2026 17:45:09 -0500 Subject: [PATCH 5/5] Updating docs --- docs/src/index.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/docs/src/index.md b/docs/src/index.md index beedac1..f6a71f7 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -46,6 +46,17 @@ Cℓ = Capse.get_Cℓ(params, Cℓ_emu) ## Detailed Guide +### Interpolation + +`Capse.jl` supports rapid and exact interpolation between the emulator's natively trained `ℓgrid` and an arbitrary user-defined `ℓgrid` via precomputed FFT plans using Chebyshev polynomials. + +```@docs +Capse.ChebyshevInterpolPlan +Capse.prepare_Cℓ_interpolation +Capse.interp_Cℓ(::AbstractVector, ::Capse.ChebyshevInterpolPlan) +Capse.interp_Cℓ(::AbstractMatrix, ::Capse.ChebyshevInterpolPlan) +``` + ### Loading Emulators `Capse.jl` supports loading pre-trained emulators from disk. Trained weights are available on [Zenodo](https://zenodo.org/record/8187935).