Towards unification

src/eft_commands.jl

@@ -3,113 +3,25 @@
 Compute the Pℓ array given the cosmological parameters array `cosmology`,
 the bias array `bs`, the growth factor `f` 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, cosmoemu::PℓNoiseEmulator)
+function get_Pℓ(cosmology::Array, D, bs::Array, cosmoemu::PℓNoiseEmulator)
 
-    P11_comp_array = get_component(cosmology, cosmoemu.Pℓ.P11)
-    Ploop_comp_array = get_component(cosmology, cosmoemu.Pℓ.Ploop)
-    Pct_comp_array = get_component(cosmology, cosmoemu.Pℓ.Pct)
-    sn_comp_array = get_component(cosmology, cosmoemu.Noise)
+    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 sum_Pℓ_components(P11_comp_array, Ploop_comp_array, Pct_comp_array,
-                             sn_comp_array, bs)
-end
-
-function get_Pℓ(cosmology::Array, biases::Array, cosmoemu::Effort.AbstractPℓEmulators)
-
-    P11_comp_array = get_component(cosmology, cosmoemu.Pℓ.P11)
-    Ploop_comp_array = get_component(cosmology, cosmoemu.Pℓ.Ploop)
-    Pct_comp_array = get_component(cosmology, cosmoemu.Pℓ.Pct)
-
-    b1, b2, b3, bs, alpha0, alpha2, alpha4, alpha6 = biases
-
-    b11 = [1, b1, b1^2]
-    bloop = [b2, b1*b2, b2^2, bs, b1*bs, b2*bs, bs^2, b3, b1*b3]
-    bct = [alpha0, alpha2, alpha4, alpha6]
-
-    P11_array = P11_comp_array*b11#Array{T}(zeros(length(P11_comp_array[1,:])))
-    Ploop_array = Ploop_comp_array*bloop#Array{T}(zeros(length(P11_comp_array[1,:])))
-    Pct_array = Pct_comp_array*bct
-
-    Pℓ = P11_array .+ Ploop_array .+ Pct_array
-
-    return Pℓ
-end
-
-function get_Pℓ(cosmology::Array, bs::Array, f, cosmoemu::AbstractBinEmulators)
-
-    mono = get_Pℓ(cosmology, bs, f, cosmoemu.MonoEmulator)
-    quad = get_Pℓ(cosmology, bs, f, cosmoemu.QuadEmulator)
-    hexa = get_Pℓ(cosmology, bs, f, cosmoemu.HexaEmulator)
-
-    return vcat(mono', quad', hexa')
-end
-
-function sum_Pℓ_components(P11_comp_array::AbstractArray{T}, Ploop_comp_array,
-    Pct_comp_array, bs, f::Number) 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 sum_Pℓ_components(P11_comp_array::AbstractArray{T}, Ploop_comp_array::AbstractArray,
-    Pct_comp_array::AbstractArray, Sn_comp_array::AbstractArray, biases::AbstractArray,) where {T}
-    b1, b2, b3, bs, alpha0, alpha2, alpha4, alpha6, sn, sn2, sn4 = biases
-
-    b11 = [1, b1, b1^2]
-    bloop = [b2, b1*b2, b2^2, bs, b1*bs, b2*bs, bs^2, b3, b1*b3]
-    bct = [alpha0, alpha2, alpha4, alpha6]
-    bsn = [sn, sn2, sn4]
-
-    P11_array = P11_comp_array*b11#Array{T}(zeros(length(P11_comp_array[1,:])))
-    Ploop_array = Ploop_comp_array*bloop#Array{T}(zeros(length(P11_comp_array[1,:])))
-    Pct_array = Pct_comp_array*bct#Array{T}(zeros(length(P11_comp_array[1,:])))
-    Sn_array = Sn_comp_array*bsn#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)
-    #bias_multiplication!(Sn_array, sn, Sn_comp_array)
-
-    Pℓ = P11_array .+ Ploop_array .+ Pct_array .+ Sn_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)

src/neural_networks.jl

@@ -5,55 +5,40 @@ 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
     TrainedEmulator::AbstractTrainedEmulators
     kgrid::Array
     InMinMax::Matrix{Float64} = zeros(8,2)
     OutMinMax::Array{Float64} = zeros(2499,2)
+    Postprocessing::Function
 end
 
 @kwdef mutable struct PctEmulator <: AbstractComponentEmulators
     TrainedEmulator::AbstractTrainedEmulators
     kgrid::Array
     InMinMax::Matrix{Float64} = zeros(8,2)
     OutMinMax::Array{Float64} = zeros(2499,2)
+    Postprocessing::Function
 end
 
 @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::AbstractComponentEmulators)
+function get_component(input_params, D, comp_emu::AbstractComponentEmulators)
     input = deepcopy(input_params)
     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)
-    As = exp(input_params[1])*1e-10
-    output .*= As
-    return reshape(output, length(comp_emu.kgrid), :)
-end
-
-function get_component(input_params, comp_emu::PloopEmulator)
-    input = deepcopy(input_params)
-    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)
-    As = exp(input_params[1])*1e-10
-    output .*= As^2
-    return reshape(output, length(comp_emu.kgrid), :)
-end
-
-function get_component(input_params, comp_emu::NoiseEmulator)
-    input = deepcopy(input_params)
-    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)
-    return reshape(output, length(comp_emu.kgrid), :)
+    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
@@ -62,6 +47,7 @@ abstract type AbstractPℓEmulators end
     P11::P11Emulator
     Ploop::PloopEmulator
     Pct::PctEmulator
+    BiasContraction::Function
 end
 
 @kwdef mutable struct PℓNoiseEmulator <: AbstractPℓEmulators

src/utils.jl

@@ -86,7 +86,8 @@ 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")
+    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)
@@ -105,23 +106,29 @@ function load_component_emulator(path::String, comp_emu; emu = SimpleChainsEmula
         TrainedEmulator = trained_emu,
         kgrid = kgrid,
         InMinMax = in_min_max,
-        OutMinMax = out_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")
+    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)
+    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)
+    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)
-    return PℓEmulator(P11=P11, Ploop=Ploop, Pct=Pct)
+    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,