From 6daf623a1328259ccf78f2e5ab3c2286451572cf Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Fri, 22 Sep 2023 07:06:57 +0200
Subject: [PATCH 01/29] Adding the AbstractCosmologicalEmulators package

---
 .github/workflows/docs.yml |  4 ++--
 Project.toml               |  9 +++++++--
 src/Effort.jl              |  3 ++-
 src/neural_networks.jl     | 27 ---------------------------
 4 files changed, 11 insertions(+), 32 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 6eee27c..5249325 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -18,11 +18,11 @@ jobs:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@latest
         with:
-          version: '1.8'
+          version: '1.9'
       - name: Install dependencies
         run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token
           DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # For authentication with SSH deploy key
-        run: julia --project=docs/ docs/make.jl
\ No newline at end of file
+        run: julia --project=docs/ docs/make.jl
diff --git a/Project.toml b/Project.toml
index 7fb6e9f..e9c3627 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,9 +1,11 @@
 name = "Effort"
 uuid = "ace19e24-02c2-4bce-bc41-b57e89c0fc53"
 authors = ["marcobonici <bonici.marco@gmail.com>"]
-version = "0.1.1"
+version = "0.1.2"
 
 [deps]
+AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 LegendrePolynomials = "3db4a2ba-fc88-11e8-3e01-49c72059a882"
@@ -11,4 +13,7 @@ LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Memoization = "6fafb56a-5788-4b4e-91ca-c0cea6611c73"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
-SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
+
+[compat]
+AbstractCosmologicalEmulators = "0.3.3"
+Adapt = "3"
diff --git a/src/Effort.jl b/src/Effort.jl
index 8febeb5..ce89c40 100644
--- a/src/Effort.jl
+++ b/src/Effort.jl
@@ -1,6 +1,8 @@
 module Effort
 
 using Base: @kwdef
+using AbstractCosmologicalEmulators
+import AbstractCosmologicalEmulators.get_emulator_description
 using DataInterpolations
 using FastGaussQuadrature
 using LegendrePolynomials
@@ -8,7 +10,6 @@ using LoopVectorization
 using Memoization
 using OrdinaryDiffEq
 using QuadGK
-using SimpleChains
 
 include("background.jl")
 include("neural_networks.jl")
diff --git a/src/neural_networks.jl b/src/neural_networks.jl
index f172451..2d52e50 100644
--- a/src/neural_networks.jl
+++ b/src/neural_networks.jl
@@ -1,26 +1,3 @@
-function _maximin_input!(x, in_MinMax)
-    for i in eachindex(x)
-        x[i] -= in_MinMax[i,1]
-        x[i] /= (in_MinMax[i,2]-in_MinMax[i,1])
-    end
-end
-
-function _inv_maximin_output!(x, out_MinMax)
-    for i in eachindex(x)
-        x[i] *= (out_MinMax[i,2]-out_MinMax[i,1])
-        x[i] += out_MinMax[i,1]
-    end
-end
-
-abstract type AbstractTrainedEmulators end
-#TODO: right now we support only SimpleChains emulators, but in the near future we may
-#support as well Flux and Lux. This is the reason behind the struct SimpleChainsEmulator
-
-@kwdef mutable struct SimpleChainsEmulator <: AbstractTrainedEmulators
-    Architecture
-    Weights
-end
-
 abstract type AbstractComponentEmulators end
 
 @kwdef mutable struct P11Emulator <: AbstractComponentEmulators
@@ -64,10 +41,6 @@ function get_component(input_params, comp_emu::PloopEmulator)
     return reshape(output, Int(length(output)/length(comp_emu.kgrid)), :)
 end
 
-function _run_emulator(input, trained_emulator::SimpleChainsEmulator)
-    return trained_emulator.Architecture(input, trained_emulator.Weights)
-end
-
 abstract type AbstractPℓEmulators end
 
 @kwdef mutable struct PℓEmulator <: AbstractPℓEmulators

From 1616bae60dca37386a710ca30ad7057edcc30883 Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Sun, 24 Sep 2023 05:59:52 +0200
Subject: [PATCH 02/29] Working with ACE

---
 src/neural_networks.jl | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/neural_networks.jl b/src/neural_networks.jl
index 2d52e50..cbbb3d9 100644
--- a/src/neural_networks.jl
+++ b/src/neural_networks.jl
@@ -23,9 +23,9 @@ end
 
 function get_component(input_params, comp_emu::AbstractComponentEmulators)
     input = deepcopy(input_params)
-    _maximin_input!(input, comp_emu.InMinMax)
-    output = Array(_run_emulator(input, comp_emu.TrainedEmulator))
-    _inv_maximin_output!(output, comp_emu.OutMinMax)
+    maximin_input!(input, comp_emu.InMinMax)
+    output = Array(run_emulator(input, comp_emu.TrainedEmulator))
+    inv_maximin_output!(output, comp_emu.OutMinMax)
     As = exp(input_params[1])*1e-10
     output .*= As
     return reshape(output, Int(length(output)/length(comp_emu.kgrid)), :)
@@ -33,9 +33,9 @@ end
 
 function get_component(input_params, comp_emu::PloopEmulator)
     input = deepcopy(input_params)
-    _maximin_input!(input, comp_emu.InMinMax)
-    output = Array(_run_emulator(input, comp_emu.TrainedEmulator))
-    _inv_maximin_output!(output, comp_emu.OutMinMax)
+    maximin_input!(input, comp_emu.InMinMax)
+    output = Array(run_emulator(input, comp_emu.TrainedEmulator))
+    inv_maximin_output!(output, comp_emu.OutMinMax)
     As = exp(input_params[1])*1e-10
     output .*= As^2
     return reshape(output, Int(length(output)/length(comp_emu.kgrid)), :)

From a483b2070501e93d949ac213817e4412ec63c68f Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Wed, 8 Nov 2023 05:49:48 +0100
Subject: [PATCH 03/29] Added window convolution

---
 src/projection.jl | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/src/projection.jl b/src/projection.jl
index 843b4ea..3eafa37 100644
--- a/src/projection.jl
+++ b/src/projection.jl
@@ -141,3 +141,22 @@ function apply_AP(k_grid_AP, k_interp, Mono_array::Array, Quad_array::Array, Hex
     int_Mono, int_Quad, int_Hexa = interp_Pℓs(Mono_array, Quad_array, Hexa_array, k_interp)
     return apply_AP(k_grid_AP, int_Mono, int_Quad, int_Hexa, q_par, q_perp)
 end
+
+function window_convolution(W,v)
+    a,b,c,d = size(W)
+    result = zeros(a,d)
+    window_convolution!(result,W,v)
+    return result
+end
+
+function window_convolution!(result, W, v)
+    @tturbo for i ∈ axes(W,1), l ∈ axes(W,4)
+        c = zero(eltype(v))
+        for k ∈ axes(W,3)
+            for j ∈ axes(W,2)
+                c += W[i,j,k,l] * v[j,k]
+            end
+        end
+        result[i,l] = c
+    end
+end

From 1b6126417d4ef224f908a2f82c20d56316159f01 Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Thu, 9 Nov 2023 05:00:16 +0100
Subject: [PATCH 04/29] Fixing typo in agn diam distane

---
 src/Effort.jl     |  2 ++
 src/background.jl | 10 +++++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/Effort.jl b/src/Effort.jl
index ce89c40..1a0cb5c 100644
--- a/src/Effort.jl
+++ b/src/Effort.jl
@@ -11,6 +11,8 @@ using Memoization
 using OrdinaryDiffEq
 using QuadGK
 
+const c_0 = 2.99792458e5
+
 include("background.jl")
 include("neural_networks.jl")
 include("eft_commands.jl")
diff --git a/src/background.jl b/src/background.jl
index c15338e..e44f49b 100644
--- a/src/background.jl
+++ b/src/background.jl
@@ -15,8 +15,16 @@ function _r̃_z(z, ΩM, w0, wa)
     return integral
 end
 
+function _r_z(z, H0, ΩM, w0, wa)
+    return c_0 * _r̃_z(z, ΩM, w0, wa) / H0
+end
+
 function _d̃A_z(z, ΩM, w0, wa)
-    return (1+z) * _r̃_z(z, ΩM, w0, wa)
+    return _r̃_z(z, ΩM, w0, wa) / (1+z)
+end
+
+function _dA_z(z, H0, ΩM, w0, wa)
+    return _r_z(z, H0, ΩM, w0, wa) / (1+z)
 end
 
 function _ρDE_z(z, w0, wa)

From b8ded333f79414fa0fb5f5161f5dcd5d734e2285 Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Thu, 9 Nov 2023 17:40:57 +0100
Subject: [PATCH 05/29] Fixed typo in default argument

---
 src/eft_commands.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/eft_commands.jl b/src/eft_commands.jl
index bf07dfe..c87da31 100644
--- a/src/eft_commands.jl
+++ b/src/eft_commands.jl
@@ -49,14 +49,14 @@ function bias_multiplication!(input_array, bias_array, Pk_input)
     end
 end
 
-function get_stoch(cϵ0, cϵ1, cϵ2, n_bar, k_grid::Array, k_nl=0.7)
+function get_stoch(cϵ0, cϵ1, cϵ2, n_bar, k_grid::Array; k_nl=0.7)
     P_stoch = zeros(3, length(k_grid))
     P_stoch[1,:] = @. 1/n_bar*(cϵ0 + cϵ1*(k_grid/k_nl)^2)
     P_stoch[2,:] = @. 1/n_bar*(cϵ2 * (k_grid/k_nl)^2)
     return P_stoch
 end
 
-function get_stoch_terms(cϵ0, cϵ1, cϵ2, n_bar, k_grid::Array, k_nl=0.7)
+function get_stoch_terms(cϵ0, cϵ1, cϵ2, n_bar, k_grid::Array; k_nl=0.7)
     P_stoch_0 = @. 1/n_bar*(cϵ0 + cϵ1*(k_grid/k_nl)^2)
     P_stoch_2 = @. 1/n_bar*(cϵ2 * (k_grid/k_nl)^2)
     return P_stoch_0, P_stoch_2
@@ -64,7 +64,7 @@ end
 
 function get_stoch_terms_binned_efficient(cϵ0, cϵ1, cϵ2, n_bar, k_edges::Array; k_nl=0.7)
     n_bins = length(k_edges)-1
-    P_stoch_0_c, _ = get_stoch_terms(cϵ0, cϵ1, cϵ2, n_bar, k_edges, k_nl)
+    P_stoch_0_c, _ = get_stoch_terms(cϵ0, cϵ1, cϵ2, n_bar, k_edges; k_nl = k_nl)
     mytype = typeof(P_stoch_0_c[1])
     P_stoch_0 = zeros(mytype, n_bins)
     P_stoch_2 = zeros(mytype, n_bins)
@@ -78,7 +78,7 @@ end
 
 function get_stoch_terms_binned(cϵ0, cϵ1, cϵ2, n_bar, k_edges::Array; k_nl=0.7, nk = 30)
     n_bins = length(k_edges)-1
-    P_stoch_0_c, _ = get_stoch_terms(cϵ0, cϵ1, cϵ2, n_bar, k_edges, k_nl)
+    P_stoch_0_c, _ = get_stoch_terms(cϵ0, cϵ1, cϵ2, n_bar, k_edges; k_nl = k_nl)
     mytype = typeof(P_stoch_0_c[1])
     P_stoch_0 = zeros(mytype, n_bins)
     P_stoch_2 = zeros(mytype, n_bins)

From f1491a1e9add78ad8e54b37109402e6d13314429 Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Thu, 9 Nov 2023 22:16:10 +0100
Subject: [PATCH 06/29] Updating AP evaluation

---
 src/projection.jl | 96 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 94 insertions(+), 2 deletions(-)

diff --git a/src/projection.jl b/src/projection.jl
index 3eafa37..9ae0d05 100644
--- a/src/projection.jl
+++ b/src/projection.jl
@@ -3,7 +3,7 @@ function _Pkμ(k, μ, Int_Mono, Int_Quad, Int_Hexa)
 end
 
 function _k_true(k_o, μ_o, q_perp, F)
-    return k_o/q_perp*sqrt(1+μ_o^2*(1/F^2-1))
+    return @. k_o/q_perp*sqrt(1+μ_o^2*(1/F^2-1))
 end
 
 function _μ_true(μ_o, F)
@@ -58,7 +58,15 @@ end
 
 function _mygemmavx(A, B, C)
     Dm = zero(eltype(C))
-    @turbo for n ∈ axes(A,1)
+    @tturbo for n ∈ axes(A,1)
+        Dm += A[n] * B[n] * C[n]
+    end
+    return Dm
+end
+
+function _mygemm(A, B, C)
+    Dm = zero(eltype(C))
+    for n ∈ axes(A,1)
         Dm += A[n] * B[n] * C[n]
     end
     return Dm
@@ -95,6 +103,80 @@ function apply_AP(k_grid, int_Mono::CubicSpline, int_Quad::CubicSpline, int_Hexa
     return result
 end
 
+function _stoch_obs(k_o, μ_o, q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2, k_nl)
+    F = q_par/q_perp
+    k_t = _k_true(k_o, μ_o, q_perp, F)
+    μ_t = _μ_true(μ_o, F)
+    return _stoch_kμ(k_t, μ_t, n_bar, cϵ0, cϵ1, cϵ2, k_nl)/(q_par*q_perp^2)
+end
+
+function _k_grid_over_nl(k_grid, k_nl)
+    return @. (k_grid/k_nl)^2
+ end
+
+ function _stoch_kμ(k_grid, μ, n_bar, cϵ0, cϵ1, cϵ2, k_nl)
+    return (cϵ0 * Pl(μ, 0) .+ Effort._k_grid_over_nl(k_grid, k_nl) .* (cϵ1 * Pl(μ, 0) +
+                        cϵ2 * Pl(μ, 2)) ) ./ n_bar
+end
+
+function get_stochs_AP(k_grid, q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2; k_nl = 0.7)
+    nk = length(k_grid)
+    n_GL_points = 5
+    #TODO: check that the extrapolation does not create problems. Maybe logextrap?
+    nodes, weights = @memoize gausslobatto(n_GL_points*2)
+    #since the integrand is symmetric, we are gonna use only half of the points
+    μ_nodes = nodes[1:n_GL_points]
+    μ_weights = weights[1:n_GL_points]
+    result = zeros(2, nk)
+
+    Pl_0 = Pl.(μ_nodes, 0)
+    Pl_2 = Pl.(μ_nodes, 2)
+
+    temp = zeros(n_GL_points)
+
+    for (k_idx, myk) in enumerate(k_grid)
+        for j in 1:n_GL_points
+            temp[j] = _stoch_obs(myk, μ_nodes[j], q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2, k_nl)
+        end
+        #we do not divided by 2 since we are using only half of the points and the result
+        #should be multiplied by 2
+        result[1, k_idx] = (2*0+1)*_mygemm(μ_weights, temp, Pl_0)
+        result[2, k_idx] = (2*2+1)*_mygemm(μ_weights, temp, Pl_2)
+    end
+    return result
+end
+
+function get_stochs_AP_new(k_grid, q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2; k_nl = 0.7, n_GL_points = 4)
+    nk = length(k_grid)
+    #TODO: check that the extrapolation does not create problems. Maybe logextrap?
+    nodes, weights = @memoize gausslobatto(n_GL_points*2)
+    #since the integrand is symmetric, we are gonna use only half of the points
+    μ_nodes = nodes[1:n_GL_points]
+    μ_weights = weights[1:n_GL_points]
+    result = zeros(2, nk)
+
+    Pl_array = zeros(2, n_GL_points)
+    Pl_array[1,:] .= Pl.(μ_nodes, 0)
+    Pl_array[2,:] .= Pl.(μ_nodes, 2)
+
+    temp = zeros(n_GL_points, nk)
+
+    for i in 1:n_GL_points
+        temp[i,:] = _stoch_obs(k_grid, μ_nodes[i], q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2,
+                               k_nl)
+    end
+
+    @turbo for i in 1:2
+        for j in 1:nk
+            for k in 1:n_GL_points
+                result[i, j] += μ_weights[k] * temp[k,j] * Pl_array[i, k]
+            end
+        end
+    end
+
+    return result
+end
+
 function q_perp_par(z, ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
     E_ref  = _E_z(z, ΩM_ref, w0_ref, wa_ref)
     E_true = _E_z(z, ΩM_true, w0_true, wa_true)
@@ -107,6 +189,16 @@ function q_perp_par(z, ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
     return q_perp, q_par
 end
 
+function q_perp_par(z, ΩM_ref, w0_ref, wa_ref, E_true, d̃A_true)
+    E_ref  = _E_z(z, ΩM_ref, w0_ref, wa_ref)
+
+    d̃A_ref  = _d̃A_z(z, ΩM_ref, w0_ref, wa_ref)
+
+    q_perp = E_ref / E_true
+    q_par  = d̃A_true / d̃A_ref
+    return q_perp, q_par
+end
+
 function apply_AP(k_grid, Mono_array::Array, Quad_array::Array, Hexa_array::Array, z, ΩM_ref,
     w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
 

From 7c4ab0b528803ddc90d9312c06f7ea70cfdbe4e9 Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Thu, 9 Nov 2023 22:28:25 +0100
Subject: [PATCH 07/29] Renaming things

---
 src/projection.jl | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/projection.jl b/src/projection.jl
index 9ae0d05..885093a 100644
--- a/src/projection.jl
+++ b/src/projection.jl
@@ -177,7 +177,7 @@ function get_stochs_AP_new(k_grid, q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2; k_nl
     return result
 end
 
-function q_perp_par(z, ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
+function q_par_perp(z, ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
     E_ref  = _E_z(z, ΩM_ref, w0_ref, wa_ref)
     E_true = _E_z(z, ΩM_true, w0_true, wa_true)
 
@@ -186,10 +186,10 @@ function q_perp_par(z, ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
 
     q_perp = E_ref/E_true
     q_par  = d̃A_true/d̃A_ref
-    return q_perp, q_par
+    return q_par, q_perp
 end
 
-function q_perp_par(z, ΩM_ref, w0_ref, wa_ref, E_true, d̃A_true)
+function q_par_perp(z, ΩM_ref, w0_ref, wa_ref, E_true, d̃A_true)
     E_ref  = _E_z(z, ΩM_ref, w0_ref, wa_ref)
 
     d̃A_ref  = _d̃A_z(z, ΩM_ref, w0_ref, wa_ref)
@@ -202,7 +202,7 @@ end
 function apply_AP(k_grid, Mono_array::Array, Quad_array::Array, Hexa_array::Array, z, ΩM_ref,
     w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
 
-    q_perp, q_par  = q_perp_par(z, ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
+    q_par, q_perp  = q_par_perp(z, ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
     int_Mono, int_Quad, int_Hexa = interp_Pℓs(Mono_array, Quad_array, Hexa_array, k_grid)
     return apply_AP(k_grid, int_Mono, int_Quad, int_Hexa, q_par, q_perp)
 end
@@ -210,7 +210,7 @@ end
 function apply_AP(k_grid_AP, k_grid, Mono_array::Array, Quad_array::Array, Hexa_array::Array, z,
     ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
 
-    q_perp, q_par  = q_perp_par(z, ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
+    q_par, q_perp  = q_par_perp(z, ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
     int_Mono, int_Quad, int_Hexa = interp_Pℓs(Mono_array, Quad_array, Hexa_array, k_grid)
 
     return apply_AP(k_grid_AP, int_Mono, int_Quad, int_Hexa, q_par, q_perp)

From 7446aef0e448e5b54cbac7b967b0f4aacf08655e Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Fri, 10 Nov 2023 20:47:23 +0100
Subject: [PATCH 08/29] Bug fixing in AP stoch

---
 src/projection.jl | 48 ++++++++++++-----------------------------------
 1 file changed, 12 insertions(+), 36 deletions(-)

diff --git a/src/projection.jl b/src/projection.jl
index 885093a..a5cb70b 100644
--- a/src/projection.jl
+++ b/src/projection.jl
@@ -18,9 +18,9 @@ function _P_obs(k_o, μ_o, q_par, q_perp, Int_Mono, Int_Quad, Int_Hexa)
 end
 
 function interp_Pℓs(Mono_array, Quad_array, Hexa_array, k_grid)
-    Int_Mono = CubicSpline(Mono_array, k_grid)
-    Int_Quad = CubicSpline(Quad_array, k_grid)
-    Int_Hexa = CubicSpline(Hexa_array, k_grid)
+    Int_Mono = CubicSpline(Mono_array, k_grid; extrapolate = true)
+    Int_Quad = CubicSpline(Quad_array, k_grid; extrapolate = true)
+    Int_Hexa = CubicSpline(Hexa_array, k_grid; extrapolate = true)
     return Int_Mono, Int_Quad, Int_Hexa
 end
 
@@ -116,37 +116,10 @@ function _k_grid_over_nl(k_grid, k_nl)
 
  function _stoch_kμ(k_grid, μ, n_bar, cϵ0, cϵ1, cϵ2, k_nl)
     return (cϵ0 * Pl(μ, 0) .+ Effort._k_grid_over_nl(k_grid, k_nl) .* (cϵ1 * Pl(μ, 0) +
-                        cϵ2 * Pl(μ, 2)) ) ./ n_bar
+            cϵ2 * Pl(μ, 2)) ) ./ n_bar
 end
 
-function get_stochs_AP(k_grid, q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2; k_nl = 0.7)
-    nk = length(k_grid)
-    n_GL_points = 5
-    #TODO: check that the extrapolation does not create problems. Maybe logextrap?
-    nodes, weights = @memoize gausslobatto(n_GL_points*2)
-    #since the integrand is symmetric, we are gonna use only half of the points
-    μ_nodes = nodes[1:n_GL_points]
-    μ_weights = weights[1:n_GL_points]
-    result = zeros(2, nk)
-
-    Pl_0 = Pl.(μ_nodes, 0)
-    Pl_2 = Pl.(μ_nodes, 2)
-
-    temp = zeros(n_GL_points)
-
-    for (k_idx, myk) in enumerate(k_grid)
-        for j in 1:n_GL_points
-            temp[j] = _stoch_obs(myk, μ_nodes[j], q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2, k_nl)
-        end
-        #we do not divided by 2 since we are using only half of the points and the result
-        #should be multiplied by 2
-        result[1, k_idx] = (2*0+1)*_mygemm(μ_weights, temp, Pl_0)
-        result[2, k_idx] = (2*2+1)*_mygemm(μ_weights, temp, Pl_2)
-    end
-    return result
-end
-
-function get_stochs_AP_new(k_grid, q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2; k_nl = 0.7, n_GL_points = 4)
+function get_stochs_AP(k_grid, q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2; k_nl = 0.7, n_GL_points = 4)
     nk = length(k_grid)
     #TODO: check that the extrapolation does not create problems. Maybe logextrap?
     nodes, weights = @memoize gausslobatto(n_GL_points*2)
@@ -166,10 +139,13 @@ function get_stochs_AP_new(k_grid, q_par, q_perp, n_bar, cϵ0, cϵ1, cϵ2; k_nl
                                k_nl)
     end
 
+    multipole_weight = [2*0+1, 2*2+1]
+
     @turbo for i in 1:2
         for j in 1:nk
             for k in 1:n_GL_points
-                result[i, j] += μ_weights[k] * temp[k,j] * Pl_array[i, k]
+                result[i, j] += μ_weights[k] * temp[k,j] * Pl_array[i, k] *
+                                multipole_weight[i]
             end
         end
     end
@@ -236,15 +212,15 @@ end
 
 function window_convolution(W,v)
     a,b,c,d = size(W)
-    result = zeros(a,d)
+    result = zeros(a,c)
     window_convolution!(result,W,v)
     return result
 end
 
 function window_convolution!(result, W, v)
-    @tturbo for i ∈ axes(W,1), l ∈ axes(W,4)
+    @tturbo for i ∈ axes(W,1), l ∈ axes(W,3)
         c = zero(eltype(v))
-        for k ∈ axes(W,3)
+        for k ∈ axes(W,4)
             for j ∈ axes(W,2)
                 c += W[i,j,k,l] * v[j,k]
             end

From 48cd1a47336d573392f07007a91b21b62d397f79 Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Fri, 10 Nov 2023 22:43:01 +0100
Subject: [PATCH 09/29] Fixed bug in window

---
 src/projection.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/projection.jl b/src/projection.jl
index a5cb70b..d3e548e 100644
--- a/src/projection.jl
+++ b/src/projection.jl
@@ -218,13 +218,13 @@ function window_convolution(W,v)
 end
 
 function window_convolution!(result, W, v)
-    @tturbo for i ∈ axes(W,1), l ∈ axes(W,3)
+    @turbo for i ∈ axes(W,1), k ∈ axes(W,3)
         c = zero(eltype(v))
-        for k ∈ axes(W,4)
+        for l ∈ axes(W,4)
             for j ∈ axes(W,2)
-                c += W[i,j,k,l] * v[j,k]
+                c += W[i,j,k,l] * v[j,l]
             end
         end
-        result[i,l] = c
+        result[i,k] = c
     end
 end

From af9575cb1b46cb92c0737603840b62e54e7dcac0 Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Sat, 11 Nov 2023 05:10:56 +0100
Subject: [PATCH 10/29] ixed typo in q_perp qpar

---
 src/projection.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/projection.jl b/src/projection.jl
index d3e548e..67b5ae3 100644
--- a/src/projection.jl
+++ b/src/projection.jl
@@ -160,8 +160,8 @@ function q_par_perp(z, ΩM_ref, w0_ref, wa_ref, ΩM_true, w0_true, wa_true)
     d̃A_ref  = _d̃A_z(z, ΩM_ref, w0_ref, wa_ref)
     d̃A_true = _d̃A_z(z, ΩM_true, w0_true, wa_true)
 
-    q_perp = E_ref/E_true
-    q_par  = d̃A_true/d̃A_ref
+    q_perp = E_true/E_ref
+    q_par  = d̃A_ref/d̃A_true
     return q_par, q_perp
 end
 
@@ -170,8 +170,8 @@ function q_par_perp(z, ΩM_ref, w0_ref, wa_ref, E_true, d̃A_true)
 
     d̃A_ref  = _d̃A_z(z, ΩM_ref, w0_ref, wa_ref)
 
-    q_perp = E_ref / E_true
-    q_par  = d̃A_true / d̃A_ref
+    q_perp = E_true/E_ref
+    q_par  = d̃A_ref/d̃A_true
     return q_perp, q_par
 end
 

From c08342f48c68e99a220e7cf8813553e69f25ff52 Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Sun, 12 Nov 2023 00:01:07 +0100
Subject: [PATCH 11/29] From cubic to quadratic spline

---
 src/projection.jl | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/projection.jl b/src/projection.jl
index 67b5ae3..a6ceca3 100644
--- a/src/projection.jl
+++ b/src/projection.jl
@@ -18,9 +18,11 @@ function _P_obs(k_o, μ_o, q_par, q_perp, Int_Mono, Int_Quad, Int_Hexa)
 end
 
 function interp_Pℓs(Mono_array, Quad_array, Hexa_array, k_grid)
-    Int_Mono = CubicSpline(Mono_array, k_grid; extrapolate = true)
-    Int_Quad = CubicSpline(Quad_array, k_grid; extrapolate = true)
-    Int_Hexa = CubicSpline(Hexa_array, k_grid; extrapolate = true)
+    #extrapolation might introduce some errors ar high k, when q << 1.
+    #maybe we should implement a log extrapolation?
+    Int_Mono = QuadraticSpline(Mono_array, k_grid; extrapolate = true)
+    Int_Quad = QuadraticSpline(Quad_array, k_grid; extrapolate = true)
+    Int_Hexa = QuadraticSpline(Hexa_array, k_grid; extrapolate = true)
     return Int_Mono, Int_Quad, Int_Hexa
 end
 
@@ -29,8 +31,8 @@ function k_projection(k_projection, Mono_array, Quad_array, Hexa_array, k_grid)
     return int_Mono.(k_projection), int_Quad.(k_projection), int_Hexa.(k_projection)
 end
 
-function apply_AP_check(k_grid, int_Mono::CubicSpline, int_Quad::CubicSpline,
-    int_Hexa::CubicSpline, q_par, q_perp)
+function apply_AP_check(k_grid, int_Mono::QuadraticSpline, int_Quad::QuadraticSpline,
+    int_Hexa::QuadraticSpline, q_par, q_perp)
     nk = length(k_grid)
     result = zeros(3, nk)
     ℓ_array = [0,2,4]
@@ -72,7 +74,7 @@ function _mygemm(A, B, C)
     return Dm
 end
 
-function apply_AP(k_grid, int_Mono::CubicSpline, int_Quad::CubicSpline, int_Hexa::CubicSpline,
+function apply_AP(k_grid, int_Mono::QuadraticSpline, int_Quad::QuadraticSpline, int_Hexa::QuadraticSpline,
     q_par, q_perp)
     nk = length(k_grid)
     n_GL_points = 5

From 535db6a328d8e0e8a48395793b7e9aa83c1e7042 Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Tue, 28 May 2024 08:33:26 -0400
Subject: [PATCH 12/29] Added precise neutrino growth

---
 src/Effort.jl     |  11 +++
 src/background.jl | 169 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 174 insertions(+), 6 deletions(-)

diff --git a/src/Effort.jl b/src/Effort.jl
index 1a0cb5c..918821e 100644
--- a/src/Effort.jl
+++ b/src/Effort.jl
@@ -13,6 +13,17 @@ using QuadGK
 
 const c_0 = 2.99792458e5
 
+function __init__()
+    min_y = get_y(0,0) #obvious, I know
+    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]
+    global F_interpolant = AkimaInterpolation(F_grid, y_grid)
+    y_grid = vcat(LinRange(min_y, 10., 10000), LinRange(10.1, max_y, 10000))
+    dFdy_grid = [_dFdy(y) for y in y_grid]
+    global dFdy_interpolant = AkimaInterpolation(dFdy_grid, y_grid)
+end
+
 include("background.jl")
 include("neural_networks.jl")
 include("eft_commands.jl")
diff --git a/src/background.jl b/src/background.jl
index e44f49b..74e491c 100644
--- a/src/background.jl
+++ b/src/background.jl
@@ -1,13 +1,94 @@
-#TODO: this part shoul be moved to a dedicate package. While necessary to a full Effort
+#TODO: this part should be moved to a dedicate package. While necessary to a full Effort
 #functionality, this could be factorized to a new module, specifically taylored to this goal
 # and maybe used in other packages
 
-function _E_z(z, ΩM, w0, wa)
+function _F(y)
+    result, _ = quadgk(x -> x^2*√(x^2+y^2)/(1+exp(x)), 0, Inf, rtol=1e-12)
+    return result
+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
+end
+
+function _ΩνE2(a, Ωγ0, mν; kB = 8.617342e-5, Tν = 0.71611*2.7255, Neff = 3.044)
+    Γν = (4/11)^(1/3)*(Neff/3)^(1/4)#0.71649^4
+    return 15/π^4*Γν^4*Ωγ0/a^4*F_interpolant(get_y(mν, a))
+end
+
+function _ΩνE2(a, Ωγ0, mν::Array; kB = 8.617342e-5, Tν = 0.71611*2.7255, Neff = 3.044)
+    Γν = (4/11)^(1/3)*(Neff/3)^(1/4)#0.71649^4
+    sum_interpolant = 0.
+    for mymν in mν
+        sum_interpolant += F_interpolant(get_y(mymν, a))
+    end
+    return 15/π^4*Γν^4*Ωγ0/a^4*sum_interpolant
+end
+
+function _dΩνE2da(a, Ωγ0, mν; kB = 8.617342e-5, Tν = 0.71611*2.7255, Neff = 3.044)
+    Γν = (4/11)^(1/3)*(Neff/3)^(1/4)#0.71649^4
+    return 15/π^4*Γν^4*Ωγ0*(-4*F_interpolant(get_y(mν, a))/a^5+
+           dFdy_interpolant(get_y(mν, a))/a^4*(mν/kB/Tν))
+end
+
+function _dΩνE2da(a, Ωγ0, mν::Array; kB = 8.617342e-5, Tν = 0.71611*2.7255, Neff = 3.044)
+    Γν = (4/11)^(1/3)*(Neff/3)^(1/4)#0.71649^4
+    sum_interpolant = 0.
+    for mymν in mν
+        sum_interpolant += -4*F_interpolant(get_y(mymν, a))/a^5+
+        dFdy_interpolant(get_y(mymν, a))/a^4*(mymν/kB/Tν)
+    end
+    return 15/π^4*Γν^4*Ωγ0*sum_interpolant
+end
+
+"""function _E_z(z, ΩM, w0, wa)
     return sqrt(ΩM*(1+z)^3+(1-ΩM)*_ρDE_z(z, w0, wa))
 end
+"""
+function _E_a(a, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    Ωγ0 = 2.469e-5/h^2
+    Ων0 = _ΩνE2(1., Ωγ0, mν)
+    ΩΛ0 = 1. - (Ωγ0 + Ωc0 + Ωb0 + Ων0)
+    return sqrt(Ωγ0*a^-4 + (Ωc0 + Ωb0)*a^-3 + ΩΛ0 * _ρDE_a(a, w0, wa)+ _ΩνE2(a, Ωγ0, mν))
+end
 
-function _H_z(z, H0, ΩM, w0, wa)
+"""function _H_z(z, H0, ΩM, w0, wa)
     return H0*_E_z(z, ΩM, w0, wa)
+end"""
+
+_H_a(a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa) = 100*h*_E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+
+function _χ_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    integral, _ = quadgk(x -> 1 /
+    _E_a(_a_z(x), Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa), 0, z, rtol=1e-6)
+    return integral*c_0/(100*h)
+end
+
+function _dEda(a, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    Ωγ0 = 2.469e-5/h^2
+    Ων0 = _ΩνE2(1., Ωγ0, mν)
+    ΩΛ0 = 1. - (Ωγ0 + Ωc0 + Ωb0 + Ων0)
+    return 0.5/_E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)*(-3(Ωc0 + Ωb0)a^-4-4Ωγ0*a^-5+
+           ΩΛ0*_dρDEda(a, w0, wa)+_dΩνE2da(a, Ωγ0, mν))
+end
+
+function _dlogEdloga(a, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    Ωγ0 = 2.469e-5/h^2
+    Ων0 = _ΩνE2(1., Ωγ0, mν)
+    ΩΛ0 = 1. - (Ωγ0 + Ωc0 + Ωb0 + Ων0)
+    return a*0.5/(_E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)^2)*(-3(Ωc0 + Ωb0)a^-4-4Ωγ0*a^-
+           5+ΩΛ0*_dρDEda(a, w0, wa)+_dΩνE2da(a, Ωγ0, mν))
+end
+
+function _ΩMa(a, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    Ωγ0 = 2.469e-5/h^2
+    Ων0 = _ΩνE2(1., Ωγ0, mν)
+    return (Ωc0 + Ωb0 + Ων0 )*a^-3 / (_E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa))^2
 end
 
 function _r̃_z(z, ΩM, w0, wa)
@@ -31,6 +112,14 @@ function _ρDE_z(z, w0, wa)
     return (1+z)^(3.0 * (1.0 + w0 + wa)) * exp(-3.0 * wa * z /(1+z))
 end
 
+function _ρDE_a(a, w0, wa)
+    return a^(-3.0 * (1.0 + w0 + wa)) * exp(3.0 * wa * (a-1))
+end
+
+function _dρDEda(a, w0, wa)
+    return 3*(-(1+w0+wa)/a+wa)*_ρDE_a(a, w0, wa)
+end
+
 function _X_z(z, ΩM, w0, wa)
     return ΩM*((1+z)^3)/((1-ΩM)*_ρDE_z(z, w0, wa))
 end
@@ -39,7 +128,7 @@ function _w_z(z, w0, wa)
     return w0+wa*z/(1+z)
 end
 
-function _growth!(du,u,p,a)
+"""function _growth!(du,u,p,a)
     ΩM = p[1]
     w0 = p[2]
     wa = p[3]
@@ -48,13 +137,29 @@ function _growth!(du,u,p,a)
     dG = u[2]
     du[1] = dG
     du[2] = -(3.5-1.5*_w_z(z, w0, wa)/(1+_X_z(z, ΩM, w0, wa)))*dG/a-1.5*(1-_w_z(z, w0,wa))/(1+_X_z(z, ΩM, w0, wa))*G/(a^2)
+end"""
+
+function _growth!(du,u,p,loga)
+    #Ωγ0 = p[1]
+    Ωc0 = p[1]
+    Ωb0 = p[2]
+    mν  = p[3]
+    h   = p[4]
+    w0  = p[5]
+    wa  = p[6]
+    a = exp(loga)
+    D = u[1]
+    dD = u[2]
+    du[1] = dD
+    du[2] = -(2+_dlogEdloga(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa))*dD+
+            1.5*_ΩMa(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)*D
 end
 
 function _a_z(z)
     return 1/(1+z)
 end
 
-function growth_solver(ΩM, w0, wa)
+"""function growth_solver(ΩM, w0, wa)
     u₀ = [1.0,0.0]
 
     aspan = (0.99e-3, 1.01)
@@ -65,9 +170,24 @@ function growth_solver(ΩM, w0, wa)
 
     sol = solve(prob, Tsit5(), abstol=1e-6, reltol=1e-6;verbose=false)
     return sol
+end"""
+
+function growth_solver(Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    amin = 1/139
+    u₀ = [amin, amin]
+
+    logaspan = (log(amin), log(1.01))
+    Ωγ0 = 2.469e-5/h^2
+
+    p = (Ωc0, Ωb0 , mν, h, w0, wa)
+
+    prob = ODEProblem(_growth!, u₀, logaspan, p)
+
+    sol = solve(prob, OrdinaryDiffEq.Tsit5(), abstol=1e-6, reltol=1e-6; verbose=false)
+    return sol
 end
 
-function _D_z(z::Array, sol::SciMLBase.ODESolution)
+"""function _D_z(z::Array, sol::SciMLBase.ODESolution)
     [u for (u,t) in sol.(_a_z.(z))] .* _a_z.(z) ./ (sol(_a_z(0.))[1,:])
 end
 
@@ -103,4 +223,41 @@ end
 function _f_z(z, ΩM, w0, wa)
     sol = growth_solver(ΩM, w0, wa)
     return _f_z(z, sol)
+end"""
+
+function _D_z(z::Array, sol::SciMLBase.ODESolution)
+    [u for (u,t) in sol.(log.(_a_z.(z)))] ./ (sol(log(_a_z(0.)))[1,:])
+end
+
+function _D_z(z, sol::SciMLBase.ODESolution)
+    return (sol(log(_a_z(z)))[1,:]/sol(log(_a_z(0.)))[1,:])[1,1]
+end
+
+function _D_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    sol = growth_solver(Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+    return _D_z(z, sol)
+end
+
+function _D_z_unnorm(z::Array, sol::SciMLBase.ODESolution)
+    [u for (u,t) in sol.(log.(_a_z.(z)))]
+end
+
+function _D_z_unnorm(z, sol::SciMLBase.ODESolution)
+    return (sol(log(_a_z(z)))[1,:])[1,1]
+end
+
+function _f_a(a, sol::SciMLBase.ODESolution)
+    D, D_prime = sol(log(a))
+    return 1 / D * D_prime
+end
+
+function _f_z(z, sol::SciMLBase.ODESolution)
+    a = _a_z.(z)
+    return _f_a(a, sol)
+end
+
+function _f_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    a = _a_z.(z)
+    sol = growth_solver(Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+    return _f_a(a, sol)
 end

From dfac5d73df94d3b777d5869d75c10667b4b5fbcd Mon Sep 17 00:00:00 2001
From: Marco Bonici <bonici.marco@gmail.com>
Date: Tue, 28 May 2024 08:46:16 -0400
Subject: [PATCH 13/29] Small change in API

---
 src/background.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/background.jl b/src/background.jl
index 74e491c..a15b48a 100644
--- a/src/background.jl
+++ b/src/background.jl
@@ -256,7 +256,7 @@ function _f_z(z, sol::SciMLBase.ODESolution)
     return _f_a(a, sol)
 end
 
-function _f_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+function _f_z(z, Ωc0, Ωb0, h, mν =0 , w0=-1., wa=0.)
     a = _a_z.(z)
     sol = growth_solver(Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
     return _f_a(a, sol)

From c5738e6cc8b35656d401129ff1c332e83e2cca37 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Sun, 21 Jul 2024 18:26:13 -0400
Subject: [PATCH 14/29] new method

---
 src/background.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/background.jl b/src/background.jl
index a15b48a..3c5850a 100644
--- a/src/background.jl
+++ b/src/background.jl
@@ -238,6 +238,11 @@ function _D_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
     return _D_z(z, sol)
 end
 
+function _D_z_unnorm(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    sol = growth_solver(Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+    return _D_z_unnorm(z, sol)
+end
+
 function _D_z_unnorm(z::Array, sol::SciMLBase.ODESolution)
     [u for (u,t) in sol.(log.(_a_z.(z)))]
 end

From 60d4d591a8f8ef4bd25dbebf05f2d3c8d53173a9 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Thu, 29 Aug 2024 12:34:56 -0400
Subject: [PATCH 15/29] Adding CI

---
 .github/workflows/test.yml | 53 ++++++++++++++++++++++++++++++++++++++
 test/Project.toml          |  6 +++++
 test/runtests.jl           | 36 ++++++++++++++++++++++++++
 3 files changed, 95 insertions(+)
 create mode 100644 .github/workflows/test.yml
 create mode 100644 test/Project.toml
 create mode 100644 test/runtests.jl

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
new file mode 100644
index 0000000..3498722
--- /dev/null
+++ b/.github/workflows/test.yml
@@ -0,0 +1,53 @@
+name: CI
+on:
+  schedule:
+    - cron: 0 0 * * *
+  pull_request:
+    branches:
+      - main
+      - develop
+      - abstractcosmoemu
+  push:
+    branches:
+        - '**'
+    tags: '*'
+jobs:
+  test:
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.7'
+          - '1.8'
+          - '1.9'
+          - '1.10'
+          - '~1.11.0-0'
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@latest
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: actions/cache@v4
+        env:
+          cache-name: cache-artifacts
+        with:
+          path: ~/.julia/artifacts
+          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
+          restore-keys: |
+            ${{ runner.os }}-test-${{ env.cache-name }}-
+            ${{ runner.os }}-test-
+            ${{ 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 }}
+        with:
+          file: lcov.info
diff --git a/test/Project.toml b/test/Project.toml
new file mode 100644
index 0000000..abaf678
--- /dev/null
+++ b/test/Project.toml
@@ -0,0 +1,6 @@
+[deps]
+AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
+NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
+SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
+Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/runtests.jl b/test/runtests.jl
new file mode 100644
index 0000000..1753f08
--- /dev/null
+++ b/test/runtests.jl
@@ -0,0 +1,36 @@
+using Test
+using NPZ
+using SimpleChains
+using Static
+using Effort
+
+mlpd = SimpleChain(
+  static(6),
+  TurboDense(tanh, 64),
+  TurboDense(tanh, 64),
+  TurboDense(tanh, 64),
+  TurboDense(tanh, 64),
+  TurboDense(tanh, 64),
+  TurboDense(identity, 40)
+)
+
+k_test = Array(LinRange(0,200, 40))
+weights = SimpleChains.init_params(mlpd)
+inminmax = rand(6,2)
+outminmax = rand(40,2)
+npzwrite("emu/l.npy", ℓ_test)
+npzwrite("emu/weights.npy", weights)
+npzwrite("emu/inminmax.npy", inminmax)
+npzwrite("emu/outminmax.npy", outminmax)
+emu = Effort.SimpleChainsEmulator(Architecture = mlpd, Weights = weights)
+
+effort = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax = inminmax,
+                                OutMinMax = outminmax)
+
+@testset "Effort tests" begin
+    cosmo = ones(6)
+    cosmo_vec = ones(6,6)
+    output = Effort.get_component(cosmo,  capse_emu)
+    output_vec = Effort.get_component(cosmo_vec, capse_emu)
+    @test isapprox(output_vec[:,1], output)
+end

From 1e63a88ba3aa741c7e0900a71245e81c7971c40a Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Thu, 29 Aug 2024 12:40:54 -0400
Subject: [PATCH 16/29] Fixing CI

---
 test/runtests.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 1753f08..146a131 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -18,7 +18,7 @@ k_test = Array(LinRange(0,200, 40))
 weights = SimpleChains.init_params(mlpd)
 inminmax = rand(6,2)
 outminmax = rand(40,2)
-npzwrite("emu/l.npy", ℓ_test)
+npzwrite("emu/l.npy", k_test)
 npzwrite("emu/weights.npy", weights)
 npzwrite("emu/inminmax.npy", inminmax)
 npzwrite("emu/outminmax.npy", outminmax)

From d759fb24cce191316922636bfa0f720f212ae3d8 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Thu, 29 Aug 2024 12:46:12 -0400
Subject: [PATCH 17/29] Fixing CI

---
 test/runtests.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 146a131..128358a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -18,10 +18,10 @@ k_test = Array(LinRange(0,200, 40))
 weights = SimpleChains.init_params(mlpd)
 inminmax = rand(6,2)
 outminmax = rand(40,2)
-npzwrite("emu/l.npy", k_test)
-npzwrite("emu/weights.npy", weights)
-npzwrite("emu/inminmax.npy", inminmax)
-npzwrite("emu/outminmax.npy", outminmax)
+#npzwrite("emu/l.npy", k_test)
+#npzwrite("emu/weights.npy", weights)
+#npzwrite("emu/inminmax.npy", inminmax)
+#npzwrite("emu/outminmax.npy", outminmax)
 emu = Effort.SimpleChainsEmulator(Architecture = mlpd, Weights = weights)
 
 effort = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax = inminmax,

From 410f7a8a511ad35241348c8508c6f93326d12dbb Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Thu, 29 Aug 2024 12:51:00 -0400
Subject: [PATCH 18/29] Fixing CI

---
 test/runtests.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 128358a..7df2b07 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -24,13 +24,13 @@ outminmax = rand(40,2)
 #npzwrite("emu/outminmax.npy", outminmax)
 emu = Effort.SimpleChainsEmulator(Architecture = mlpd, Weights = weights)
 
-effort = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax = inminmax,
+effort_emu = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax = inminmax,
                                 OutMinMax = outminmax)
 
 @testset "Effort tests" begin
     cosmo = ones(6)
     cosmo_vec = ones(6,6)
-    output = Effort.get_component(cosmo,  capse_emu)
-    output_vec = Effort.get_component(cosmo_vec, capse_emu)
+    output = Effort.get_component(cosmo,  effort_emu)
+    output_vec = Effort.get_component(cosmo_vec, effort_emu)
     @test isapprox(output_vec[:,1], output)
 end

From aca5b0902fd193f82f8dbcccc564054f6de92848 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Thu, 29 Aug 2024 13:43:47 -0400
Subject: [PATCH 19/29] Updating CI

---
 test/runtests.jl | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 7df2b07..dd02706 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -18,19 +18,12 @@ k_test = Array(LinRange(0,200, 40))
 weights = SimpleChains.init_params(mlpd)
 inminmax = rand(6,2)
 outminmax = rand(40,2)
-#npzwrite("emu/l.npy", k_test)
-#npzwrite("emu/weights.npy", weights)
-#npzwrite("emu/inminmax.npy", inminmax)
-#npzwrite("emu/outminmax.npy", outminmax)
+a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa = [1., 1e-5, 0.25, 0.05, 0., 0.67, -1, 0.]
 emu = Effort.SimpleChainsEmulator(Architecture = mlpd, Weights = weights)
 
 effort_emu = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax = inminmax,
                                 OutMinMax = outminmax)
 
 @testset "Effort tests" begin
-    cosmo = ones(6)
-    cosmo_vec = ones(6,6)
-    output = Effort.get_component(cosmo,  effort_emu)
-    output_vec = Effort.get_component(cosmo_vec, effort_emu)
-    @test isapprox(output_vec[:,1], output)
+    @test isapprox(Effort._H_a(a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa), h*100)
 end

From 94d9ae8e094433ceec7d704904c752c995600aa5 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Thu, 29 Aug 2024 14:30:03 -0400
Subject: [PATCH 20/29] Updating readme; removed old lines

---
 README.md         | 13 +++++----
 src/background.jl | 69 +++--------------------------------------------
 test/runtests.jl  |  4 ++-
 3 files changed, 13 insertions(+), 73 deletions(-)

diff --git a/README.md b/README.md
index ccfeca0..68a2e65 100644
--- a/README.md
+++ b/README.md
@@ -1,13 +1,12 @@
 # Effort.jl
-[![](https://img.shields.io/badge/docs-stable-blue.svg)](https://cosmologicalemulators.github.io/Effort.jl/stable)
-[![](https://img.shields.io/badge/docs-dev-blue.svg)](https://cosmologicalemulators.github.io/Effort.jl/dev)
-![size](https://img.shields.io/github/repo-size/CosmologicalEmulators/Effort.jl)
-[![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle)
-[![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor's%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
+
+| **Documentation** | **Build Status** | **Code style** |
+|:--------:|:----------------:|:----------------:|
+| [![](https://img.shields.io/badge/docs-dev-blue.svg)](https://cosmologicalemulators.github.io/Effort.jl/dev) [![](https://img.shields.io/badge/docs-stable-blue.svg)](https://cosmologicalemulators.github.io/Effort.jl/stable) | [![Build status (Github Actions)](https://github.com/CosmologicalEmulators/Effort.jl/workflows/CI/badge.svg)](https://github.com/CosmologicalEmulators/Effort.jl/actions) [![codecov](https://codecov.io/gh/CosmologicalEmulators/Effort.jl/branch/abstractcosmoemu/graph/badge.svg?token=0PYHCWVL67)](https://codecov.io/gh/CosmologicalEmulators/Effort.jl) | [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle) [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor's%20Guide-blueviolet)](https://github.com/SciML/ColPrac) |
 
 Repository containing the EFfective Field theORy surrogaTe.
 
 ### Authors
 
-- Marco Bonici, INAF - Institute of Space Astrophysics and Cosmic Physics (IASF), Milano
-- Guido D'Amico, Università Degli Studi di Parma
+- Marco Bonici, PostDoctoral Researcher at Waterloo Centre for Astrophysics
+- Guido D'Amico, Associate Professor at Università Degli Studi di Parma
diff --git a/src/background.jl b/src/background.jl
index 3c5850a..3ea8ef0 100644
--- a/src/background.jl
+++ b/src/background.jl
@@ -57,9 +57,10 @@ function _E_a(a, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
     return sqrt(Ωγ0*a^-4 + (Ωc0 + Ωb0)*a^-3 + ΩΛ0 * _ρDE_a(a, w0, wa)+ _ΩνE2(a, Ωγ0, mν))
 end
 
-"""function _H_z(z, H0, ΩM, w0, wa)
-    return H0*_E_z(z, ΩM, w0, wa)
-end"""
+function _E_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    a = _a_z.(z)
+    return _E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+end
 
 _H_a(a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa) = 100*h*_E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
 
@@ -128,17 +129,6 @@ function _w_z(z, w0, wa)
     return w0+wa*z/(1+z)
 end
 
-"""function _growth!(du,u,p,a)
-    ΩM = p[1]
-    w0 = p[2]
-    wa = p[3]
-    z = 1.0 / a - 1.0
-    G = u[1]
-    dG = u[2]
-    du[1] = dG
-    du[2] = -(3.5-1.5*_w_z(z, w0, wa)/(1+_X_z(z, ΩM, w0, wa)))*dG/a-1.5*(1-_w_z(z, w0,wa))/(1+_X_z(z, ΩM, w0, wa))*G/(a^2)
-end"""
-
 function _growth!(du,u,p,loga)
     #Ωγ0 = p[1]
     Ωc0 = p[1]
@@ -159,19 +149,6 @@ function _a_z(z)
     return 1/(1+z)
 end
 
-"""function growth_solver(ΩM, w0, wa)
-    u₀ = [1.0,0.0]
-
-    aspan = (0.99e-3, 1.01)
-
-    p = [ΩM, w0, wa]
-
-    prob = ODEProblem(_growth!, u₀, aspan, p)
-
-    sol = solve(prob, Tsit5(), abstol=1e-6, reltol=1e-6;verbose=false)
-    return sol
-end"""
-
 function growth_solver(Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
     amin = 1/139
     u₀ = [amin, amin]
@@ -187,44 +164,6 @@ function growth_solver(Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
     return sol
 end
 
-"""function _D_z(z::Array, sol::SciMLBase.ODESolution)
-    [u for (u,t) in sol.(_a_z.(z))] .* _a_z.(z) ./ (sol(_a_z(0.))[1,:])
-end
-
-function _D_z(z, sol::SciMLBase.ODESolution)
-    return (Effort._a_z(z) .* sol(Effort._a_z(z))[1,:]/sol(Effort._a_z(0.))[1,:])[1,1]
-end
-
-function _D_z(z, ΩM, w0, wa)
-    sol = growth_solver(ΩM, w0, wa)
-    return _D_z(z, sol)
-end
-
-function _f_a(a, sol::SciMLBase.ODESolution)
-    G, G_prime = sol(a)
-    D = G * a
-    D_prime = G_prime * a + G
-    return a / D * D_prime
-end
-
-function _f_a(a::Array, sol::SciMLBase.ODESolution)
-    G = [u for (u,t) in sol.(a)]
-    G_prime = [t for (u,t) in sol.(a)]
-    D = G .* a
-    D_prime = G_prime .* a .+ G
-    return a ./ D .* D_prime
-end
-
-function _f_z(z, sol::SciMLBase.ODESolution)
-    a = _a_z.(z)
-    return _f_a(a, sol)
-end
-
-function _f_z(z, ΩM, w0, wa)
-    sol = growth_solver(ΩM, w0, wa)
-    return _f_z(z, sol)
-end"""
-
 function _D_z(z::Array, sol::SciMLBase.ODESolution)
     [u for (u,t) in sol.(log.(_a_z.(z)))] ./ (sol(log(_a_z(0.)))[1,:])
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index dd02706..5805615 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -18,7 +18,8 @@ k_test = Array(LinRange(0,200, 40))
 weights = SimpleChains.init_params(mlpd)
 inminmax = rand(6,2)
 outminmax = rand(40,2)
-a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa = [1., 1e-5, 0.25, 0.05, 0., 0.67, -1, 0.]
+z, a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa = [0., 1., 1e-5, 0.25, 0.05, 0., 0.67, -1, 0.]
+ΩM = Ωc0 + Ωb0
 emu = Effort.SimpleChainsEmulator(Architecture = mlpd, Weights = weights)
 
 effort_emu = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax = inminmax,
@@ -26,4 +27,5 @@ effort_emu = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax =
 
 @testset "Effort tests" begin
     @test isapprox(Effort._H_a(a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa), h*100)
+    @test isapprox(Effort._E_a(a, Ωc0, Ωb0, h), 1.)
 end

From 05b8d2c8251b78b6bd4e98d455557d8061b86b0e Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Thu, 29 Aug 2024 15:25:41 -0400
Subject: [PATCH 21/29] Towards diff D(z)

---
 src/background.jl | 16 ++++++++++++++++
 test/Project.toml |  3 +++
 test/runtests.jl  | 11 +++++++++++
 3 files changed, 30 insertions(+)

diff --git a/src/background.jl b/src/background.jl
index 3ea8ef0..08e3f8c 100644
--- a/src/background.jl
+++ b/src/background.jl
@@ -164,6 +164,22 @@ function growth_solver(Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
     return sol
 end
 
+function growth_solver(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    amin = 1/139
+    loga = vcat(log.(_a_z.(z)), 0.)
+    u₀ = [amin, amin]
+
+    logaspan = (log(amin), log(1.01))
+    Ωγ0 = 2.469e-5/h^2
+
+    p = [Ωc0, Ωb0 , mν, h, w0, wa]
+
+    prob = ODEProblem(_growth!, u₀, logaspan, p)
+
+    sol = solve(prob, OrdinaryDiffEq.Tsit5(), abstol=1e-6, reltol=1e-6; saveat=loga)[1:2,:]
+    return sol
+end
+
 function _D_z(z::Array, sol::SciMLBase.ODESolution)
     [u for (u,t) in sol.(log.(_a_z.(z)))] ./ (sol(log(_a_z(0.)))[1,:])
 end
diff --git a/test/Project.toml b/test/Project.toml
index abaf678..48d9612 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,6 +1,9 @@
 [deps]
 AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
diff --git a/test/runtests.jl b/test/runtests.jl
index 5805615..14d8c55 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -3,6 +3,9 @@ using NPZ
 using SimpleChains
 using Static
 using Effort
+using ForwardDiff
+using Zygote
+using SciMLSensitivity
 
 mlpd = SimpleChain(
   static(6),
@@ -25,7 +28,15 @@ emu = Effort.SimpleChainsEmulator(Architecture = mlpd, Weights = weights)
 effort_emu = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax = inminmax,
                                 OutMinMax = outminmax)
 
+x = [Ωc0, Ωb0, h]
+
+function pippo(z, x)
+    Ωc0, Ωb0, h = x
+    sum(Effort.growth_solver(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.))
+end
+
 @testset "Effort tests" begin
     @test isapprox(Effort._H_a(a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa), h*100)
     @test isapprox(Effort._E_a(a, Ωc0, Ωb0, h), 1.)
+    @test isapprox(Zygote.gradient(x->pippo(z, x), x)[1], ForwardDiff.gradient(x->pippo(z, x), x), rtol=1e-6)
 end

From fec4ea33cadc0f825f2f68c552d30e0f8607dfe9 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Thu, 29 Aug 2024 16:03:56 -0400
Subject: [PATCH 22/29] Added tests

---
 src/background.jl | 6 ++++++
 test/runtests.jl  | 4 +++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/background.jl b/src/background.jl
index 08e3f8c..370e61e 100644
--- a/src/background.jl
+++ b/src/background.jl
@@ -189,6 +189,12 @@ function _D_z(z, sol::SciMLBase.ODESolution)
 end
 
 function _D_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+    sol = growth_solver(z, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+    D_z = reverse(sol[1,1:end-1])./sol[1,end]
+    return D_z
+end
+
+function _D_z_old(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
     sol = growth_solver(Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
     return _D_z(z, sol)
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 14d8c55..1b688a6 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -21,7 +21,8 @@ k_test = Array(LinRange(0,200, 40))
 weights = SimpleChains.init_params(mlpd)
 inminmax = rand(6,2)
 outminmax = rand(40,2)
-z, a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa = [0., 1., 1e-5, 0.25, 0.05, 0., 0.67, -1, 0.]
+a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa = [1., 1e-5, 0.25, 0.05, 0., 0.67, -1, 0.]
+z = [0.4, 0.5]
 ΩM = Ωc0 + Ωb0
 emu = Effort.SimpleChainsEmulator(Architecture = mlpd, Weights = weights)
 
@@ -39,4 +40,5 @@ end
     @test isapprox(Effort._H_a(a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa), h*100)
     @test isapprox(Effort._E_a(a, Ωc0, Ωb0, h), 1.)
     @test isapprox(Zygote.gradient(x->pippo(z, x), x)[1], ForwardDiff.gradient(x->pippo(z, x), x), rtol=1e-6)
+    @test isapprox(Effort._D_z_old(z, Ωc0, Ωb0, h), Effort._D_z(z, Ωc0, Ωb0, h), rtol=1e-9)
 end

From f19d878be76e6bffb444f2c729b4918d2f94ec27 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Thu, 29 Aug 2024 16:24:39 -0400
Subject: [PATCH 23/29] More tests

---
 .github/workflows/test.yml |  1 -
 src/background.jl          | 27 +++++++++++++++++----------
 test/Project.toml          |  1 +
 test/runtests.jl           |  1 +
 4 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 3498722..28528da 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -23,7 +23,6 @@ jobs:
           - '1.8'
           - '1.9'
           - '1.10'
-          - '~1.11.0-0'
         os:
           - ubuntu-latest
         arch:
diff --git a/src/background.jl b/src/background.jl
index 370e61e..6a050e3 100644
--- a/src/background.jl
+++ b/src/background.jl
@@ -180,11 +180,11 @@ function growth_solver(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
     return sol
 end
 
-function _D_z(z::Array, sol::SciMLBase.ODESolution)
+function _D_z_old(z::Array, sol::SciMLBase.ODESolution)
     [u for (u,t) in sol.(log.(_a_z.(z)))] ./ (sol(log(_a_z(0.)))[1,:])
 end
 
-function _D_z(z, sol::SciMLBase.ODESolution)
+function _D_z_old(z, sol::SciMLBase.ODESolution)
     return (sol(log(_a_z(z)))[1,:]/sol(log(_a_z(0.)))[1,:])[1,1]
 end
 
@@ -196,7 +196,7 @@ end
 
 function _D_z_old(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
     sol = growth_solver(Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
-    return _D_z(z, sol)
+    return _D_z_old(z, sol)
 end
 
 function _D_z_unnorm(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
@@ -212,18 +212,25 @@ function _D_z_unnorm(z, sol::SciMLBase.ODESolution)
     return (sol(log(_a_z(z)))[1,:])[1,1]
 end
 
-function _f_a(a, sol::SciMLBase.ODESolution)
-    D, D_prime = sol(log(a))
-    return 1 / D * D_prime
+function _f_a_old(a, sol::SciMLBase.ODESolution)
+    D, D_prime = sol.(log.(a))
+    return @. 1 / D * D_prime
 end
 
-function _f_z(z, sol::SciMLBase.ODESolution)
+function _f_z_old(z, sol::SciMLBase.ODESolution)
     a = _a_z.(z)
-    return _f_a(a, sol)
+    return _f_a_old(a, sol)
 end
 
-function _f_z(z, Ωc0, Ωb0, h, mν =0 , w0=-1., wa=0.)
+function _f_z_old(z, Ωc0, Ωb0, h; mν =0 , w0=-1., wa=0.)
     a = _a_z.(z)
     sol = growth_solver(Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
-    return _f_a(a, sol)
+    return _f_a_old(a, sol)
+end
+
+function _f_z(z, Ωc0, Ωb0, h; mν =0 , w0=-1., wa=0.)
+    sol = growth_solver(z, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+    D = sol[1,1:end-1]
+    D_prime = sol[2,1:end-1]
+    return @. 1 / D * D_prime
 end
diff --git a/test/Project.toml b/test/Project.toml
index 48d9612..e1708be 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,5 +1,6 @@
 [deps]
 AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
+Effort = "ace19e24-02c2-4bce-bc41-b57e89c0fc53"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
diff --git a/test/runtests.jl b/test/runtests.jl
index 1b688a6..3b75d67 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -41,4 +41,5 @@ end
     @test isapprox(Effort._E_a(a, Ωc0, Ωb0, h), 1.)
     @test isapprox(Zygote.gradient(x->pippo(z, x), x)[1], ForwardDiff.gradient(x->pippo(z, x), x), rtol=1e-6)
     @test isapprox(Effort._D_z_old(z, Ωc0, Ωb0, h), Effort._D_z(z, Ωc0, Ωb0, h), rtol=1e-9)
+    @test isapprox(Effort._f_z_old(0.4, Ωc0, Ωb0, h), Effort._f_z(0.4, Ωc0, Ωb0, h)[1], rtol=1e-9)
 end

From 4c94322b2ca7540683d20adc5cfd6a03107b4926 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Thu, 29 Aug 2024 16:39:43 -0400
Subject: [PATCH 24/29] Added finiteDiff test

---
 Project.toml      | 2 +-
 test/Project.toml | 2 +-
 test/runtests.jl  | 6 ++++--
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/Project.toml b/Project.toml
index e9c3627..64e9b05 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Effort"
 uuid = "ace19e24-02c2-4bce-bc41-b57e89c0fc53"
 authors = ["marcobonici <bonici.marco@gmail.com>"]
-version = "0.1.2"
+version = "0.2.0"
 
 [deps]
 AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
diff --git a/test/Project.toml b/test/Project.toml
index e1708be..ea13341 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,6 +1,6 @@
 [deps]
 AbstractCosmologicalEmulators = "c83c1981-e5c4-4837-9eb8-c9b1572acfc6"
-Effort = "ace19e24-02c2-4bce-bc41-b57e89c0fc53"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
diff --git a/test/runtests.jl b/test/runtests.jl
index 3b75d67..0f87d82 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -5,6 +5,7 @@ using Static
 using Effort
 using ForwardDiff
 using Zygote
+using FiniteDiff
 using SciMLSensitivity
 
 mlpd = SimpleChain(
@@ -33,13 +34,14 @@ x = [Ωc0, Ωb0, h]
 
 function pippo(z, x)
     Ωc0, Ωb0, h = x
-    sum(Effort.growth_solver(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.))
+    sum(Effort._D_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.))
 end
 
 @testset "Effort tests" begin
     @test isapprox(Effort._H_a(a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa), h*100)
     @test isapprox(Effort._E_a(a, Ωc0, Ωb0, h), 1.)
-    @test isapprox(Zygote.gradient(x->pippo(z, x), x)[1], ForwardDiff.gradient(x->pippo(z, x), x), rtol=1e-6)
     @test isapprox(Effort._D_z_old(z, Ωc0, Ωb0, h), Effort._D_z(z, Ωc0, Ωb0, h), rtol=1e-9)
     @test isapprox(Effort._f_z_old(0.4, Ωc0, Ωb0, h), Effort._f_z(0.4, Ωc0, Ωb0, h)[1], rtol=1e-9)
+    @test isapprox(Zygote.gradient(x->pippo(z, x), x)[1], ForwardDiff.gradient(x->pippo(z, x), x), rtol=1e-5)
+    @test isapprox(FiniteDiff.finite_difference_gradient(x->pippo(z, x), x), ForwardDiff.gradient(x->pippo(z, x), x), rtol=1e-5)
 end

From 13766bdac34076db25e4366de0adf299d39b8f00 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Fri, 30 Aug 2024 10:02:41 -0400
Subject: [PATCH 25/29] Updating comov distance

---
 src/background.jl | 91 +++++++++++++++++++++++------------------------
 test/runtests.jl  | 18 +++++-----
 2 files changed, 54 insertions(+), 55 deletions(-)

diff --git a/src/background.jl b/src/background.jl
index 6a050e3..97bfb8e 100644
--- a/src/background.jl
+++ b/src/background.jl
@@ -50,63 +50,63 @@ end
     return sqrt(ΩM*(1+z)^3+(1-ΩM)*_ρDE_z(z, w0, wa))
 end
 """
-function _E_a(a, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+function _E_a(a, Ωm0, h; mν =0., w0=-1., wa=0.)
     Ωγ0 = 2.469e-5/h^2
     Ων0 = _ΩνE2(1., Ωγ0, mν)
-    ΩΛ0 = 1. - (Ωγ0 + Ωc0 + Ωb0 + Ων0)
-    return sqrt(Ωγ0*a^-4 + (Ωc0 + Ωb0)*a^-3 + ΩΛ0 * _ρDE_a(a, w0, wa)+ _ΩνE2(a, Ωγ0, mν))
+    ΩΛ0 = 1. - (Ωγ0 + Ωm0 + Ων0)
+    return sqrt(Ωγ0*a^-4 + Ωm0*a^-3 + ΩΛ0 * _ρDE_a(a, w0, wa)+ _ΩνE2(a, Ωγ0, mν))
 end
 
-function _E_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+function _E_z(z, Ωm0, h; mν =0., w0=-1., wa=0.)
     a = _a_z.(z)
-    return _E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+    return _E_a(a, Ωm0, h; mν =mν, w0=w0, wa=wa)
 end
 
-_H_a(a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa) = 100*h*_E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+_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, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+function _χ_z(z, Ωm0, h; mν =0., w0=-1., wa=0.)
     integral, _ = quadgk(x -> 1 /
-    _E_a(_a_z(x), Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa), 0, z, rtol=1e-6)
+    _E_a(_a_z(x), Ωm0, h; mν =mν, w0=w0, wa=wa), 0, z, rtol=1e-6)
     return integral*c_0/(100*h)
 end
 
-function _dEda(a, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+function _dEda(a, Ωm0, h; mν =0., w0=-1., wa=0.)
     Ωγ0 = 2.469e-5/h^2
     Ων0 = _ΩνE2(1., Ωγ0, mν)
-    ΩΛ0 = 1. - (Ωγ0 + Ωc0 + Ωb0 + Ων0)
-    return 0.5/_E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)*(-3(Ωc0 + Ωb0)a^-4-4Ωγ0*a^-5+
+    ΩΛ0 = 1. - (Ωγ0 + Ωm0 + Ων0)
+    return 0.5/_E_a(a, Ωm0, h; mν =mν, w0=w0, wa=wa)*(-3(Ωm0)a^-4-4Ωγ0*a^-5+
            ΩΛ0*_dρDEda(a, w0, wa)+_dΩνE2da(a, Ωγ0, mν))
 end
 
-function _dlogEdloga(a, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+function _dlogEdloga(a, Ωm0, h; mν =0., w0=-1., wa=0.)
     Ωγ0 = 2.469e-5/h^2
     Ων0 = _ΩνE2(1., Ωγ0, mν)
-    ΩΛ0 = 1. - (Ωγ0 + Ωc0 + Ωb0 + Ων0)
-    return a*0.5/(_E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)^2)*(-3(Ωc0 + Ωb0)a^-4-4Ωγ0*a^-
+    ΩΛ0 = 1. - (Ωγ0 + Ωm0 + Ων0)
+    return a*0.5/(_E_a(a, Ωm0, h; mν =mν, w0=w0, wa=wa)^2)*(-3(Ωm0)a^-4-4Ωγ0*a^-
            5+ΩΛ0*_dρDEda(a, w0, wa)+_dΩνE2da(a, Ωγ0, mν))
 end
 
-function _ΩMa(a, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+function _ΩMa(a, Ωm0, h; mν =0., w0=-1., wa=0.)
     Ωγ0 = 2.469e-5/h^2
     Ων0 = _ΩνE2(1., Ωγ0, mν)
-    return (Ωc0 + Ωb0 + Ων0 )*a^-3 / (_E_a(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa))^2
+    return (Ωm0 + Ων0 )*a^-3 / (_E_a(a, Ωm0, h; mν =mν, w0=w0, wa=wa))^2
 end
 
-function _r̃_z(z, ΩM, w0, wa)
-    integral, _ = quadgk(x -> 1 / _E_z(x, ΩM, w0, wa), 0, z, rtol=1e-10)
+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
 end
 
-function _r_z(z, H0, ΩM, w0, wa)
-    return c_0 * _r̃_z(z, ΩM, w0, wa) / H0
+function _r_z(z, ΩM, h; mν =0., w0=-1., wa=0.)
+    return c_0 * _r̃_z(z, ΩM, h; mν =mν, w0=w0, wa=wa) / (100*h)
 end
 
-function _d̃A_z(z, ΩM, w0, wa)
-    return _r̃_z(z, ΩM, w0, wa) / (1+z)
+function _d̃A_z(z, ΩM, h; mν =0., w0=-1., wa=0.)
+    return _r̃_z(z, ΩM, h; mν =mν, w0=w0, wa=wa) / (1+z)
 end
 
-function _dA_z(z, H0, ΩM, w0, wa)
-    return _r_z(z, H0, ΩM, w0, wa) / (1+z)
+function _dA_z(z, ΩM, h; mν =0., w0=-1., wa=0.)
+    return _r_z(z, ΩM, h; mν =mν, w0=w0, wa=wa) / (1+z)
 end
 
 function _ρDE_z(z, w0, wa)
@@ -131,32 +131,31 @@ end
 
 function _growth!(du,u,p,loga)
     #Ωγ0 = p[1]
-    Ωc0 = p[1]
-    Ωb0 = p[2]
-    mν  = p[3]
-    h   = p[4]
-    w0  = p[5]
-    wa  = p[6]
+    Ωm0 = p[1]
+    mν  = p[2]
+    h   = p[3]
+    w0  = p[4]
+    wa  = p[5]
     a = exp(loga)
     D = u[1]
     dD = u[2]
     du[1] = dD
-    du[2] = -(2+_dlogEdloga(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa))*dD+
-            1.5*_ΩMa(a, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)*D
+    du[2] = -(2+_dlogEdloga(a, Ωm0, h; mν =mν, w0=w0, wa=wa))*dD+
+            1.5*_ΩMa(a, Ωm0, h; mν =mν, w0=w0, wa=wa)*D
 end
 
 function _a_z(z)
     return 1/(1+z)
 end
 
-function growth_solver(Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+function growth_solver(Ωm0, h; mν =0., w0=-1., wa=0.)
     amin = 1/139
     u₀ = [amin, amin]
 
     logaspan = (log(amin), log(1.01))
     Ωγ0 = 2.469e-5/h^2
 
-    p = (Ωc0, Ωb0 , mν, h, w0, wa)
+    p = (Ωm0, mν, h, w0, wa)
 
     prob = ODEProblem(_growth!, u₀, logaspan, p)
 
@@ -164,7 +163,7 @@ function growth_solver(Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
     return sol
 end
 
-function growth_solver(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
+function growth_solver(z, Ωm0, h; mν =0., w0=-1., wa=0.)
     amin = 1/139
     loga = vcat(log.(_a_z.(z)), 0.)
     u₀ = [amin, amin]
@@ -172,7 +171,7 @@ function growth_solver(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
     logaspan = (log(amin), log(1.01))
     Ωγ0 = 2.469e-5/h^2
 
-    p = [Ωc0, Ωb0 , mν, h, w0, wa]
+    p = [Ωm0, mν, h, w0, wa]
 
     prob = ODEProblem(_growth!, u₀, logaspan, p)
 
@@ -188,19 +187,19 @@ function _D_z_old(z, sol::SciMLBase.ODESolution)
     return (sol(log(_a_z(z)))[1,:]/sol(log(_a_z(0.)))[1,:])[1,1]
 end
 
-function _D_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
-    sol = growth_solver(z, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+function _D_z(z, Ωm0, h; mν =0., w0=-1., wa=0.)
+    sol = growth_solver(z,Ωm0, h; mν =mν, w0=w0, wa=wa)
     D_z = reverse(sol[1,1:end-1])./sol[1,end]
     return D_z
 end
 
-function _D_z_old(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
-    sol = growth_solver(Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+function _D_z_old(z, Ωm0, h; mν =0., w0=-1., wa=0.)
+    sol = growth_solver(Ωm0, h; mν =mν, w0=w0, wa=wa)
     return _D_z_old(z, sol)
 end
 
-function _D_z_unnorm(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.)
-    sol = growth_solver(Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+function _D_z_unnorm(z, Ωm0, h; mν =0., w0=-1., wa=0.)
+    sol = growth_solver(Ωm0, h; mν =mν, w0=w0, wa=wa)
     return _D_z_unnorm(z, sol)
 end
 
@@ -222,14 +221,14 @@ function _f_z_old(z, sol::SciMLBase.ODESolution)
     return _f_a_old(a, sol)
 end
 
-function _f_z_old(z, Ωc0, Ωb0, h; mν =0 , w0=-1., wa=0.)
+function _f_z_old(z, Ωm0, h; mν =0 , w0=-1., wa=0.)
     a = _a_z.(z)
-    sol = growth_solver(Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+    sol = growth_solver(Ωm0, h; mν =mν, w0=w0, wa=wa)
     return _f_a_old(a, sol)
 end
 
-function _f_z(z, Ωc0, Ωb0, h; mν =0 , w0=-1., wa=0.)
-    sol = growth_solver(z, Ωc0, Ωb0, h; mν =mν, w0=w0, wa=wa)
+function _f_z(z, Ωm0, h; mν =0 , w0=-1., wa=0.)
+    sol = growth_solver(z, Ωm0, h; mν =mν, w0=w0, wa=wa)
     D = sol[1,1:end-1]
     D_prime = sol[2,1:end-1]
     return @. 1 / D * D_prime
diff --git a/test/runtests.jl b/test/runtests.jl
index 0f87d82..1588457 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -22,26 +22,26 @@ k_test = Array(LinRange(0,200, 40))
 weights = SimpleChains.init_params(mlpd)
 inminmax = rand(6,2)
 outminmax = rand(40,2)
-a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa = [1., 1e-5, 0.25, 0.05, 0., 0.67, -1, 0.]
+a, Ωγ0, Ωm0, mν, h, w0, wa = [1., 1e-5, 0.3, 0.06, 0.67, -1.1, 0.2]
 z = [0.4, 0.5]
-ΩM = Ωc0 + Ωb0
+
 emu = Effort.SimpleChainsEmulator(Architecture = mlpd, Weights = weights)
 
 effort_emu = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax = inminmax,
                                 OutMinMax = outminmax)
 
-x = [Ωc0, Ωb0, h]
+x = [Ωm0, h, mν, w0, wa]
 
 function pippo(z, x)
-    Ωc0, Ωb0, h = x
-    sum(Effort._D_z(z, Ωc0, Ωb0, h; mν =0., w0=-1., wa=0.))
+    Ωm0, h, mν, w0, wa = x
+    sum(Effort._D_z(z, Ωm0, h; mν =mν, w0=w0, wa=wa))
 end
 
 @testset "Effort tests" begin
-    @test isapprox(Effort._H_a(a, Ωγ0, Ωc0, Ωb0, mν, h, w0, wa), h*100)
-    @test isapprox(Effort._E_a(a, Ωc0, Ωb0, h), 1.)
-    @test isapprox(Effort._D_z_old(z, Ωc0, Ωb0, h), Effort._D_z(z, Ωc0, Ωb0, h), rtol=1e-9)
-    @test isapprox(Effort._f_z_old(0.4, Ωc0, Ωb0, h), Effort._f_z(0.4, Ωc0, Ωb0, h)[1], rtol=1e-9)
+    @test isapprox(Effort._H_a(a, Ωγ0, Ωm0, mν, h, w0, wa), h*100)
+    @test isapprox(Effort._E_a(a, Ωm0, h), 1.)
+    @test isapprox(Effort._D_z_old(z, Ωm0, h), Effort._D_z(z, Ωm0, h), rtol=1e-9)
+    @test isapprox(Effort._f_z_old(0.4, Ωm0, h), Effort._f_z(0.4, Ωm0, h)[1], rtol=1e-9)
     @test isapprox(Zygote.gradient(x->pippo(z, x), x)[1], ForwardDiff.gradient(x->pippo(z, x), x), rtol=1e-5)
     @test isapprox(FiniteDiff.finite_difference_gradient(x->pippo(z, x), x), ForwardDiff.gradient(x->pippo(z, x), x), rtol=1e-5)
 end

From 94705d6d42d7de9b1fbd058e43ca8d0341b42821 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Fri, 30 Aug 2024 10:12:20 -0400
Subject: [PATCH 26/29] Added test for growth rate diff

---
 test/runtests.jl | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 1588457..82ca89f 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -32,16 +32,23 @@ effort_emu = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax =
 
 x = [Ωm0, h, mν, w0, wa]
 
-function pippo(z, x)
+function D_z_x(z, x)
     Ωm0, h, mν, w0, wa = x
     sum(Effort._D_z(z, Ωm0, h; mν =mν, w0=w0, wa=wa))
 end
 
+function f_z_x(z, x)
+    Ωm0, h, mν, w0, wa = x
+    sum(Effort._f_z(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.)
     @test isapprox(Effort._D_z_old(z, Ωm0, h), Effort._D_z(z, Ωm0, h), rtol=1e-9)
     @test isapprox(Effort._f_z_old(0.4, Ωm0, h), Effort._f_z(0.4, Ωm0, h)[1], rtol=1e-9)
-    @test isapprox(Zygote.gradient(x->pippo(z, x), x)[1], ForwardDiff.gradient(x->pippo(z, x), x), rtol=1e-5)
-    @test isapprox(FiniteDiff.finite_difference_gradient(x->pippo(z, x), x), ForwardDiff.gradient(x->pippo(z, x), x), rtol=1e-5)
+    @test isapprox(Zygote.gradient(x->D_z_x(z, x), x)[1], ForwardDiff.gradient(x->D_z_x(z, x), x), rtol=1e-5)
+    @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-5)
 end

From 51086ecda32d556bb155e19fd2116b6e75f60318 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Fri, 30 Aug 2024 10:22:00 -0400
Subject: [PATCH 27/29] Updating tolerance finitediff test

---
 test/runtests.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 82ca89f..9ca8a72 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -50,5 +50,5 @@ end
     @test isapprox(Zygote.gradient(x->D_z_x(z, x), x)[1], ForwardDiff.gradient(x->D_z_x(z, x), x), rtol=1e-5)
     @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-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)
 end

From 8cc6877c6adc01335d8ff0ee23eb1e56eaa2b829 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Fri, 30 Aug 2024 10:52:59 -0400
Subject: [PATCH 28/29] Removed useless string

---
 src/background.jl | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/src/background.jl b/src/background.jl
index 97bfb8e..b3f9d86 100644
--- a/src/background.jl
+++ b/src/background.jl
@@ -46,10 +46,6 @@ function _dΩνE2da(a, Ωγ0, mν::Array; kB = 8.617342e-5, Tν = 0.71611*2.7255
     return 15/π^4*Γν^4*Ωγ0*sum_interpolant
 end
 
-"""function _E_z(z, ΩM, w0, wa)
-    return sqrt(ΩM*(1+z)^3+(1-ΩM)*_ρDE_z(z, w0, wa))
-end
-"""
 function _E_a(a, Ωm0, h; mν =0., w0=-1., wa=0.)
     Ωγ0 = 2.469e-5/h^2
     Ων0 = _ΩνE2(1., Ωγ0, mν)

From 9a392b41f4d32d448f6ba2ecc8176287dc077a74 Mon Sep 17 00:00:00 2001
From: marcobonici <bonici.marco@gmail.com>
Date: Fri, 30 Aug 2024 11:18:53 -0400
Subject: [PATCH 29/29] Updating docs

---
 .github/workflows/docs.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 5249325..61f428e 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -15,10 +15,10 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@latest
         with:
-          version: '1.9'
+          version: '1.10'
       - name: Install dependencies
         run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy