diff --git a/project/ana_cov_comp/paramfiles/cov_dr6_v4_20231128.dict b/project/ana_cov_comp/paramfiles/cov_dr6_v4_20231128.dict index caff04cc..a559a782 100644 --- a/project/ana_cov_comp/paramfiles/cov_dr6_v4_20231128.dict +++ b/project/ana_cov_comp/paramfiles/cov_dr6_v4_20231128.dict @@ -1,4 +1,6 @@ -data_dir = '/scratch/gpfs/zatkins/data/simonsobs/PSpipe/project/ana_cov_comp/cov_dr6_v4_20231128/' +iam = 'zatkins_on_della' # 'xzackli_on_perlmutter' + +data_dir = {'zatkins_on_della': '/scratch/gpfs/zatkins/data/simonsobs/PSpipe/project/ana_cov_comp/cov_dr6_v4_20231128/', 'xzackli_on_perlmutter': '/pscratch/sd/x/xzackli/so/data/dr6v4/'}[iam] surveys = ["dr6"] arrays_dr6 = {'pa4': ['f220'], 'pa5': ['f090', 'f150'], 'pa6': ['f090', 'f150']} @@ -47,6 +49,8 @@ couplings_dir = data_dir + 'couplings' sq_win_alms_dir = data_dir + 'sq_win_alms' covariances_dir = data_dir + 'covariances' split_covariances_dir = data_dir + 'split_covariances' +noise_model_savgol_w = 75 +noise_model_savgol_k = 4 use_toeplitz_cov = True @@ -78,31 +82,29 @@ edge_skip_rescale = 1 cross_link_threshold = 0.97 n_med_ivar = 3 -window_dir = '/pscratch/sd/x/xzackli/so/data/windows/' - -window_T_dr6_pa4_f150 = window_dir + "window_dr6_pa4_f150_baseline.fits" -window_pol_dr6_pa4_f150 = window_dir + "window_dr6_pa4_f150_baseline.fits" -window_kspace_dr6_pa4_f150 = window_dir + "kspace_mask_dr6_pa4_f150.fits" +window_T_dr6_pa4_f150 = data_dir + "windows/window_dr6_pa4_f150_baseline.fits" +window_pol_dr6_pa4_f150 = data_dir + "windows/window_dr6_pa4_f150_baseline.fits" +window_kspace_dr6_pa4_f150 = data_dir + "windows/kspace_mask_dr6_pa4_f150.fits" -window_T_dr6_pa4_f220 = window_dir + "window_dr6_pa4_f220_baseline.fits" -window_pol_dr6_pa4_f220 = window_dir + "window_dr6_pa4_f220_baseline.fits" -window_kspace_dr6_pa4_f220 = window_dir + "kspace_mask_dr6_pa4_f220.fits" +window_T_dr6_pa4_f220 = data_dir + "windows/window_dr6_pa4_f220_baseline.fits" +window_pol_dr6_pa4_f220 = data_dir + "windows/window_dr6_pa4_f220_baseline.fits" +window_kspace_dr6_pa4_f220 = data_dir + "windows/kspace_mask_dr6_pa4_f220.fits" -window_T_dr6_pa5_f090 = window_dir + "window_dr6_pa5_f090_baseline.fits" -window_pol_dr6_pa5_f090 = window_dir + "window_dr6_pa5_f090_baseline.fits" -window_kspace_dr6_pa5_f090 = window_dir + "kspace_mask_dr6_pa5_f090.fits" +window_T_dr6_pa5_f090 = data_dir + "windows/window_dr6_pa5_f090_baseline.fits" +window_pol_dr6_pa5_f090 = data_dir + "windows/window_dr6_pa5_f090_baseline.fits" +window_kspace_dr6_pa5_f090 = data_dir + "windows/kspace_mask_dr6_pa5_f090.fits" -window_T_dr6_pa5_f150 = window_dir + "window_dr6_pa5_f150_baseline.fits" -window_pol_dr6_pa5_f150 = window_dir + "window_dr6_pa5_f150_baseline.fits" -window_kspace_dr6_pa5_f150 = window_dir + "kspace_mask_dr6_pa5_f150.fits" +window_T_dr6_pa5_f150 = data_dir + "windows/window_dr6_pa5_f150_baseline.fits" +window_pol_dr6_pa5_f150 = data_dir + "windows/window_dr6_pa5_f150_baseline.fits" +window_kspace_dr6_pa5_f150 = data_dir + "windows/kspace_mask_dr6_pa5_f150.fits" -window_T_dr6_pa6_f090 = window_dir + "window_dr6_pa6_f090_baseline.fits" -window_pol_dr6_pa6_f090 = window_dir + "window_dr6_pa6_f090_baseline.fits" -window_kspace_dr6_pa6_f090 = window_dir + "kspace_mask_dr6_pa6_f090.fits" +window_T_dr6_pa6_f090 = data_dir + "windows/window_dr6_pa6_f090_baseline.fits" +window_pol_dr6_pa6_f090 = data_dir + "windows/window_dr6_pa6_f090_baseline.fits" +window_kspace_dr6_pa6_f090 = data_dir + "windows/kspace_mask_dr6_pa6_f090.fits" -window_T_dr6_pa6_f150 = window_dir + "window_dr6_pa6_f150_baseline.fits" -window_pol_dr6_pa6_f150 = window_dir + "window_dr6_pa6_f150_baseline.fits" -window_kspace_dr6_pa6_f150 = window_dir + "kspace_mask_dr6_pa6_f150.fits" +window_T_dr6_pa6_f150 = data_dir + "windows/window_dr6_pa6_f150_baseline.fits" +window_pol_dr6_pa6_f150 = data_dir + "windows/window_dr6_pa6_f150_baseline.fits" +window_kspace_dr6_pa6_f150 = data_dir + "windows/kspace_mask_dr6_pa6_f150.fits" window_T_dr6_pa4_f220_alias = "w_dr6_pa4_f220" window_pol_dr6_pa4_f220_alias = "w_dr6_pa4_f220" @@ -116,7 +118,7 @@ window_T_dr6_pa6_f150_alias = "w_dr6_pa6_f150" window_pol_dr6_pa6_f150_alias = "w_dr6_pa6_f150" # maps -map_dir = '/pscratch/sd/x/xzackli/so/data/dr6v4/maps/' +map_dir = {'zatkins_on_della': '/scratch/gpfs/ACT/dr6v4/maps/dr6v4_20230316/release/', 'xzackli_on_perlmutter': '/pscratch/sd/x/xzackli/so/data/dr6v4/maps/'}[iam] n_splits_dr6 = 4 src_free_maps_dr6 = True diff --git a/project/ana_cov_comp/python/get_2pt_ewin_alms.py b/project/ana_cov_comp/python/get_2pt_ewin_alms.py index f9e987f8..339d7cac 100644 --- a/project/ana_cov_comp/python/get_2pt_ewin_alms.py +++ b/project/ana_cov_comp/python/get_2pt_ewin_alms.py @@ -45,7 +45,7 @@ if field_info not in field_infos: field_infos.append(field_info) else: - raise ValueError(f'{field_info} is not unique') + raise ValueError(f'{field_info=} is not unique') ewin_info_s = psc.get_ewin_info_from_field_info(field_info, d, mode='w') if ewin_info_s not in ewin_infos: diff --git a/project/ana_cov_comp/python/get_4pt_coupling_matrices_recipe.py b/project/ana_cov_comp/python/get_4pt_coupling_matrices_recipe.py index 06e4f17f..be68f684 100644 --- a/project/ana_cov_comp/python/get_4pt_coupling_matrices_recipe.py +++ b/project/ana_cov_comp/python/get_4pt_coupling_matrices_recipe.py @@ -44,7 +44,7 @@ # we define the canon by the windows order. we first build the fields, # then use a mapping from fields to windows to build the canonical # windows -recipe = {'S_only':[], 'SxN':[], 'N_only':[]} +recipe = {'S_only': [], 'SxN': [], 'N_only': []} field_infos = [] ewin_infos = [] for sv1 in surveys: @@ -128,7 +128,7 @@ 'spintype': spintype, 'w1': f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy', 'w2': f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy', - 'lmax':lmax, + 'lmax': lmax, 'input_alm': True, 'coupling_fn': coupling_fn }) @@ -186,7 +186,7 @@ 'spintype': spintype, 'w1': f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy', 'w2': f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy', - 'lmax':lmax, + 'lmax': lmax, 'input_alm': True, 'coupling_fn': coupling_fn }) @@ -245,7 +245,7 @@ 'spintype': spintype, 'w1': f'{ewin_alms_dir}/{ewin_name1}x{ewin_name2}_alm.npy', 'w2': f'{ewin_alms_dir}/{ewin_name3}x{ewin_name4}_alm.npy', - 'lmax':lmax, + 'lmax': lmax, 'input_alm': True, 'coupling_fn': coupling_fn }) diff --git a/project/ana_cov_comp/python/get_alms.py b/project/ana_cov_comp/python/get_alms.py deleted file mode 100644 index dc2fe6d3..00000000 --- a/project/ana_cov_comp/python/get_alms.py +++ /dev/null @@ -1,120 +0,0 @@ -""" -This script compute all alms write them to disk. -It uses the window function provided in the dictionnary file. -Optionally, it applies a calibration to the maps, a kspace filter and deconvolve the CAR pixel window function. -""" - -import sys -import time - -import numpy as np -from pixell import enmap -from pspipe_utils import kspace, log, misc, pspipe_list -from pspy import pspy_utils, so_dict, so_map, so_mpi, sph_tools - -d = so_dict.so_dict() -d.read_from_file(sys.argv[1]) -log = log.get_logger(**d) - -surveys = d["surveys"] -lmax = d["lmax"] -deconvolve_pixwin = d["deconvolve_pixwin"] -niter = d["niter"] -apply_kspace_filter = d["apply_kspace_filter"] - - -window_dir = d["window_dir"] -alms_dir = d["alms_dir"] -pspy_utils.create_directory(alms_dir) - -n_ar, sv_list, ar_list = pspipe_list.get_arrays_list(d) - -log.info(f"number of arrays for the mpi loop : {n_ar}") -so_mpi.init(True) -subtasks = so_mpi.taskrange(imin=0, imax=n_ar-1) - -if len(sys.argv) == 4: - log.info(f"computing only the alms : {int(sys.argv[2])}:{int(sys.argv[3])}") - subtasks = subtasks[int(sys.argv[2]):int(sys.argv[3])] -log.info(subtasks) - -for task in subtasks: - task = int(task) - sv, ar = sv_list[task], ar_list[task] - - log.info(f"[{task}] Computing alm for '{sv}' survey and '{ar}' array") - - win_T = so_map.read_map(d[f"window_T_{sv}_{ar}"]) - win_pol = so_map.read_map(d[f"window_pol_{sv}_{ar}"]) - - window_tuple = (win_T, win_pol) - - if win_T.pixel == "CAR": - # if None, don't apply window in fourier operations, - # will result in exception if ks_f["weighted"] is True - if d[f"window_kspace_{sv}_{ar}"] is not None: - win_kspace = so_map.read_map(d[f"window_kspace_{sv}_{ar}"]) - else: - win_kspace = None - - if apply_kspace_filter: - ks_f = d[f"k_filter_{sv}"] - filter = kspace.get_kspace_filter(win_T, ks_f) - - inv_pixwin_lxly = None - if deconvolve_pixwin: - if d[f"pixwin_{sv}"]["pix"] == "CAR": - # compute the CAR pixel function in fourier space - wy, wx = enmap.calc_window(win_T.data.shape, order=d[f"pixwin_{sv}"]["order"]) - inv_pixwin_lxly = (wy[:,None] * wx[None,:]) ** (-1) - - - maps = d[f"maps_{sv}_{ar}"] - cal, pol_eff = d[f"cal_{sv}_{ar}"], d[f"pol_eff_{sv}_{ar}"] - - t0 = time.time() - for k, map in enumerate(maps): - - if win_T.pixel == "CAR": - split = so_map.read_map(map, geometry=win_T.data.geometry) - - ### NOTE: comment-out because we don't care about ps, and - ### *don't* want ps mask! - # if d[f"src_free_maps_{sv}"] == True: - # ps_map_name = map.replace("srcfree.fits", "model.fits") - # if ps_map_name == map: - # raise ValueError("No model map is provided! Check map names!") - # ps_map = so_map.read_map(ps_map_name) - # ps_mask = so_map.read_map(d[f"ps_mask_{sv}_{ar}"]) - # ps_map.data *= ps_mask.data - # split.data += ps_map.data - - if apply_kspace_filter: - log.info(f"[{task}] apply kspace filter on {map}") - split = kspace.filter_map(split, - filter, - win_kspace, - inv_pixwin=inv_pixwin_lxly, - weighted_filter=ks_f["weighted"], - use_ducc_rfft=True) - - else: - log.info(f"[{task}] WARNING: no kspace filter is applied") - if deconvolve_pixwin: - split = so_map.fourier_convolution(split, - inv_pixwin_lxly, - window=win_kspace, - use_ducc_rfft=True) - - elif win_T.pixel == "HEALPIX": - split = so_map.read_map(map) - - split = split.calibrate(cal=cal, pol_eff=pol_eff) - - if d["remove_mean"] == True: - split = split.subtract_mean(window_tuple) - - master_alms = sph_tools.get_alms(split, window_tuple, niter, lmax) - np.save(f"{alms_dir}/alms_{sv}_{ar}_{k}.npy", master_alms) - - log.info(f"[{task}] execution time {time.time() - t0} seconds") diff --git a/project/ana_cov_comp/python/get_covariance_blocks.py b/project/ana_cov_comp/python/get_covariance_blocks.py deleted file mode 100644 index f8ea1e0f..00000000 --- a/project/ana_cov_comp/python/get_covariance_blocks.py +++ /dev/null @@ -1,156 +0,0 @@ -""" -This script compute the analytical covariance matrix elements. -""" -import sys - -import numpy as np -from pspipe_utils import best_fits, log, pspipe_list -from pspy import pspy_utils, so_cov, so_dict, so_map, so_mcm, so_mpi - -d = so_dict.so_dict() -d.read_from_file(sys.argv[1]) - -log = log.get_logger(**d) - -mcms_dir = d["mcms_dir"] -spectra_dir = d["spectra_dir"] -noise_dir = d["noise_model_dir"] -bestfit_dir = d["best_fits_dir"] -sq_win_alms_dir = d["sq_win_alms_dir"] - -cov_dir = d["covariances_dir"] - -pspy_utils.create_directory(cov_dir) - -surveys = d["surveys"] -binning_file = d["binning_file"] -lmax = d["lmax"] -niter = d["niter"] -binned_mcm = d["binned_mcm"] -apply_kspace_filter = d["apply_kspace_filter"] -cov_T_E_only = d["cov_T_E_only"] - - -spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"] -spin_pairs = ["spin0xspin0", "spin0xspin2", "spin2xspin0", "spin2xspin2"] - -# fast_coupling is designed to be fast but is not for general usage -# In particular we assume that the same window function is used in T and Pol -fast_coupling = True - -arrays, n_splits, bl_dict = {}, {}, {} -for sv in surveys: - arrays[sv] = d[f"arrays_{sv}"] - for ar in arrays[sv]: - l_beam, bl_dict[sv, ar] = pspy_utils.read_beam_file(d[f"beam_{sv}_{ar}"]) - id_beam = np.where((l_beam >= 2) & (l_beam < lmax)) - bl_dict[sv, ar] = bl_dict[sv, ar][id_beam] - n_splits[sv] = d[f"n_splits_{sv}"] - assert n_splits[sv] == len(d.get(f"maps_{sv}_{ar}", [0]*n_splits[sv])), "the number of splits does not correspond to the number of maps" - - if fast_coupling: - # This loop check that this is what was specified in the dictfile - assert d[f"window_T_{sv}_{ar}"] == d[f"window_pol_{sv}_{ar}"], "T and pol windows have to be the same" - -l_cmb, cmb_dict = best_fits.cmb_dict_from_file(bestfit_dir + "/cmb.dat", lmax, spectra) - -array_list = [f"{sv}_{ar}" for sv in surveys for ar in arrays[sv]] -l_fg, fg_dict = best_fits.fg_dict_from_files(bestfit_dir + "/fg_{}x{}.dat", array_list, lmax, spectra) - -f_name_noise = noise_dir + "/mean_{}x{}_{}_noise.dat" -l_noise, nl_dict = best_fits.noise_dict_from_files(f_name_noise, surveys, arrays, lmax, spectra, n_splits=n_splits) - -spec_name_list = pspipe_list.get_spec_name_list(d) -l, ps_all, nl_all = best_fits.get_all_best_fit(spec_name_list, - l_cmb, - cmb_dict, - fg_dict, - spectra, - nl_dict=nl_dict, - bl_dict=bl_dict) - -ncovs, na_list, nb_list, nc_list, nd_list = pspipe_list.get_covariances_list(d) - -if d["use_toeplitz_cov"] == True: - log.info("we will use the toeplitz approximation") - l_exact, l_band, l_toep = 800, 2000, 2750 -else: - l_exact, l_band, l_toep = None, None, None - -log.info(f"number of covariance matrices to compute : {ncovs}") -so_mpi.init(True) -subtasks = so_mpi.taskrange(imin=0, imax=ncovs - 1) - -if len(sys.argv) == 4: - log.info(f"computing only the covariance matrices : {int(sys.argv[2])}:{int(sys.argv[3])}") - subtasks = subtasks[int(sys.argv[2]):int(sys.argv[3])] -log.info(subtasks) - -for task in subtasks: - task = int(task) - na, nb, nc, nd = na_list[task], nb_list[task], nc_list[task], nd_list[task] - na_r, nb_r, nc_r, nd_r = na.replace("&", "_"), nb.replace("&", "_"), nc.replace("&", "_"), nd.replace("&", "_") - - log.info(f"[{task}] cov element ({na_r} x {nb_r}, {nc_r} x {nd_r})") - - if fast_coupling: - - coupling = so_cov.fast_cov_coupling_spin0and2(sq_win_alms_dir, - [na_r, nb_r, nc_r, nd_r], - lmax, - l_exact=l_exact, - l_band=l_band, - l_toep=l_toep) - - else: - win = {} - win["Ta"] = so_map.read_map(d[f"window_T_{na_r}"]) - win["Tb"] = so_map.read_map(d[f"window_T_{nb_r}"]) - win["Tc"] = so_map.read_map(d[f"window_T_{nc_r}"]) - win["Td"] = so_map.read_map(d[f"window_T_{nd_r}"]) - win["Pa"] = so_map.read_map(d[f"window_pol_{na_r}"]) - win["Pb"] = so_map.read_map(d[f"window_pol_{nb_r}"]) - win["Pc"] = so_map.read_map(d[f"window_pol_{nc_r}"]) - win["Pd"] = so_map.read_map(d[f"window_pol_{nd_r}"]) - - coupling = so_cov.cov_coupling_spin0and2_simple(win, - lmax, - niter=niter, - l_exact=l_exact, - l_band=l_band, - l_toep=l_toep) - - - - try: mbb_inv_ab, Bbl_ab = so_mcm.read_coupling(prefix=f"{mcms_dir}/{na_r}x{nb_r}", spin_pairs=spin_pairs) - except: mbb_inv_ab, Bbl_ab = so_mcm.read_coupling(prefix=f"{mcms_dir}/{nb_r}x{na_r}", spin_pairs=spin_pairs) - - try: mbb_inv_cd, Bbl_cd = so_mcm.read_coupling(prefix=f"{mcms_dir}/{nc_r}x{nd_r}", spin_pairs=spin_pairs) - except: mbb_inv_cd, Bbl_cd = so_mcm.read_coupling(prefix=f"{mcms_dir}/{nd_r}x{nc_r}", spin_pairs=spin_pairs) - - - analytic_cov = so_cov.generalized_cov_spin0and2(coupling, - [na, nb, nc, nd], - n_splits, - ps_all, - nl_all, - lmax, - binning_file, - mbb_inv_ab, - mbb_inv_cd, - binned_mcm=binned_mcm, - cov_T_E_only=cov_T_E_only, - dtype=np.float32) - - if apply_kspace_filter == True: - # Some heuristic correction for the number of modes lost due to the transfer function - # This should be tested against simulation and revisited - - one_d_tf_ab = np.loadtxt(f"{spectra_dir}/one_dimension_kspace_tf_{na_r}x{nb_r}.dat") - one_d_tf_cd = np.loadtxt(f"{spectra_dir}/one_dimension_kspace_tf_{nc_r}x{nd_r}.dat") - one_d_tf = np.minimum(one_d_tf_ab, one_d_tf_cd) - # sqrt(tf) is an approx for the number of modes masked in a given map so (2l+1)*fsky*sqrt(tf) - # is our proxy for the number of modes - analytic_cov /= np.outer(one_d_tf ** (1.0 / 4.0), one_d_tf ** (1.0 / 4.0)) - - np.save(f"{cov_dir}/analytic_cov_{na_r}x{nb_r}_{nc_r}x{nd_r}.npy", analytic_cov) diff --git a/project/ana_cov_comp/python/get_mcm_and_bbl.py b/project/ana_cov_comp/python/get_mcm_and_bbl.py deleted file mode 100644 index 9e7e3a2d..00000000 --- a/project/ana_cov_comp/python/get_mcm_and_bbl.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -This script computes the mode coupling matrices and the binning matrices Bbl -for the different surveys and arrays. -""" - -import sys - -from pspipe_utils import log, pspipe_list -from pspy import pspy_utils, so_dict, so_map, so_mcm, so_mpi - -d = so_dict.so_dict() -d.read_from_file(sys.argv[1]) -log = log.get_logger(**d) - -mcm_dir = d["mcms_dir"] -pspy_utils.create_directory(mcm_dir) - -surveys = d["surveys"] -lmax = d["lmax"] -binned_mcm = d["binned_mcm"] - -if d["use_toeplitz_mcm"] == True: - log.info("we will use the toeplitz approximation") - l_exact, l_band, l_toep = 800, 2000, 2750 -else: - l_exact, l_band, l_toep = None, None, None - -n_mcms, sv1_list, ar1_list, sv2_list, ar2_list = pspipe_list.get_spectra_list(d) - -log.info(f"number of mcm matrices to compute : {n_mcms}") -so_mpi.init(True) -subtasks = so_mpi.taskrange(imin=0, imax=n_mcms - 1) - -if len(sys.argv) == 4: - log.info(f"computing only the mcm matrices : {int(sys.argv[2])}:{int(sys.argv[3])}") - subtasks = subtasks[int(sys.argv[2]):int(sys.argv[3])] - -for task in subtasks: - task = int(task) - sv1, ar1, sv2, ar2 = sv1_list[task], ar1_list[task], sv2_list[task], ar2_list[task] - - log.info(f"[{task:02d}] mcm matrix for {sv1}_{ar1} x {sv2}_{ar2}") - - l, bl1 = pspy_utils.read_beam_file(d[f"beam_{sv1}_{ar1}"]) - win1_T = so_map.read_map(d[f"window_T_{sv1}_{ar1}"]) - win1_pol = so_map.read_map(d[f"window_pol_{sv1}_{ar1}"]) - - l, bl2 = pspy_utils.read_beam_file(d[f"beam_{sv2}_{ar2}"]) - win2_T = so_map.read_map(d[f"window_T_{sv2}_{ar2}"]) - win2_pol = so_map.read_map(d[f"window_pol_{sv2}_{ar2}"]) - - mbb_inv, Bbl = so_mcm.mcm_and_bbl_spin0and2(win1=(win1_T, win1_pol), - win2=(win2_T, win2_pol), - bl1=(bl1, bl1), - bl2=(bl2, bl2), - binning_file=d["binning_file"], - niter=d["niter"], - lmax=d["lmax"], - type=d["type"], - l_exact=l_exact, - l_band=l_band, - l_toep=l_toep, - binned_mcm=binned_mcm, - save_file=f"{mcm_dir}/{sv1}_{ar1}x{sv2}_{ar2}") diff --git a/project/ana_cov_comp/python/get_noise_model.py b/project/ana_cov_comp/python/get_noise_model.py index 2747e02f..8299e8f5 100644 --- a/project/ana_cov_comp/python/get_noise_model.py +++ b/project/ana_cov_comp/python/get_noise_model.py @@ -1,115 +1,289 @@ -#TO BE TESTED MORE +""" +This script computes the noise model for the covariance from the measured alms, +using auto - crosses. For INKA covariances, we only need the pseudospectra. +""" +import sys -import matplotlib +from pspipe_utils import log +from pspy import so_dict, pspy_utils -matplotlib.use("Agg") -import sys +from pixell import curvedsky import numpy as np -import pylab as plt -import scipy.interpolate -from pspipe_utils import log -from pspy import pspy_utils, so_dict, so_spectra - - -def interpolate_dict(lb, cb, lth, spectra, force_positive=True, l_inf_lmin_equal_lmin=True, discard_cross=True): - cl_dict = {} - for spec in spectra: - cl = scipy.interpolate.interp1d(lb, cb[spec], fill_value = "extrapolate") - cl_dict[spec] = cl(lth) - if l_inf_lmin_equal_lmin: - id = np.where(lth <= np.min(lb)) - cl_dict[spec][id]= cb[spec][0] - if force_positive: - cl_dict[spec] = np.abs(cl_dict[spec]) - if discard_cross: - if spec not in ["TT", "EE", "BB"]: - cl_dict[spec] = np.zeros(len(lth)) - return cl_dict +from scipy.signal import savgol_filter +import matplotlib.pyplot as plt + +from itertools import combinations, combinations_with_replacement as cwr, product +import os d = so_dict.so_dict() d.read_from_file(sys.argv[1]) -log = log.get_logger(**d) -spectra_dir = d["spectra_dir"] +log = log.get_logger(**d) -ps_model_dir = d["noise_model_dir"] -plot_dir = d["plot_dir"] + "/noise_model" +alms_dir = d['alms_dir'] +savgol_w = d['noise_model_savgol_w'] +savgol_k = d['noise_model_savgol_k'] -pspy_utils.create_directory(ps_model_dir) +noise_model_dir = d["noise_model_dir"] +plot_dir = os.path.join(d["plot_dir"], "noise_model") +pspy_utils.create_directory(noise_model_dir) pspy_utils.create_directory(plot_dir) -spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"] -surveys = d["surveys"] -lmax = d["lmax"] -type = d["type"] -binning_file = d["binning_file"] +surveys = d['surveys'] +arrays = {sv: d[f'arrays_{sv}'] for sv in surveys} -lth = np.arange(2, lmax+2) +# we will make noise models for each survey and array, +# so "everything" happens inside this loop +for sv1 in surveys: + for ar1 in arrays[sv1]: + print(f'Doing {sv1}, {ar1}') -for sv in surveys: - arrays = d[f"arrays_{sv}"] - for id_ar1, ar1 in enumerate(arrays): - for id_ar2, ar2 in enumerate(arrays): - if id_ar1 > id_ar2: continue + # load alms for this survey and array + # we are assuming: + # - noise is uncorrelated between surveys + # - noise is uncorrelated between arrays within surveys + # - alms have shape (npol=3, nalm) + alms = [] + nchan = len(arrays[sv1][ar1]) + + for i, chan1 in enumerate(arrays[sv1][ar1]): + if i == 0: + nsplit = len(d[f'maps_{sv1}_{ar1}_{chan1}']) + else: + _nsplit = len(d[f'maps_{sv1}_{ar1}_{chan1}']) + assert _nsplit == nsplit, \ + f'sv={sv1}, ar={ar1}, chan={chan1}, nsplit={_nsplit}, expected {nsplit}' + + for split1 in range(nsplit): + alms.append(np.load(f'{alms_dir}/alms_{sv1}_{ar1}_{chan1}_{split1}.npy')) + assert alms[-1].ndim == 2, \ + f'sv={sv1}, ar={ar1}, chan={chan1}, split={split1} alm.ndim={alms[-1].ndim}, expected 2' + assert alms[-1].shape[0] == 3, \ + f'sv={sv1}, ar={ar1}, chan={chan1}, split={split1} alm.shape[0]={alms[-1].shape[0]}, expected 3' + + alms = np.asarray(alms).reshape(nchan, nsplit, 3, -1) # will fail if shape mismatches - log.info(f"Computing noise for '{sv}' survey and '{ar1}x{ar2}' arrays") + # get the signal model for this survey and array from + # the average of split cross spectra. + # we are assuming: + # - noise is uncorrelated between splits within arrays + nell = curvedsky.nalm2lmax(alms.shape[-1]) + 1 - l, bl_ar1 = pspy_utils.read_beam_file(d[f"beam_{sv}_{ar1}"]) - l, bl_ar2 = pspy_utils.read_beam_file(d[f"beam_{sv}_{ar2}"]) + signal_model = 0 + count = 0 + for split1, split2 in combinations(range(nsplit), r=2): + alms1, alms2 = alms[:, split1], alms[:, split2] + + _signal_model = np.zeros((nchan, 3, nchan, 3, nell), dtype=alms.real.dtype) + + # signal model is not guaranteed to be symmetric cause of split crosses + for preidx1, preidx2 in product(np.ndindex((nchan, 3)), repeat=2): + _signal_model[(*preidx1, *preidx2)] = curvedsky.alm2cl(alms1[preidx1], alms2[preidx2]) + + signal_model += _signal_model + count += 1 + assert count == nsplit * (nsplit - 1) / 2, \ + f'Calculated {count=} but expected {nsplit * (nsplit - 1) / 2}' + signal_model /= count + + # signal model is not guaranteed to be symmetric cause of split crosses, + # but we know it should be. in effect, we get extra independent samples + # of the off diagonals by virtue of there being 2 of them. alternatively, + # could have iterated split1, split2 over permutations in above loop + signal_model = (signal_model + np.moveaxis(signal_model, (0, 1), (2, 3))) / 2 + # each noise model will have the following shape: + # (nchan, npol=3, nchan, npol=3, nell) + # we are assuming: + # - noise may be correlated between channels and pols + # - beams and masks don't vary over split + for split1 in range(nsplit): + print(f'Doing {sv1}, {ar1}, set{split1}') - lb, nbs_ar1xar1 = so_spectra.read_ps(f"{spectra_dir}/{type}_{sv}_{ar1}x{sv}_{ar1}_noise.dat", spectra=spectra) - lb, nbs_ar1xar2 = so_spectra.read_ps(f"{spectra_dir}/{type}_{sv}_{ar1}x{sv}_{ar2}_noise.dat", spectra=spectra) - lb, nbs_ar2xar2 = so_spectra.read_ps(f"{spectra_dir}/{type}_{sv}_{ar2}x{sv}_{ar2}_noise.dat", spectra=spectra) + alms1 = alms[:, split1] + total_model = np.zeros((nchan, 3, nchan, 3, nell), dtype=alms.real.dtype) + + # signal model is guaranteed to be symmetric cause of split auto + for preidx1, preidx2 in cwr(np.ndindex((nchan, 3)), r=2): + total_model[(*preidx1, *preidx2)] = curvedsky.alm2cl(alms1[preidx1], alms1[preidx2]) + if preidx1 != preidx2: + total_model[(*preidx2, *preidx1)] = total_model[(*preidx1, *preidx2)] - lb, bb_ar1 = pspy_utils.naive_binning(l, bl_ar1, binning_file, lmax) - lb, bb_ar2 = pspy_utils.naive_binning(l, bl_ar2, binning_file, lmax) + inp_noise_model = total_model - signal_model + out_noise_model = np.zeros_like(inp_noise_model) - Rb = {} - for spec in spectra: - nbs_ar1xar1[spec] *= bb_ar1 * bb_ar1 - nbs_ar1xar2[spec] *= bb_ar1 * bb_ar2 - nbs_ar2xar2[spec] *= bb_ar2 * bb_ar2 - Rb[spec] = nbs_ar1xar2[spec] / np.sqrt(np.abs(nbs_ar1xar1[spec] * nbs_ar2xar2[spec])) + inp_corr_model = np.zeros_like(inp_noise_model) + out_corr_model = np.zeros_like(inp_noise_model) + + mask_model = np.zeros_like(inp_noise_model, dtype=bool) + + # huzzah! now we have a noisy thing that we want to smooth, + # nonparametrically. we need to fit the diagonals first, because + # we will fit the off-diagonal corrs, and need to convert them back + # to cross-spectra using the fit diagonals; i.e., the fit diagonals + # all need to exist before doing the off-diagonals. + for preidx1 in np.ndindex(nchan, 3): + + # only if we are TT, we start at l=0, otherwise l=2 + if preidx1[1] == 0: + start = 0 + else: + assert np.all(inp_noise_model[(*preidx1, *preidx1)][:2] == 0), \ + f'{preidx1=}; expected 0 for l=0,1' + start = 2 + + y = inp_noise_model[(*preidx1, *preidx1)][start:].copy() + + # might start <= 0, so first we clip and fit without log, and + # add fitted pts back where it is <= 0. in principle + # this could still fail, so if it doesn't, then this was a + # fluctuation consistent with noise + mask = y <= 0 + if np.any(mask): + print(f'{preidx1=} has {np.sum(mask)} values <= 0') + np.clip(y, 0, None, out=y) + fit_y = savgol_filter(y, savgol_w, savgol_k) + y[mask] = fit_y[mask] + assert not np.any(y <= 0), \ + print(f'{preidx1=} still has {np.sum(y <= 0)} values <= 0 after correcting') + mask_model[(*preidx1, *preidx1)][start:] = mask + + # now fit in log space + fit_y = savgol_filter(np.log(y), savgol_w, savgol_k) + fit_y = np.exp(fit_y) + out_noise_model[(*preidx1, *preidx1)][start:] = fit_y + + for preidx1, preidx2 in combinations(np.ndindex((nchan, 3)), r=2): + + # only if we are TT, we start at l=0, otherwise l=2 + if preidx1[1] == 0 and preidx2[1] == 0: + start = 0 + else: + assert np.all(inp_noise_model[(*preidx1, *preidx2)][:2] == 0), \ + f'{preidx1=}, {preidx2=}; expected 0 for l=0,1' + start = 2 + + y = inp_noise_model[(*preidx1, *preidx2)][start:].copy() + y /= np.sqrt(out_noise_model[(*preidx1, *preidx1)][start:]) + y /= np.sqrt(out_noise_model[(*preidx2, *preidx2)][start:]) + + inp_corr_model[(*preidx1, *preidx2)][start:] = y + inp_corr_model[(*preidx2, *preidx1)][start:] = y + + # might start >= |1|, so first we clip and fit without arctanh, and + # add fitted pts back where it is >= |1|. in principle + # this could still fail, so if it doesn't, then this was a + # fluctuation consistent with noise + mask = np.logical_or(y <= -1, 1 <= y) + if np.any(mask): + print(f'{preidx1=}, {preidx2=} has {np.sum(mask)} values >= |1|') + np.clip(y, -1, 1, out=y) + fit_y = savgol_filter(y, savgol_w, savgol_k) + y[mask] = fit_y[mask] + assert not np.any(np.logical_or(y <= -1, 1 <= y)), \ + print(f'{preidx1=}, {preidx2=} still has {np.logical_or(y <= -1, 1 <= y)} values >= |1| after correcting') + mask_model[(*preidx1, *preidx2)][start:] = mask + mask_model[(*preidx2, *preidx1)][start:] = mask + + # now fit in arctanh space + fit_y = savgol_filter(np.arctanh(y), savgol_w, savgol_k) + fit_y = np.tanh(fit_y) + out_corr_model[(*preidx1, *preidx2)][start:] = fit_y + out_corr_model[(*preidx2, *preidx1)][start:] = fit_y + + fit_y *= np.sqrt(out_noise_model[(*preidx1, *preidx1)][start:]) + fit_y *= np.sqrt(out_noise_model[(*preidx2, *preidx2)][start:]) + out_noise_model[(*preidx1, *preidx2)][start:] = fit_y + out_noise_model[(*preidx2, *preidx1)][start:] = fit_y + + # plot and save + for preidx1, preidx2 in cwr(np.ndindex((nchan, 3)), r=2): + + if preidx1 == preidx2: + # ps + raw_y = inp_noise_model[(*preidx1, *preidx2)] + fit_y = out_noise_model[(*preidx1, *preidx2)] + red_mask = mask_model[(*preidx1, *preidx2)] + l = np.arange(raw_y.size) + fig, axs = plt.subplots(ncols=2, nrows=2, sharex='col', sharey='row', + figsize=(12, 8), dpi=100, height_ratios=(2, 1), + layout='constrained') + axs[0, 0].loglog(l, raw_y, alpha=0.3, label='raw') + axs[0, 0].loglog(l, fit_y, alpha=1, label='fit') + axs[0, 0].scatter(l[red_mask], fit_y[red_mask], alpha=1, c='r') + axs[0, 0].set_ylabel('$N_{\ell}$', fontsize=12) + axs[0, 0].legend() + axs[0, 0].grid() + + axs[1, 0].semilogx(l, (fit_y / raw_y), alpha=1, color='k') + axs[1, 0].scatter(l[red_mask], (fit_y / raw_y)[red_mask], alpha=1, c='r') + axs[1, 0].set_xlabel('$\ell$', fontsize=12) + axs[1, 0].set_ylabel('fit / raw', fontsize=12) + axs[1, 0].grid() + + axs[0, 1].semilogy(l, raw_y, alpha=0.3, label='raw') + axs[0, 1].semilogy(l, fit_y, alpha=1, label='fit') + axs[0, 1].scatter(l[red_mask], fit_y[red_mask], alpha=1, c='r') + axs[0, 1].legend() + axs[0, 1].grid() + + axs[1, 1].plot(l, (fit_y / raw_y), alpha=1, color='k') + axs[1, 1].scatter(l[red_mask], (fit_y / raw_y)[red_mask], alpha=1, c='r') + axs[1, 1].set_xlabel('$\ell$', fontsize=12) + axs[1, 1].grid() + + chan1, chan2 = arrays[sv1][ar1][preidx1[0]], arrays[sv1][ar1][preidx2[0]] + pol1, pol2 = 'TEB'[preidx1[1]], 'TEB'[preidx2[1]] + fig.suptitle(f'{sv1}_{ar1}_set{split1} {chan1}_{pol1}x{chan2}_{pol2}') + + plt_fn = f'{plot_dir}/{sv1}_{ar1}_set{split1}_{chan1}_{pol1}x{chan2}_{pol2}_ps.png' + plt.savefig(plt_fn, bbox_inches='tight') + plt.close() - if ar1 == ar2: - nlth = interpolate_dict(lb, nbs_ar1xar1, lth, spectra) - else: - Rlth = interpolate_dict(lb, Rb, lth, spectra) - nlth_ar1xar1 = interpolate_dict(lb, nbs_ar1xar1, lth, spectra) - nlth_ar2xar2 = interpolate_dict(lb, nbs_ar2xar2, lth, spectra) - nlth = {} - for spec in spectra: - nlth[spec] = Rlth[spec] * np.sqrt(nlth_ar1xar1[spec] * nlth_ar2xar2[spec]) - - - for spec in spectra: - plt.figure(figsize=(12,12)) - plt.plot(lth, - nlth[spec], - label="interpolate", - color="lightblue") - if ar1 == ar2: - nbs= nbs_ar1xar1[spec] else: - nbs= nbs_ar1xar2[spec] - - plt.plot(lb, - nbs, - ".", - label = f"{sv} {ar1}x{ar2}", - color="red") - plt.legend(fontsize=20) - plt.savefig(f"{plot_dir}/noise_interpolate_{ar1}x{ar2}_{sv}_{spec}.png", bbox_inches="tight") - plt.clf() - plt.close() - - spec_name_noise_mean = f"mean_{ar1}x{ar2}_{sv}_noise" - so_spectra.write_ps(ps_model_dir + f"/{spec_name_noise_mean}.dat", lth, nlth, type, spectra=spectra) - - if ar2 != ar1: - spec_name_noise_mean = f"mean_{ar2}x{ar1}_{sv}_noise" - TE, ET, TB, BT, EB, BE = nlth["ET"], nlth["TE"], nlth["BT"], nlth["TB"], nlth["BE"], nlth["EB"] - nlth["TE"], nlth["ET"], nlth["TB"], nlth["BT"], nlth["EB"], nlth["BE"] = TE, ET, TB, BT, EB, BE - so_spectra.write_ps(ps_model_dir + f"/{spec_name_noise_mean}.dat", lth, nlth, type, spectra=spectra) + # corrs, ps + for i, (raw_y, fit_y) in enumerate(((inp_corr_model[(*preidx1, *preidx2)], out_corr_model[(*preidx1, *preidx2)]), + (inp_noise_model[(*preidx1, *preidx2)], out_noise_model[(*preidx1, *preidx2)]))): + red_mask = mask_model[(*preidx1, *preidx2)] + l = np.arange(raw_y.size) + fig, axs = plt.subplots(ncols=2, nrows=2, sharex='col', sharey='row', + figsize=(12, 8), dpi=100, height_ratios=(2, 1), + layout='constrained') + axs[0, 0].semilogx(l, raw_y, alpha=0.3, label='raw') + axs[0, 0].semilogx(l, fit_y, alpha=1, label='fit') + axs[0, 0].scatter(l[red_mask], fit_y[red_mask], alpha=1, c='r') + axs[0, 0].set_ylabel(['$r_{\ell}$', '$N_{\ell}$'][i], fontsize=12) + axs[0, 0].legend() + axs[0, 0].grid() + + axs[1, 0].semilogx(l, (fit_y - raw_y), alpha=1, color='k') + axs[1, 0].scatter(l[red_mask], (fit_y - raw_y)[red_mask], alpha=1, c='r') + axs[1, 0].set_xlabel('$\ell$', fontsize=12) + axs[1, 0].set_ylabel('fit - raw', fontsize=12) + axs[1, 0].grid() + + axs[0, 1].plot(l, raw_y, alpha=0.3, label='raw') + axs[0, 1].plot(l, fit_y, alpha=1, label='fit') + axs[0, 1].scatter(l[red_mask], fit_y[red_mask], alpha=1, c='r') + + axs[0, 1].legend() + axs[0, 1].grid() + + axs[1, 1].plot(l, (fit_y - raw_y), alpha=1, color='k') + axs[1, 1].scatter(l[red_mask], (fit_y - raw_y)[red_mask], alpha=1, c='r') + axs[1, 1].set_xlabel('$\ell$', fontsize=12) + axs[1, 1].grid() + + chan1, chan2 = arrays[sv1][ar1][preidx1[0]], arrays[sv1][ar1][preidx2[0]] + pol1, pol2 = 'TEB'[preidx1[1]], 'TEB'[preidx2[1]] + fig.suptitle(f'{sv1}_{ar1}_set{split1} {chan1}_{pol1}x{chan2}_{pol2}') + + plt_fn = f"{plot_dir}/{sv1}_{ar1}_set{split1}_{chan1}_{pol1}x{chan2}_{pol2}_{['corr', 'ps'][i]}.png" + plt.savefig(plt_fn, bbox_inches='tight') + plt.close() + + noise_model_fn = f'{noise_model_dir}/{sv1}_{ar1}_set{split1}_noise_model.npy' + np.save(noise_model_fn, out_noise_model) + + alms = None # help with memory \ No newline at end of file diff --git a/project/ana_cov_comp/python/get_per_split_noise.py b/project/ana_cov_comp/python/get_per_split_noise.py deleted file mode 100644 index 85d1aa0c..00000000 --- a/project/ana_cov_comp/python/get_per_split_noise.py +++ /dev/null @@ -1,165 +0,0 @@ -""" -This script is used to compute the per-split noise -from pre-computed power spectra. `write_split_spectra` has to be -set to `True` to run this script. -""" -from pspy import so_dict, pspy_utils, so_spectra -import sys -from pspipe_utils import pspipe_list -import scipy.interpolate -import numpy as np -import matplotlib.pyplot as plt - -def interpolate_dict(lb, cb, lth, spectra, force_positive=True, l_inf_lmin_equal_lmin=True, discard_cross=True): - cl_dict = {} - for spec in spectra: - cl = scipy.interpolate.interp1d(lb, cb[spec], fill_value="extrapolate") - cl_dict[spec] = cl(lth) - if l_inf_lmin_equal_lmin: - id = np.where(lth <= np.min(lb)) - cl_dict[spec][id]= cb[spec][0] - if force_positive: - cl_dict[spec] = np.abs(cl_dict[spec]) - if discard_cross: - if spec not in ["TT", "EE", "BB"]: - cl_dict[spec] = np.zeros(len(lth)) - return cl_dict - -d = so_dict.so_dict() -d.read_from_file(sys.argv[1]) - -spectra_dir = d["spectra_dir"] -ps_model_dir = d["noise_model_dir"] - -split_noise_dir = d["split_noise_dir"] -plot_dir = d["plot_dir"] + "/split_noise" - -pspy_utils.create_directory(split_noise_dir) -pspy_utils.create_directory(plot_dir) - -spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"] -surveys = d["surveys"] -lmax = d["lmax"] -type = d["type"] -binning_file = d["binning_file"] - -lth = np.arange(2, lmax+2) - -arrays = {sv: d[f"arrays_{sv}"] for sv in surveys} -n_splits = {sv: d[f"n_splits_{sv}"] for sv in surveys} - -spec_name_list = pspipe_list.get_spec_name_list(d) - -split_noise_dict = {} -for spec_name in spec_name_list: - sv1ar1, sv2ar2 = spec_name.split("x") - sv1, ar1 = sv1ar1.split("&") - sv2, ar2 = sv2ar2.split("&") - - if sv1 != sv2: continue - - split_ps_dict = {} - for i in range(n_splits[sv1]): - for j in range(n_splits[sv2]): - if i > j and sv1ar1 == sv2ar2: continue - - lb, ps = so_spectra.read_ps(f"{spectra_dir}/Dl_{sv1}_{ar1}x{sv2}_{ar2}_{i}{j}.dat", - spectra=spectra) - split_ps_dict[i, j] = ps - - for i in range(n_splits[sv1]): - - cross_ps_dict = {spec: [] for spec in spectra} - - for id1, id2 in list(split_ps_dict.keys()): - if id1 == id2: continue - - for spec in spectra: - cross_ps_dict[spec].append(split_ps_dict[id1, id2][spec]) - - for spec in spectra: - cross_ps_dict[spec] = np.mean(cross_ps_dict[spec], axis=0) - - split_noise = {spec: split_ps_dict[i, i][spec] - cross_ps_dict[spec] for spec in spectra} - - split_noise_dict[sv1, ar1, sv2, ar2, i] = split_noise - - spec_name = f"Dl_{sv1}_{ar1}x{sv2}_{ar2}_{i}{i}_noise.dat" - so_spectra.write_ps(f"{split_noise_dir}/{spec_name}", lb, split_noise, type, spectra=spectra) - -for (sv1, ar1, sv2, ar2, split), split_noise in split_noise_dict.items(): - - # Noise model - l, bl1 = pspy_utils.read_beam_file(d[f"beam_{sv1}_{ar1}"]) - l, bl2 = pspy_utils.read_beam_file(d[f"beam_{sv2}_{ar2}"]) - - lb, bb1 = pspy_utils.naive_binning(l, bl1, binning_file, lmax) - lb, bb2 = pspy_utils.naive_binning(l, bl2, binning_file, lmax) - - if ar1 == ar2: - for spec in spectra: - split_noise[spec] = split_noise[spec] * bb1 * bb2 - nlth = interpolate_dict(lb, split_noise, lth, spectra) - - else: - nb_ar1xar1 = {spec: split_noise_dict[sv1, ar1, sv1, ar1, split][spec] * bb1 * bb1 for spec in spectra} - nb_ar2xar2 = {spec: split_noise_dict[sv2, ar2, sv2, ar2, split][spec] * bb2 * bb2 for spec in spectra} - nb_ar1xar2 = {spec: split_noise_dict[sv1, ar1, sv2, ar2, split][spec] * bb1 * bb2 for spec in spectra} - - Rb = {spec : nb_ar1xar2[spec] / np.sqrt(np.abs(nb_ar1xar1[spec] * nb_ar2xar2[spec])) for spec in spectra} - - Rlth = interpolate_dict(lb, Rb, lth, spectra) - nlth_ar1xar1 = interpolate_dict(lb, nb_ar1xar1, lth, spectra) - nlth_ar2xar2 = interpolate_dict(lb, nb_ar2xar2, lth, spectra) - - nlth = {spec: Rlth[spec] * np.sqrt(nlth_ar1xar1[spec] * nlth_ar2xar2[spec]) for spec in spectra} - - spec_model_name = f"Dl_{ar1}_{split}x{ar2}_{split}_{sv1}_noise_model.dat" - so_spectra.write_ps(f"{split_noise_dir}/{spec_model_name}", lth, nlth, type, spectra=spectra) - - for split2 in range(n_splits[sv1]): - if split == split2: continue - nl_xsplit = {spec: np.zeros_like(nlth[spec]) for spec in spectra} - spec_name = f"Dl_{ar1}_{split}x{ar2}_{split2}_{sv1}_noise_model.dat" - so_spectra.write_ps(f"{split_noise_dir}/{spec_name}", lth, nl_xsplit, type, spectra=spectra) - -# Plots -for spec_name in spec_name_list: - sv1ar1, sv2ar2 = spec_name.split("x") - sv1, ar1 = sv1ar1.split("&") - sv2, ar2 = sv2ar2.split("&") - - if sv1 != sv2: continue - - noise_model_file = f"{ps_model_dir}/mean_{ar1}x{ar2}_{sv1}_noise.dat" - ell, nl_mean = so_spectra.read_ps(noise_model_file, spectra=spectra) - - for spec in ["TT", "EE", "BB"]: - plt.figure(figsize = (13, 8)) - for i in range(n_splits[sv1]): - split_noise_model_file = f"{split_noise_dir}/Dl_{ar1}_{i}x{ar2}_{i}_{sv1}_noise_model.dat" - ell, nl = so_spectra.read_ps(split_noise_model_file, spectra=spectra) - plt.plot(ell, nl[spec], label=f"Split {i}") - - plt.plot(ell, nl_mean[spec] * n_splits[sv1], label="Mean noise", color="k", ls="--") - plt.title(f"{sv1ar1}x{sv2ar2}", fontsize=16.5) - plt.xlabel(r"$\ell$", fontsize=15) - plt.ylabel(f"$N_\ell^{{{spec}}}$", fontsize=15) - plt.yscale("log") - plt.legend() - plt.tight_layout() - plt.savefig(f"{plot_dir}/noise_{sv1ar1}x{sv2ar2}_{spec}.png", dpi=300) - - if sv1ar1 == sv2ar2: - plt.figure(figsize=(13,8)) - plt.axhline(1, color="k", ls="--") - for i in range(n_splits[sv1]): - split_noise_model_file = f"{split_noise_dir}/Dl_{ar1}_{i}x{ar2}_{i}_{sv1}_noise_model.dat" - ell, nl = so_spectra.read_ps(split_noise_model_file, spectra = spectra) - plt.plot(ell, nl[spec]/nl_mean[spec]/n_splits[sv1], label = f"Split {i}") - plt.title(f"{ar1}x{ar2}", fontsize = 16.5) - plt.xlabel(r"$\ell$", fontsize = 15) - plt.ylabel(r"$N_\ell^{%s}/N_\ell^{%s,\mathrm{mean}}$" % (spec, spec), fontsize = 15) - plt.legend() - plt.tight_layout() - plt.savefig(f"{plot_dir}/noise_ratio_{sv1ar1}x{sv2ar2}_{spec}.png", dpi = 300) diff --git a/project/ana_cov_comp/python/get_spectra_from_alms.py b/project/ana_cov_comp/python/get_spectra_from_alms.py deleted file mode 100644 index fafcb02b..00000000 --- a/project/ana_cov_comp/python/get_spectra_from_alms.py +++ /dev/null @@ -1,188 +0,0 @@ -""" -This script compute all power spectra and write them to disk. -It uses the window function provided in the dictionnary file. -Optionally, it applies a calibration to the maps, a kspace filter and deconvolve the pixel window function. -The spectra are then combined in mean auto, cross and noise power spectrum and written to disk. -If write_all_spectra=True, each individual spectrum is also written to disk. -""" - -import sys - -import healpy as hp -import numpy as np -from pixell import enmap -from pspipe_utils import kspace, log, pspipe_list, transfer_function -from pspy import pspy_utils, so_dict, so_map, so_mcm, so_mpi, so_spectra - -d = so_dict.so_dict() -d.read_from_file(sys.argv[1]) -log = log.get_logger(**d) - -surveys = d["surveys"] -lmax = d["lmax"] -type = d["type"] -binning_file = d["binning_file"] -write_all_spectra = d["write_splits_spectra"] -deconvolve_pixwin = d["deconvolve_pixwin"] -binned_mcm = d["binned_mcm"] -apply_kspace_filter = d["apply_kspace_filter"] -cov_T_E_only = d["cov_T_E_only"] - -mcm_dir = d["mcms_dir"] -spec_dir = d["spectra_dir"] -alms_dir = d["alms_dir"] - -pspy_utils.create_directory(spec_dir) - -master_alms, nsplit, arrays, templates, filter_dicts = {}, {}, {}, {}, {} -# read the alms - -for sv in surveys: - arrays[sv] = d[f"arrays_{sv}"] - if apply_kspace_filter == True: filter_dicts[sv] = d[f"k_filter_{sv}"] - templates[sv] = so_map.read_map(d[f"window_T_{sv}_{arrays[sv][0]}"]) - - for ar in arrays[sv]: - maps = d[f"maps_{sv}_{ar}"] - nsplit[sv] = len(maps) - log.info(f"{nsplit[sv]} split of survey: {sv}, array {ar}") - for k, map in enumerate(maps): - master_alms[sv, ar, k] = np.load(f"{alms_dir}/alms_{sv}_{ar}_{k}.npy") - log.debug(f"{alms_dir}/alms_{sv}_{ar}_{k}.npy") - -# compute the transfer functions -_, _, lb, _ = pspy_utils.read_binning_file(binning_file, lmax) -n_bins = len(lb) -spec_name_list = pspipe_list.get_spec_name_list(d, char="_") - -if apply_kspace_filter: - kspace_tf_path = d["kspace_tf_path"] - if kspace_tf_path == "analytical": - kspace_transfer_matrix = kspace.build_analytic_kspace_filter_matrices(surveys, - arrays, - templates, - filter_dicts, - binning_file, - lmax) - else: - kspace_transfer_matrix = {} - for spec_name in spec_name_list: - kspace_transfer_matrix[spec_name] = np.load(f"{kspace_tf_path}/kspace_matrix_{spec_name}.npy", allow_pickle=True) - - - # this will be used in the covariance computation - for spec_name in spec_name_list: - one_d_tf = kspace_transfer_matrix[spec_name].diagonal() - if cov_T_E_only == True: one_d_tf = one_d_tf[:4 * n_bins] - np.savetxt(f"{spec_dir}/one_dimension_kspace_tf_{spec_name}.dat", one_d_tf) - -# compute the power spectra - -spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"] - -n_spec, sv1_list, ar1_list, sv2_list, ar2_list = pspipe_list.get_spectra_list(d) -log.info(f"number of spectra to compute : {n_spec}") -so_mpi.init(True) -subtasks = so_mpi.taskrange(imin=0, imax=n_spec - 1) - -for task in subtasks: - task = int(task) - sv1, ar1, sv2, ar2 = sv1_list[task], ar1_list[task], sv2_list[task], ar2_list[task] - - log.info(f"[{task}] Computing spectra for {sv1}_{ar1} x {sv2}_{ar2}") - - xtra_pw1, xtra_pw2, mm_tf1, mm_tf2 = None, None, None, None - if deconvolve_pixwin: - if d[f"pixwin_{sv1}"]["pix"] == "HEALPIX": - pixwin_l = hp.pixwin(d[f"pixwin_{sv1}"]["nside"]) - lb, xtra_pw1 = pspy_utils.naive_binning(np.arange(len(pixwin_l)), pixwin_l, binning_file, lmax) - if d[f"pixwin_{sv2}"]["pix"] == "HEALPIX": - pixwin_l = hp.pixwin(d[f"pixwin_{sv2}"]["nside"]) - lb, xtra_pw2 = pspy_utils.naive_binning(np.arange(len(pixwin_l)), pixwin_l, binning_file, lmax) - - - if d[f"deconvolve_map_maker_tf_{sv1}"]: - mm_tf1 = so_spectra.read_ps(d[f"mm_tf_{sv1}_{ar1}.dat"], spectra=spectra) - if d[f"deconvolve_map_maker_tf_{sv2}"]: - mm_tf2 = so_spectra.read_ps(d[f"mm_tf_{sv2}_{ar2}.dat"], spectra=spectra) - - - ps_dict = {} - for spec in spectra: - ps_dict[spec, "auto"] = [] - ps_dict[spec, "cross"] = [] - - nsplits_1 = nsplit[sv1] - nsplits_2 = nsplit[sv2] - - spin_pairs = ["spin0xspin0", "spin0xspin2", "spin2xspin0", "spin2xspin2"] - mbb_inv, Bbl = so_mcm.read_coupling(prefix=f"{mcm_dir}/{sv1}_{ar1}x{sv2}_{ar2}", - spin_pairs=spin_pairs) - - - for s1 in range(nsplits_1): - for s2 in range(nsplits_2): - if (sv1 == sv2) & (ar1 == ar2) & (s1>s2) : continue - - l, ps_master = so_spectra.get_spectra_pixell(master_alms[sv1, ar1, s1], - master_alms[sv2, ar2, s2], - spectra=spectra) - - spec_name=f"{type}_{sv1}_{ar1}x{sv2}_{ar2}_{s1}{s2}" - - lb, ps = so_spectra.bin_spectra(l, - ps_master, - binning_file, - lmax, - type=type, - mbb_inv=mbb_inv, - spectra=spectra, - binned_mcm=binned_mcm) - if apply_kspace_filter: - lb, ps = kspace.deconvolve_kspace_filter_matrix(lb, - ps, - kspace_transfer_matrix[f"{sv1}_{ar1}x{sv2}_{ar2}"], - spectra) - - - lb, ps = transfer_function.deconvolve_xtra_tf(lb, - ps, - spectra, - xtra_pw1=xtra_pw1, - xtra_pw2=xtra_pw2, - mm_tf1=mm_tf1, - mm_tf2=mm_tf2) - - if write_all_spectra: - so_spectra.write_ps(spec_dir + f"/{spec_name}.dat", lb, ps, type, spectra=spectra) - - for count, spec in enumerate(spectra): - if (s1 == s2) & (sv1 == sv2): - if count == 0: log.debug(f"[{task}] auto {sv1}_{ar1} x {sv2}_{ar2} {s1}{s2}") - ps_dict[spec, "auto"] += [ps[spec]] - else: - if count == 0: log.debug(f"[{task}] cross {sv1}_{ar1} x {sv2}_{ar2} {s1}{s2}") - ps_dict[spec, "cross"] += [ps[spec]] - - ps_dict_auto_mean = {} - ps_dict_cross_mean = {} - ps_dict_noise_mean = {} - - for spec in spectra: - ps_dict_cross_mean[spec] = np.mean(ps_dict[spec, "cross"], axis=0) - spec_name_cross = f"{type}_{sv1}_{ar1}x{sv2}_{ar2}_cross" - - if ar1 == ar2 and sv1 == sv2: - # Average TE / ET so that for same array same season TE = ET - ps_dict_cross_mean[spec] = (np.mean(ps_dict[spec, "cross"], axis=0) + np.mean(ps_dict[spec[::-1], "cross"], axis=0)) / 2. - - if sv1 == sv2: - ps_dict_auto_mean[spec] = np.mean(ps_dict[spec, "auto"], axis=0) - spec_name_auto = f"{type}_{sv1}_{ar1}x{sv2}_{ar2}_auto" - ps_dict_noise_mean[spec] = (ps_dict_auto_mean[spec] - ps_dict_cross_mean[spec]) / nsplit[sv1] - spec_name_noise = f"{type}_{sv1}_{ar1}x{sv2}_{ar2}_noise" - - so_spectra.write_ps(spec_dir + f"/{spec_name_cross}.dat", lb, ps_dict_cross_mean, type, spectra=spectra) - if sv1 == sv2: - so_spectra.write_ps(spec_dir + f"/{spec_name_auto}.dat", lb, ps_dict_auto_mean, type, spectra=spectra) - so_spectra.write_ps(spec_dir + f"/{spec_name_noise}.dat", lb, ps_dict_noise_mean, type, spectra=spectra) diff --git a/project/ana_cov_comp/python/get_split_covariance_aniso.py b/project/ana_cov_comp/python/get_split_covariance_aniso.py deleted file mode 100644 index 750b437b..00000000 --- a/project/ana_cov_comp/python/get_split_covariance_aniso.py +++ /dev/null @@ -1,171 +0,0 @@ -""" -This script compute the analytical covariance matrix elements -between split power spectra -""" -import sys -import numpy as np -from pspipe_utils import best_fits, log -from pspy import pspy_utils, so_cov, so_dict, so_map, so_mcm, so_mpi, so_spectra -from itertools import combinations_with_replacement as cwr -from itertools import combinations -from pspipe_utils import best_fits, pspipe_list - -d = so_dict.so_dict() -d.read_from_file(sys.argv[1]) - -log = log.get_logger(**d) - -mcms_dir = d["mcms_dir"] -noise_dir = d["split_noise_dir"] -bestfit_dir = d["best_fits_dir"] - -cov_dir = d["split_covariances_dir"] - -pspy_utils.create_directory(cov_dir) - -surveys = d["surveys"] -binning_file = d["binning_file"] -lmax = d["lmax"] -niter = d["niter"] -binned_mcm = d["binned_mcm"] -apply_kspace_filter = d["apply_kspace_filter"] -cov_T_E_only = d["cov_T_E_only"] - -spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"] -spin_pairs = ["spin0xspin0", "spin0xspin2", "spin2xspin0", "spin2xspin2"] - - -# l_cmb, cmb_dict = best_fits.cmb_dict_from_file(bestfit_dir + "/cmb.dat", lmax, spectra) - -arrays = {sv: d[f"arrays_{sv}"] for sv in surveys} -n_splits = {sv: d[f"n_splits_{sv}"] for sv in surveys} -array_list = [f"{sv}_{ar}" for sv in surveys for ar in arrays[sv]] - -spec_name_list = pspipe_list.get_spec_name_list(d) -sv_array_list = [f"{sv}_{ar}" for sv in surveys for ar in arrays[sv]] -l_fg, fg_dict = best_fits.fg_dict_from_files(bestfit_dir + "/fg_{}x{}.dat", sv_array_list, lmax, spectra) - -l_cmb, cmb_dict = best_fits.cmb_dict_from_file(bestfit_dir + "/cmb.dat", lmax, spectra) -l, ps_all = best_fits.get_all_best_fit(spec_name_list, - l_cmb, - cmb_dict, - fg_dict, - spectra) - -noise_dict = {} -f_name_noise = f"{noise_dir}/Dl_" + "{}x{}_{}_noise_model.dat" -for sv in surveys: - for ar in arrays[sv]: - split_list = {sv: [f"{ar}_{i}" for i in range(n_splits[sv])]} - _, nlth = best_fits.noise_dict_from_files(f_name_noise, [sv], split_list, lmax=d["lmax"],spectra=spectra) - noise_dict.update(nlth) - -cov_name = [] -for sv in surveys: - for ar in arrays[sv]: - for id1, xsplit1 in enumerate(combinations(range(n_splits[sv]), 2)): - for id2, xsplit2 in enumerate(combinations(range(n_splits[sv]), 2)): - if id1 > id2: continue - - - for pol in ("T",): # TODO: add back E - - na = f"{sv}&{ar}&{xsplit1[0]}&{pol}" - nb = f"{sv}&{ar}&{xsplit1[1]}&{pol}" - nc = f"{sv}&{ar}&{xsplit2[0]}&{pol}" - nd = f"{sv}&{ar}&{xsplit2[1]}&{pol}" - - cov_name.append((na, nb, nc, nd)) - -ncovs = len(cov_name) - - -log.info(f"number of covariance matrices to compute : {ncovs}") - -so_mpi.init(True) -subtasks = so_mpi.taskrange(imin=0, imax=ncovs - 1) - -if len(sys.argv) == 4: - log.info(f"computing only the covariance matrices : {int(sys.argv[2])}:{int(sys.argv[3])}") - subtasks = subtasks[int(sys.argv[2]):int(sys.argv[3])] -log.info(subtasks) - -for task in subtasks: - na, nb, nc, nd = cov_name[task] - - sv_a, ar_a, split_a, pol_a = na.split("&") - sv_b, ar_b, split_b, pol_b = nb.split("&") - sv_c, ar_c, split_c, pol_c = nc.split("&") - sv_d, ar_d, split_d, pol_d = nd.split("&") - - na = f"{sv_a}&{ar_a}" - nb = f"{sv_b}&{ar_b}" - nc = f"{sv_c}&{ar_c}" - nd = f"{sv_d}&{ar_d}" - - na_r, nb_r, nc_r, nd_r = na.replace("&", "_"), nb.replace("&", "_"), nc.replace("&", "_"), nd.replace("&", "_") - - try: mbb_inv_ab, Bbl_ab = so_mcm.read_coupling(prefix=f"{mcms_dir}/{na_r}x{nb_r}", spin_pairs=spin_pairs) - except: mbb_inv_ab, Bbl_ab = so_mcm.read_coupling(prefix=f"{mcms_dir}/{nb_r}x{na_r}", spin_pairs=spin_pairs) - - try: mbb_inv_cd, Bbl_cd = so_mcm.read_coupling(prefix=f"{mcms_dir}/{nc_r}x{nd_r}", spin_pairs=spin_pairs) - except: mbb_inv_cd, Bbl_cd = so_mcm.read_coupling(prefix=f"{mcms_dir}/{nd_r}x{nc_r}", spin_pairs=spin_pairs) - - log.info(f"[task] cov element ({na_r} x {nb_r}, {nc_r} x {nd_r}) {split_a}{split_b}{split_c}{split_d} {pol_a}{pol_b}{pol_c}{pol_d}") - - splits = [split_a, split_b, split_c, split_d] - survey_id = [pol_a + "a", pol_b + "b", pol_c + "c", pol_d + "d"] - splitname2splitindex = dict(zip(splits, range(len(splits)))) - - survey_combo = sv_a, sv_b, sv_c, sv_d - array_combo = ar_a, ar_b, ar_c, ar_d - - win, var = {}, {} - white_noise = {} - for splitname1, id1, sv1, ar1 in zip(splits, survey_id, survey_combo, array_combo): - # log.info(f"{name1}, {id1}, {sv1}, {ar1}") - split_index = splitname2splitindex[splitname1] - ivar_fname = d[f"ivar_{sv1}_{ar1}"][split_index] - ivar1 = so_map.read_map(ivar_fname).data - var1 = np.reciprocal(ivar1,where=(ivar1!=0)) - maskpol = "T" if id1[0] == "T" else "pol" - win[id1] = so_map.read_map(d[f"window_{maskpol}_{sv1}_{ar1}"]) - white_noise[id1] = so_cov.measure_white_noise_level(var1, win[id1].data) - var[id1] = so_cov.make_weighted_variance_map(var1, win[id1]) - - - log.info(white_noise) - - - Clth_dict = {} - Rl_dict = {} - - for splitname1, id1, sv1, ar1 in zip(splits, survey_id, survey_combo, array_combo): - s1 = splitname1.replace("set", "") - X1 = id1[0] - ndl = noise_dict[sv1, (f"{ar1}_{s1}"), (f"{ar1}_{s1}")][X1 + X1] - ell = np.arange(2, lmax) - dlfac = (ell * (ell + 1) / (2 * np.pi)) - nl = ndl[:(lmax-2)] / dlfac - - Rl_dict[id1] = np.sqrt(nl / white_noise[id1]) - - for name2, id2, sv2, ar2 in zip(splits, survey_id, survey_combo, array_combo): - s2 = name2.replace("set", "") - X2 = id2[0] - # k = (f"{sv1}&{ar1}&{s1}", f"{sv2}&{ar2}&{s2}", X1+X2) - dlth = ps_all[f"{sv1}&{ar1}", f"{sv2}&{ar2}", X1+X2] - Clth_dict[id1 + id2] = dlth / dlfac[:(lmax-2)] - - - couplings = so_cov.generate_aniso_couplings_TTTT(survey_id, splits, win, var, lmax) - - cov_tttt = so_cov.coupled_cov_aniso_same_pol(survey_id, Clth_dict, Rl_dict, couplings) - - np.save(f"{cov_dir}/analytic_aniso_coupled_cov_{na_r}_{split_a}x{nb_r}_{split_b}_{nc_r}_{split_c}x{nd_r}_{split_d}.npy", - cov_tttt) - - full_analytic_cov = mbb_inv_ab["spin0xspin0"] @ cov_tttt @ mbb_inv_cd["spin0xspin0"].T - analytic_cov = so_cov.bin_mat(full_analytic_cov, binning_file, lmax) - np.save(f"{cov_dir}/analytic_aniso_cov_{na_r}_{split_a}x{nb_r}_{split_b}_{nc_r}_{split_c}x{nd_r}_{split_d}.npy", - analytic_cov) diff --git a/project/ana_cov_comp/python/get_sq_windows_alms.py b/project/ana_cov_comp/python/get_sq_windows_alms.py deleted file mode 100644 index 09ffe47a..00000000 --- a/project/ana_cov_comp/python/get_sq_windows_alms.py +++ /dev/null @@ -1,65 +0,0 @@ -""" -This script compute all alms squared windows, it's a necessary step of covariance computation. -""" -import sys - -import numpy as np -from pspipe_utils import log, pspipe_list -from pspy import pspy_utils, so_dict, so_map, so_mpi, sph_tools - - -def mult(map_a, map_b, log): - res_a = 1 / map_a.data.pixsize() - res_b = 1 / map_b.data.pixsize() - - if res_a == res_b: - prod = map_a.copy() - prod.data *= map_b.data - elif res_a < res_b: - log.info("resample map a") - prod = map_b.copy() - map_a_proj = so_map.car2car(map_a, map_b) - prod.data *= map_a_proj.data - elif res_b < res_a: - log.info("resample map b") - prod = map_a.copy() - map_b_proj = so_map.car2car(map_b, map_a) - prod.data *= map_b_proj.data - - return prod - - -d = so_dict.so_dict() -d.read_from_file(sys.argv[1]) -log = log.get_logger(**d) - -surveys = d["surveys"] -lmax = d["lmax"] -niter = d["niter"] - -sq_win_alms_dir = d["sq_win_alms_dir"] - -pspy_utils.create_directory(sq_win_alms_dir) - -n_sq_alms, sv1_list, ar1_list, sv2_list, ar2_list = pspipe_list.get_spectra_list(d) - - -log.info(f"number of sq win alms to compute : {n_sq_alms}") -so_mpi.init(True) -subtasks = so_mpi.taskrange(imin=0, imax=n_sq_alms - 1) -print(subtasks) -for task in subtasks: - task = int(task) - sv1, ar1, sv2, ar2 = sv1_list[task], ar1_list[task], sv2_list[task], ar2_list[task] - - log.info(f"[{task:02d}] Computing map2alm for {sv1}_{ar1}x{sv2}_{ar2}...") - - win_T1 = so_map.read_map(d["window_T_%s_%s" % (sv1, ar1)]) - win_T2 = so_map.read_map(d["window_T_%s_%s" % (sv2, ar2)]) - - sq_win = mult(win_T1, win_T2, log) - # sq_win = win_T1.copy() - # sq_win.data[:] *= win_T2.data[:] - sqwin_alm = sph_tools.map2alm(sq_win, niter=niter, lmax=lmax) - - np.save(f"{sq_win_alms_dir}/alms_{sv1}_{ar1}x{sv2}_{ar2}.npy", sqwin_alm) diff --git a/project/ana_cov_comp/python/get_window_dr6.py b/project/ana_cov_comp/python/get_window_dr6.py deleted file mode 100644 index f2cf0303..00000000 --- a/project/ana_cov_comp/python/get_window_dr6.py +++ /dev/null @@ -1,245 +0,0 @@ -""" -This script create the window functions used in the PS computation -They consist of a point source mask, a galactic mask, a mask based on the amount of cross linking in the data, and a coordinate mask, note that -we also produce a window that include the pixel weighting. -The different masks are apodized. -We also produce a kspace-mask that will later be used for the kspace filtering operation, in order to remove the edges of the survey and avoid nasty pixels. -""" - -import sys - -import numpy as np -from pspipe_utils import log, pspipe_list -from pspy import pspy_utils, so_dict, so_map, so_mpi, so_window -from pixell import enmap - -def create_crosslink_mask(xlink_map, cross_link_threshold): - """ - Create a mask to remove pixels with a small amount of x-linking - We compute this using product from the map maker which assess the amount - of scan direction that hits each pixels in the map - the product have 3 component and we compute sqrt(Q **2 + U ** 2) / I by analogy with the polarisation fraction - A high value of this quantity means low level of xlinking, we mask all pixels above a given threshold - note that the mask is designed on a downgraded version of the maps, this is to avoid small scale structure in the mask - Parameters - ---------- - xlink_map: so_map - 3 component so_map assessing the direction of scan hitting each pixels - cross_link_threshold: float - a threshold above which region of the sky are considered not x-linked - """ - - xlink = so_map.read_map(xlink_map) - xlink_lowres = xlink.downgrade(32) - with np.errstate(invalid="ignore"): - x_mask = ( - np.sqrt(xlink_lowres.data[1] ** 2 + xlink_lowres.data[2] ** 2) / xlink_lowres.data[0] - ) - x_mask[np.isnan(x_mask)] = 1 - x_mask[x_mask >= cross_link_threshold] = 1 - x_mask[x_mask < cross_link_threshold] = 0 - x_mask = 1 - x_mask - - xlink_lowres.data[0] = x_mask - xlink = so_map.car2car(xlink_lowres, xlink) - x_mask = xlink.data[0].copy() - id = np.where(x_mask > 0.9) - x_mask[:] = 0 - x_mask[id] = 1 - return x_mask - -def apply_coordinate_mask(mask, coord): - """ - Apply a mask based on coordinate - - Parameters - ---------- - mask: so_map - the mask on which the coordinate mask will be applied - coord: list of list offloat - format is [[dec0, ra0], [dec1, ra1]] - we create a rectangle mask from this coordinate - in the convention assumed in this code - dec0, ra0 are the coordiante of the top left corner - dec1, ra1 are the coordinate of the bottom right corner - """ - - dec_ra = np.deg2rad(coord) - pix1 = mask.data.sky2pix(dec_ra[0]) - pix2 = mask.data.sky2pix(dec_ra[1]) - min_pix = np.min([pix1, pix2], axis=0).astype(int) - max_pix = np.max([pix1, pix2], axis=0).astype(int) - mask.data[min_pix[0] : max_pix[0], min_pix[1] : max_pix[1]] = 0 - - return mask - -d = so_dict.so_dict() -d.read_from_file(sys.argv[1]) -log = log.get_logger(**d) - - -surveys = d["surveys"] -# the apodisation length of the point source mask in degree -apod_pts_source_degree = d["apod_pts_source_degree"] -# the apodisation length of the final survey mask -apod_survey_degree = d["apod_survey_degree"] -# the threshold on the amount of cross linking to keep the data in -cross_link_threshold = d["cross_link_threshold"] -# pixel weight with inverse variance above n_ivar * median are set to ivar -# this ensure that the window is not totally dominated by few pixels with too much weight -n_med_ivar = d["n_med_ivar"] - -# we will skip the edges of the survey where the noise is very difficult to model, the default is to skip 0.5 degree for -# constructing the kspace mask and 2 degree for the final survey mask, this parameter can be used to rescale the default -rescale = d["edge_skip_rescale"] - -window_dir = d["window_dir"] - -pspy_utils.create_directory(window_dir) - -patch = None -if "patch" in d: - patch = so_map.read_map(d["patch"]) - -# here we list the different windows that need to be computed, we will then do a MPI loops over this list -n_wins, sv_list, ar_list = pspipe_list.get_arrays_list(d) - -log.info(f"number of windows to compute : {n_wins}") - -my_masks = {} - -so_mpi.init(True) -subtasks = so_mpi.taskrange(imin=0, imax=n_wins - 1) - -if len(sys.argv) == 4: - log.info(f"computing only the windows : {int(sys.argv[2])}:{int(sys.argv[3])}") - subtasks = subtasks[int(sys.argv[2]):int(sys.argv[3])] - -for task in subtasks: - task = int(task) - sv, ar = sv_list[task], ar_list[task] - - log.info(f"[{task}] create windows for '{sv}' survey and '{ar}' array...") - - gal_mask = so_map.read_map(d[f"gal_mask_{sv}_{ar}"]) - - my_masks["baseline"] = gal_mask.copy() - my_masks["baseline"].data[:] = 1 - - maps = d[f"maps_{sv}_{ar}"] - - ivar_all = gal_mask.copy() - ivar_all.data[:] = 0 - - # the first step is to iterate on the maps of a given array to identify pixels with zero values - # we will form a first survey mask as the union of all these split masks - # we also compute the average ivar map - - log.info(f"[{task}] create survey mask from ivar maps") - - for k, map in enumerate(maps): - if d[f"src_free_maps_{sv}"] == True: - index = map.find("map_srcfree.fits") - else: - index = map.find("map.fits") - - ivar_map = map[:index] + "ivar.fits" - ivar_map = so_map.read_map(ivar_map) - my_masks["baseline"].data[ivar_map.data[:] == 0.0] = 0.0 - ivar_all.data[:] += ivar_map.data[:] - - ivar_all.data[:] /= np.max(ivar_all.data[:]) - - # optionnaly apply an extra coordinate mask - if d[f"coord_mask_{sv}_{ar}"] is not None: - log.info(f"[{task}] apply coord mask") - my_masks["baseline"] = apply_coordinate_mask(my_masks["baseline"], d[f"coord_mask_{sv}_{ar}"]) - - - log.info(f"[{task}] compute distance to the edges and remove {0.5*rescale:.2f} degree from the edges") - # compute the distance to the nearest 0 - dist = so_window.get_distance(my_masks["baseline"], rmax = 4 * rescale * np.pi / 180) - # here we remove pixels near the edges in order to avoid pixels with very low hit count - # note the hardcoded 0.5 degree value that can be rescale with the edge_skip_rescale argument in the dictfile. - my_masks["baseline"].data[dist.data < 0.5 * rescale] = 0 - - # apply the galactic mask - log.info(f"[{task}] apply galactic mask") - my_masks["baseline"].data *= gal_mask.data - - # with this we can create the k space mask this will only be used for applying the kspace filter - log.info(f"[{task}] appodize kspace mask with {rescale:.2f} apod and write it to disk") - my_masks["kspace"] = my_masks["baseline"].copy() - - # we apodize this k space mask with a 1 degree apodisation - - my_masks["kspace"] = so_window.create_apodization(my_masks["kspace"], "C1", 1 * rescale, use_rmax=True) - my_masks["kspace"].data = my_masks["kspace"].data.astype(np.float32) - my_masks["kspace"].write_map(f"{window_dir}/kspace_mask_{sv}_{ar}.fits") - - # compare to the kspace mask we will skip for the nominal mask - # an additional 2 degrees to avoid ringing from the filter - - dist = so_window.get_distance(my_masks["baseline"], rmax = 4 * rescale * np.pi / 180) - my_masks["baseline"].data[dist.data < 2 * rescale] = 0 - - # now we create a xlink mask based on the xlink threshold - - log.info(f"[{task}] create xlink mask") - - my_masks["xlink"] = my_masks["baseline"].copy() - for k, map in enumerate(maps): - if d[f"src_free_maps_{sv}"] == True: - index = map.find("map_srcfree.fits") - else: - index = map.find("map.fits") - - xlink_map = map[:index] + "xlink.fits" - x_mask = create_crosslink_mask(xlink_map, cross_link_threshold) - my_masks["xlink"].data *= x_mask - - - for mask_type in ["baseline", "xlink"]: - - # optionnaly apply a patch mask - - if patch is not None: - log.info(f"[{task}] apply patch mask") - - my_masks[mask_type].data *= patch.data - - - log.info(f"[{task}] apodize mask with {apod_survey_degree:.2f} apod") - - # apodisation of the final mask - my_masks[mask_type] = so_window.create_apodization(my_masks[mask_type], "C1", apod_survey_degree, use_rmax=True) - - # create a point source mask and apodise it - - log.info(f"[{task}] include ps mask") - - ps_mask = so_map.read_map(d[f"ps_mask_{sv}_{ar}"]) - ps_mask = so_window.create_apodization(ps_mask, "C1", apod_pts_source_degree, use_rmax=True) - my_masks[mask_type].data *= ps_mask.data - - my_masks[mask_type].data = my_masks[mask_type].data.astype(np.float32) - my_masks[mask_type].write_map(f"{window_dir}/window_{sv}_{ar}_{mask_type}.fits") - - # we also make a version of the windows taking into account the ivar of the maps - log.info(f"[{task}] include ivar ") - - mask_type_w = mask_type + "_ivar" - - my_masks[mask_type_w] = my_masks[mask_type].copy() - id = np.where(ivar_all.data[:] * my_masks[mask_type_w].data[:] != 0) - med = np.median(ivar_all.data[id]) - ivar_all.data[ivar_all.data[:] > n_med_ivar * med] = n_med_ivar * med - my_masks[mask_type_w].data[:] *= ivar_all.data[:] - - my_masks[mask_type_w].data = my_masks[mask_type_w].data.astype(np.float32) - my_masks[mask_type_w].write_map(f"{window_dir}/window_w_{sv}_{ar}_{mask_type}.fits") - - for mask_type, mask in my_masks.items(): - log.info(f"[{task}] downgrade and plot {mask_type} ") - mask = mask.downgrade(4) - mask.plot(file_name=f"{window_dir}/{mask_type}_mask_{sv}_{ar}")