make sacc file compatible with bbpower

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,27 +107,36 @@ 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"]
             ):
 
                 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")

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")