From 43e252378a21bf21b674463aa997cd007ac8cb0f Mon Sep 17 00:00:00 2001 From: Louis Thibaut Date: Wed, 6 Nov 2024 13:17:42 +0100 Subject: [PATCH] add soft --- .../custom_nulls/code/get_calibs_inout.py | 222 +++++++++++++ .../custom_nulls/code/get_calibs_surveys.py | 209 +++++++++++++ .../code/get_polar_eff_LCDM_cal.py | 213 +++++++++++++ .../code/get_polar_eff_LCDM_cal_inout.py | 218 +++++++++++++ .../custom_nulls/code/get_window_surveys.py | 224 ++++++++++++++ .../code/mc_mnms_get_custom_nlms.py | 127 ++++++++ .../mc_mnms_get_spectra_from_custom_nlms.py | 292 ++++++++++++++++++ 7 files changed, 1505 insertions(+) create mode 100644 project/data_analysis/python/null_tests/custom_nulls/code/get_calibs_inout.py create mode 100644 project/data_analysis/python/null_tests/custom_nulls/code/get_calibs_surveys.py create mode 100644 project/data_analysis/python/null_tests/custom_nulls/code/get_polar_eff_LCDM_cal.py create mode 100644 project/data_analysis/python/null_tests/custom_nulls/code/get_polar_eff_LCDM_cal_inout.py create mode 100644 project/data_analysis/python/null_tests/custom_nulls/code/get_window_surveys.py create mode 100644 project/data_analysis/python/null_tests/custom_nulls/code/mc_mnms_get_custom_nlms.py create mode 100644 project/data_analysis/python/null_tests/custom_nulls/code/mc_mnms_get_spectra_from_custom_nlms.py diff --git a/project/data_analysis/python/null_tests/custom_nulls/code/get_calibs_inout.py b/project/data_analysis/python/null_tests/custom_nulls/code/get_calibs_inout.py new file mode 100644 index 00000000..0eebd184 --- /dev/null +++ b/project/data_analysis/python/null_tests/custom_nulls/code/get_calibs_inout.py @@ -0,0 +1,222 @@ +import matplotlib +matplotlib.use("Agg") +from pspy import pspy_utils, so_dict, so_spectra, so_cov +from pspipe_utils import consistency, log +import numpy as np +import pylab as plt +import pickle +import sys + + +def get_proj_pattern(test, map_set, ref_map_set): + + if test == "AxA-AxP": + proj_pattern = np.array([1, -1, 0]) + # [1, -1, 0]] x [ AxA, AxP, PxP].T = AxA - AxP + name = f"{map_set}x{map_set}-{map_set}x{ref_map_set}" + elif test == "AxA-PxP": + proj_pattern = np.array([1, 0, -1]) + # [1, 0, -1]] x [ AxA, AxP, PxP].T = AxA - PxP + name = f"{map_set}x{map_set}-{ref_map_set}x{ref_map_set}" + elif test == "PxP-AxP": + proj_pattern = np.array([0, -1, 1]) + # [[0, -1, 1]] x [ AxA, AxP, PxP].T = PxP - AxP + name = f"{ref_map_set}x{ref_map_set}-{map_set}x{ref_map_set}" + + return name, proj_pattern + +d = so_dict.so_dict() +d.read_from_file(sys.argv[1]) +log = log.get_logger(**d) + +planck_corr = True +subtract_bf_fg = True + + +output_dir = "calibration_results" +spec_dir = "spectra" +bestfir_dir = "best_fits" +cov_dir = "covariances" + +if planck_corr: + spec_dir = "spectra_leak_corr_planck_bias_corr" + output_dir += "_planck_bias_corrected" + +if subtract_bf_fg: + output_dir += "_fg_sub" + + + +_, _, lb, _ = pspy_utils.read_binning_file(d["binning_file"], d["lmax"]) +n_bins = len(lb) + +spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"] + +if d["cov_T_E_only"] == True: + modes = ["TT", "TE", "ET", "EE"] +else: + modes = spectra + + +# Create output dirs + +residual_output_dir = f"{output_dir}/residuals" +plot_output_dir = f"{output_dir}/plots" +chains_dir = f"{output_dir}/chains" + +pspy_utils.create_directory(output_dir) +pspy_utils.create_directory(residual_output_dir) +pspy_utils.create_directory(chains_dir) +pspy_utils.create_directory(plot_output_dir) + + +surveys = d["surveys"] +try: + surveys.remove("Planck") +except: + pass + +results_dict = {} +ref_map_sets = {} +for sv in surveys: + + # Define the multipole range used to obtain + # the calibration amplitudes + multipole_range = {f"{sv}_pa4_f150_in": [1250, 1800], + f"{sv}_pa4_f220_in": [1250, 2000], + f"{sv}_pa5_f090_in": [1000, 2000], + f"{sv}_pa5_f150_in": [800, 2000], + f"{sv}_pa6_f090_in": [1000, 2000], + f"{sv}_pa6_f150_in": [600, 2000], + f"{sv}_pa4_f150_out": [1250, 1800], + f"{sv}_pa4_f220_out": [1250, 2000], + f"{sv}_pa5_f090_out": [1000, 2000], + f"{sv}_pa5_f150_out": [800, 2000], + f"{sv}_pa6_f090_out": [1000, 2000], + f"{sv}_pa6_f150_out": [600, 2000]} + + # Define the reference arrays + ref_map_sets[f"{sv}_pa4_f150_in"] = "Planck_f143" + ref_map_sets[f"{sv}_pa4_f220_in"] = "Planck_f217" + ref_map_sets[f"{sv}_pa5_f090_in"] = "Planck_f143" + ref_map_sets[f"{sv}_pa5_f150_in"] = "Planck_f143" + ref_map_sets[f"{sv}_pa6_f090_in"] = "Planck_f143" + ref_map_sets[f"{sv}_pa6_f150_in"] = "Planck_f143" + + ref_map_sets[f"{sv}_pa4_f150_out"] = "Planck_f143" + ref_map_sets[f"{sv}_pa4_f220_out"] = "Planck_f217" + ref_map_sets[f"{sv}_pa5_f090_out"] = "Planck_f143" + ref_map_sets[f"{sv}_pa5_f150_out"] = "Planck_f143" + ref_map_sets[f"{sv}_pa6_f090_out"] = "Planck_f143" + ref_map_sets[f"{sv}_pa6_f150_out"] = "Planck_f143" + + y_lims = {"TT": (-100000, 75000),} + + tests = ["AxA-AxP", "AxA-PxP", "PxP-AxP"] + + for test in tests: + for ar in d[f"arrays_{sv}"]: + + map_set = f"{sv}_{ar}" + ref_map_set = ref_map_sets[map_set] + + name, proj_pattern = get_proj_pattern(test, map_set, ref_map_set) + + spectra_for_cal = [(map_set, map_set, "TT"), + (map_set, ref_map_set, "TT"), + (ref_map_set, ref_map_set, "TT")] + + # Load spectra and cov + ps_dict = {} + cov_dict = {} + for i, (ms1, ms2, m1) in enumerate(spectra_for_cal): + _, ps = so_spectra.read_ps(f"{spec_dir}/Dl_{ms1}x{ms2}_cross.dat", spectra=spectra) + + ps_dict[ms1, ms2, m1] = ps[m1] + + if (m1 == "TT") & (subtract_bf_fg): + log.info(f"remove fg {m1} {ms1} x {ms2}") + l_fg, bf_fg = so_spectra.read_ps(f"{bestfir_dir}/fg_{ms1}x{ms2}.dat", spectra=spectra) + _, bf_fg_TT_binned = pspy_utils.naive_binning(l_fg, bf_fg["TT"], d["binning_file"], d["lmax"]) + ps_dict[ms1, ms2, m1] -= bf_fg_TT_binned + + + ps_dict[ms1, ms2, m1] = ps[m1] + for j, (ms3, ms4, m2) in enumerate(spectra_for_cal): + if j < i: continue + cov = np.load(f"{cov_dir}/analytic_cov_{ms1}x{ms2}_{ms3}x{ms4}.npy") + cov = so_cov.selectblock(cov, modes, n_bins = n_bins, block = m1 + m2) + cov_dict[(ms1, ms2, m1), (ms3, ms4, m2)] = cov + + # Concatenation + spec_vec, full_cov = consistency.append_spectra_and_cov(ps_dict, cov_dict, spectra_for_cal) + + # Save and plot residuals before calibration + res_spectrum, res_cov = consistency.project_spectra_vec_and_cov(spec_vec, full_cov, proj_pattern) + np.savetxt(f"{residual_output_dir}/residual_{name}_before.dat", np.array([lb, res_spectrum]).T) + np.savetxt(f"{residual_output_dir}/residual_cov_{name}.dat", res_cov) + + lmin, lmax = multipole_range[map_set] + id = np.where((lb >= lmin) & (lb <= lmax))[0] + consistency.plot_residual(lb, res_spectrum, {"analytical": res_cov}, "TT", f"{map_set} {test}", + f"{plot_output_dir}/residual_{name}_before", + lrange=id, l_pow=1, ylims=y_lims["TT"]) + + # Calibrate the spectra + cal_mean, cal_std = consistency.get_calibration_amplitudes(spec_vec, full_cov, + proj_pattern, "TT", id, + f"{chains_dir}/{name}") + + + results_dict[name] = {"multipole_range": multipole_range[map_set], + "ref_map_set": ref_map_set, + "calibs": [cal_mean, cal_std]} + + calib_vec = np.array([cal_mean**2, cal_mean, 1]) + res_spectrum, res_cov = consistency.project_spectra_vec_and_cov(spec_vec, full_cov, + proj_pattern, + calib_vec = calib_vec) + + np.savetxt(f"{residual_output_dir}/residual_{name}_after.dat", np.array([lb, res_spectrum]).T) + consistency.plot_residual(lb, res_spectrum, {"analytical": res_cov}, "TT", f"{map_set} {test}", + f"{plot_output_dir}/residual_{name}_after", + lrange=id, l_pow=1, ylims=y_lims["TT"]) + + + + + +for sv in surveys: + + # plot the cal factors + color_list = ["blue", "red", "green"] + + for i, ar in enumerate(d[f"arrays_{sv}"]): + map_set = f"{sv}_{ar}" + ref_map_set = ref_map_sets[map_set] + print(f"**************") + print(f"calibration {map_set} with {ref_map_set}") + + for j, test in enumerate(tests): + name, _ = get_proj_pattern(test, map_set, ref_map_set) + cal, std = results_dict[name]["calibs"] + print(f"{test}, cal: {cal}, sigma cal: {std}") + + plt.errorbar(i + 0.9 + j * 0.1, cal, std, label = test, + color = color_list[j], marker = ".", + ls = "None", + markersize=6.5, + markeredgewidth=2) + + if i == 0: + plt.legend(fontsize = 15) + + x = np.arange(1, len(d[f"arrays_{sv}"]) + 1) + plt.xticks(x, d[f"arrays_{sv}"]) + plt.ylim(0.967, 1.06) + plt.tight_layout() + plt.savefig(f"{plot_output_dir}/calibs_summary_{sv}.pdf", bbox_inches="tight") + plt.clf() + plt.close() + + pickle.dump(results_dict, open(f"{output_dir}/calibs_dict_{sv}.pkl", "wb")) diff --git a/project/data_analysis/python/null_tests/custom_nulls/code/get_calibs_surveys.py b/project/data_analysis/python/null_tests/custom_nulls/code/get_calibs_surveys.py new file mode 100644 index 00000000..ea332360 --- /dev/null +++ b/project/data_analysis/python/null_tests/custom_nulls/code/get_calibs_surveys.py @@ -0,0 +1,209 @@ +import matplotlib +matplotlib.use("Agg") +from pspy import pspy_utils, so_dict, so_spectra, so_cov +from pspipe_utils import consistency, log +import numpy as np +import pylab as plt +import pickle +import sys + + +def get_proj_pattern(test, map_set, ref_map_set): + + if test == "AxA-AxP": + proj_pattern = np.array([1, -1, 0]) + # [1, -1, 0]] x [ AxA, AxP, PxP].T = AxA - AxP + name = f"{map_set}x{map_set}-{map_set}x{ref_map_set}" + elif test == "AxA-PxP": + proj_pattern = np.array([1, 0, -1]) + # [1, 0, -1]] x [ AxA, AxP, PxP].T = AxA - PxP + name = f"{map_set}x{map_set}-{ref_map_set}x{ref_map_set}" + elif test == "PxP-AxP": + proj_pattern = np.array([0, -1, 1]) + # [[0, -1, 1]] x [ AxA, AxP, PxP].T = PxP - AxP + name = f"{ref_map_set}x{ref_map_set}-{map_set}x{ref_map_set}" + + return name, proj_pattern + +d = so_dict.so_dict() +d.read_from_file(sys.argv[1]) +log = log.get_logger(**d) + +planck_corr = True +subtract_bf_fg = True + + +output_dir = "calibration_results" +spec_dir = "spectra" +bestfir_dir = "best_fits" +cov_dir = "covariances" + +if planck_corr: + spec_dir = "spectra_leak_corr_planck_bias_corr" + output_dir += "_planck_bias_corrected" + +if subtract_bf_fg: + output_dir += "_fg_sub" + + + +_, _, lb, _ = pspy_utils.read_binning_file(d["binning_file"], d["lmax"]) +n_bins = len(lb) + +spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"] + +if d["cov_T_E_only"] == True: + modes = ["TT", "TE", "ET", "EE"] +else: + modes = spectra + + +# Create output dirs + +residual_output_dir = f"{output_dir}/residuals" +plot_output_dir = f"{output_dir}/plots" +chains_dir = f"{output_dir}/chains" + +pspy_utils.create_directory(output_dir) +pspy_utils.create_directory(residual_output_dir) +pspy_utils.create_directory(chains_dir) +pspy_utils.create_directory(plot_output_dir) + + +surveys = d["surveys"] +try: + surveys.remove("Planck") +except: + pass + +results_dict = {} +ref_map_sets = {} +for sv in surveys: + + # Define the multipole range used to obtain + # the calibration amplitudes + multipole_range = {f"{sv}_pa4_f150": [1250, 1800], + f"{sv}_pa4_f220": [1250, 2000], + f"{sv}_pa5_f090": [1000, 2000], + f"{sv}_pa5_f150": [800, 2000], + f"{sv}_pa6_f090": [1000, 2000], + f"{sv}_pa6_f150": [600, 2000]} + + # Define the reference arrays + ref_map_sets[f"{sv}_pa4_f150"] = "Planck_f143" + ref_map_sets[f"{sv}_pa4_f220"] = "Planck_f217" + ref_map_sets[f"{sv}_pa5_f090"] = "Planck_f143" + ref_map_sets[f"{sv}_pa5_f150"] = "Planck_f143" + ref_map_sets[f"{sv}_pa6_f090"] = "Planck_f143" + ref_map_sets[f"{sv}_pa6_f150"] = "Planck_f143" + + y_lims = {"TT": (-100000, 75000),} + + tests = ["AxA-AxP", "AxA-PxP", "PxP-AxP"] + + for test in tests: + for ar in d[f"arrays_{sv}"]: + + map_set = f"{sv}_{ar}" + ref_map_set = ref_map_sets[map_set] + + name, proj_pattern = get_proj_pattern(test, map_set, ref_map_set) + + spectra_for_cal = [(map_set, map_set, "TT"), + (map_set, ref_map_set, "TT"), + (ref_map_set, ref_map_set, "TT")] + + # Load spectra and cov + ps_dict = {} + cov_dict = {} + for i, (ms1, ms2, m1) in enumerate(spectra_for_cal): + _, ps = so_spectra.read_ps(f"{spec_dir}/Dl_{ms1}x{ms2}_cross.dat", spectra=spectra) + + ps_dict[ms1, ms2, m1] = ps[m1] + + if (m1 == "TT") & (subtract_bf_fg): + log.info(f"remove fg {m1} {ms1} x {ms2}") + l_fg, bf_fg = so_spectra.read_ps(f"{bestfir_dir}/fg_{ms1}x{ms2}.dat", spectra=spectra) + _, bf_fg_TT_binned = pspy_utils.naive_binning(l_fg, bf_fg["TT"], d["binning_file"], d["lmax"]) + ps_dict[ms1, ms2, m1] -= bf_fg_TT_binned + + + ps_dict[ms1, ms2, m1] = ps[m1] + for j, (ms3, ms4, m2) in enumerate(spectra_for_cal): + if j < i: continue + cov = np.load(f"{cov_dir}/analytic_cov_{ms1}x{ms2}_{ms3}x{ms4}.npy") + cov = so_cov.selectblock(cov, modes, n_bins = n_bins, block = m1 + m2) + cov_dict[(ms1, ms2, m1), (ms3, ms4, m2)] = cov + + # Concatenation + spec_vec, full_cov = consistency.append_spectra_and_cov(ps_dict, cov_dict, spectra_for_cal) + + # Save and plot residuals before calibration + res_spectrum, res_cov = consistency.project_spectra_vec_and_cov(spec_vec, full_cov, proj_pattern) + np.savetxt(f"{residual_output_dir}/residual_{name}_before.dat", np.array([lb, res_spectrum]).T) + np.savetxt(f"{residual_output_dir}/residual_cov_{name}.dat", res_cov) + + lmin, lmax = multipole_range[map_set] + id = np.where((lb >= lmin) & (lb <= lmax))[0] + consistency.plot_residual(lb, res_spectrum, {"analytical": res_cov}, "TT", f"{map_set} {test}", + f"{plot_output_dir}/residual_{name}_before", + lrange=id, l_pow=1, ylims=y_lims["TT"]) + + # Calibrate the spectra + cal_mean, cal_std = consistency.get_calibration_amplitudes(spec_vec, full_cov, + proj_pattern, "TT", id, + f"{chains_dir}/{name}") + + + results_dict[name] = {"multipole_range": multipole_range[map_set], + "ref_map_set": ref_map_set, + "calibs": [cal_mean, cal_std]} + + calib_vec = np.array([cal_mean**2, cal_mean, 1]) + res_spectrum, res_cov = consistency.project_spectra_vec_and_cov(spec_vec, full_cov, + proj_pattern, + calib_vec = calib_vec) + + np.savetxt(f"{residual_output_dir}/residual_{name}_after.dat", np.array([lb, res_spectrum]).T) + consistency.plot_residual(lb, res_spectrum, {"analytical": res_cov}, "TT", f"{map_set} {test}", + f"{plot_output_dir}/residual_{name}_after", + lrange=id, l_pow=1, ylims=y_lims["TT"]) + + + + + +for sv in surveys: + + # plot the cal factors + color_list = ["blue", "red", "green"] + + for i, ar in enumerate(d[f"arrays_{sv}"]): + map_set = f"{sv}_{ar}" + ref_map_set = ref_map_sets[map_set] + print(f"**************") + print(f"calibration {map_set} with {ref_map_set}") + + for j, test in enumerate(tests): + name, _ = get_proj_pattern(test, map_set, ref_map_set) + cal, std = results_dict[name]["calibs"] + print(f"{test}, cal: {cal}, sigma cal: {std}") + + plt.errorbar(i + 0.9 + j * 0.1, cal, std, label = test, + color = color_list[j], marker = ".", + ls = "None", + markersize=6.5, + markeredgewidth=2) + + if i == 0: + plt.legend(fontsize = 15) + + x = np.arange(1, len(d[f"arrays_{sv}"]) + 1) + plt.xticks(x, d[f"arrays_{sv}"]) + plt.ylim(0.967, 1.06) + plt.tight_layout() + plt.savefig(f"{plot_output_dir}/calibs_summary_{sv}.pdf", bbox_inches="tight") + plt.clf() + plt.close() + + pickle.dump(results_dict, open(f"{output_dir}/calibs_dict_{sv}.pkl", "wb")) diff --git a/project/data_analysis/python/null_tests/custom_nulls/code/get_polar_eff_LCDM_cal.py b/project/data_analysis/python/null_tests/custom_nulls/code/get_polar_eff_LCDM_cal.py new file mode 100644 index 00000000..51ab4837 --- /dev/null +++ b/project/data_analysis/python/null_tests/custom_nulls/code/get_polar_eff_LCDM_cal.py @@ -0,0 +1,213 @@ +""" +This script is used to calibrate the polarization efficiencies +for each array with respect to a LCDM bestfit (default: cmb.dat) +Note that I assume that the spectra are uncalibrated and apply the calibration in this script +""" +from pspy import so_dict, pspy_utils, so_spectra, so_cov +from getdist.mcsamples import loadMCSamples +from pspipe_utils import best_fits +import matplotlib.pyplot as plt +import getdist.plots as gdplt +from cobaya.run import run +import numpy as np +import sys + +d = so_dict.so_dict() +d.read_from_file(sys.argv[1]) + +# Set up directories +ps_dir = "spectra_leak_corr" +planck_corr = False + +if planck_corr: + ps_dir = "spectra_leak_corr_planck_bias_corr" + + + +cov_dir = "covariances" +bestfit_dir = "best_fits" +mcm_dir = "mcms" +output_dir = "pol_eff_results" +pspy_utils.create_directory(output_dir) + +ps_filename = f"{bestfit_dir}/cmb.dat" + +# Load paramfiles info +surveys = d["surveys"] +arrays = {sv: d[f"arrays_{sv}"] for sv in surveys} +lmax = d["lmax"] + +spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"] + +# Load ps theory file +l_th, ps_th = so_spectra.read_ps(ps_filename, spectra=spectra) +l_th = l_th.astype(int) + + +# Load foreground dict and dust prior, then normalize (with polarized dust normalized to 1) +do_bandpass_integration = d["do_bandpass_integration"] +fg_components = d["fg_components"] +fg_params = d["fg_params"] + +dust_priors = { + "EE": {"loc": fg_params["a_gee"], "scale": 0.008}, + "TE": {"loc": fg_params["a_gte"], "scale": 0.015} +} + +fg_components["ee"] = ["dust"] +fg_components["te"] = ["dust"] +fg_params["a_gee"] = 1. +fg_params["a_gte"] = 1. + +print(dust_priors) + +passbands = {} +for sv in surveys: + for ar in arrays[sv]: + freq_info = d[f"freq_info_{sv}_{ar}"] + if do_bandpass_integration: + nu_ghz, pb = np.loadtxt(freq_info["passband"]).T + else: + nu_ghz, pb = np.array([freq_info["freq_tag"]]), np.array([1.]) + + passbands[f"{sv}_{ar}"] = [nu_ghz, pb] + +fg_dict = best_fits.get_foreground_dict(l_th, passbands, fg_components, fg_params, d["fg_norm"]) + +# Load priors on dust amplitudes +# from Planck 353 GHz spectra +# computed in ACT DR6 windows + + +# Calibration range +lmin_cal = 500 +lmax_cal = {} + +my_surveys = d["surveys"] +try: + my_surveys.remove("Planck") +except: + pass + +for sv in my_surveys: + lmax_cal[f"{sv}_pa4_f220"] = 2000 + lmax_cal[f"{sv}_pa5_f090"] = 1200 + lmax_cal[f"{sv}_pa5_f150"] = 2000 + lmax_cal[f"{sv}_pa6_f090"] = 1200 + lmax_cal[f"{sv}_pa6_f150"] = 2000 + +lmax_cal["Planck_f100"] = 1200 +lmax_cal["Planck_f143"] = 2000 +lmax_cal["Planck_f217"] = 2000 + + +# Define useful functions +def get_model(cmb_th, fg_th, Bbl, dust_amp, pol_eff, mode): + ps_theory = (cmb_th + dust_amp * fg_th) * pol_eff ** mode.count("E") + return Bbl @ ps_theory + +pol_eff_mean, pol_eff_std, dust_mean, dust_std = {}, {}, {}, {} +for sv in surveys: + for ar in arrays[sv]: + for spectrum in ["EE", "TE"]: + + print(spectrum, sv, ar, lmin_cal, lmax_cal[f"{sv}_{ar}"]) + + # Load ps and cov + spec_name = f"{sv}_{ar}x{sv}_{ar}" + lb, ps = so_spectra.read_ps(f"{ps_dir}/Dl_{spec_name}_cross.dat", spectra=spectra) + cov = np.load(f"{cov_dir}/analytic_cov_{spec_name}_{spec_name}.npy") + leakage_cov = np.load(f"{cov_dir}/leakage_cov_{spec_name}_{spec_name}.npy") + + cal = d[f"cal_{sv}_{ar}"] + + + cov += leakage_cov + # Select the spectrum + ps = ps[spectrum] * cal ** 2 + n_bins = len(lb) + + cov = so_cov.selectblock(cov, spectra, n_bins=n_bins, block=spectrum+spectrum) + + # Multipole cuts + id = np.where((lb >= lmin_cal) & (lb <= lmax_cal[f"{sv}_{ar}"]))[0] + ps = ps[id] + cov = cov[np.ix_(id, id)] + invcov = np.linalg.inv(cov) + + # Load Bbl + spin_pair = "spin2xspin2" if spectrum == "EE" else ("spin0xspin2" if spectrum == "TE" else None) + if spin_pair is None: + raise ValueError("spectrum must be set to either 'EE' or 'TE'") + Bbl = np.load(f"{mcm_dir}/{spec_name}_Bbl_{spin_pair}.npy") + Bbl = Bbl[:n_bins, :lmax] + + # Get theory + cmb_th = ps_th[spectrum][:lmax] + fg_th = fg_dict[spectrum.lower(), "dust", f"{sv}_{ar}", f"{sv}_{ar}"][:lmax] + + # Define loglike + def loglike(pol_eff, dust_amp): + theory = get_model(cmb_th, fg_th, Bbl, dust_amp, pol_eff, mode=spectrum) + theory = theory[id] + residual = ps - theory + chi2 = residual @ invcov @ residual + return -0.5 * chi2 + + loc, scale = dust_priors[spectrum]["loc"], dust_priors[spectrum]["scale"] + + # Prepare MCMC sampling + info = { + "likelihood": { + "pol_eff": loglike + }, + "params": { + "pol_eff": { + "prior": { + "min": 0.5, + "max": 1.5 + }, + "latex": r"\epsilon_\mathrm{pol}^{%s}" % f"{sv}_{ar}".replace("_", "\_") + }, + "dust_amp": { + "prior": { + "dist": "norm", + "loc": loc, + "scale": scale + }, + "latex": r"A_\mathrm{dust}^{%s}" % spectrum + }, + }, + "sampler": { + "mcmc": { + "max_tries": 10**6, + "Rminus1_stop": 0.01, + "Rminus1_cl_stop": 0.01 + } + }, + "output": f"{output_dir}/chain_{spectrum}_{sv}_{ar}", + "force": True, + } + + updated_info, sampler = run(info) + + samples = loadMCSamples(f"{output_dir}/chain_{spectrum}_{sv}_{ar}", settings={"ignore_rows": 0.5}) + pol_eff_mean[sv, ar, spectrum] = samples.mean("pol_eff") + pol_eff_std[sv, ar, spectrum] = np.sqrt(samples.cov(["pol_eff"])[0, 0]) + dust_mean[sv, ar, spectrum] = samples.mean("dust_amp") + dust_std[sv, ar, spectrum] = np.sqrt(samples.cov(["dust_amp"])[0, 0]) + + + gdplot = gdplt.get_subplot_plotter() + gdplot.triangle_plot(samples, ["pol_eff", "dust_amp"], filled=True, title_limit=1) + plt.savefig(f"{output_dir}/posterior_dist_{spectrum}_{sv}_{ar}.png", dpi=300, bbox_inches="tight") + + +for sv in surveys: + for ar in arrays[sv]: + print(f"**************") + for spectrum in ["EE", "TE"]: + p_eff, std = pol_eff_mean[sv, ar, spectrum], pol_eff_std[sv, ar, spectrum] + print(f"{sv} {ar} {spectrum} {p_eff} {std}") + + print(f"**************") diff --git a/project/data_analysis/python/null_tests/custom_nulls/code/get_polar_eff_LCDM_cal_inout.py b/project/data_analysis/python/null_tests/custom_nulls/code/get_polar_eff_LCDM_cal_inout.py new file mode 100644 index 00000000..10dc8468 --- /dev/null +++ b/project/data_analysis/python/null_tests/custom_nulls/code/get_polar_eff_LCDM_cal_inout.py @@ -0,0 +1,218 @@ +""" +This script is used to calibrate the polarization efficiencies +for each array with respect to a LCDM bestfit (default: cmb.dat) +Note that I assume that the spectra are uncalibrated and apply the calibration in this script +""" +from pspy import so_dict, pspy_utils, so_spectra, so_cov +from getdist.mcsamples import loadMCSamples +from pspipe_utils import best_fits +import matplotlib.pyplot as plt +import getdist.plots as gdplt +from cobaya.run import run +import numpy as np +import sys + +d = so_dict.so_dict() +d.read_from_file(sys.argv[1]) + +# Set up directories +ps_dir = "spectra_leak_corr" +planck_corr = False + +if planck_corr: + ps_dir = "spectra_leak_corr_planck_bias_corr" + + + +cov_dir = "covariances" +bestfit_dir = "best_fits" +mcm_dir = "mcms" +output_dir = "pol_eff_results" +pspy_utils.create_directory(output_dir) + +ps_filename = f"{bestfit_dir}/cmb.dat" + +# Load paramfiles info +surveys = d["surveys"] +arrays = {sv: d[f"arrays_{sv}"] for sv in surveys} +lmax = d["lmax"] + +spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"] + +# Load ps theory file +l_th, ps_th = so_spectra.read_ps(ps_filename, spectra=spectra) +l_th = l_th.astype(int) + + +# Load foreground dict and dust prior, then normalize (with polarized dust normalized to 1) +do_bandpass_integration = d["do_bandpass_integration"] +fg_components = d["fg_components"] +fg_params = d["fg_params"] + +dust_priors = { + "EE": {"loc": fg_params["a_gee"], "scale": 0.008}, + "TE": {"loc": fg_params["a_gte"], "scale": 0.015} +} + +fg_components["ee"] = ["dust"] +fg_components["te"] = ["dust"] +fg_params["a_gee"] = 1. +fg_params["a_gte"] = 1. + +print(dust_priors) + +passbands = {} +for sv in surveys: + for ar in arrays[sv]: + freq_info = d[f"freq_info_{sv}_{ar}"] + if do_bandpass_integration: + nu_ghz, pb = np.loadtxt(freq_info["passband"]).T + else: + nu_ghz, pb = np.array([freq_info["freq_tag"]]), np.array([1.]) + + passbands[f"{sv}_{ar}"] = [nu_ghz, pb] + +fg_dict = best_fits.get_foreground_dict(l_th, passbands, fg_components, fg_params, d["fg_norm"]) + +# Load priors on dust amplitudes +# from Planck 353 GHz spectra +# computed in ACT DR6 windows + + +# Calibration range +lmin_cal = 500 +lmax_cal = {} + +my_surveys = d["surveys"] +try: + my_surveys.remove("Planck") +except: + pass + +for sv in my_surveys: + lmax_cal[f"{sv}_pa4_f220_in"] = 2000 + lmax_cal[f"{sv}_pa5_f090_in"] = 1200 + lmax_cal[f"{sv}_pa5_f150_in"] = 2000 + lmax_cal[f"{sv}_pa6_f090_in"] = 1200 + lmax_cal[f"{sv}_pa6_f150_in"] = 2000 + lmax_cal[f"{sv}_pa4_f220_out"] = 2000 + lmax_cal[f"{sv}_pa5_f090_out"] = 1200 + lmax_cal[f"{sv}_pa5_f150_out"] = 2000 + lmax_cal[f"{sv}_pa6_f090_out"] = 1200 + lmax_cal[f"{sv}_pa6_f150_out"] = 2000 + +lmax_cal["Planck_f100"] = 1200 +lmax_cal["Planck_f143"] = 2000 +lmax_cal["Planck_f217"] = 2000 + + +# Define useful functions +def get_model(cmb_th, fg_th, Bbl, dust_amp, pol_eff, mode): + ps_theory = (cmb_th + dust_amp * fg_th) * pol_eff ** mode.count("E") + return Bbl @ ps_theory + +pol_eff_mean, pol_eff_std, dust_mean, dust_std = {}, {}, {}, {} +for sv in surveys: + for ar in arrays[sv]: + for spectrum in ["EE", "TE"]: + + print(spectrum, sv, ar, lmin_cal, lmax_cal[f"{sv}_{ar}"]) + + # Load ps and cov + spec_name = f"{sv}_{ar}x{sv}_{ar}" + lb, ps = so_spectra.read_ps(f"{ps_dir}/Dl_{spec_name}_cross.dat", spectra=spectra) + cov = np.load(f"{cov_dir}/analytic_cov_{spec_name}_{spec_name}.npy") + leakage_cov = np.load(f"{cov_dir}/leakage_cov_{spec_name}_{spec_name}.npy") + + cal = d[f"cal_{sv}_{ar}"] + + + cov += leakage_cov + # Select the spectrum + ps = ps[spectrum] * cal ** 2 + n_bins = len(lb) + + cov = so_cov.selectblock(cov, spectra, n_bins=n_bins, block=spectrum+spectrum) + + # Multipole cuts + id = np.where((lb >= lmin_cal) & (lb <= lmax_cal[f"{sv}_{ar}"]))[0] + ps = ps[id] + cov = cov[np.ix_(id, id)] + invcov = np.linalg.inv(cov) + + # Load Bbl + spin_pair = "spin2xspin2" if spectrum == "EE" else ("spin0xspin2" if spectrum == "TE" else None) + if spin_pair is None: + raise ValueError("spectrum must be set to either 'EE' or 'TE'") + Bbl = np.load(f"{mcm_dir}/{spec_name}_Bbl_{spin_pair}.npy") + Bbl = Bbl[:n_bins, :lmax] + + # Get theory + cmb_th = ps_th[spectrum][:lmax] + fg_th = fg_dict[spectrum.lower(), "dust", f"{sv}_{ar}", f"{sv}_{ar}"][:lmax] + + # Define loglike + def loglike(pol_eff, dust_amp): + theory = get_model(cmb_th, fg_th, Bbl, dust_amp, pol_eff, mode=spectrum) + theory = theory[id] + residual = ps - theory + chi2 = residual @ invcov @ residual + return -0.5 * chi2 + + loc, scale = dust_priors[spectrum]["loc"], dust_priors[spectrum]["scale"] + + # Prepare MCMC sampling + info = { + "likelihood": { + "pol_eff": loglike + }, + "params": { + "pol_eff": { + "prior": { + "min": 0.5, + "max": 1.5 + }, + "latex": r"\epsilon_\mathrm{pol}^{%s}" % f"{sv}_{ar}".replace("_", "\_") + }, + "dust_amp": { + "prior": { + "dist": "norm", + "loc": loc, + "scale": scale + }, + "latex": r"A_\mathrm{dust}^{%s}" % spectrum + }, + }, + "sampler": { + "mcmc": { + "max_tries": 10**6, + "Rminus1_stop": 0.01, + "Rminus1_cl_stop": 0.01 + } + }, + "output": f"{output_dir}/chain_{spectrum}_{sv}_{ar}", + "force": True, + } + + updated_info, sampler = run(info) + + samples = loadMCSamples(f"{output_dir}/chain_{spectrum}_{sv}_{ar}", settings={"ignore_rows": 0.5}) + pol_eff_mean[sv, ar, spectrum] = samples.mean("pol_eff") + pol_eff_std[sv, ar, spectrum] = np.sqrt(samples.cov(["pol_eff"])[0, 0]) + dust_mean[sv, ar, spectrum] = samples.mean("dust_amp") + dust_std[sv, ar, spectrum] = np.sqrt(samples.cov(["dust_amp"])[0, 0]) + + + gdplot = gdplt.get_subplot_plotter() + gdplot.triangle_plot(samples, ["pol_eff", "dust_amp"], filled=True, title_limit=1) + plt.savefig(f"{output_dir}/posterior_dist_{spectrum}_{sv}_{ar}.png", dpi=300, bbox_inches="tight") + + +for sv in surveys: + for ar in arrays[sv]: + print(f"**************") + for spectrum in ["EE", "TE"]: + p_eff, std = pol_eff_mean[sv, ar, spectrum], pol_eff_std[sv, ar, spectrum] + print(f"{sv} {ar} {spectrum} {p_eff} {std}") + + print(f"**************") diff --git a/project/data_analysis/python/null_tests/custom_nulls/code/get_window_surveys.py b/project/data_analysis/python/null_tests/custom_nulls/code/get_window_surveys.py new file mode 100644 index 00000000..128dcdea --- /dev/null +++ b/project/data_analysis/python/null_tests/custom_nulls/code/get_window_surveys.py @@ -0,0 +1,224 @@ +""" +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 pixell import enmap +from pspipe_utils import log, pspipe_list +from pspy import pspy_utils, so_dict, so_map, so_mpi, so_window + + +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 + + +d = so_dict.so_dict() +d.read_from_file(sys.argv[1]) +log = log.get_logger(**d) + + +# 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 = "windows" + +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 +# we also force surveys to be dr6 (otherwise product will be missing) + +surveys = d["surveys"] +try: + surveys.remove("Planck") +except: + pass + + +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) + +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[:]) + + if d[f"extra_mask_{sv}_{ar}"] is not None: + log.info(f"[{task}] apply extra mask") + extra_mask = so_map.read_map(d[f"extra_mask_{sv}_{ar}"]) + my_masks["baseline"].data[:] *= extra_mask.data[:] + + 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}") diff --git a/project/data_analysis/python/null_tests/custom_nulls/code/mc_mnms_get_custom_nlms.py b/project/data_analysis/python/null_tests/custom_nulls/code/mc_mnms_get_custom_nlms.py new file mode 100644 index 00000000..e0b384e9 --- /dev/null +++ b/project/data_analysis/python/null_tests/custom_nulls/code/mc_mnms_get_custom_nlms.py @@ -0,0 +1,127 @@ +description = """ +generates noise alms for custom nulls +""" +from pspy import pspy_utils, so_dict, so_mpi +from mnms import noise_models as nm +from pixell import curvedsky +from pspipe_utils import log +import numpy as np +import time +import sys +import argparse + +parser = argparse.ArgumentParser(description=description, + formatter_class=argparse.ArgumentDefaultsHelpFormatter) +parser.add_argument("--paramfile", type=str, + help="Filename (full or relative path) of paramfile to use") +parser.add_argument("--iStart", type=int, default=None) +parser.add_argument("--iStop", type=int, default=None) +parser.add_argument("--survey", type=str, default=None) +args = parser.parse_args() + +d = so_dict.so_dict() +d.read_from_file(args.paramfile) +log = log.get_logger(**d) + + +lmax = d["lmax"] +survey = args.survey +iStart = args.iStart +iStop = args.iStop + +all_surveys = ["el1", "el2", "el3", "pwv1", "pwv2", "t1", "t2", "inout"] +sim_num_start = {} +for count, sv in enumerate(all_surveys): + sim_num_start[sv] = count * 1000 + +log.info(f"generating sims from {iStart} to {iStop} for survey: {survey}, seed start at {sim_num_start[survey]}") + +# Set the lmax corresponding to a pre-computed noise model +lmax_noise_sim = 10800 # Should be moved to the paramfile ? + +# Aliases for arrays +arrays_alias = { + "pa4": {"pa4a": "pa4_f150", "pa4b": "pa4_f220"}, + "pa5": {"pa5a": "pa5_f090", "pa5b": "pa5_f150"}, + "pa6": {"pa6a": "pa6_f090", "pa6b": "pa6_f150"} +} + +wafers = ["pa4", "pa5", "pa6"] +# Load the noise models +noise_models = {} +for wafer_name in wafers: + if survey in ["el1", "el2", "el3"]: + noise_models[wafer_name] = nm.BaseNoiseModel.from_config(f"act_dr6v4_el_split", + d[f"noise_sim_type_{wafer_name}"], + *arrays_alias[wafer_name].keys(), + el_split=[survey]) + if survey in ["pwv1", "pwv2"]: + noise_models[wafer_name] = nm.BaseNoiseModel.from_config(f"act_dr6v4_pwv_split", + d[f"noise_sim_type_{wafer_name}"], + *arrays_alias[wafer_name].keys(), + pwv_split=[survey]) + if survey in ["t1", "t2"]: + arrays_alias["pa4"] = {"pa4b": "pa4_f220"} + noise_models[wafer_name] = nm.BaseNoiseModel.from_config(f"act_dr6v4_t_split", + d[f"noise_sim_type_{wafer_name}"], + *arrays_alias[wafer_name].keys(), + t_split=[survey]) + + + if survey == "inout": + noise_models[wafer_name] = nm.BaseNoiseModel.from_config("act_dr6v4_inout_split", + d[f"noise_sim_type_{wafer_name}"], + *arrays_alias[wafer_name].keys(), + inout_split=["inout1", "inout2"]) + +if survey == "inout": + arrays_alias["pa4"] = {"qid1": "pa4_f150_in", "qid2": "pa4_f220_in", + "qid3": "pa4_f150_out", "qid4": "pa4_f220_out"} + arrays_alias["pa5"] = {"qid1": "pa5_f090_in", "qid2": "pa5_f150_in", + "qid3": "pa5_f090_out", "qid4": "pa5_f150_out"} + arrays_alias["pa6"] = {"qid1": "pa6_f090_in", "qid2": "pa6_f150_in", + "qid3": "pa6_f090_out", "qid4": "pa6_f150_out"} + + +# Create output dir +nlms_dir = f"mnms_noise_alms_{survey}" +pspy_utils.create_directory(nlms_dir) + +n_splits = 2 if survey in ["t1", "t2"] else 4 +mpi_list = [] +for id_split in range(n_splits): + mpi_list.append(id_split) + +# we will use mpi over the number of splits +so_mpi.init(True) +subtasks = so_mpi.taskrange(imin=0, imax=len(mpi_list)-1) +for id_mpi in subtasks: + k = mpi_list[id_mpi] + t1 = time.time() + for wafer_name in wafers: + for iii in range(iStart, iStop + 1): + t2 = time.time() + + sim_arrays = noise_models[wafer_name].get_sim(split_num=k, + sim_num=sim_num_start[survey] + iii, + lmax=lmax_noise_sim, + alm=False, + keep_model=True) + + log.info(f"[Sim n° {iii}] {survey} {wafer_name} split {k} noise realization generated in {time.time()-t2:.2f} s") + + for i, (qid, ar) in enumerate(arrays_alias[wafer_name].items()): + + cal, pol_eff = d[f"cal_{survey}_{ar}"], d[f"pol_eff_{survey}_{ar}"] + + sim_arrays[i, 0, :][0] *= cal + sim_arrays[i, 0, :][1] *= cal / pol_eff + sim_arrays[i, 0, :][2] *= cal / pol_eff + + noise_alms = curvedsky.map2alm(sim_arrays[i, 0, :], lmax=lmax) + + np.save(f"{nlms_dir}/mnms_nlms_{survey}_{ar}_set{k}_{iii:05d}.npy", noise_alms) + + noise_models[wafer_name].cache_clear() + + log.info(f"split {k} - {iStop-iStart+1} sims generated in {time.time()-t1:.2f} s") diff --git a/project/data_analysis/python/null_tests/custom_nulls/code/mc_mnms_get_spectra_from_custom_nlms.py b/project/data_analysis/python/null_tests/custom_nulls/code/mc_mnms_get_spectra_from_custom_nlms.py new file mode 100644 index 00000000..718a94c7 --- /dev/null +++ b/project/data_analysis/python/null_tests/custom_nulls/code/mc_mnms_get_spectra_from_custom_nlms.py @@ -0,0 +1,292 @@ +""" +generates simulation for customs nulls +""" + +import sys +import time + +import healpy as hp +import numpy as np +from pixell import curvedsky, enmap +from pspipe_utils import kspace, log, misc, pspipe_list, simulation, transfer_function +from pspy import pspy_utils, so_dict, so_map, so_mcm, so_mpi, so_spectra, 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"] +niter = d["niter"] +type = d["type"] +binning_file = d["binning_file"] +write_all_spectra = d["write_splits_spectra"] +sim_alm_dtype = d["sim_alm_dtype"] +binned_mcm = d["binned_mcm"] +apply_kspace_filter = d["apply_kspace_filter"] +if sim_alm_dtype in ["complex64", "complex128"]: sim_alm_dtype = getattr(np, sim_alm_dtype) +else: raise ValueError(f"Unsupported sim_alm_dtype {sim_alm_dtype}") +dtype = np.float32 if sim_alm_dtype == "complex64" else np.float64 + +window_dir = "windows" +mcm_dir = "mcms" +spec_dir = "sim_spectra" +bestfit_dir = "best_fits" + +nlms_dir = "/pscratch/sd/x/xgarrido/sims/nlms/" + +pspy_utils.create_directory(spec_dir) + +spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"] +spin_pairs = ["spin0xspin0", "spin0xspin2", "spin2xspin0", "spin2xspin2"] + +# prepare the tempalte and the filter +arrays, templates, filters, n_splits, filter_dicts, pixwin, inv_pixwin = {}, {}, {}, {}, {}, {}, {} +spec_name_list = pspipe_list.get_spec_name_list(d, delimiter="_") + +for sv in surveys: + arrays[sv] = d[f"arrays_{sv}"] + n_splits[sv] = len(d[f"maps_{sv}_{arrays[sv][0]}"]) + log.info(f"Running with {n_splits[sv]} splits for survey {sv}") + template_name = d[f"maps_{sv}_{arrays[sv][0]}"][0] + templates[sv] = so_map.read_map(template_name) + + if d[f"pixwin_{sv}"]["pix"] == "CAR": + wy, wx = enmap.calc_window(templates[sv].data.shape, + order=d[f"pixwin_{sv}"]["order"]) + pixwin[sv] = (wy[:, None] * wx[None, :]) + inv_pixwin[sv] = pixwin[sv] ** (-1) + elif d[f"pixwin_{sv}"]["pix"] == "HEALPIX": + pw_l = hp.pixwin(d[f"pixwin_{sv}"]["nside"]) + pixwin[sv] = pw_l + inv_pixwin[sv] = pw_l ** (-1) + + if apply_kspace_filter: + filter_dicts[sv] = d[f"k_filter_{sv}"] + filters[sv] = kspace.get_kspace_filter(templates[sv], filter_dicts[sv], dtype=np.float32) + +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 = {} + TE_corr = {} + for spec_name in spec_name_list: + kspace_transfer_matrix[spec_name] = np.load(f"{kspace_tf_path}/kspace_matrix_{spec_name}.npy") + _, TE_corr[spec_name] = so_spectra.read_ps(f"{kspace_tf_path}/TE_correction_{spec_name}.dat", spectra=spectra) + + +f_name_cmb = bestfit_dir + "/cmb.dat" +ps_mat = simulation.cmb_matrix_from_file(f_name_cmb, lmax, spectra) + +f_name_fg = bestfit_dir + "/fg_{}x{}.dat" +array_list = [f"{sv}_{ar}" for sv in surveys for ar in arrays[sv]] +l, fg_mat = simulation.foreground_matrix_from_files(f_name_fg, array_list, lmax, spectra) + +# we will use mpi over the number of simulations +so_mpi.init(True) +subtasks = so_mpi.taskrange(imin=d["iStart"], imax=d["iStop"]) + +for iii in subtasks: + t0 = time.time() + + # generate cmb alms and foreground alms + # cmb alms will be of shape (3, lm) 3 standing for T,E,B + + # Set seed if needed + if d["seed_sims"]: + np.random.seed(iii) + alms_cmb = curvedsky.rand_alm(ps_mat, lmax=lmax, dtype="complex64") + fglms = simulation.generate_fg_alms(fg_mat, array_list, lmax, dtype="complex64") + + log.info(f"[Sim n° {iii}] Generate signal alms in {time.time()-t0:.2f} s") + master_alms = {} + + for sv in surveys: + + t1 = time.time() + + signal_alms = {} + for ar in arrays[sv]: + + signal_alms[ar] = alms_cmb + fglms[f"{sv}_{ar}"] + l, bl = misc.read_beams(d[f"beam_T_{sv}_{ar}"], d[f"beam_pol_{sv}_{ar}"]) + signal_alms[ar] = misc.apply_beams(signal_alms[ar], bl) + + # since the mnms noise sim include a pixwin, we convolve the signal ones + signal = sph_tools.alm2map(signal_alms[ar], templates[sv]) + win_kspace = so_map.read_map(d[f"window_kspace_{sv}_{ar}"]) + # Convolve the signal map with the pixwin only for CAR pixellization + if d[f"pixwin_{sv}"]["pix"] == "CAR": + signal = signal.convolve_with_pixwin(niter=niter, pixwin=pixwin[sv], window=win_kspace, use_ducc_rfft=True) + signal_alms[ar] = sph_tools.map2alm(signal, niter, lmax) + # Convolve the signal with the pixwin in harm. space for Planck sims + if d[f"pixwin_{sv}"]["pix"] == "HEALPIX": + signal_alms[ar] = curvedsky.almxfl(signal_alms[ar], pixwin[sv]) + + log.info(f"[Sim n° {iii}] Convolve beam and pixwin in {time.time()-t1:.2f} s") + + wafers = sorted({ar.split("_")[0] for ar in arrays[sv]}) + + for k in range(n_splits[sv]): + for ar in arrays[sv]: + + t3 = time.time() + + 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) + del win_T, win_pol + + log.info(f"[Sim n° {iii}] Read window in {time.time()-t3:.2f} s") + + + if any("_out" in my_array for my_array in arrays[sv]): + noise_alm_file_name = f"{nlms_dir}/mnms_noise_alms_inout/mnms_nlms_inout_{ar}_set{k}_{iii:05d}.npy" + else: + noise_alm_file_name = f"{nlms_dir}/mnms_noise_alms_{sv}/mnms_nlms_{sv}_{ar}_set{k}_{iii:05d}.npy" + + log.info(f"[Sim n° {iii}] read nlm noise file {noise_alm_file_name}") + noise_alms = np.load(noise_alm_file_name) + + t4 = time.time() + + split = sph_tools.alm2map(signal_alms[ar] + noise_alms, templates[sv]) + + log.info(f"[Sim n° {iii}] alm2map for split {k} and array {ar} done in {time.time()-t4:.2f} s") + + t5 = time.time() + if (window_tuple[0].pixel == "CAR") & (apply_kspace_filter): + + win_kspace = so_map.read_map(d[f"window_kspace_{sv}_{ar}"]) + + inv_pwin = inv_pixwin[sv] if d[f"pixwin_{sv}"]["pix"] == "CAR" else None + + split = kspace.filter_map(split, + filters[sv], + win_kspace, + inv_pixwin = inv_pwin, + weighted_filter=filter_dicts[sv]["weighted"], + use_ducc_rfft=True) + + + del win_kspace + + log.info(f"[Sim n° {iii}] Filter split {k} and array {ar} in {time.time()-t5:.2f} s") + + #cal, pol_eff = d[f"cal_{sv}_{ar}"], d[f"pol_eff_{sv}_{ar}"] + #split = split.calibrate(cal=cal, pol_eff=pol_eff) + + if d["remove_mean"] == True: + split = split.subtract_mean(window_tuple) + + t6 = time.time() + master_alms[sv, ar, k] = sph_tools.get_alms(split, window_tuple, niter, lmax, dtype=sim_alm_dtype) + log.info(f"[Sim n° {iii}] Final map2alm done in {time.time()-t6:.2f} s") + + ps_dict = {} + + t7 = time.time() + + n_spec, sv1_list, ar1_list, sv2_list, ar2_list = pspipe_list.get_spectra_list(d) + + for i_spec in range(n_spec): + + sv1, ar1, sv2, ar2 = sv1_list[i_spec], ar1_list[i_spec], sv2_list[i_spec], ar2_list[i_spec] + + for spec in spectra: + ps_dict[spec, "auto"] = [] + ps_dict[spec, "cross"] = [] + + mbb_inv, Bbl = so_mcm.read_coupling(prefix=f"{mcm_dir}/{sv1}_{ar1}x{sv2}_{ar2}", + spin_pairs=spin_pairs) + + for s1 in range(n_splits[sv1]): + for s2 in range(n_splits[sv2]): + 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 kspace_tf_path == "analytical": + xtra_corr = None + else: + xtra_corr = TE_corr[f"{sv1}_{ar1}x{sv2}_{ar2}"] + + lb, ps = kspace.deconvolve_kspace_filter_matrix(lb, + ps, + kspace_transfer_matrix[f"{sv1}_{ar1}x{sv2}_{ar2}"], + spectra, + xtra_corr=xtra_corr) + + if d[f"pixwin_{sv1}"]["pix"] == "HEALPIX": + _, xtra_pw1 = pspy_utils.naive_binning(np.arange(len(pixwin[sv1])), pixwin[sv1], binning_file, lmax) + else: + xtra_pw1 = None + if d[f"pixwin_{sv2}"]["pix"] == "HEALPIX": + _, xtra_pw2 = pspy_utils.naive_binning(np.arange(len(pixwin[sv2])), pixwin[sv2], binning_file, lmax) + else: + xtra_pw2 = None + lb, ps = transfer_function.deconvolve_xtra_tf(lb, + ps, + spectra, + xtra_pw1=xtra_pw1, + xtra_pw2=xtra_pw2) + + if write_all_spectra: + so_spectra.write_ps(spec_dir + f"/{spec_name}_%05d.dat" % (iii), lb, ps, type, spectra=spectra) + + for count, spec in enumerate(spectra): + if (s1 == s2) & (sv1 == sv2): + if count == 0: log.info(f"[Sim n° {iii}] auto {sv1}_{ar1} X {sv2}_{ar2} {s1}{s2}") + ps_dict[spec, "auto"] += [ps[spec]] + else: + if count == 0: log.info(f"[Sim n° {iii}] 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) + + 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) + ps_dict_noise_mean[spec] = (ps_dict_auto_mean[spec] - ps_dict_cross_mean[spec]) / n_splits[sv1] + + spec_name_cross = f"{type}_{sv1}_{ar1}x{sv2}_{ar2}_cross_{iii:05d}" + so_spectra.write_ps(spec_dir + f"/{spec_name_cross}.dat", lb, ps_dict_cross_mean, type, spectra=spectra) + + if sv1 == sv2: + spec_name_auto = f"{type}_{sv1}_{ar1}x{sv2}_{ar2}_auto_{iii:05d}" + so_spectra.write_ps(spec_dir + f"/{spec_name_auto}.dat", lb, ps_dict_auto_mean, type, spectra=spectra) + spec_name_noise = f"{type}_{sv1}_{ar1}x{sv2}_{ar2}_noise_{iii:05d}" + so_spectra.write_ps(spec_dir + f"/{spec_name_noise}.dat", lb, ps_dict_noise_mean, type, spectra=spectra) + + log.info(f"[Sim n° {iii}] Spectra computation done in {time.time()-t7:.2f} s") + log.info(f"[Sim n° {iii}] Done in {time.time()-t0:.2f} s")