Abstractcosmoemu

.github/workflows/test.yml

@@ -19,8 +19,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.7'
-          - '1.8'
           - '1.9'
           - '1.10'
         os:

Project.toml

@@ -6,9 +6,12 @@ version = "0.2.0"
 [deps]
 AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
+Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
 LegendrePolynomials = "3db4a2ba-fc88-11e8-3e01-49c72059a882"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Memoization = "6fafb56a-5788-4b4e-91ca-c0cea6611c73"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"

src/Effort.jl

@@ -3,18 +3,20 @@ module Effort
 using Base: @kwdef
 using AbstractCosmologicalEmulators
 import AbstractCosmologicalEmulators.get_emulator_description
+using ChainRulesCore
 using DataInterpolations
 using FastGaussQuadrature
 using LegendrePolynomials
 using LoopVectorization
 using Memoization
 using OrdinaryDiffEq
-using QuadGK
+using Integrals
+using LinearAlgebra
 
 const c_0 = 2.99792458e5
 
 function __init__()
-    min_y = get_y(0,0) #obvious, I know
+    min_y = get_y(0,0) #obvious, I knowadd OrdinaryDiffEqTsit5
     max_y = get_y(1,10)
     y_grid = vcat(LinRange(min_y, 100, 100), LinRange(100.1, max_y, 1000))
     F_grid = [_F(y) for y in y_grid]
@@ -28,5 +30,7 @@ include("background.jl")
 include("neural_networks.jl")
 include("eft_commands.jl")
 include("projection.jl")
+include("utils.jl")
+include("chainrules.jl")
 
 end # module

src/background.jl

@@ -3,17 +3,23 @@
 # and maybe used in other packages
 
 function _F(y)
-    result, _ = quadgk(x -> x^2*√(x^2+y^2)/(1+exp(x)), 0, Inf, rtol=1e-12)
-    return result
+    f(x, y) = x^2*√(x^2+y^2)/(1+exp(x))
+    domain = (zero(eltype(Inf)), Inf) # (lb, ub)
+    prob = IntegralProblem(f, domain, y; reltol=1e-12)
+    sol = solve(prob, QuadGKJL())[1]
+    return sol
 end
 
 function get_y(mν, a; kB = 8.617342e-5, Tν = 0.71611*2.7255)
     return mν * a / (kB * Tν)
 end
 
 function _dFdy(y)
-    result, _ = quadgk(x -> x^2/((1+exp(x))*√(x^2+y^2)), 0, Inf, rtol=1e-12)
-    return result*y
+    f(x, y) = x^2/((1+exp(x))*√(x^2+y^2))
+    domain = (zero(eltype(Inf)), Inf) # (lb, ub)
+    prob = IntegralProblem(f, domain, y; reltol=1e-12)
+    sol = solve(prob, QuadGKJL())[1]
+    return sol*y
 end
 
 function _ΩνE2(a, Ωγ0, mν; kB = 8.617342e-5, Tν = 0.71611*2.7255, Neff = 3.044)
@@ -61,9 +67,12 @@
 _H_a(a, Ωγ0, Ωm0, mν, h, w0, wa) = 100*h*_E_a(a, Ωm0, h; mν =mν, w0=w0, wa=wa)
 
 function _χ_z(z, Ωm0, h; mν =0., w0=-1., wa=0.)
-    integral, _ = quadgk(x -> 1 /
-    _E_a(_a_z(x), Ωm0, h; mν =mν, w0=w0, wa=wa), 0, z, rtol=1e-6)
-    return integral*c_0/(100*h)
+    p = [Ωm0, h, mν , w0, wa]
+    f(x, p) = 1 / _E_a(_a_z(x), p[1], p[2]; mν =p[3], w0=p[4], wa=p[5])
+    domain = (zero(eltype(z)), z) # (lb, ub)
+    prob = IntegralProblem(f, domain, y; reltol=1e-12)
+    sol = solve(prob, QuadGKJL())[1]
+    return sol*c_0/(100*h)
 end
 
 function _dEda(a, Ωm0, h; mν =0., w0=-1., wa=0.)
@@ -88,9 +97,23 @@
     return (Ωm0 + Ων0 )*a^-3 / (_E_a(a, Ωm0, h; mν =mν, w0=w0, wa=wa))^2
 end
 
+function _r̃_z_check(z, ΩM, h; mν =0., w0=-1., wa=0.)
+    p = [ΩM, h, mν , w0, wa]
+    f(x, p) = 1 / _E_a(_a_z(x), p[1], p[2]; mν =p[3], w0=p[4], wa=p[5])
+    domain = (zero(eltype(z)), z) # (lb, ub)
+    prob = IntegralProblem(f, domain, p; reltol=1e-12)
+    sol = solve(prob, QuadGKJL())[1]
+    return sol
+end
+
 function _r̃_z(z, ΩM, h; mν =0., w0=-1., wa=0.)
-    integral, _ = quadgk(x -> 1 / _E_z(x, ΩM, h; mν =mν, w0=w0, wa=wa), 0, z, rtol=1e-10)
-    return integral
+    z_array, weigths_array = _transformed_weights(FastGaussQuadrature.gausslegendre, 9, 0, z)
+    integrand_array = @. 1. / _E_a(_a_z(z_array), ΩM, h; mν =mν, w0=w0, wa=wa)
+    return dot(weigths_array, integrand_array)
+end
+
+function _r_z_check(z, ΩM, h; mν =0., w0=-1., wa=0.)
+    return c_0 * _r̃_z_check(z, ΩM, h; mν =mν, w0=w0, wa=wa) / (100*h)
 end
 
 function _r_z(z, ΩM, h; mν =0., w0=-1., wa=0.)
@@ -126,7 +149,6 @@
 end
 
 function _growth!(du,u,p,loga)
-    #Ωγ0 = p[1]
     Ωm0 = p[1]
     mν  = p[2]
     h   = p[3]

src/chainrules.jl

@@ -0,0 +1,2 @@
+@non_differentiable LinRange(a,b,n)
+@non_differentiable _transformed_weights(quadrature_rule, order, a,b)

src/utils.jl

@@ -0,0 +1,6 @@
+function _transformed_weights(quadrature_rule, order, a,b)
+    x, w = quadrature_rule(order)
+    x = (b-a)/2. .* x .+ (b+a)/2.
+    w = (b-a)/2. .* w
+    return x, w
+end

test/runtests.jl

@@ -42,6 +42,16 @@ function f_z_x(z, x)
     sum(Effort._f_z(z, Ωm0, h; mν =mν, w0=w0, wa=wa))
 end
 
+function r_z_x(z, x)
+    Ωm0, h, mν, w0, wa = x
+    sum(Effort._r_z(z, Ωm0, h; mν =mν, w0=w0, wa=wa))
+end
+
+function r_z_check_x(z, x)
+    Ωm0, h, mν, w0, wa = x
+    sum(Effort._r_z_check(z, Ωm0, h; mν =mν, w0=w0, wa=wa))
+end
+
 @testset "Effort tests" begin
     @test isapprox(Effort._H_a(a, Ωγ0, Ωm0, mν, h, w0, wa), h*100)
     @test isapprox(Effort._E_a(a, Ωm0, h), 1.)
@@ -51,4 +61,8 @@ end
     @test isapprox(FiniteDiff.finite_difference_gradient(x->D_z_x(z, x), x), ForwardDiff.gradient(x->D_z_x(z, x), x), rtol=1e-5)
     @test isapprox(Zygote.gradient(x->f_z_x(z, x), x)[1], ForwardDiff.gradient(x->f_z_x(z, x), x), rtol=1e-5)
     @test isapprox(FiniteDiff.finite_difference_gradient(x->f_z_x(z, x), x), ForwardDiff.gradient(x->f_z_x(z, x), x), rtol=1e-4)
+    @test isapprox(FiniteDiff.finite_difference_gradient(x->r_z_x(3., x), x), ForwardDiff.gradient(x->r_z_x(3., x), x), rtol=1e-7)
+    @test isapprox(Zygote.gradient(x->r_z_x(3., x), x)[1], ForwardDiff.gradient(x->r_z_x(3., x), x), rtol=1e-6)
+    @test isapprox(Zygote.gradient(x->r_z_x(3., x), x)[1], Zygote.gradient(x->r_z_check_x(3., x), x)[1], rtol=1e-7)
+    @test isapprox(Effort._r_z(3., Ωm0, h; mν =mν, w0=w0, wa=wa), Effort._r_z_check(3., Ωm0, h; mν =mν, w0=w0, wa=wa), rtol=1e-6)
 end