diff --git a/pipeline/create_sacc_file.py b/pipeline/create_sacc_file.py index bb56876..b2d894f 100644 --- a/pipeline/create_sacc_file.py +++ b/pipeline/create_sacc_file.py @@ -1,5 +1,6 @@ import argparse from soopercool import BBmeta +from soopercool import mpi_utils as mpi import sacc from itertools import product import numpy as np @@ -33,6 +34,7 @@ def main(args): """ meta = BBmeta(args.globals) + verbose = args.verbose out_dir = meta.output_directory sacc_dir = f"{out_dir}/saccs" @@ -43,7 +45,16 @@ def main(args): binning = np.load(meta.binning_file) nmt_binning = nmt.NmtBin.from_edges(binning["bin_low"], binning["bin_high"] + 1) + ls = np.arange(nmt_binning.lmax+1) lb = nmt_binning.get_effective_ells() + n_bins = len(lb) + lwin = np.zeros((len(ls), n_bins)) + + for id_bin in range(n_bins): + weights = np.array(nmt_binning.get_weight_list(id_bin)) + multipoles = np.array(nmt_binning.get_ell_list(id_bin)) + for il, l in enumerate(multipoles): + lwin[l, id_bin] = weights[il] field_pairs = [m1+m2 for m1, m2 in product("TEB", repeat=2)] @@ -96,13 +107,22 @@ def main(args): full_cov = np.triu(full_cov) full_cov += full_cov.T - np.diag(full_cov.diagonal()) - for id_sim in range(Nsims): + use_mpi4py = args.sims + mpi.init(use_mpi4py) + + for id_sim in mpi.taskrange(Nsims - 1): sim_label = f"_{id_sim:04d}" if Nsims > 1 else "" + s_wins = sacc.BandpowerWindow(ls, lwin) + s = sacc.Sacc() for ms in map_sets: + f = float(meta.freq_tag_from_map_set(ms)) + if verbose: + print(f"# {id_sim} | {ms}") + for spin, qty in zip( [0, 2], ["cmb_temperature", "cmb_polarization"] @@ -110,13 +130,13 @@ def main(args): s.add_tracer(**{ "tracer_type": "NuMap", - "name": f"{ms}_s{spin}", + "name": f"{ms}", "quantity": qty, "spin": spin, - "nu": [meta.freq_tag_from_map_set(ms)], + "nu": [f-1., f, f+1], # Delta bandpasses. TODO: generalize "ell": lb, - "beam": np.ones_like(lb), # TODO, - "bandpass": [1.] # TODO + "beam": np.ones_like(lb), # Unit beam. TODO: generalize + "bandpass": [0., 1., 0.] # Deltas. TODO: generalize }) for i, (ms1, ms2) in enumerate(ps_names): @@ -127,15 +147,13 @@ def main(args): for fp in field_pairs: f1, f2 = fp - spin1 = 0 if f1 == "T" else 2 - spin2 = 0 if f2 == "T" else 2 s.add_ell_cl(**{ "data_type": f"cl_{data_types[f1]}{data_types[f2]}", - "tracer1": f"{ms1}_s{spin1}", - "tracer2": f"{ms2}_s{spin2}", + "tracer1": f"{ms1}", + "tracer2": f"{ms2}", "ell": lb, "x": cells[fp], - "window": np.ones_like(lb) # TODO + "window": s_wins }) s.add_covariance(full_cov) @@ -152,6 +170,8 @@ def main(args): ) parser.add_argument("--globals", type=str, help="Path to the yaml file") + parser.add_argument("--verbose", action="store_true", help="Verbose mode.") + mode = parser.add_mutually_exclusive_group() mode.add_argument("--sims", action="store_true") mode.add_argument("--data", action="store_true") diff --git a/pipeline/sacc_plotter.py b/pipeline/sacc_plotter.py index 61e1ac6..aec46b7 100644 --- a/pipeline/sacc_plotter.py +++ b/pipeline/sacc_plotter.py @@ -1,5 +1,6 @@ import matplotlib.pyplot as plt from soopercool import BBmeta +from soopercool import mpi_utils as mpi import argparse from itertools import product import sacc @@ -85,7 +86,7 @@ def plot_spectrum(lb, cb, cb_err, title, ylabel, xlim, markeredgewidth=1.75) sub.set_xlim(*xlim) - sub.set_ylim(-4.5, 4.5) + sub.set_ylim(-10, 10) if save_file: plt.savefig(save_file, bbox_inches="tight") @@ -101,6 +102,7 @@ def main(args): spectra. """ meta = BBmeta(args.globals) + verbose = args.verbose out_dir = meta.output_directory sacc_dir = f"{out_dir}/saccs" @@ -117,7 +119,6 @@ def main(args): field_pairs = [m1+m2 for m1, m2 in product("TEB", repeat=2)] ps_names = meta.get_ps_names_list(type="all", coadd=True) - spins = {"T": 0, "E": 2, "B": 2} types = {"T": "0", "E": "e", "B": "b"} if args.data: @@ -188,8 +189,11 @@ def main(args): plot_data[ms1, ms2, fp]["y_th"] = y_th plot_data[ms1, ms2, fp]["th_binned"] = th_binned - for id_sim in range(Nsims): - sim_label = f"_{id_sim:04d}" if Nsims > 1 else "" + use_mpi4py = args.sims + mpi.init(use_mpi4py) + + for id_sim in mpi.taskrange(Nsims - 1): + sim_label = f"_{id_sim:04d}" if args.sims else "" s = sacc.Sacc.load_fits(f"{sacc_dir}/cl_and_cov_sacc{sim_label}.fits") @@ -202,8 +206,8 @@ def main(args): ell, cl, cov = s.get_ell_cl( f"cl_{types[f1]}{types[f2]}", - f"{ms1}_s{spins[f1]}", - f"{ms2}_s{spins[f2]}", + ms1, + ms2, return_cov=True) idx_bad = idx_bad_tf[ftag1, ftag2][fp] @@ -220,14 +224,18 @@ def main(args): plot_data[ms1, ms2, fp]["ylabel"] = fp for ms1, ms2 in ps_names: + if verbose: + print(f"# {id_sim} | {ms1} x {ms2}") for fp in field_pairs: plot_data[ms1, ms2, fp]["y"] = np.mean( plot_data[ms1, ms2, fp]["y"], axis=0 ) plot_data[ms1, ms2, fp]["err"] /= np.sqrt(Nsims) - - plot_name = f"{ms1}_{ms2}_{fp}.pdf" + if args.sims: + plot_name = f"plot_cells_sims_{ms1}_{ms2}_{fp}.pdf" + if args.data: + plot_name = f"plot_cells_data_{ms1}_{ms2}_{fp}.pdf" plot_spectrum( plot_data[ms1, ms2, fp]["x"], @@ -248,6 +256,8 @@ def main(args): parser = argparse.ArgumentParser(description='Sacc plotter') parser.add_argument("--globals", type=str, help="Path to the global configuration file") + parser.add_argument("--verbose", action="store_true", help="Verbose mode.") + mode = parser.add_mutually_exclusive_group() mode.add_argument("--sims", action="store_true") mode.add_argument("--data", action="store_true")