Merge pull request #79 from CosmologicalEmulators/develop

Develop

Author: marcobonici

.github/workflows/ci.yml

@@ -46,7 +46,7 @@ jobs:
         with:
           directories: src,ext
       - uses: codecov/codecov-action@v5
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
         with:
-          file: lcov.info
+          files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: false

Project.toml

@@ -1,6 +1,6 @@
 name = "AbstractCosmologicalEmulators"
 uuid = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
-version = "0.9.2"
+version = "0.9.3"
 authors = ["Marco Bonici <bonici.marco@gmail.com>"]
 
 [deps]
@@ -60,9 +60,10 @@ Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ADTypes", "DataInterpolations", "DifferentiationInterface", "FastChebInterp", "FastGaussQuadrature", "FiniteDifferences", "ForwardDiff", "Integrals", "JSON", "LinearAlgebra", "Lux", "Mooncake", "NPZ", "OrdinaryDiffEqTsit5", "Random", "SimpleChains", "Test", "Zygote"]
+test = ["ADTypes", "DataInterpolations", "DifferentiationInterface", "FastChebInterp", "FastGaussQuadrature", "FiniteDifferences", "ForwardDiff", "Integrals", "JSON", "LinearAlgebra", "Lux", "Mooncake", "NPZ", "OrdinaryDiffEqTsit5", "Random", "SciMLSensitivity", "SimpleChains", "Test", "Zygote"]

ext/BackgroundCosmologyExt/background.jl

@@ -116,30 +116,30 @@ function _Ωma(a, cosmo::w0waCDMCosmology)
     return _Ωma(a, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0)
 end
 
-function r̃_z(z::Number, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0)
-    z_array, weigths_array = _transformed_weights(FastGaussQuadrature.gausslegendre, 9, 0, z)
+function r̃_z(z::Number, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0, order=9)
+    z_array, weigths_array = _transformed_weights(FastGaussQuadrature.gausslegendre, order, 0, z)
     integrand_array = 1.0 ./ E_a(_a_z(z_array), Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0)
     return dot(weigths_array, integrand_array)
 end
 
-function r̃_z(z::AbstractArray, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0)
-    return [r̃_z(zi, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0) for zi in z]
+function r̃_z(z::AbstractArray, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0, order=9)
+    return [r̃_z(zi, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0, order=order) for zi in z]
 end
 
-function r̃_z(z, cosmo::w0waCDMCosmology)
+function r̃_z(z, cosmo::w0waCDMCosmology; order=9)
     Ωcb0 = (cosmo.ωb + cosmo.ωc) / cosmo.h^2
     Ωk0 = cosmo.ωk / cosmo.h^2
-    return r̃_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0)
+    return r̃_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0, order=order)
 end
 
-function r_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0)
-    return c_0 * r̃_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0) / (100 * h)
+function r_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0, order=9)
+    return c_0 * r̃_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0, order=order) / (100 * h)
 end
 
-function r_z(z, cosmo::w0waCDMCosmology)
+function r_z(z, cosmo::w0waCDMCosmology; order=9)
     Ωcb0 = (cosmo.ωb + cosmo.ωc) / cosmo.h^2
     Ωk0 = cosmo.ωk / cosmo.h^2
-    return r_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0)
+    return r_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0, order=order)
 end
 
 function S_of_K(Ω::Number, r)
@@ -154,54 +154,54 @@ function S_of_K(Ω::Number, r)
     end
 end
 
-function d̃M_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0)
-    return S_of_K(Ωk0, r̃_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0))
+function d̃M_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0, order=9)
+    return S_of_K(Ωk0, r̃_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0, order=order))
 end
 
-function d̃M_z(z, cosmo::w0waCDMCosmology)
+function d̃M_z(z, cosmo::w0waCDMCosmology; order=9)
     Ωcb0 = (cosmo.ωb + cosmo.ωc) / cosmo.h^2
     Ωk0 = cosmo.ωk / cosmo.h^2
-    return d̃M_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0)
+    return d̃M_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0, order=order)
 end
 
-function dM_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0)
-    return c_0 * d̃M_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0) / (100 * h)
+function dM_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0, order=9)
+    return c_0 * d̃M_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0, order=order) / (100 * h)
 end
 
-function dM_z(z, cosmo::w0waCDMCosmology)
+function dM_z(z, cosmo::w0waCDMCosmology; order=9)
     Ωcb0 = (cosmo.ωb + cosmo.ωc) / cosmo.h^2
     Ωk0 = cosmo.ωk / cosmo.h^2
-    return dM_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0)
+    return dM_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0, order=order)
 end
 
-function d̃A_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0)
-    return d̃M_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0) ./ (1 .+ z)
+function d̃A_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0, order=9)
+    return d̃M_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0, order=order) ./ (1 .+ z)
 end
 
-function d̃A_z(z, cosmo::w0waCDMCosmology)
+function d̃A_z(z, cosmo::w0waCDMCosmology; order=9)
     Ωcb0 = (cosmo.ωb + cosmo.ωc) / cosmo.h^2
     Ωk0 = cosmo.ωk / cosmo.h^2
-    return d̃A_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0)
+    return d̃A_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0, order=order)
 end
 
-function dA_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0)
-    return dM_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0) ./ (1 .+ z)
+function dA_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0, order=9)
+    return dM_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0, order=order) ./ (1 .+ z)
 end
 
-function dA_z(z, cosmo::w0waCDMCosmology)
+function dA_z(z, cosmo::w0waCDMCosmology; order=9)
     Ωcb0 = (cosmo.ωb + cosmo.ωc) / cosmo.h^2
     Ωk0 = cosmo.ωk / cosmo.h^2
-    return dA_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0)
+    return dA_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0, order=order)
 end
 
-function dL_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0)
-    return dM_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0) .* (1 .+ z)
+function dL_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0, order=9)
+    return dM_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa, Ωk0=Ωk0, order=order) .* (1 .+ z)
 end
 
-function dL_z(z, cosmo::w0waCDMCosmology)
+function dL_z(z, cosmo::w0waCDMCosmology; order=9)
     Ωcb0 = (cosmo.ωb + cosmo.ωc) / cosmo.h^2
     Ωk0 = cosmo.ωk / cosmo.h^2
-    return dL_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0)
+    return dL_z(z, Ωcb0, cosmo.h; mν=cosmo.mν, w0=cosmo.w0, wa=cosmo.wa, Ωk0=Ωk0, order=order)
 end
 
 function _growth!(du, u, p, loga)

test/runtests.jl

@@ -15,6 +15,7 @@ using DataInterpolations
 using LinearAlgebra
 using FastGaussQuadrature
 using FiniteDifferences
+using SciMLSensitivity
 using AbstractCosmologicalEmulators
 
 @testset "AbstractEmulators test" begin

test/test_background.jl

@@ -335,6 +335,12 @@ if !isnothing(ext)
             @test isapprox(dL_z(2.0, Ωcb0_class, h_class; mν=mν_class, w0=w0_class, wa=wa_class), 17445.761528257153, rtol=1e-4)
             @test isapprox(dA_z(2.0, Ωcb0_class, h_class; mν=mν_class, w0=w0_class, wa=wa_class), 1938.4179475841324, rtol=1e-4)
         end
+
+        @testset "z =1200.0" begin
+            @test isapprox(r_z(1200.0, Ωcb0_class, h_class; mν=mν_class, w0=w0_class, wa=wa_class, order=120), 14405.271466533255, rtol=2e-3)
+            @test isapprox(dL_z(1200.0, Ωcb0_class, h_class; mν=mν_class, w0=w0_class, wa=wa_class, order=120), 17300731.03130642, rtol=2e-3)
+            @test isapprox(dA_z(1200.0, Ωcb0_class, h_class; mν=mν_class, w0=w0_class, wa=wa_class, order=120), 11.9943975574798, rtol=2e-3)
+        end
     end
 
     @testset "CLASS comparison tests - cosmology 2" begin
@@ -375,6 +381,12 @@ if !isnothing(ext)
             @test isapprox(dL_z(2.0, Ωcb0_class2, h_class2; mν=mν_class2, w0=w0_class2, wa=wa_class2), 15764.581310383495, rtol=1e-4)
             @test isapprox(dA_z(2.0, Ωcb0_class2, h_class2; mν=mν_class2, w0=w0_class2, wa=wa_class2), 1751.6201455981693, rtol=1e-4)
         end
+
+        @testset "z = 1200.0" begin
+            @test isapprox(r_z(1200.0, Ωcb0_class2, h_class2; mν=mν_class2, w0=w0_class2, wa=wa_class2, order=120), 12906.201419127055, rtol=2e-3)
+            @test isapprox(dL_z(1200.0, Ωcb0_class2, h_class2; mν=mν_class2, w0=w0_class2, wa=wa_class2, order=120), 15500347.904371575, rtol=2e-3)
+            @test isapprox(dA_z(1200.0, Ωcb0_class2, h_class2; mν=mν_class2, w0=w0_class2, wa=wa_class2, order=120), 10.746212672045832, rtol=2e-3)
+        end
     end
 
     @testset "CLASS comparison tests - cosmology 3" begin
@@ -416,6 +428,13 @@ if !isnothing(ext)
             @test isapprox(dL_z(2.0, Ωcb0_class3, h_class3; mν=mν_class3, w0=w0_class3, wa=wa_class3, Ωk0=Ωk0_class3), 11240.362895970827594, rtol=1e-4)
             @test isapprox(dA_z(2.0, Ωcb0_class3, h_class3; mν=mν_class3, w0=w0_class3, wa=wa_class3, Ωk0=Ωk0_class3), 1248.929210663426375, rtol=1e-4)
         end
+
+        # z = 2.0 values
+        @testset "z = 1200.0" begin
+            @test isapprox(r_z(1200.0, Ωcb0_class3, h_class3; mν=mν_class3, w0=w0_class3, wa=wa_class3, Ωk0=Ωk0_class3, order=120), 10695.833449050964, rtol=2e-3)
+            @test isapprox(dL_z(1200.0, Ωcb0_class3, h_class3; mν=mν_class3, w0=w0_class3, wa=wa_class3, Ωk0=Ωk0_class3, order=120), 15749661.605665019, rtol=2e-3)
+            @test isapprox(dA_z(1200.0, Ωcb0_class3, h_class3; mν=mν_class3, w0=w0_class3, wa=wa_class3, Ωk0=Ωk0_class3, order=120), 10.919058989604848, rtol=2e-3)
+        end
     end
 
     @testset "Missing coverage tests" begin

test/test_lux_emulator_autodiff.jl

@@ -25,8 +25,10 @@ using Mooncake
     n_output = 10
 
     # Create a Lux neural network
+    # We use a closure for tanh to avoid internal mapping to tanh_fast
+    # which can cause issues with Mooncake in some environments
     lux_model = Chain(
-        Dense(n_input => 8, tanh),
+        Dense(n_input => 8, x -> tanh(x)),
         Dense(8 => n_output)
     )