Merge pull request #65 from xzackli/easy_features

Place to add (and test) quick features

examples/plot_deriv_cl.jl

@@ -1,9 +1,96 @@
 using Bolt
 using ForwardDiff
-# using PyPlot
 using Plots
 using BenchmarkTools
+using Plots.PlotMeasures
 
+
+# General derivates of  all other parameters
+function clp(𝕡::DT,ells) where DT
+    bg = Background(𝕡; x_grid=-20.0:0.01:0.0, nq=15)
+    𝕣 = Bolt.RECFAST(bg=bg, Yp=𝕡.Y_p, OmegaB=𝕡.Ω_b, OmegaG=𝕡.Ω_r)
+    k_grid = quadratic_k(0.1bg.H₀, 1000bg.H₀, 100)
+    ih = IonizationHistory(𝕣, 𝕡, bg)
+    sf = source_grid(𝕡, bg, ih, k_grid, BasicNewtonian())
+    return cltt(ells, 𝕡, bg, ih, sf)
+end
+
+
+pt = CosmoParams() #dummy struct
+ells = 10:20:1200
+# do a loop over everyone to see who is slow...this is bad for performance but who cares we are just checking times inside...
+fds = zeros(length(fieldnames(typeof(pt))),length(ells))
+fws = zeros(length(fieldnames(typeof(pt))),length(ells))
+
+for (i,name) in enumerate(fieldnames(typeof(pt)))
+    println("i = ", i, " $name = ",getfield(pt,name))
+    xv = getfield(pt,name)#pt.name
+    function clx(x::DT,name,ells) where DT
+        𝕡n = CosmoParams{DT}(;Dict((fieldnames(typeof(pt))[i])=>x)...)
+        return clp(𝕡n,ells)
+    end
+    f(x) = clx(x,name,ells)
+    fwddiff_∂cl = ForwardDiff.derivative(f,xv)
+    println(fwddiff_∂cl[1])
+    Δ = 1e-2*xv 
+    finitediff_∂cl = (f(xv + Δ) .- f(xv - Δ)) ./ 2Δ
+    fws[i,:] = fwddiff_∂cl
+    fds[i,:] = finitediff_∂cl
+end
+
+#multiply the mnu by its value to get something of OOM as the others...
+#just an ugliness of the Mpc units for the eV neutrino mass...
+fws[end,:] .*= 𝕡.Σm_ν
+fds[end,:] .*= 𝕡.Σm_ν
+
+# save the results
+writedlm("ClTT_AD.dat",fws)
+writedlm("ClTT_FD1e-2.dat",fds)
+
+#make some plots
+pnames = string.( fieldnames(typeof(pt)) )
+for i in 1:length(fieldnames(typeof(pt)))
+    plot(title=raw"$ p = $"*pnames[i])
+    plot!(ells, fws[i,:] .* ells.^2,label="AD")
+    plot!(ells, fds[i,:] .* ells.^2,ls=:dash,label="FD, "*raw"$\frac{\Delta p}{p} = 10^{-2}$")
+    xlabel!(raw"$\ell$")
+    ylabel!(raw"$\ell^{2} \frac{\partial C_{\ell}^{TT}}{\partial p }$",left_margin=5mm)
+    savefig("AD_vs_FD1e-2_all_params_"*pnames[i]*".pdf")
+end
+
+
+
+#-----FIXME move the below, which just tests scalar polzn spectra, and has nothing to do with derivs
+
+𝕡 = CosmoParams()
+bg = Background(𝕡; x_grid=-20.0:0.01:0.0, nq=15)
+𝕣 = Bolt.RECFAST(bg=bg, Yp=𝕡.Y_p, OmegaB=𝕡.Ω_b, OmegaG=𝕡.Ω_r)
+ih = IonizationHistory(𝕣, 𝕡, bg)
+k_grid = quadratic_k(0.1bg.H₀, 1000bg.H₀, 100)
+sf_t = source_grid(𝕡, bg, ih, k_grid, BasicNewtonian())
+sf_e = source_grid_P(𝕡, bg, ih, k_grid, BasicNewtonian())
+
+ℓs = 10:20:1200
+Cℓtt = cltt(ℓs, 𝕡, bg, ih, sf_t)
+Cℓte = clte(ℓs, 𝕡, bg, ih, sf_t,sf_e)
+Cℓee = clee(ℓs, 𝕡, bg, ih, sf_e)
+
+ℓfac = ℓs.*(ℓs.+1)
+plot(ℓs, @. ( ℓfac * Cℓtt))
+xlabel!(raw"$\ell$")
+ylabel!(raw"$\ell(\ell+1)C^{TT}_{\ell}$")
+# savefig("./test/cltt.png")
+plot(ℓs, @. (ℓfac *Cℓte))
+xlabel!(raw"$\ell$")
+ylabel!(raw"$\ell(\ell+1)C^{TE}_{\ell}$")
+# savefig("./test/clte.png")
+plot(ℓs, @. (ℓfac * Cℓee))
+xlabel!(raw"$\ell$")
+ylabel!(raw"$\ell(\ell+1)C^{EE}_{\ell}$")
+# savefig("./test/clee.png")
+
+
+#--- old code
 # Cₗ as a function of baryon density
 #Some noise at lowest few ells...9/17/21
 
@@ -22,20 +109,18 @@ f(Ω_b) = clb(Ω_b, ells)
 @time cl = f(0.046)
 @time ∂cl = ForwardDiff.derivative(f, 0.046)  # you can just ForwardDiff the whole thing
 
-##
-# clf()
-# plt.figure(figsize=(10,5))
-plot(ells, cl .* ells.^2, label=raw"$C_{\ell}$",lw=0,marker=:circle)
-plot!(ells, ∂cl .* ells.^2 / 10,
-    label=raw"$\partial C_{\ell}/\partial\Omega_b / 10$",lw=0,marker=:circle)
+@profview f(0.046)
+#58.130749 seconds (74.01 M allocations: 7.458 GiB, 3.10% gc time, 2.92% compilation time)
+@profview ForwardDiff.derivative(f, 0.046)
+#219.854089 seconds (139.67 M allocations: 13.256 GiB, 1.45% gc time, 8.75% compilation time)
 
-# uncomment to see finite differences
-# Δ = 1e-3
-@time finitediff_∂cl = (f(0.046 + Δ) .- f(0.046 - Δ)) ./ 2Δ
-plot!(ells, finitediff_∂cl .* ells.^2 / 10, ls=:dash,
-    label=raw"$\Delta C_{\ell}/\Delta\Omega_b / 10$")
 
-ylabel!(raw"$\ell^2 C_{\ell}^{TT}$")
-xlabel!(raw"$\ell$")
-ylims!(-0.3, 0.5)
-# savefig("docs/assets/example_spectrum.png")
+ells = 10:20:1200
+fh(h) = clh(h, ells)
+@time clhh = fh(0.7)
+#^ 59.244405 seconds (72.37 M allocations: 7.376 GiB, 3.04% gc time, 0.28% compilation time)
+@time ∂clhh = ForwardDiff.derivative(fh, 0.7)  # you can just ForwardDiff the whole thing
+#^ 217.447492 seconds (109.88 M allocations: 12.040 GiB, 1.53% gc time, 9.47% compilation time)
+
+@profview fh(0.7)
+@profview ForwardDiff.derivative(fh, 0.7)

scripts/plot_perts_x.jl

@@ -1,4 +1,4 @@
-using Revise
+# using Revise
 using Bolt
 using Plots
 using Printf
@@ -36,7 +36,7 @@ kclass = retnf[2][1] #read class k mode from file (in h/Mpc)
 k = (bg.H₀*3e5/100)*kclass #get k in our units
 class_pxs = transpose(reduce(hcat,ret[1:end .!= 2]))
 class_pxsnf = transpose(reduce(hcat,retnf[1:end .!= 2]))
-dipole_fac = kclass/𝕡.h #for later normalization
+dipole_fac = 2kclass #for later normalization
 
 #see above plot but for this particular mode at this cosmology the k condition happens later
 this_rsa_switch = x_grid[argmin(abs.(k .* bg.η.(x_grid) .- 45))]
@@ -49,7 +49,7 @@ println("k = ", kclass," log10k = ", log10(kclass), " h/Mpc")
 ℓᵧ=50
 ℓ_ν=50
 ℓ_mν=20
-reltol=1e-5 #cheaper  rtol
+reltol=1e-8 #cheaper  rtol
 pertlen = 2(ℓᵧ+1)+(ℓ_ν+1)+(ℓ_mν+1)*n_q+5
 results=zeros(pertlen,length(x_grid))
 ℳρ,ℳσ = zeros(length(x_grid)),zeros(length(x_grid)) #arrays for the massive neutrino integrated perts
@@ -77,35 +77,39 @@ results_with_rsa = boltsolve_rsa(hierarchy; reltol=reltol)
 #photon Θ0 monopole
 plot(class_pxs[1,:],log10.(abs.(class_pxsnf[2,:])),
       label=raw"$\Theta_{0,\rm{CLASS}}$",legend=:topleft)
-plot!(x_grid, log10.(abs.(results_with_rsa[1,:]* 𝕡.h*4)),
-      label=raw"$4 h \Theta_{0,\rm{Bolt}}$",ls=:dash)
+plot!(x_grid, log10.(abs.(results_with_rsa[1,:]*4)),
+      label=raw"$4 \Theta_{0,\rm{Bolt,rsa}}$",ls=:dash)
+plot!(x_grid, log10.(abs.(results[1,:]*4)),
+      label=raw"$4 \Theta_{0,\rm{Bolt}}$",ls=:dash)
 vline!([this_rsa_switch],label="RSA switch",ls=:dot)
 xlims!(-8,0)
 xlabel!(raw"$x$")
 ylabel!(raw"$\delta_{i}(x)$")
 #photon Θ1 dipole
 plot(class_pxs[1,:],log10.(abs.(class_pxsnf[10,:])),
-     label=raw"$\Theta_{0,\rm{CLASS}}$",
+     label=raw"$\Theta_{1,\rm{CLASS}}$",
      legend=:topleft)
 plot!(x_grid, log10.(abs.(results_with_rsa[2,:] * dipole_fac)),
-      label=raw"$4 h \Theta_{0,\rm{Bolt}}$",ls=:dash)
+      label=raw"$4 \Theta_{1,\rm{Bolt,rsa}}$",ls=:dash)
+plot!(x_grid, log10.(abs.(results[2,:] * dipole_fac)),
+      label=raw"$4 \Theta_{1,\rm{Bolt}}$",ls=:dash)
 vline!([this_rsa_switch],label="RSA switch",ls=:dot)
 xlims!(-5,1)
 xlims!(-7,0)
 
 #neutrino 𝒩0 monopole
 plot(class_pxs[1,:],log10.(abs.(class_pxsnf[5,:])),
       label=raw"$\mathcal{N}_{0,\rm{CLASS}}$",legend=:topleft)
-plot!(x_grid, log10.(abs.(results_with_rsa[1+2(ℓᵧ+1),:]* 𝕡.h*4)),
-      label=raw"$4 h \mathcal{N}_{0,\rm{Bolt}}$",ls=:dash)
+plot!(x_grid, log10.(abs.(results_with_rsa[1+2(ℓᵧ+1),:]*4)),
+      label=raw"$4 \mathcal{N}_{0,\rm{Bolt}}$",ls=:dash)
 vline!([this_rsa_switch],label="RSA switch",ls=:dot)
 ylims!(-.2,1)
 xlims!(-7,0)
 #neutrino 𝒩1 dipole
 plot(class_pxs[1,:],log10.(abs.(class_pxs[13,:])),
       label=raw"$\mathcal{N}_{0,\rm{CLASS}}$",legend=:topleft)
 plot!(x_grid, log10.(abs.(results_with_rsa[2(ℓᵧ+1)+2,:]* dipole_fac)),
-      label=raw"$4 h \mathcal{N}_{0,\rm{Bolt}}$",ls=:dash)
+      label=raw"$4 \mathcal{N}_{0,\rm{Bolt}}$",ls=:dash)
 vline!([this_rsa_switch],label="RSA switch",ls=:dot)
 xlims!(-5,1)
 
@@ -124,15 +128,15 @@ ours_2_Mpcm1 = 1/(2.1331e-35 *3e5) #unit conversion 1/([km/s/Mpc]*[c/km/s])
 plot(class_pxsnf[1,:],log10.(abs.(class_pxsnf[4,:])),
      label=raw"$\delta_{c,\rm{CLASS}}$",
      legend=:topleft)
-plot!(x_grid,log10.(results[2(ℓᵧ+1)+(ℓ_ν+1)+(ℓ_mν+1)*n_q+2,:]* 𝕡.h),
+plot!(x_grid,log10.(results[2(ℓᵧ+1)+(ℓ_ν+1)+(ℓ_mν+1)*n_q+2,:]),
       label=raw"$h \delta_{\rm{Bolt}}$",ls=:dash)
 
 #baryon δ_b
 plot(class_pxsnf[1,:],log10.(abs.(class_pxsnf[3,:])),
     label=raw"$\delta_{b,\rm{CLASS}}$",legend=:topleft)
-plot!(x_grid,log10.(abs.(results[2(ℓᵧ+1)+(ℓ_ν+1)+(ℓ_mν+1)*n_q+4,:]* 𝕡.h)),
+plot!(x_grid,log10.(abs.(results[2(ℓᵧ+1)+(ℓ_ν+1)+(ℓ_mν+1)*n_q+4,:])),
     label=raw"$h \delta_{b,\rm{Bolt}}$",ls=:dash)
-xlims!(-8,-4)
+xlims!(-8,0)
 
 #massive neutrino monopole ℳ0
 plot(class_pxsnf[1,:],log10.(abs.(class_pxsnf[6,:])),
@@ -141,7 +145,7 @@ plot(class_pxsnf[1,:],log10.(abs.(class_pxsnf[6,:])),
 plot!(class_pxs[1,:],log10.(abs.(class_pxs[6,:])),
     label=raw"$m\nu_{0,\rm{CLASS,f}}$",
     ls=:dot)
-plot!(x_grid, log10.(abs.(ℳρ* 𝕡.h)),
+plot!(x_grid, log10.(abs.(ℳρ)),
     label=raw"$h m\nu_{0,\rm{Bolt}}$",ls=:dash)
     #ls=:dot)
 vline!([xhor],ls=:dot,c=:black,label=raw"$k/(2\pi a H h)=1$")

src/Bolt.jl

@@ -6,8 +6,7 @@ export IonizationHistory, AbstractIonizationHistory, IonizationIntegrator
 export Peebles, PeeblesI
 export ρ_σ,ρP_0,f0,dlnf0dlnq,θ,oldH_a #FIXME: quick hack to look at perts
 export Hierarchy, boltsolve, BasicNewtonian,unpack,rsa_perts!,boltsolve_rsa
-export IE,initial_conditions,unpack,ie_unpack
-export source_grid, quadratic_k, cltt,log10_k,plin
+export source_grid, quadratic_k, cltt,log10_k,plin, source_grid_P, clte,clee
 export z2a, a2z, x2a, a2x, z2x, x2z, to_ui, from_ui, dxdq
 
 using Parameters
@@ -49,6 +48,8 @@ unnatural(x, y) = UnitfulCosmo.unmpc(x, y)
 # all unit conversions. should distribute these in-situ someday. Mpc units
 const km_s_Mpc_100 = ustrip(natural(100.0u"km/s/Mpc"))  # [Mpc^-1]
 const G_natural = ustrip(natural(float(NewtonianConstantOfGravitation))) # [Mpc^2]
+const mass_natural = ustrip(UnitfulCosmo.mpc(1.0u"eV")) # [Mpc^-1] #FIXME? Might want to put neutrino conversion elsewhere, but need to convert eV
+
 
 abstract type AbstractCosmoParams{T} end
 
@@ -61,7 +62,7 @@ abstract type AbstractCosmoParams{T} end
     n = 1.0  # scalar spectral index
     Y_p = 0.24  # primordial helium fraction
     N_ν = 3.046 #effective number of relativisic species (PDG25 value)
-    Σm_ν = 0.06 #sum of neutrino masses (eV), Planck 15 default ΛCDM value
+    Σm_ν = 0.06*mass_natural #sum of neutrino masses (eV), Planck 15 default ΛCDM value
 end
 
 include("util.jl")

src/perturbations.jl

@@ -362,7 +362,7 @@ function source_function(du, u, hierarchy::Hierarchy{T, BasicNewtonian}, x) wher
     _, σℳ′ = ρ_σ(ℳ′[0:nq-1], ℳ′[2*nq:3*nq-1], bg, a, par)
     Ψ = -Φ - 12H₀² / k^2 / a^2 * (par.Ω_r * Θ[2]
                                   + Ω_ν * 𝒩[2] #add rel quadrupole
-                                  + σℳ / bg.ρ_crit) #why am I doing this? - because H0 pulls out a factor of rho crit - just unit conversion
+                                  + σℳ / bg.ρ_crit /4) #why am I doing this? - because H0 pulls out a factor of rho crit - just unit conversion
                                                                    #this introduces a factor of bg density I cancel using the integrated bg mnu density now
 
    Ψ′ = -Φ′ - 12H₀² / k^2 / a^2 * (par.Ω_r * (Θ′[2] - 2 * Θ[2])
@@ -381,3 +381,24 @@ function source_function(du, u, hierarchy::Hierarchy{T, BasicNewtonian}, x) wher
         ℋₓ^2 * (g̃ₓ′′ * Π + 2g̃ₓ′ * Π′ + g̃ₓ * Π′′))
     return term1 + term2 + term3
 end
+
+# The polarization source function (jms 6/6/22 UNTESTED!) SZ eqn 12d (in our units)
+function source_function_P(du, u, hierarchy::Hierarchy{T, BasicNewtonian}, x) where T
+    k, ℓᵧ, par, bg, ih,nq = hierarchy.k, hierarchy.ℓᵧ, hierarchy.par, hierarchy.bg, hierarchy.ih,hierarchy.nq
+    H₀² = bg.H₀^2
+    ℋₓ, ℋₓ′, ℋₓ′′ = bg.ℋ(x), bg.ℋ′(x), bg.ℋ′′(x)
+    τₓ, τₓ′, τₓ′′ = ih.τ(x), ih.τ′(x), ih.τ′′(x)
+    g̃ₓ, g̃ₓ′, g̃ₓ′′ = ih.g̃(x), ih.g̃′(x), ih.g̃′′(x)
+    a = x2a(x)
+    ρ0ℳ = bg.ρ₀ℳ(x) #get current value of massive neutrino backround density from spline
+    Tν =  (par.N_ν/3)^(1/4) *(4/11)^(1/3) * (15/ π^2 *ρ_crit(par) *par.Ω_r)^(1/4)
+    Ω_ν =  7*(2/3)*par.N_ν/8 *(4/11)^(4/3) *par.Ω_r
+    logqmin,logqmax=log10(Tν/30),log10(Tν*30)
+    q_pts = xq2q.(bg.quad_pts,logqmin,logqmax)
+    Θ, Θᵖ, 𝒩, ℳ, Φ, δ, v, δ_b, v_b = unpack(u, hierarchy)  # the Θ, Θᵖ are mutable views (see unpack)
+    Θ′, Θᵖ′, 𝒩′, ℳ′, Φ′, δ′, v′, δ_b′, v_b′ = unpack(du, hierarchy)
+
+    y = k*(bg.η(bg.x_grid[end]) - bg.η(x))
+    Π = Θ[2] + Θᵖ[2] + Θᵖ[0]
+    return (3/(4y^2)) * g̃ₓ * Π 
+end