minor rSZ cleanup

examples/make_map_raw_car_martine.jl

@@ -0,0 +1,204 @@
+# modified from Zack Li
+
+using Pixell, WCS, XGPaint
+using Cosmology
+using Interpolations
+import XGPaint: AbstractProfile
+using HDF5
+import JSON
+using JLD2
+using ThreadsX
+using FileIO
+using Healpix
+
+print("Threads: ", Threads.nthreads(), "\n")
+modeltype::String = "break"
+healpix = true
+
+
+if healpix
+    nside = 1024
+    m = HealpixMap{Float64,RingOrder}(nside)
+else
+    shape, wcs = fullsky_geometry(0.5 * Pixell.arcminute)
+    # for a box:
+    # box = [30   -30;           # RA
+    #        -25     25] * Pixell.degree  # DEC
+    # shape, wcs = geometry(CarClenshawCurtis{Float64}, box, 0.5 * Pixell.arcminute)
+    m = Enmap(zeros(shape), wcs)
+    # precomputed sky angles
+    α_map, δ_map = posmap(shape, wcs)
+    psa = (sin_α=sin.(α_map), cos_α=cos.(α_map), sin_δ=sin.(δ_map), cos_δ=cos.(δ_map))
+end
+
+##
+fid = h5open("/mnt/raid-cita/mlokken/buzzard/catalogs/halos/buzzard_halos.hdf5", "r")
+# fid = h5open("/mnt/scratch-lustre/mlokken/pkpatch/halos_hdf5.h5", "r")
+ra, dec = deg2rad.(fid["ra"]), deg2rad.(fid["dec"])
+redshift = collect(fid["z"])
+halo_mass = collect(fid["m200c"])
+# choose the cutoff for the pressure profiles here
+cutoff = 4 # the standard
+apply_beam = true
+##
+perm = sortperm(dec, alg=ThreadsX.MergeSort)
+ra = ra[perm]
+dec = dec[perm]
+redshift = redshift[perm]
+halo_mass = halo_mass[perm]
+
+
+
+##
+print("Precomputing the model profile grid.\n")
+
+# set up a profile to paint
+if modeltype=="break"
+    p = XGPaint.BreakModel(Omega_c=0.24, Omega_b=0.046, h=0.7, alpha_break=1.5) # Buzzard cosmology
+elseif modeltype=="battaglia"
+    p = XGPaint.BattagliaProfile(Omega_c=0.24, Omega_b=0.046, h=0.7) # Buzzard cosmology
+end
+
+# beam stuff
+N_logθ = 512
+rft = RadialFourierTransform(n=N_logθ, pad=256)
+beamstr = ""
+
+if apply_beam
+    ells = rft.l
+    fwhm = deg2rad(1.6/60.)  # ACT beam
+    σ² = fwhm^2 / (8log(2))
+    lbeam = @. exp(-ells * (ells+1) * σ² / 2)
+    beamstr = "_1p6arcmin"
+end
+model_file::String = "cached_$(modeltype)_$(beamstr).jld2"
+if isfile(model_file)
+    print("Found cached Battaglia profile model. Loading from disk.\n")
+    model = load(model_file)
+    prof_logθs, prof_redshift, prof_logMs, prof_y = model["prof_logθs"], 
+        model["prof_redshift"], model["prof_logMs"], model["prof_y"]
+else
+    print("Didn't find a cached profile model. Computing and saving.\n")
+    logθ_min, logθ_max = log(minimum(rft.r)), log(maximum(rft.r))
+    @time prof_logθs, prof_redshift, prof_logMs, prof_y = profile_grid(p; 
+        N_logθ=length(rft.r), logθ_min=logθ_min, logθ_max=logθ_max)
+
+    if apply_beam
+        # now apply the beam
+        @time XGPaint.transform_profile_grid!(prof_y, rft, lbeam)
+        prof_y = prof_y[begin+rft.pad:end-rft.pad, :, :]
+        prof_logθs = prof_logθs[begin+rft.pad:end-rft.pad]
+        XGPaint.cleanup_negatives!(prof_y)
+    end
+    jldsave(model_file; prof_logθs, prof_redshift, prof_logMs, prof_y)
+
+end
+
+if ~healpix # do this for CAR
+    itp = Interpolations.interpolate(log.(prof_y), BSpline(Cubic(Line(OnGrid()))))
+    sitp = scale(itp, prof_logθs, prof_redshift, prof_logMs);
+end
+
+
+##
+function paint_map!(m::Enmap, p::XGPaint.AbstractProfile, psa, sitp, masses, 
+                    redshifts, αs, δs, irange; mult=4)
+    for i in irange
+        α₀ = αs[i]
+        δ₀ = δs[i]
+        mh = masses[i]
+        z = redshifts[i]
+        θmax = XGPaint.θmax(p, mh * XGPaint.M_sun, z, mult=mult)
+        profile_paint!(m, α₀, δ₀, psa, sitp, z, mh, θmax)
+    end
+end
+
+function paint_map!(m::HealpixMap, p::XGPaint.AbstractProfile{T}, masses, 
+    redshifts, αs, δs, irange; mult=4) where {T}
+    for i in irange
+        α₀ = αs[i]
+        δ₀ = δs[i]
+        mh = masses[i]
+        z = redshifts[i]
+        θmax = XGPaint.θmax(p, mh * XGPaint.M_sun, z, mult=mult)
+        θmin = 0.0 # don't know what this should be # error
+        beam_interp = XGPaint.realspacegaussbeam(T,fwhm)[0] #error
+        w = XGPaint.HealpixPaintingWorkspace(nside, θmin, θmax, beam_interp)
+        profile_paint!(m, α₀, δ₀, w, z, mh, θmax)    
+    end
+end
+
+function chunked_paint!(m::Enmap, p::XGPaint.AbstractProfile, psa, sitp, masses, 
+                        redshifts, αs, δs; mult=4)
+    m .= 0.0
+
+    N_sources = length(masses)
+    chunksize = ceil(Int, N_sources / (2Threads.nthreads()))
+    chunks = chunk(N_sources, chunksize);
+
+    Threads.@threads for i in 1:Threads.nthreads()
+        chunk_i = 2i
+        i1, i2 = chunks[chunk_i]
+        paint_map!(m, p, psa, sitp, masses, redshifts, αs, δs, i1:i2, mult=mult)
+    end
+
+    Threads.@threads for i in 1:Threads.nthreads()
+        chunk_i = 2i - 1
+        i1, i2 = chunks[chunk_i]
+        paint_map!(m, p, psa, sitp, masses, redshifts, αs, δs, i1:i2, mult=mult)
+    end
+end
+
+function chunked_paint!(m::HealpixMap, p::XGPaint.AbstractProfile, masses, 
+                        redshifts, αs, δs; mult=4)
+    m .= 0.0
+
+    N_sources = length(masses)
+    chunksize = ceil(Int, N_sources / (2Threads.nthreads()))
+    chunks = chunk(N_sources, chunksize);
+
+    Threads.@threads for i in 1:Threads.nthreads()
+        chunk_i = 2i
+        i1, i2 = chunks[chunk_i]
+        paint_map!(m, p, masses, redshifts, αs, δs, i1:i2, mult=mult)
+    end
+
+    Threads.@threads for i in 1:Threads.nthreads()
+        chunk_i = 2i - 1
+        i1, i2 = chunks[chunk_i]
+        paint_map!(m, p,  masses, redshifts, αs, δs, i1:i2, mult=mult)
+    end
+end
+##
+# cut = eachindex(halo_mass)[begin:end-5]
+
+print(maximum(ra), minimum(ra), maximum(dec), minimum(dec))
+print("Painting map.\n")
+if healpix
+    @time chunked_paint!(m, p, halo_mass, redshift, ra, dec, mult=cutoff) # for a Healpix
+else
+    @time chunked_paint!(m, p, psa, sitp, halo_mass, redshift, ra, dec, mult=cutoff) # for an Enmap
+end
+
+#
+write_map(
+    "/mnt/raid-cita/mlokken/buzzard/ymaps/ymap_buzzard_$(modeltype)_bbps_car$(beamstr)_cutoff$(cutoff).fits",
+    m)
+
+using PyPlot
+plt.clf()
+plt.figure()
+plt.imshow(log10.(m.data'))
+plt.axis("off")
+plt.savefig("test_Healpix.png", bbox_inches="tight",pad_inches = 0)
+plt.gcf()
+
+
+
+# ##
+# m .= 0
+# profile_paint!(m, 0., 0., psa, sitp, 0.1, 1e14, π/90)
+
+# ##
+# using Plots
+# Plots.plot(log10.(m))

src/XGPaint.jl

@@ -18,9 +18,13 @@ using JLD2, FileIO
 import ThreadsX
 import Distributions
 import DataInterpolations  # used for uneven spaced interpolators
-
 using DelimitedFiles
 
+import PhysicalConstants.CODATA2018 as constants
+const M_sun = 1.98847e30u"kg"
+const T_cmb =  2.725 * u"K"
+const P_e_factor = constants.σ_e / (constants.m_e * constants.c_0^2)
+
 include("./util.jl")
 include("./model.jl")
 include("./profiles.jl")
@@ -35,5 +39,6 @@ export get_cosmology, read_halo_catalog_hdf5, sort_halo_catalog
 export Radio_Sehgal2009, CIB_Planck2013, CIB_Scarfy, CO_CROWNED, LRG_Yuan23
 export paint!, generate_sources, process_sources, profile_grid, profile_paint!, profileworkspace, paint_szp!, profile_grid_szp, profile_paint_szp!, paint_rsz!, profile_grid_rsz, profile_paint_rsz!
 export build_interpolator, Battaglia16ThermalSZProfile, Battaglia16RelativisticSZProfile, build_interpolator_szp, build_interpolator_rsz
+export Battaglia16SZPackProfile
 
 end # module

src/profiles.jl

@@ -1,13 +1,4 @@
 
-import PhysicalConstants.CODATA2018 as constants
-using Unitful
-const M_sun = 1.98847e30u"kg"
-const T_cmb =  2.725 * u"K"
-const P_e_factor = constants.σ_e / (constants.m_e * constants.c_0^2)
-using Cosmology
-using QuadGK
-
-
 
 # RECTANGULAR WORKSPACES
 
@@ -63,9 +54,6 @@ function Battaglia16ThermalSZProfile(; Omega_c::T=0.2589, Omega_b::T=0.0486, h::
     return Battaglia16ThermalSZProfile(f_b, cosmo)
 end
 
-abstract type AbstractPaintingProblem{T} end
-
-
 function BreakModel(; Omega_c::T=0.2589, Omega_b::T=0.0486, h::T=0.6774, alpha_break::T=1.5, M_break::T=2.0*10^14) where {T <: Real}
     #alpha_break = 1.486 from Shivam P paper by Nate's sleuthing
     OmegaM=Omega_b+Omega_c

src/profiles_rsz.jl

@@ -1,13 +1,4 @@
 
-import PhysicalConstants.CODATA2018 as constants
-using Unitful
-const M_sun = 1.98847e30u"kg"
-const P_e_factor = constants.σ_e / (constants.m_e * constants.c_0^2)
-const T_cmb =  2.725 * u"K"
-using Cosmology
-using QuadGK
-
-
 struct Battaglia16RelativisticSZProfile{T,C} <: AbstractGNFW{T}
     f_b::T  # Omega_b / Omega_c = 0.0486 / 0.2589
     cosmo::C
@@ -89,7 +80,6 @@ function rSZ(𝕡, M_200, z, r)
     I = (X^3/(ℯ^X-1)) * (2*(2π)^4*(constants.k_B*T_cmb)^3)/((constants.h*constants.c_0)^2) * n 
     T = I/abs((2 * constants.h^2 * ω^4 * ℯ^X)/(constants.k_B * constants.c_0^2 * T_cmb * (ℯ^X - 1)^2))
 
-    #return T
     return abs(T)
 end