From 5f71e0f8abae3a749f96ce1b3dd3596073ccf6f4 Mon Sep 17 00:00:00 2001 From: Shamik Ghosh Date: Fri, 21 Mar 2025 13:37:08 -0700 Subject: [PATCH 1/4] First working version for Planck --- .../phase2/planck/rotate_and_combine_maps.py | 225 ++++++++++++++++++ .../planck/rotate_and_combine_maps.slurm | 17 ++ 2 files changed, 242 insertions(+) create mode 100644 chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py create mode 100755 chile_optimization_sims/phase2/planck/rotate_and_combine_maps.slurm diff --git a/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py b/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py new file mode 100644 index 0000000..11656f6 --- /dev/null +++ b/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py @@ -0,0 +1,225 @@ +# 2025-03-20: Original version by Shamik Ghosh + +import glob +import os +import sys + +import astropy.units as u +import healpy as hp +from mpi4py import MPI +import numpy as np +import os + +comm = MPI.COMM_WORLD +ntask = comm.size +rank = comm.rank + +if rank == 0: + print(f"Running with {ntask} MPI tasks") +prefix = f"{rank:04} : " + +npipe_noise_dir = '/global/cfs/cdirs/cmb/data/planck2020/npipe/npipe6v20_sim/residual/' +beam_dir = '/global/cfs/cdirs/cmb/data/planck2020/npipe/npipe6v20/quickpol/' +fgdir = "/global/cfs/cdirs/cmb/gsharing/panexp_v1_planck/galactic_foregrounds_mediumcomplexity/" +cmbdir = "/global/cfs/cdirs/cmb/data/generic/cmb/ffp10/mc/scalar/" +tensordir = "/global/cfs/cdirs/cmb/data/generic/cmb/ffp10/mc/tensor/" +outputdir = "/pscratch/sd/s/shamikg/chile_optimization_sims/planck/" + +os.makedirs(outputdir, exist_ok=True) +os.makedirs(f"{outputdir}cmb/", exist_ok=True) +os.makedirs(f"{outputdir}noise/", exist_ok=True) +os.makedirs(f"{outputdir}total/", exist_ok=True) + +npipe_noise_ini = 200 + +freqs = [ + "030", + "044", + "070", + "100", + "143", + "217", + "353", + "545", + "857", +] + +rot = hp.Rotator(coord=('G', 'C')) + +ijob = -1 +for band in freqs: + + print(prefix + f"band = {band} GHz") + + # Get the foregrounds + + nside = 2048 + lmin = 0 + lmax = 2 * nside + + fname_Bl = f"{beam_dir}/Bl_TEB_npipe6v20_{band}GHzx{band}GHz.fits" + if int(band) < 100: + fname_fg = f"{fgdir}" \ + f"sobs_mbs-s0017-20250208_LFI_mission_{band}_galactic_foregrounds_mediumcomplexity_healpix.fits" + else: + fname_fg = f"{fgdir}" \ + f"sobs_mbs-s0017-20250208_HFI_mission_{band}_galactic_foregrounds_mediumcomplexity_healpix.fits" + + fg = None + Bl = None + for mc in range(100): + + ijob += 1 + if ijob % ntask != rank: + continue + + # Check if the last maps for this MC were already written + if int(band) < 100: + fname_total = f"{outputdir}total/planck_total_{band}_mc_{mc:04}.fits" + if os.path.isfile(fname_total): + print(prefix + f" {fname_total} exists, skipping mc = {mc}") + continue + else: + print(prefix + f" {fname_total} does not exist.") + elif int(band) >= 100: + fname_total5 = f"{outputdir}total/planck_total_{band}_mc_{mc:04}.fits" + if os.path.isfile(fname_total5): + print(prefix + f" {fname_total5} exists, skipping mc = {mc}") + continue + + # Proceed with the simulation + + if fg is None: + # First time foregrounds are needed + print(prefix + f" Reading {fname_fg}") + fg = hp.read_map(fname_fg, field=None) * 1e-6 # to K_CMB + print(prefix + f" Expanding foreground in alm") + fg_alms = hp.map2alm(fg, lmax=lmax, use_pixel_weights=True) + + + # Get the beam + if Bl is None: + print(prefix + f" Reading {fname_Bl}") + Bl = hp.read_cl(fname_Bl) + + + # Process all survey lengths for this band and this MC + + # Load the CMB realization + + print(prefix + f" MC = {mc}") + fname_cmb = f"{cmbdir}/ffp10noaber_lensed_scl_cmb_000_alm_mc_{mc:04}.fits" + print(prefix + f" Reading {fname_cmb}") + alms = hp.read_alm(fname_cmb, hdu=(1, 2, 3)) + lmax_in = hp.Alm.getlmax(alms[0].size) + if lmax != lmax_in: + alms = hp.resize_alm(alms, lmax_in, lmax_in, lmax, lmax) + + # Add primordial B-modes to even-numbered MC + + # if mc % 2 == 0: + # fname_tensor = f"{tensordir}/ffp10_ten_cmb_000_alm_mc_{mc:04}.fits" + # print(prefix + f" Reading {fname_tensor}") + # tensor_alms = hp.read_alm(fname_tensor, hdu=(1, 2, 3)) + # lmax_in = hp.Alm.getlmax(tensor_alms[0].size) + # if lmax != lmax_in: + # tensor_alms = hp.resize_alm(tensor_alms, lmax_in, lmax_in, lmax, lmax) + # r_in = 0.01 # FFP10 simulations have r=0.01 + # r_out = 0.003 # We want r = 0.003 + # scale = np.sqrt(r_out / r_in) # Power spectrum ratios in map domain + # for i in range(3): + # alms[i] += scale * tensor_alms[i] + # r = r_out + # else: + fname_tensor = None + r = 0 + + for i in range(3): + hp.almxfl(alms[i], Bl[i], inplace=True) + + # Make a pure CMB map for reference + + print(prefix + " Synthesizing CMB map") + cmb = hp.alm2map(alms, nside) # K_CMB + fname_cmb = f"{outputdir}cmb/planck_cmb_{band}_mc_{mc:04}.fits" + print(prefix + f" Writing {fname_cmb}") + os.makedirs(os.path.dirname(fname_cmb), exist_ok=True) + hp.write_map( + fname_cmb, + cmb.astype(np.float32), + dtype=np.float32, + coord="C", + column_units="K_CMB", + extra_header=[ + ("BEAM", f"NPIPE TEB {band}GHzx{band}GHz"), + ("CMB", fname_cmb), + ("TENSOR", fname_tensor), + ("R", r, "tensor-scalar ratio"), + ], + overwrite=True, + ) + + # Combine the foregrounds with the CMB and translate into a map + # UPDATE 2024-11-20: Beam smoothing and pixel window smoothing are now not + # performed on the total signal alms. + for i in range(3): + alms[i] += fg_alms[i] + print(prefix + " Synthesizing sky map") + sky = hp.alm2map(alms, nside) # K_CMB + + + # Read and rotate noise maps + noise_mc = npipe_noise_ini+mc + fname_npipe = f"{npipe_noise_dir}{noise_mc:04}residual_npipe6v20_{band}_{noise_mc:04}.fits" + print(prefix + f" Reading {fname_npipe}") + noise_npipe = hp.read_map(fname_npipe, field=None) + noise_npipe[noise_npipe == hp.UNSEEN] = 0. + noise_npipe = rot.rotate_map_pixel(noise_npipe) + + fname_noise = f"{outputdir}noise/planck_noise_{band}_mc_{mc:04}.fits" + print(prefix + f" Writing {fname_noise}") + os.makedirs(os.path.dirname(fname_noise), exist_ok=True) + + args = { + "dtype" : np.float32, + "coord" : "C", + "column_units" : "K_CMB", + "extra_header" : [ + ("NOISE", f"Planck NPIPE {band}"), + ("NPIPE MC", noise_mc), + ("NPIPE FILE", os.path.basename(fname_npipe)), + ], + "overwrite" : True, + } + hp.write_map(fname_noise, noise_npipe.astype(np.float32), **args) + + + # coadd all maps and save to file + total = sky + noise_npipe + fname_total = f"{outputdir}total/planck_total_{band}_mc_{mc:04}.fits" + os.makedirs(os.path.dirname(fname_total), exist_ok=True) + args = { + "dtype" : np.float32, + "coord" : "C", + "column_units" : "K_CMB", + "extra_header" : [ + ("BEAM", f"NPIPE TEB {band}GHzx{band}GHz"), + ("FOREGRND", os.path.basename(fname_fg)), + ("CMB", os.path.basename(fname_cmb)), + ("TENSOR", os.path.basename(fname_tensor)), + ("R", r, "tensor-scalar ratio"), + ("LMIN_CMB", lmin, "High-pass cut-off"), + ("NOISE", os.path.basename(fname_noise)), + ], + "overwrite" : True, + } + print(prefix + f" Writing {fname_total}") + hp.write_map(fname_total, total.astype(np.float32), **args) + +comm.Barrier() +if rank == 0: + print("All done!") +comm.Barrier() + + + diff --git a/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.slurm b/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.slurm new file mode 100755 index 0000000..9e950b7 --- /dev/null +++ b/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.slurm @@ -0,0 +1,17 @@ +#!/bin/bash +#SBATCH --qos=debug +#SBATCH --time=00:30:00 +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=16 +#SBATCH --cpus-per-task=16 +#SBATCH --job-name=planck_rotate_combine +#SBATCH --licenses=SCRATCH +#SBATCH --constraint=cpu +#SBATCH --account=mp107a + +export OMP_PROC_BIND=spread +export OMP_PLACES=threads +export OMP_NUM_THREADS=8 + + +srun python -u rotate_and_combine_maps.py \ No newline at end of file From 866c7f7277f70fbb303f65daa9a038ad82780258 Mon Sep 17 00:00:00 2001 From: Shamik Ghosh Date: Fri, 21 Mar 2025 13:56:35 -0700 Subject: [PATCH 2/4] Fix npipe noise path --- .../phase2/planck/rotate_and_combine_maps.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py b/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py index 11656f6..5ee9a1f 100644 --- a/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py +++ b/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py @@ -18,7 +18,7 @@ print(f"Running with {ntask} MPI tasks") prefix = f"{rank:04} : " -npipe_noise_dir = '/global/cfs/cdirs/cmb/data/planck2020/npipe/npipe6v20_sim/residual/' +npipe_noise_dir = '/global/cfs/cdirs/cmb/data/planck2020/npipe/npipe6v20_sim/' beam_dir = '/global/cfs/cdirs/cmb/data/planck2020/npipe/npipe6v20/quickpol/' fgdir = "/global/cfs/cdirs/cmb/gsharing/panexp_v1_planck/galactic_foregrounds_mediumcomplexity/" cmbdir = "/global/cfs/cdirs/cmb/data/generic/cmb/ffp10/mc/scalar/" @@ -170,7 +170,7 @@ # Read and rotate noise maps noise_mc = npipe_noise_ini+mc - fname_npipe = f"{npipe_noise_dir}{noise_mc:04}residual_npipe6v20_{band}_{noise_mc:04}.fits" + fname_npipe = f"{npipe_noise_dir}{noise_mc:04}/residual/residual_npipe6v20_{band}_{noise_mc:04}.fits" print(prefix + f" Reading {fname_npipe}") noise_npipe = hp.read_map(fname_npipe, field=None) noise_npipe[noise_npipe == hp.UNSEEN] = 0. From b7405394ea84013090dec0cff86567902ee93487 Mon Sep 17 00:00:00 2001 From: Shamik Ghosh Date: Fri, 21 Mar 2025 15:59:20 -0700 Subject: [PATCH 3/4] Fix bugs with fits header --- .../phase2/planck/rotate_and_combine_maps.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py b/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py index 5ee9a1f..ae6a01d 100644 --- a/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py +++ b/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py @@ -52,16 +52,18 @@ print(prefix + f"band = {band} GHz") # Get the foregrounds - - nside = 2048 - lmin = 0 - lmax = 2 * nside fname_Bl = f"{beam_dir}/Bl_TEB_npipe6v20_{band}GHzx{band}GHz.fits" if int(band) < 100: + nside = 1024 + lmin = 0 + lmax = 2 * nside fname_fg = f"{fgdir}" \ f"sobs_mbs-s0017-20250208_LFI_mission_{band}_galactic_foregrounds_mediumcomplexity_healpix.fits" else: + nside = 2048 + lmin = 0 + lmax = 2 * nside fname_fg = f"{fgdir}" \ f"sobs_mbs-s0017-20250208_HFI_mission_{band}_galactic_foregrounds_mediumcomplexity_healpix.fits" @@ -174,7 +176,8 @@ print(prefix + f" Reading {fname_npipe}") noise_npipe = hp.read_map(fname_npipe, field=None) noise_npipe[noise_npipe == hp.UNSEEN] = 0. - noise_npipe = rot.rotate_map_pixel(noise_npipe) + noise_npipe = np.array(rot.rotate_map_pixel(noise_npipe)) + print(noise_npipe.shape) fname_noise = f"{outputdir}noise/planck_noise_{band}_mc_{mc:04}.fits" print(prefix + f" Writing {fname_noise}") @@ -206,13 +209,18 @@ ("BEAM", f"NPIPE TEB {band}GHzx{band}GHz"), ("FOREGRND", os.path.basename(fname_fg)), ("CMB", os.path.basename(fname_cmb)), - ("TENSOR", os.path.basename(fname_tensor)), ("R", r, "tensor-scalar ratio"), ("LMIN_CMB", lmin, "High-pass cut-off"), ("NOISE", os.path.basename(fname_noise)), ], "overwrite" : True, } + + if fname_tensor is not None: + args["extra_header"].append(("TENSOR", os.path.basename(fname_tensor))) + else: + args["extra_header"].append(("TENSOR", "None")) + print(prefix + f" Writing {fname_total}") hp.write_map(fname_total, total.astype(np.float32), **args) From 1d0e7f10871b796eabfd301d544943eed13f947c Mon Sep 17 00:00:00 2001 From: Shamik Ghosh Date: Fri, 21 Mar 2025 16:44:40 -0700 Subject: [PATCH 4/4] Defer unpolarizaed channels --- .../phase2/planck/rotate_and_combine_maps.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py b/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py index ae6a01d..613be74 100644 --- a/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py +++ b/chile_optimization_sims/phase2/planck/rotate_and_combine_maps.py @@ -1,3 +1,5 @@ +# TO-DO: Add 545 and 857 channels which are not polarized and might have different units + # 2025-03-20: Original version by Shamik Ghosh import glob @@ -40,8 +42,8 @@ "143", "217", "353", - "545", - "857", + # "545", + # "857", ] rot = hp.Rotator(coord=('G', 'C')) @@ -53,7 +55,11 @@ # Get the foregrounds - fname_Bl = f"{beam_dir}/Bl_TEB_npipe6v20_{band}GHzx{band}GHz.fits" + if int(band) < 500: + fname_Bl = f"{beam_dir}/Bl_TEB_npipe6v20_{band}GHzx{band}GHz.fits" + else: + fname_Bl = f"{beam_dir}/Bl_npipe6v20_{band}GHzx{band}GHz.fits" + if int(band) < 100: nside = 1024 lmin = 0