cosmo type redone

src/boltzmann.jl

@@ -100,7 +100,7 @@ function _inv_y_transformation(emulator::Emulator, point)
 end
 
 function lin_Pk0(mode::Val{:EmuPk}, cpar::CosmoPar, settings::Settings)
-    cosmotype = settings.cosmo_type
+    cosmo_type = _get_cosmo_type(cpar)
     wc = cpar.Ωc*cpar.h^2
     wb = cpar.Ωb*cpar.h^2
     ln1010As = log((10^10)*cpar.As)
@@ -109,7 +109,7 @@ function lin_Pk0(mode::Val{:EmuPk}, cpar::CosmoPar, settings::Settings)
 
     lks_emul = emulator.training_karr
     nk = length(lks_emul)
-    pk0s_t = zeros(cosmotype, nk)
+    pk0s_t = zeros(cosmo_type, nk)
     @inbounds for i in 1:nk
         kernel = _get_kernel(emulator.trans_cosmos, params_t, emulator.hypers[i, :])
         pk0s_t[i] = dot(vec(kernel), vec(emulator.alphas[i,:]))
@@ -135,7 +135,7 @@ function lin_Pk0(mode::Val{:Custom}, cpar::CosmoPar, settings::Settings; kwargs.
 end
 
 lin_Pk0(mode::Symbol, cpar::CosmoPar, settings::Settings; kwargs...) = lin_Pk0(Val(mode), cpar, settings; kwargs...)
-lin_Pk0(@nospecialize(mode), cpar::CosmoPar, settings::Settings; kwargs...) = error("Tk mode $(typeof(i)) not supported.")
+lin_Pk0(@nospecialize(mode), cpar::CosmoPar, settings::Settings; kwargs...) = error("Tk mode $(typeof(mode)) not supported.")
 
 function _get_kernel(arr1, arr2, hyper)
     arr1_w = @.(arr1/exp(hyper[2:6]))

src/core.jl

@@ -11,7 +11,6 @@ Kwargs:
 - `nk::Int=500`: number of nodes in the k-scale array used to compute matter power spectrum grid.
 - `nℓ::Int=300`: number of nodes in the multipoles array.
 - `using_As::Bool=false`: `True` if using the `As` parameter.
-- `cosmo_type::Type=Float64` : type of cosmological parameters. 
 - `tk_mode::String=:EisHu` : choice of transfer function.
 - `Dz_mode::String=:RK2` : choice of method to compute the linear growth factor.
 - `Pk_mode::String=:linear` : choice of method to apply non-linear corrections to the matter power spectrum.
@@ -59,14 +58,12 @@ mutable struct Settings
     dlogk
 
     using_As::Bool
-
-    cosmo_type::DataType
     tk_mode::Symbol
     Dz_mode::Symbol
     Pk_mode::Symbol
 end
 
-Settings(;kwargs...) = begin
+function Settings(;kwargs...)
     nz = get(kwargs, :nz, 300)
     nz_chi = get(kwargs, :nz_chi, 1000)
     nz_t = get(kwargs, :nz_t, 350)
@@ -83,7 +80,6 @@ Settings(;kwargs...) = begin
     ℓs = range(0, stop=2000, length=nℓ)
 
     using_As = get(kwargs, :using_As, false)
-    cosmo_type = get(kwargs, :cosmo_type, Float64)
     tk_mode = get(kwargs, :tk_mode, :EisHu)
     Dz_mode = get(kwargs, :Dz_mode, :RK2)
     Pk_mode = get(kwargs, :Pk_mode, :linear)
@@ -94,10 +90,10 @@ Settings(;kwargs...) = begin
         using_As = false
     end   
 
-    Settings(nz, nz_chi, nz_t, nk, nℓ,
-             xs, zs, zs_chi, zs_t, ks, ℓs, logk,  dlogk,
-             using_As,
-             cosmo_type, tk_mode, Dz_mode, Pk_mode)
+    return Settings(nz, nz_chi, nz_t, nk, nℓ,
+        xs, zs, zs_chi, zs_t, ks, ℓs, logk,  dlogk,
+        using_As,
+        tk_mode, Dz_mode, Pk_mode)
 end
 
 """
@@ -158,7 +154,7 @@ mutable struct CosmoPar{T}
     ΩΛ::T
 end
 
-CosmoPar(;kwargs...) = begin
+function CosmoPar(;kwargs...)
     kwargs = Dict(kwargs)
 
     Ωm = get(kwargs, :Ωm, 0.3)
@@ -179,19 +175,20 @@ CosmoPar(;kwargs...) = begin
     Ωr = get(kwargs, :Ωr, Ωg*f_rel)
     Ωc = Ωm-Ωb
     ΩΛ = 1-Ωm-Ωr
-    CosmoPar{cosmo_type}(Ωm, Ωb, h, ns, As, σ8,
-                         θCMB, Y_p, N_ν, Σm_ν,
-                         Ωg, Ωr, Ωc, ΩΛ)
+    return CosmoPar{cosmo_type}(Ωm, Ωb, h, ns, As, σ8,
+        θCMB, Y_p, N_ν, Σm_ν,
+        Ωg, Ωr, Ωc, ΩΛ)
 end
 
 
 function _get_cosmo_type(x::CosmoPar{T}) where{T}
     return T
 end
 
-struct Cosmology
+mutable struct Cosmology
     settings::Settings
     cpar::CosmoPar
+    cosmo_type::DataType
     chi::AbstractInterpolation
     z_of_chi::AbstractInterpolation
     t_of_z::AbstractInterpolation
@@ -246,9 +243,9 @@ mutable struct Cosmology
 end     
 ``` 
 """
-Cosmology(cpar::CosmoPar, settings::Settings; kwargs...) = begin
+function Cosmology(cpar::CosmoPar, settings::Settings; kwargs...)
     # Load settings
-    cosmo_type = settings.cosmo_type
+    cosmo_type = _get_cosmo_type(cpar)
     zs_chi, nz_chi = settings.zs_chi, settings.nz_chi
     zs, nz = settings.zs, settings.nz
     logk, nk = settings.logk, settings.nk
@@ -278,13 +275,23 @@ Cosmology(cpar::CosmoPar, settings::Settings; kwargs...) = begin
         Pki = linear_interpolation((logk, zs), log.(Pks);
                                    extrapolation_bc=Line())
     elseif settings.Pk_mode == :Halofit
-        Pki = get_PKnonlin(cpar, zs, ks, pk0, Dzs;
-                          cosmo_type=cosmo_type)
+        Pki = get_PKnonlin(cpar, zs, ks, pk0, Dzs)
     else 
         @error("Pk mode not implemented")
     end
-    Cosmology(settings, cpar, chii, zi, ti, chis[end],
-              chi_LSS, Dzi, fs8zi, pki, Pki)
+    return Cosmology(
+        settings,
+        cpar,
+        cosmo_type,
+        chii,
+        zi,
+        ti,
+        chis[end],
+        chi_LSS,
+        Dzi,
+        fs8zi,
+        pki,
+        Pki)
 end
 
 """
@@ -297,19 +304,16 @@ Returns:
 - `Cosmology` : cosmology structure.
 
 """
-Cosmology(;kwargs...) = begin
+function Cosmology(;kwargs...)
     kwargs=Dict(kwargs)
     if :As ∈ keys(kwargs)
         using_As = true
     else
         using_As = false
     end     
     cpar = CosmoPar(;kwargs...)
-    cosmo_type = _get_cosmo_type(cpar)
-    settings = Settings(;cosmo_type=cosmo_type,
-                         using_As=using_As,
-                         kwargs...)                  
-    Cosmology(cpar, settings)
+    settings = Settings(; using_As=using_As, kwargs...)                  
+    return Cosmology(cpar, settings)
 end
 
 function Ez(cpar::CosmoPar, z)

src/growth.jl

@@ -8,6 +8,7 @@ get_growth(mode::Symbol, cpar::CosmoPar, settings::Settings; kwargs...) = get_gr
 get_growth(@nospecialize(mode), cpar::CosmoPar, settings::Settings; kwargs...) = error("Dz mode $(typeof(i)) not supported.")
 
 function get_growth(mode::Val{:RK2}, cpar::CosmoPar, settings::Settings; kwargs...)
+    cosmo_type = _get_cosmo_type(cpar)
     x = settings.xs
     z = settings.zs
     a = @.(1/(1+z))
@@ -16,8 +17,8 @@ function get_growth(mode::Val{:RK2}, cpar::CosmoPar, settings::Settings; kwargs.
     aa = reverse(a)
     e = Ez(cpar, z)
     ee = reverse(e)
-    dd = zeros(settings.cosmo_type, settings.nz)
-    yy = zeros(settings.cosmo_type, settings.nz)
+    dd = zeros(cosmo_type, settings.nz)
+    yy = zeros(cosmo_type, settings.nz)
     dd[1] = aa[1]
     yy[1] = aa[1]^3*ee[end]
 

src/halofit.jl

@@ -24,7 +24,8 @@ function _get_σ2(lks, ks, pks, R, kind)
     return integral
 end
 
-function get_PKnonlin(cpar::CosmoPar, z, k, PkLz0, Dzs; cosmo_type::DataType=Real)
+function get_PKnonlin(cpar::CosmoPar, z, k, PkLz0, Dzs)
+    cosmo_type = _get_cosmo_type(cpar)
     nk = length(k)
     nz = length(z)
     logk = log.(k)

src/spectra.jl

@@ -15,7 +15,7 @@ function Cℓintegrand(cosmo::Cosmology,
                      t2::AbstractInterpolation,
                      ℓ::Number)
     sett = cosmo.settings
-    chis = zeros(sett.cosmo_type, sett.nk)
+    chis = zeros(cosmo.cosmo_type, sett.nk)
     chis[1:sett.nk] = (ℓ+0.5) ./ sett.ks
     chis .*= (chis .< cosmo.chi_max)
     z = cosmo.z_of_chi(chis)

src/theory.jl

@@ -1,12 +1,12 @@
-function Theory(cosmology::Cosmology,
+function Theory(cosmo::Cosmology,
                 names, types, pairs,
                 idx, files;
                 Nuisances=Dict())
 
     nui_type =  eltype(valtype(Nuisances))
     if !(nui_type <: Float64) & (nui_type != Any)
         if nui_type != Real
-            cosmology.settings.cosmo_type = nui_type
+            cosmo.cosmo_type = nui_type
         end
     end
 
@@ -21,7 +21,7 @@ function Theory(cosmology::Cosmology,
             nz = get(Nuisances, string(name, "_", "nz"), nz_mean)
             zs = get(Nuisances, string(name, "_", "zs"), zs_mean)
             dz = get(Nuisances, string(name, "_", "dz"), 0.0)
-            tracer = NumberCountsTracer(cosmology, zs .+ dz, nz;
+            tracer = NumberCountsTracer(cosmo, zs .+ dz, nz;
                                         b=b)
         elseif t_type == "galaxy_shear"
             zs_mean, nz_mean = files[string("nz_", name)]
@@ -31,11 +31,11 @@ function Theory(cosmology::Cosmology,
             nz = get(Nuisances, string(name, "_", "nz"), nz_mean)
             zs = get(Nuisances, string(name, "_", "zs"), zs_mean)
             dz = get(Nuisances, string(name, "_", "dz"), 0.0)
-            tracer = WeakLensingTracer(cosmology, zs .+ dz, nz;
+            tracer = WeakLensingTracer(cosmo, zs .+ dz, nz;
                                        m=m, IA_params=IA_params)
 
         elseif t_type == "cmb_convergence"
-            tracer = CMBLensingTracer(cosmology)
+            tracer = CMBLensingTracer(cosmo)
 
         else
             @error("Tracer not implemented")
@@ -46,20 +46,20 @@ function Theory(cosmology::Cosmology,
 
     npairs = length(pairs)
     total_len = last(idx)
-    cls = zeros(cosmology.settings.cosmo_type, total_len)
+    cls = zeros(cosmo.cosmo_type, total_len)
     @inbounds Threads.@threads :static for i in 1:npairs
         name1, name2 = pairs[i]
         ls = files[string("ls_", name1, "_", name2)]
         tracer1 = tracers[name1]
         tracer2 = tracers[name2]
-        cls[idx[i]+1:idx[i+1]] = angularCℓs(cosmology, tracer1, tracer2, ls)
+        cls[idx[i]+1:idx[i+1]] = angularCℓs(cosmo, tracer1, tracer2, ls)
     end
 
     return cls
 end
 
 """
-    Theory(cosmology::Cosmology,
+    Theory(cosmo::Cosmology,
            instructions::Instructions, files;
            Nuisances=Dict())
 
@@ -68,7 +68,7 @@ a `Meta` objectm, a `files` npz file and \
 a dictionary of nuisance parameters.
 
 Arguments:
-- `cosmology::Cosmology` : `Cosmology` object.
+- `cosmo::Cosmology` : `Cosmology` object.
 - `meta::Meta` : `Meta` object.
 - `files` : `files` `npz` file.
 - `Nuisances::Dict` : dictonary of nuisace parameters. 
@@ -88,7 +88,7 @@ end
 ```
 - files: npz file
 """
-function Theory(cosmology::Cosmology,
+function Theory(cosmo::Cosmology,
                 instructions::Instructions, files;
                 Nuisances=Dict())
 
@@ -97,7 +97,7 @@ function Theory(cosmology::Cosmology,
     pairs = instructions.pairs
     idx = instructions.idx
 
-    return Theory(cosmology::Cosmology,
+    return Theory(cosmo::Cosmology,
                   names, types, pairs,
                   idx, files;
                   Nuisances=Nuisances)

src/tracers.jl

@@ -27,7 +27,7 @@ Kwargs:
 Returns:
 - `NumberCountsTracer::NumberCountsTracer` : Number counts tracer structure.
 """
-NumberCountsTracer(cosmo::Cosmology, z_n, nz; b=1.0) = begin
+function NumberCountsTracer(cosmo::Cosmology, z_n, nz; b=1.0)
     nz_int = linear_interpolation(z_n, nz, extrapolation_bc=0)
     res = cosmo.settings.nz_t
     z_w = range(0.00001, stop=z_n[end], length=res)
@@ -40,7 +40,7 @@ NumberCountsTracer(cosmo::Cosmology, z_n, nz; b=1.0) = begin
     w_arr = @. (nz_w*hz/nz_norm)
     wint = linear_interpolation(chi, b .* w_arr, extrapolation_bc=Line())
     F::Function = ℓ -> 1
-    NumberCountsTracer(wint, F)
+    return NumberCountsTracer(wint, F)
 end
 
 """
@@ -62,10 +62,9 @@ struct WeakLensingTracer <: Tracer
     F::Function
 end
 
-WeakLensingTracer(cosmo::Cosmology, z_n, nz; IA_params = [0.0, 0.0], m=0.0, kwargs...) = begin
+function WeakLensingTracer(cosmo::Cosmology, z_n, nz; IA_params = [0.0, 0.0], m=0.0, kwargs...)
     nz_int = linear_interpolation(z_n, nz, extrapolation_bc=0) 
-    cosmo_type = cosmo.settings.cosmo_type
-    res =cosmo.settings.nz_t
+    res = cosmo.settings.nz_t
     z_w = range(0.00001, stop=z_n[end], length=res)
     dz_w = (z_w[end]-z_w[1])/res
     nz_w = nz_int(z_w)
@@ -79,7 +78,7 @@ WeakLensingTracer(cosmo::Cosmology, z_n, nz; IA_params = [0.0, 0.0], m=0.0, kwar
     #      at all zs.
     # Calculate integral at each chi
     w_itg(chii) = @.(nz_w*(1-chii/chi))
-    w_arr = zeros(cosmo_type, res)
+    w_arr = zeros(cosmo.cosmo_type, res)
     @inbounds for i in 1:res-3
         w_arr[i] = integrate(z_w[i:res], w_itg(chi[i])[i:res], SimpsonEven())
     end
@@ -102,7 +101,7 @@ WeakLensingTracer(cosmo::Cosmology, z_n, nz; IA_params = [0.0, 0.0], m=0.0, kwar
     b = m+1.0
     wint = linear_interpolation(chi, b.*w_arr, extrapolation_bc=Line())
     F::Function = ℓ -> @.(sqrt((ℓ+2)*(ℓ+1)*ℓ*(ℓ-1))/(ℓ+0.5)^2)
-    WeakLensingTracer(wint, F)
+    return WeakLensingTracer(wint, F)
 end
 
 """
@@ -127,7 +126,7 @@ Arguments:
 Returns:
 - `CMBLensingTracer::CMBLensingTracer` : CMB lensing tracer structure.
 """
-CMBLensingTracer(cosmo::Cosmology) = begin
+function CMBLensingTracer(cosmo::Cosmology)
     # chi array
     chis = range(0.0, stop=cosmo.chi_max, length=cosmo.settings.nz_t)
     zs = cosmo.z_of_chi(chis)
@@ -140,7 +139,7 @@ CMBLensingTracer(cosmo::Cosmology) = begin
     # Interpolate
     wint = linear_interpolation(chis, w_arr, extrapolation_bc=0.0)
     F::Function = ℓ -> @.(((ℓ+1)*ℓ)/(ℓ+0.5)^2) 
-    CMBLensingTracer(wint, F)
+    return CMBLensingTracer(wint, F)
 end
 
 """