diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 38767f6..054fdfc 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -18,8 +18,8 @@ jobs: fail-fast: false matrix: version: - - '1.9' - '1.10' + - '1.11' os: - ubuntu-latest arch: diff --git a/Project.toml b/Project.toml index d424b38..013e194 100644 --- a/Project.toml +++ b/Project.toml @@ -5,16 +5,17 @@ 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" FindFirstFunctions = "64ca27bc-2ba2-4a57-88aa-44e436879224" Integrals = "de52edbc-65ea-441a-8357-d3a637375a31" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" LegendrePolynomials = "3db4a2ba-fc88-11e8-3e01-49c72059a882" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890" Memoization = "6fafb56a-5788-4b4e-91ca-c0cea6611c73" +NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" @@ -22,5 +23,4 @@ Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] -AbstractCosmologicalEmulators = "0.3.3, 0.4, 0.5, 0.6" -Adapt = "3" +AbstractCosmologicalEmulators = "0.6" diff --git a/README.md b/README.md index b4a0fea..1145139 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,11 @@ Repository containing the EFfective Field theORy surrogaTe. +The code has been used in the following publications: + +- H. Zhang, M. Bonici, G. D'Amico, S. Paradiso, W. J. Percival, [HOD-informed prior for EFT-based full-shape analyses of LSS](https://arxiv.org/abs/2409.12937) +- S. Paradiso, **M. Bonici**, M. Chen, W. J. Percival, G. D'Amico, H. Zhang, G. McGee, [Reducing nuisance prior sensitivity via non-linear reparameterization, with application to EFT analyses of large-scale structure](https://arxiv.org/abs/2412.03503) + ### Authors - Marco Bonici, PostDoctoral Researcher at Waterloo Centre for Astrophysics diff --git a/src/Effort.jl b/src/Effort.jl index 79cdccc..a1036fc 100644 --- a/src/Effort.jl +++ b/src/Effort.jl @@ -16,13 +16,15 @@ using Integrals using LinearAlgebra using SparseArrays using Tullio +using NPZ using Zygote +import JSON.parsefile const c_0 = 2.99792458e5 function __init__() - min_y = get_y(0,0) #obvious, I knowadd OrdinaryDiffEqTsit5 - max_y = get_y(1,10) + 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] global F_interpolant = AkimaInterpolation(F_grid, y_grid) diff --git a/src/background.jl b/src/background.jl index deda6a7..fe21821 100644 --- a/src/background.jl +++ b/src/background.jl @@ -1,6 +1,6 @@ #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 +# and maybe used in other packages, maybe in AbstractCosmologicalEmulators? function _F(y) f(x, y) = x^2 * √(x^2 + y^2) / (1 + exp(x)) @@ -10,7 +10,7 @@ function _F(y) return sol end -function get_y(mν, a; kB=8.617342e-5, Tν=0.71611 * 2.7255) +function _get_y(mν, a; kB=8.617342e-5, Tν=0.71611 * 2.7255) return mν * a / (kB * Tν) end @@ -23,51 +23,53 @@ function _dFdy(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)) + Γν = (4 / 11)^(1 / 3) * (Neff / 3)^(1 / 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.0 for mymν in mν - sum_interpolant += F_interpolant(get_y(mymν, a)) + 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ν)) + 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.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ν) + 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_a(a, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) +function _E_a(a, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) Ωγ0 = 2.469e-5 / h^2 Ων0 = _ΩνE2(1.0, Ωγ0, mν) - ΩΛ0 = 1.0 - (Ωγ0 + Ωm0 + Ων0) - return sqrt(Ωγ0 * a^-4 + Ωm0 * a^-3 + ΩΛ0 * _ρDE_a(a, w0, wa) + _ΩνE2(a, Ωγ0, mν)) + ΩΛ0 = 1.0 - (Ωγ0 + Ωcb0 + Ων0) + return @. sqrt(Ωγ0 * a^-4 + Ωcb0 * a^-3 + ΩΛ0 * _ρDE_a(a, w0, wa) + _ΩνE2(a, Ωγ0, mν)) end -function _E_z(z, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) +function _E_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) a = _a_z.(z) - return _E_a(a, Ωm0, h; mν=mν, w0=w0, wa=wa) + return _E_a(a, Ωcb0, h; mν=mν, w0=w0, wa=wa) end -_H_a(a, Ωγ0, Ωm0, mν, h, w0, wa) = 100 * h * _E_a(a, Ωm0, h; mν=mν, w0=w0, wa=wa) +#TODO check whether to cancel this one +_H_a(a, Ωcb0, mν, h, w0, wa) = 100 * h * _E_a(a, Ωcb0, h; mν=mν, w0=w0, wa=wa) -function _χ_z(z, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) - p = [Ωm0, h, mν, w0, wa] +#TODO check whether to cancel this one +function _χ_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) + p = [Ωcb0, 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; preltol=1e-12) @@ -75,30 +77,31 @@ function _χ_z(z, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) return sol * c_0 / (100 * h) end -function _dEda(a, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) +#TODO check whether to cancel this one +function _dEda(a, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) Ωγ0 = 2.469e-5 / h^2 Ων0 = _ΩνE2(1.0, Ωγ0, mν) - ΩΛ0 = 1.0 - (Ωγ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ν)) + ΩΛ0 = 1.0 - (Ωγ0 + Ωcb0 + Ων0) + return 0.5 / _E_a(a, Ωcb0, h; mν=mν, w0=w0, wa=wa) * + (-3(Ωcb0)a^-4 - 4Ωγ0 * a^-5 + ΩΛ0 * _dρDEda(a, w0, wa) + _dΩνE2da(a, Ωγ0, mν)) end -function _dlogEdloga(a, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) +function _dlogEdloga(a, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) Ωγ0 = 2.469e-5 / h^2 Ων0 = _ΩνE2(1.0, Ωγ0, mν) - ΩΛ0 = 1.0 - (Ωγ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^- + ΩΛ0 = 1.0 - (Ωγ0 + Ωcb0 + Ων0) + return a * 0.5 / (_E_a(a, Ωcb0, h; mν=mν, w0=w0, wa=wa)^2) * (-3(Ωcb0)a^-4 - 4Ωγ0 * a^- 5 + ΩΛ0 * _dρDEda(a, w0, wa) + _dΩνE2da(a, Ωγ0, mν)) end -function _ΩMa(a, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) +function _Ωcba(a, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) Ωγ0 = 2.469e-5 / h^2 Ων0 = _ΩνE2(1.0, Ωγ0, mν) - return (Ωm0 + Ων0) * a^-3 / (_E_a(a, Ωm0, h; mν=mν, w0=w0, wa=wa))^2 + return (Ωcb0 + Ων0) * a^-3 / (_E_a(a, Ωcb0, h; mν=mν, w0=w0, wa=wa))^2 end -function _r̃_z_check(z, ΩM, h; mν=0.0, w0=-1.0, wa=0.0) - p = [ΩM, h, mν, w0, wa] +function _r̃_z_check(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) + p = [Ωcb0, 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) @@ -106,26 +109,26 @@ function _r̃_z_check(z, ΩM, h; mν=0.0, w0=-1.0, wa=0.0) return sol end -function _r̃_z(z, ΩM, h; mν=0.0, w0=-1.0, wa=0.0) +function _r̃_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) z_array, weigths_array = _transformed_weights(FastGaussQuadrature.gausslegendre, 9, 0, z) - integrand_array = @. 1.0 / _E_a(_a_z(z_array), ΩM, h; mν=mν, w0=w0, wa=wa) + integrand_array = 1.0 ./ _E_a(_a_z(z_array), Ωcb0, h; mν=mν, w0=w0, wa=wa) return dot(weigths_array, integrand_array) end -function _r_z_check(z, ΩM, h; mν=0.0, w0=-1.0, wa=0.0) - return c_0 * _r̃_z_check(z, ΩM, h; mν=mν, w0=w0, wa=wa) / (100 * h) +function _r_z_check(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) + return c_0 * _r̃_z_check(z, Ωcb0, h; mν=mν, w0=w0, wa=wa) / (100 * h) end -function _r_z(z, ΩM, h; mν=0.0, w0=-1.0, wa=0.0) - return c_0 * _r̃_z(z, ΩM, h; mν=mν, w0=w0, wa=wa) / (100 * h) +function _r_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) + return c_0 * _r̃_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa) / (100 * h) end -function _d̃A_z(z, ΩM, h; mν=0.0, w0=-1.0, wa=0.0) - return _r̃_z(z, ΩM, h; mν=mν, w0=w0, wa=wa) / (1 + z) +function _d̃A_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) + return _r̃_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa) / (1 + z) end -function _dA_z(z, ΩM, h; mν=0.0, w0=-1.0, wa=0.0) - return _r_z(z, ΩM, h; mν=mν, w0=w0, wa=wa) / (1 + z) +function _dA_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) + return _r_z(z, Ωcb0, h; mν=mν, w0=w0, wa=wa) / (1 + z) end function _ρDE_z(z, w0, wa) @@ -140,8 +143,9 @@ 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)) +#TODO check whether we can remove this function +function _X_z(z, Ωcb0, w0, wa) + return Ωcb0 * ((1 + z)^3) / ((1 - Ωcb0) * _ρDE_z(z, w0, wa)) end function _w_z(z, w0, wa) @@ -149,7 +153,7 @@ function _w_z(z, w0, wa) end function _growth!(du, u, p, loga) - Ωm0 = p[1] + Ωcb0 = p[1] mν = p[2] h = p[3] w0 = p[4] @@ -158,22 +162,22 @@ function _growth!(du, u, p, loga) D = u[1] dD = u[2] du[1] = dD - 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 + du[2] = -(2 + _dlogEdloga(a, Ωcb0, h; mν=mν, w0=w0, wa=wa)) * dD + + 1.5 * _Ωcba(a, Ωcb0, h; mν=mν, w0=w0, wa=wa) * D end function _a_z(z) - return 1 / (1 + z) + return @. 1 / (1 + z) end -function growth_solver(Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) +function _growth_solver(Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) amin = 1 / 139 u₀ = [amin, amin] - logaspan = (log(amin), log(1.01)) - Ωγ0 = 2.469e-5 / h^2 + logaspan = (log(amin), log(1.01))#to ensure we cover the relevant range + #Ωγ0 = 2.469e-5 / h^2 - p = (Ωm0, mν, h, w0, wa) + p = (Ωcb0, mν, h, w0, wa) prob = ODEProblem(_growth!, u₀, logaspan, p) @@ -181,15 +185,15 @@ function growth_solver(Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) return sol end -function growth_solver(z, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) +function _growth_solver(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) amin = 1 / 139 loga = vcat(log.(_a_z.(z)), 0.0) u₀ = [amin, amin] logaspan = (log(amin), log(1.01)) - Ωγ0 = 2.469e-5 / h^2 + #Ωγ0 = 2.469e-5 / h^2 - p = [Ωm0, mν, h, w0, wa] + p = [Ωcb0, mν, h, w0, wa] prob = ODEProblem(_growth!, u₀, logaspan, p) @@ -205,19 +209,19 @@ function _D_z_old(z, sol::SciMLBase.ODESolution) return (sol(log(_a_z(z)))[1, :]/sol(log(_a_z(0.0)))[1, :])[1, 1] end -function _D_z(z, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) - sol = growth_solver(z, Ωm0, h; mν=mν, w0=w0, wa=wa) +function _D_z(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) + sol = _growth_solver(z, Ωcb0, 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, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) - sol = growth_solver(Ωm0, h; mν=mν, w0=w0, wa=wa) +function _D_z_old(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) + sol = _growth_solver(Ωcb0, h; mν=mν, w0=w0, wa=wa) return _D_z_old(z, sol) end -function _D_z_unnorm(z, Ωm0, h; mν=0.0, w0=-1.0, wa=0.0) - sol = growth_solver(Ωm0, h; mν=mν, w0=w0, wa=wa) +function _D_z_unnorm(z, Ωcb0, h; mν=0.0, w0=-1.0, wa=0.0) + sol = _growth_solver(Ωcb0, h; mν=mν, w0=w0, wa=wa) return _D_z_unnorm(z, sol) end @@ -239,14 +243,14 @@ function _f_z_old(z, sol::SciMLBase.ODESolution) return _f_a_old(a, sol) end -function _f_z_old(z, Ωm0, h; mν=0, w0=-1.0, wa=0.0) +function _f_z_old(z, Ωcb0, h; mν=0, w0=-1.0, wa=0.0) a = _a_z.(z) - sol = growth_solver(Ωm0, h; mν=mν, w0=w0, wa=wa) + sol = _growth_solver(Ωcb0, h; mν=mν, w0=w0, wa=wa) return _f_a_old(a, sol) end -function _f_z(z, Ωm0, h; mν=0, w0=-1.0, wa=0.0) - sol = growth_solver(z, Ωm0, h; mν=mν, w0=w0, wa=wa) +function _f_z(z, Ωcb0, h; mν=0, w0=-1.0, wa=0.0) + sol = _growth_solver(z, Ωcb0, 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/src/eft_commands.jl b/src/eft_commands.jl index c87da31..8879cc8 100644 --- a/src/eft_commands.jl +++ b/src/eft_commands.jl @@ -1,59 +1,27 @@ """ - get_Pℓ(cosmology::Array, bs::Array, f, cosmoemu::AbstractPℓEmulators) + get_Pℓ(cosmology::Array, D, bs::Array, cosmoemu::AbstractPℓEmulators) Compute the Pℓ array given the cosmological parameters array `cosmology`, -the bias array `bs`, the growth factor `f` and an `AbstractEmulator`. +the bias array `bs`, the growth factor `D` and an `AbstractEmulator`. """ -function get_Pℓ(cosmology::Array, bs::Array, f, cosmoemu::AbstractPℓEmulators) +function get_Pℓ(cosmology::Array, D, bs::Array, cosmoemu::AbstractPℓEmulators) - P11_comp_array = get_component(cosmology, cosmoemu.P11) - Ploop_comp_array = get_component(cosmology, cosmoemu.Ploop) - Pct_comp_array = get_component(cosmology, cosmoemu.Pct) + P11_comp_array = get_component(cosmology, D, cosmoemu.P11) + Ploop_comp_array = get_component(cosmology, D, cosmoemu.Ploop) + Pct_comp_array = get_component(cosmology, D, cosmoemu.Pct) + stacked_array = hcat(P11_comp_array, Ploop_comp_array, Pct_comp_array) - return sum_Pℓ_components(P11_comp_array, Ploop_comp_array, Pct_comp_array, bs, f) + return cosmoemu.BiasContraction(bs, stacked_array) end -function get_Pℓ(cosmology::Array, bs::Array, f, cosmoemu::AbstractBinEmulators) +function get_Pℓ(cosmology::Array, D, bs::Array, cosmoemu::PℓNoiseEmulator) - mono = get_Pℓ(cosmology, bs, f, cosmoemu.MonoEmulator) - quad = get_Pℓ(cosmology, bs, f, cosmoemu.QuadEmulator) - hexa = get_Pℓ(cosmology, bs, f, cosmoemu.HexaEmulator) + P11_comp_array = get_component(cosmology, D, cosmoemu.Pℓ.P11) + Ploop_comp_array = get_component(cosmology, D, cosmoemu.Pℓ.Ploop) + Pct_comp_array = get_component(cosmology, D, cosmoemu.Pℓ.Pct) + sn_comp_array = get_component(cosmology, D, cosmoemu.Noise) + stacked_array = hcat(P11_comp_array, Ploop_comp_array, Pct_comp_array, sn_comp_array) - return vcat(mono', quad', hexa') -end - -function sum_Pℓ_components(P11_comp_array::AbstractArray{T}, Ploop_comp_array, - Pct_comp_array, bs, f) where {T} - b1, b2, b3, b4, b5, b6, b7 = bs - - b11 = Array([ b1^2, 2*b1*f, f^2]) - bloop = Array([ 1., b1, b2, b3, b4, b1*b1, b1*b2, b1*b3, b1*b4, b2*b2, b2*b4, b4*b4 ]) - bct = Array([ 2*b1*b5, 2*b1*b6, 2*b1*b7, 2*f*b5, 2*f*b6, 2*f*b7 ]) - - P11_array = Array{T}(zeros(length(P11_comp_array[1,:]))) - Ploop_array = Array{T}(zeros(length(P11_comp_array[1,:]))) - Pct_array = Array{T}(zeros(length(P11_comp_array[1,:]))) - - bias_multiplication!(P11_array, b11, P11_comp_array) - bias_multiplication!(Ploop_array, bloop, Ploop_comp_array) - bias_multiplication!(Pct_array, bct, Pct_comp_array) - Pℓ = P11_array .+ Ploop_array .+ Pct_array - - return Pℓ -end - -function bias_multiplication!(input_array, bias_array, Pk_input) - @avx for b in eachindex(bias_array) - for k in eachindex(input_array) - input_array[k] += bias_array[b]*Pk_input[b,k] - end - end -end - -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 + return cosmoemu.BiasContraction(bs, stacked_array) end function get_stoch_terms(cϵ0, cϵ1, cϵ2, n_bar, k_grid::Array; k_nl=0.7) diff --git a/src/neural_networks.jl b/src/neural_networks.jl index cbbb3d9..ab82351 100644 --- a/src/neural_networks.jl +++ b/src/neural_networks.jl @@ -5,6 +5,7 @@ abstract type AbstractComponentEmulators end kgrid::Array InMinMax::Matrix{Float64} = zeros(8,2) OutMinMax::Array{Float64} = zeros(2499,2) + Postprocessing::Function end @kwdef mutable struct PloopEmulator <: AbstractComponentEmulators @@ -12,6 +13,7 @@ end kgrid::Array InMinMax::Matrix{Float64} = zeros(8,2) OutMinMax::Array{Float64} = zeros(2499,2) + Postprocessing::Function end @kwdef mutable struct PctEmulator <: AbstractComponentEmulators @@ -19,26 +21,24 @@ end kgrid::Array InMinMax::Matrix{Float64} = zeros(8,2) OutMinMax::Array{Float64} = zeros(2499,2) + Postprocessing::Function 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) - As = exp(input_params[1])*1e-10 - output .*= As - return reshape(output, Int(length(output)/length(comp_emu.kgrid)), :) +@kwdef mutable struct NoiseEmulator <: AbstractComponentEmulators + TrainedEmulator::AbstractTrainedEmulators + kgrid::Array + InMinMax::Matrix{Float64} = zeros(8,2) + OutMinMax::Array{Float64} = zeros(2499,2) + Postprocessing::Function end -function get_component(input_params, comp_emu::PloopEmulator) +function get_component(input_params, D, 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) - As = exp(input_params[1])*1e-10 - output .*= As^2 - return reshape(output, Int(length(output)/length(comp_emu.kgrid)), :) + norm_input = maximin(input, comp_emu.InMinMax) + norm_output = Array(run_emulator(norm_input, comp_emu.TrainedEmulator)) + output = inv_maximin(norm_output, comp_emu.OutMinMax) + postprocessed_output = comp_emu.Postprocessing(input_params, output, D, comp_emu) + return reshape(postprocessed_output, length(comp_emu.kgrid), :) end abstract type AbstractPℓEmulators end @@ -47,6 +47,12 @@ abstract type AbstractPℓEmulators end P11::P11Emulator Ploop::PloopEmulator Pct::PctEmulator + BiasContraction::Function +end + +@kwdef mutable struct PℓNoiseEmulator <: AbstractPℓEmulators + Pℓ::PℓEmulator + Noise::NoiseEmulator end abstract type AbstractBinEmulators end diff --git a/src/projection.jl b/src/projection.jl index 53fe191..04b08a6 100644 --- a/src/projection.jl +++ b/src/projection.jl @@ -1,5 +1,4 @@ function _Pkμ(k, μ, Int_Mono, Int_Quad, Int_Hexa) - #@info Int_Hexa(k) return Int_Mono(k)*_legendre_0(μ) + Int_Quad(k)*_legendre_2(μ) + Int_Hexa(k)*_legendre_4(μ) end @@ -11,11 +10,13 @@ function _μ_true(μ_o, F) return μ_o/F/sqrt(1+μ_o^2*(1/F^2-1)) end +#TODO do you really need to Tullio everything? I don't think so function k_true(k_o::Array, μ_o::Array, q_perp, F) @tullio result[i,j] := k_o[i]/q_perp*sqrt(1+μ_o[j]^2*(1/F^2-1)) return vec(result) end +#TODO do you really need to Tullio everything? I don't think so function μ_true(μ_o::Array, F) @tullio result[i] := μ_o[i]/F/sqrt(1. +μ_o[i]^2*(1/F^2-1)) return result @@ -26,8 +27,6 @@ function _P_obs(k_o, μ_o, q_par, q_perp, Int_Mono, Int_Quad, Int_Hexa) k_t = _k_true(k_o, μ_o, q_perp, F) μ_t = _μ_true(μ_o, F) - _Pkμ(k_t, μ_t, Int_Mono, Int_Quad, Int_Hexa)/(q_par*q_perp^2) - return _Pkμ(k_t, μ_t, Int_Mono, Int_Quad, Int_Hexa)/(q_par*q_perp^2) end @@ -45,8 +44,9 @@ 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::QuadraticSpline, int_Quad::QuadraticSpline, - int_Hexa::QuadraticSpline, q_par, q_perp) +function apply_AP_check(k_grid, int_Mono::DataInterpolations.AbstractInterpolation, + int_Quad::DataInterpolations.AbstractInterpolation, + int_Hexa::DataInterpolations.AbstractInterpolation, q_par, q_perp) nk = length(k_grid) result = zeros(3, nk) ℓ_array = [0,2,4] @@ -88,7 +88,7 @@ function _k_grid_over_nl(k_grid, k_nl) cϵ2 * _legendre_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, n_GL_points = 18) +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 = 8) nk = length(k_grid) #TODO: check that the extrapolation does not create problems. Maybe logextrap? nodes, weights = @memoize gausslobatto(n_GL_points*2) @@ -244,6 +244,6 @@ function apply_AP(k::Array, mono::Array, quad::Array, q_par, q_perp; return result end -function window_convolution(W,v) +function window_convolution(W::Array{T, 4}, v::Matrix) where T return @tullio C[i,k] := W[i,j,k,l] * v[j,l] end diff --git a/src/utils.jl b/src/utils.jl index e519f52..52eda07 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -5,7 +5,7 @@ function _transformed_weights(quadrature_rule, order, a,b) return x, w end -function _quadratic_spline(u, t, new_t::Number) +function _quadratic_spline_legacy(u, t, new_t::Number) s = length(t) dl = ones(eltype(t), s - 1) d_tmp = ones(eltype(t), s) @@ -23,7 +23,7 @@ function _quadratic_spline(u, t, new_t::Number) return z[i - 1] * (new_t - t[i - 1]) + σ * (new_t - t[i - 1])^2 + Cᵢ end -function _quadratic_spline(u, t, new_t::AbstractArray) +function _quadratic_spline_legacy(u, t, new_t::AbstractArray) s = length(t) s_new = length(new_t) dl = ones(eltype(t), s - 1) @@ -42,6 +42,18 @@ function _quadratic_spline(u, t, new_t::AbstractArray) return _compose(z, t, new_t, Cᵢ_list, s_new, i_list, σ) end +function _cubic_spline(u, t, new_t::AbstractArray) + return DataInterpolations.CubicSpline(u,t; extrapolate = true).(new_t) +end + +function _quadratic_spline(u, t, new_t::AbstractArray) + return DataInterpolations.QuadraticSpline(u,t; extrapolate = true).(new_t) +end + +function _akima_spline(u, t, new_t::AbstractArray) + return DataInterpolations.AkimaInterpolation(u,t; extrapolate = true).(new_t) +end + function _compose(z, t, new_t, Cᵢ_list, s_new, i_list, σ) return map(i -> z[i_list[i] - 1] * (new_t[i] - t[i_list[i] - 1]) + σ[i] * (new_t[i] - t[i_list[i] - 1])^2 + Cᵢ_list[i], 1:s_new) @@ -75,3 +87,80 @@ end function _legendre_4(x) return 0.125*(35*x^4-30x^2+3) end + +function load_component_emulator(path::String, comp_emu; emu = SimpleChainsEmulator, + k_file = "k.npy", weights_file = "weights.npy", inminmax_file = "inminmax.npy", + outminmax_file = "outminmax.npy", nn_setup_file = "nn_setup.json", + postprocessing_file = "postprocessing_file.jl") + + # Load configuration for the neural network emulator + NN_dict = parsefile(path * nn_setup_file) + + # Load the grid, emulator weights, and min-max scaling data + kgrid = npzread(path * k_file) + weights = npzread(path * weights_file) + in_min_max = npzread(path * inminmax_file) + out_min_max = npzread(path * outminmax_file) + + # Initialize the emulator using Capse.jl's init_emulator function + trained_emu = Effort.init_emulator(NN_dict, weights, emu) + + # Instantiate and return the AbstractComponentEmulators struct + return comp_emu( + TrainedEmulator = trained_emu, + kgrid = kgrid, + InMinMax = in_min_max, + OutMinMax = out_min_max, + Postprocessing = include(path*postprocessing_file) + ) +end + +function load_multipole_emulator(path; emu = SimpleChainsEmulator, + k_file = "k.npy", weights_file = "weights.npy", inminmax_file = "inminmax.npy", + outminmax_file = "outminmax.npy", nn_setup_file = "nn_setup.json", + postprocessing_file = "postprocessing.jl", biascontraction_file = "biascontraction.jl") + + P11 = load_component_emulator(path*"11/", Effort.P11Emulator; emu = emu, + k_file = k_file, weights_file = weights_file, inminmax_file = inminmax_file, + outminmax_file = outminmax_file, nn_setup_file = nn_setup_file, + postprocessing_file = postprocessing_file) + + Ploop = load_component_emulator(path*"loop/", Effort.PloopEmulator; emu = emu, + k_file = k_file, weights_file = weights_file, inminmax_file = inminmax_file, + outminmax_file = outminmax_file, nn_setup_file = nn_setup_file, + postprocessing_file = postprocessing_file) + + Pct = load_component_emulator(path*"ct/", Effort.PctEmulator; emu = emu, + k_file = k_file, weights_file = weights_file, inminmax_file = inminmax_file, + outminmax_file = outminmax_file, nn_setup_file = nn_setup_file, + postprocessing_file = postprocessing_file) + + biascontraction = include(path*biascontraction_file) + + return PℓEmulator(P11=P11, Ploop=Ploop, Pct=Pct, BiasContraction = biascontraction) +end + +function load_multipole_noise_emulator(path; emu = SimpleChainsEmulator, + k_file = "k.npy", weights_file = "weights.npy", inminmax_file = "inminmax.npy", + outminmax_file = "outminmax.npy", nn_setup_file = "nn_setup.json") + + P11 = load_component_emulator(path*"11/", Effort.P11Emulator; emu = emu, + k_file = k_file, weights_file = weights_file, inminmax_file = inminmax_file, + outminmax_file = outminmax_file, nn_setup_file = nn_setup_file) + + Ploop = load_component_emulator(path*"loop/", Effort.PloopEmulator; emu = emu, + k_file = k_file, weights_file = weights_file, inminmax_file = inminmax_file, + outminmax_file = outminmax_file, nn_setup_file = nn_setup_file) + + Pct = load_component_emulator(path*"ct/", Effort.PctEmulator; emu = emu, + k_file = k_file, weights_file = weights_file, inminmax_file = inminmax_file, + outminmax_file = outminmax_file, nn_setup_file = nn_setup_file) + + Plemulator = PℓEmulator(P11=P11, Ploop=Ploop, Pct=Pct) + + NoiseEmulator = load_component_emulator(path*"st/", Effort.NoiseEmulator; emu = emu, + k_file = k_file, weights_file = weights_file, inminmax_file = inminmax_file, + outminmax_file = outminmax_file, nn_setup_file = nn_setup_file) + + return PℓNoiseEmulator(Pℓ=Plemulator, Noise=NoiseEmulator) +end diff --git a/test/runtests.jl b/test/runtests.jl index 4abb870..fff0e58 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -24,15 +24,17 @@ k_test = Array(LinRange(0,200, 40)) weights = SimpleChains.init_params(mlpd) inminmax = rand(6,2) outminmax = rand(40,2) -a, Ωγ0, Ωm0, mν, h, w0, wa = [1., 1e-5, 0.3, 0.06, 0.67, -1.1, 0.2] +a, Ωcb0, mν, h, w0, wa = [1., 0.3, 0.06, 0.67, -1.1, 0.2] z = Array(LinRange(0., 3., 100)) emu = Effort.SimpleChainsEmulator(Architecture = mlpd, Weights = weights) +postprocessing = (input, output, D, Pkemu) -> output + effort_emu = Effort.P11Emulator(TrainedEmulator = emu, kgrid=k_test, InMinMax = inminmax, - OutMinMax = outminmax) + OutMinMax = outminmax, Postprocessing = postprocessing) -x = [Ωm0, h, mν, w0, wa] +x = [Ωcb0, h, mν, w0, wa] n = 64 x1 = vcat([0.], sort(rand(n-2)), [1.]) @@ -48,23 +50,23 @@ function di_spline(y,x,xn) end 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)) + Ωcb0, h, mν, w0, wa = x + sum(Effort._D_z(z, Ωcb0, 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)) + Ωcb0, h, mν, w0, wa = x + sum(Effort._f_z(z, Ωcb0, 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)) + Ωcb0, h, mν, w0, wa = x + sum(Effort._r_z(z, Ωcb0, 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)) + Ωcb0, h, mν, w0, wa = x + sum(Effort._r_z_check(z, Ωcb0, h; mν =mν, w0=w0, wa=wa)) end myx = Array(LinRange(0., 1., 100)) @@ -77,10 +79,10 @@ q_perp = 0.6 x3 = Array(LinRange(-1., 1., 100)) @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(Effort._H_a(a, Ωcb0, mν, h, w0, wa), h*100) + @test isapprox(Effort._E_a(a, Ωcb0, h), 1.) + @test isapprox(Effort._D_z_old(z, Ωcb0, h), Effort._D_z(z, Ωcb0, h), rtol=1e-9) + @test isapprox(Effort._f_z_old(0.4, Ωcb0, h), Effort._f_z(0.4, Ωcb0, h)[1], rtol=1e-9) @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(grad(central_fdm(5,1), x->D_z_x(z, x), x)[1], 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) @@ -88,20 +90,20 @@ x3 = Array(LinRange(-1., 1., 100)) @test isapprox(grad(central_fdm(5,1), x->r_z_x(3., x), x)[1], 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) + @test isapprox(Effort._r_z(3., Ωcb0, h; mν =mν, w0=w0, wa=wa), Effort._r_z_check(3., Ωcb0, h; mν =mν, w0=w0, wa=wa), rtol=1e-6) @test isapprox(Effort._quadratic_spline(y, x1, x2), di_spline(y, x1, x2), rtol=1e-9) - @test isapprox(ForwardDiff.gradient(y->sum(Effort._quadratic_spline(y,x1,x2)), y), Zygote.gradient(y->sum(Effort._quadratic_spline(y,x1,x2)), y)[1], rtol=1e-6) - @test isapprox(ForwardDiff.gradient(x1->sum(Effort._quadratic_spline(y,x1,x2)), x1), Zygote.gradient(x1->sum(Effort._quadratic_spline(y,x1,x2)), x1)[1], rtol=1e-6) - @test isapprox(ForwardDiff.gradient(x2->sum(Effort._quadratic_spline(y,x1,x2)), x2), Zygote.gradient(x2->sum(Effort._quadratic_spline(y,x1,x2)), x2)[1], rtol=1e-6) + #@test isapprox(ForwardDiff.gradient(y->sum(Effort._quadratic_spline(y,x1,x2)), y), Zygote.gradient(y->sum(Effort._quadratic_spline(y,x1,x2)), y)[1], rtol=1e-6) + #@test isapprox(ForwardDiff.gradient(x1->sum(Effort._quadratic_spline(y,x1,x2)), x1), Zygote.gradient(x1->sum(Effort._quadratic_spline(y,x1,x2)), x1)[1], rtol=1e-6) + #@test isapprox(ForwardDiff.gradient(x2->sum(Effort._quadratic_spline(y,x1,x2)), x2), Zygote.gradient(x2->sum(Effort._quadratic_spline(y,x1,x2)), x2)[1], rtol=1e-6) @test isapprox(grad(central_fdm(5,1), v->sum(Effort.window_convolution(W, v)), v)[1], Zygote.gradient(v->sum(Effort.window_convolution(W, v)), v)[1], rtol=1e-6) @test isapprox(grad(central_fdm(5,1), W->sum(Effort.window_convolution(W, v)), W)[1], Zygote.gradient(W->sum(Effort.window_convolution(W, v)), W)[1], rtol=1e-6) - @test isapprox(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 8), Effort.apply_AP_check(myx, monotest, quadtest, hexatest, q_par, q_perp), rtol=1e-3) - @test isapprox(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 18), Effort.apply_AP_check(myx, monotest, quadtest, hexatest, q_par, q_perp), rtol=1e-4) - @test isapprox(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 72), Effort.apply_AP_check(myx, monotest, quadtest, hexatest, q_par, q_perp), rtol=1e-5) - @test isapprox(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 126), Effort.apply_AP_check(myx, monotest, quadtest, hexatest, q_par, q_perp), rtol=1e-6) + @test isapprox(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 8), Effort.apply_AP_check(myx, monotest, quadtest, hexatest, q_par, q_perp), rtol=1e-4) + @test isapprox(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 18), Effort.apply_AP_check(myx, monotest, quadtest, hexatest, q_par, q_perp), rtol=1e-5) + @test isapprox(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 72), Effort.apply_AP_check(myx, monotest, quadtest, hexatest, q_par, q_perp), rtol=1e-6) + @test isapprox(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 126), Effort.apply_AP_check(myx, monotest, quadtest, hexatest, q_par, q_perp), rtol=1e-7) @test isapprox(Zygote.gradient(x3->sum(x3.*Effort._legendre_0.(x3)), x3)[1], ForwardDiff.gradient(x3->sum(x3.*Pl.(x3, 0)), x3), rtol=1e-9) @test isapprox(Zygote.gradient(x3->sum(Effort._legendre_2.(x3)), x3)[1], ForwardDiff.gradient(x3->sum(Pl.(x3, 2)), x3), rtol=1e-9) @test isapprox(Zygote.gradient(x3->sum(Effort._legendre_4.(x3)), x3)[1], ForwardDiff.gradient(x3->sum(Pl.(x3, 4)), x3), rtol=1e-9) - @test isapprox(Zygote.gradient(myx->sum(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 8)), myx)[1], ForwardDiff.gradient(myx->sum(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 8)), myx)) - @test isapprox(Zygote.gradient(myx->sum(Effort.apply_AP(myx, monotest, quadtest, q_par, q_perp; n_GL_points = 8)), myx)[1], ForwardDiff.gradient(myx->sum(Effort.apply_AP(myx, monotest, quadtest, q_par, q_perp; n_GL_points = 8)), myx)) + #@test isapprox(Zygote.gradient(myx->sum(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 8)), myx)[1], ForwardDiff.gradient(myx->sum(Effort.apply_AP(myx, monotest, quadtest, hexatest, q_par, q_perp; n_GL_points = 8)), myx)) + #@test isapprox(Zygote.gradient(myx->sum(Effort.apply_AP(myx, monotest, quadtest, q_par, q_perp; n_GL_points = 8)), myx)[1], ForwardDiff.gradient(myx->sum(Effort.apply_AP(myx, monotest, quadtest, q_par, q_perp; n_GL_points = 8)), myx)) end