get_spectra_from_maps: only load necessary alms. get_pseudo2datavec: dont do Bbl (saves disk and Bbl not really needed). get_mcm_bbl_and_pseudosignal: per ell pseudosignal spectra for INKA also computed and saved to disk

Author: zatkins2

project/SO/pISO/python/get_mcm_and_bbl.py …O/python/get_mcm_bbl_and_pseudosignal.pyproject/SO/pISO/python/get_mcm_and_bbl.py renamed to project/SO/pISO/python/get_mcm_bbl_and_pseudosignal.py project/SO/pISO/python/get_mcm_and_bbl.py renamed to project/SO/pISO/python/get_mcm_bbl_and_pseudosignal.py

@@ -29,6 +29,12 @@
 mcm_dir = d['mcm_dir']
 pspy_utils.create_directory(mcm_dir)
 
+bestfit_dir = d["best_fits_dir"]
+pseudo_dir = opj(bestfit_dir, 'pseudo')
+pspy_utils.create_directory(pseudo_dir)
+
+spectra = ["TT", "TE", "TB", "ET", "BT", "EE", "EB", "BE", "BB"]
+
 surveys = d["surveys"]
 niter = d['niter']
 lmax = d["lmax"]
@@ -141,10 +147,32 @@
     nbins = len(bin_hi)
     Pbl = so_spectra.get_binning_matrix(bin_lo, bin_hi, lmax, type)
 
-    # loop over map pairs for the post-processing and saving
     for t, task in enumerate(subtasks):
         sv1, m1, sv2, m2 = sv1_list[task], m1_list[task], sv2_list[task], m2_list[task]
-    
+        spec_name = f"{sv1}_{m1}x{sv2}_{m2}"
+
+        # we need to get the best-fit pseudosignal spectra for the covariance. we do
+        # that here to avoid recalculating all the unbinned mcms again in a
+        # different script. NOTE: we need beamed Cls
+        l, signal_dict, save_type = so_spectra.read_ps(opj(bestfit_dir, f'cmb_and_fg_{spec_name}.dat'), spectra=spectra, return_type=True)
+        assert l[0] == 2, f'Bestfit spectra assumed to start at l=2, got l={l[0]}'
+
+        if save_type == 'Dl':
+            fac = 2*np.pi / (l*(l + 1))
+        else:
+            assert save_type == 'Cl', f'save_type must be Dl or Cl, got {save_type}'
+            fac = 1
+
+        # trim to match mcm
+        for k in signal_dict.keys():
+            signal_dict[k] = signal_dict[k][:lmax-2] * fac[:lmax-2]
+        l = l[:lmax-2]
+
+        # the fully realized mcm matrix would be a lot of memory
+        pseudosignal_dict = so_spectra.spin2spin_array_matmul_spec_dict(mcm[t], signal_dict)
+        so_spectra.write_ps(opj(pseudo_dir, f'pseudo_signal_{spec_name}.dat'), l, pseudosignal_dict, 'Cl', spectra=spectra)
+
+        # now do the binning 
         if binned_mcm:
             mxx = np.zeros((5, nbins, nbins)) # b x b
             Bbl = np.zeros((5, nbins, lmax))
@@ -194,6 +222,5 @@
 
         log.info(f"[{task:02d}] Saving mcm matrix for {sv1}_{m1} x {sv2}_{m2}")
 
-        prefix = opj(f"{mcm_dir}", f"{sv1}_{m1}x{sv2}_{m2}")
-        np.save(prefix + "_mode_coupling_inv.npy", mbl_inv)
-        np.save(prefix + "_Bbl.npy", Bbl)
+        np.save(opj(f"{mcm_dir}", spec_name + "_mode_coupling_inv.npy") , mbl_inv)
+        np.save(opj(f"{mcm_dir}", spec_name + "_Bbl.npy"), Bbl)

…t/SO/pISO/python/get_pspipe_operators.py …ect/SO/pISO/python/get_pseudo2datavec.pyproject/SO/pISO/python/get_pspipe_operators.py renamed to project/SO/pISO/python/get_pseudo2datavec.py project/SO/pISO/python/get_pspipe_operators.py renamed to project/SO/pISO/python/get_pseudo2datavec.py

@@ -92,15 +92,13 @@
 
     log.info(f'[Rank {so_mpi.rank}] Calculating operators for {spec_name}')
 
-    # get the mbl_inv, Bbl matrices for this array cross
+    # get the mbl_inv for this array cross
     prefix = opj(f"{mcm_dir}", spec_name)
     mbl_inv = np.load(prefix + "_mode_coupling_inv.npy")
-    Bbl = np.load(prefix + "_Bbl.npy")
 
-    # need to splice mbl_inv and Bbl into spectra-ordered arrays
+    # need to splice mbl_inv into spectra-ordered arrays
     # TODO: consider disk-space, memory (could be sparse)
     pseudo2datavec = so_mcm.get_spec2spec_array_from_spin2spin_array(mbl_inv, dense=True)
-    th2datavec = so_mcm.get_spec2spec_array_from_spin2spin_array(Bbl, spin0=not binned_mcm, dense=True)
 
     # get the inv_kspace matrix for this array cross, if necessary
     if apply_kspace_filter:
@@ -115,14 +113,12 @@
     if d[f"pixwin_{sv2}"]["pix"] == "HEALPIX" and deconvolve_pixwin:
         pseudo2datavec /= np.tile(pixwins[sv2], 9)[:, None] # apply on the left
 
-    # save 
+    # save and plot
     np.save(opj(f'{mcm_dir}', f'pseudo2datavec_{spec_name}.npy'), pseudo2datavec)
-    np.save(opj(f'{mcm_dir}', f'th2datavec_{spec_name}.npy'), th2datavec)
-
-    for name, var in (('pseudo2datavec', pseudo2datavec), ('th2datavec', th2datavec)):
-        plt.figure(figsize=(10, 8))
-        plt.imshow(np.log(np.abs(var[:, ::4])), aspect=25)
-        plt.colorbar()
-        plt.title(f'{name} {spec_name}')
-        plt.savefig(opj(f'{plot_dir}', f'{name}_{spec_name}'), bbox_inches='tight')
-        plt.close()
+
+    plt.figure(figsize=(10, 8))
+    plt.imshow(np.log(np.abs(pseudo2datavec)), aspect=100)
+    plt.colorbar()
+    plt.title(f'pseudo2datavec {spec_name}')
+    plt.savefig(opj(f'{plot_dir}', f'pseudo2datavec_{spec_name}'), bbox_inches='tight')
+    plt.close()

project/SO/pISO/python/get_spectra_from_maps.py

@@ -364,7 +364,7 @@
 
                         # TODO: would be cleaner if could just make ps_map with
                         # cutoff instead of relying on the mask
-                        ps_mask = so_map.read_map(d[f"ps_mask_{sv}_{m}"])
+                        ps_mask = so_map.read_map(d[f"ps_mask_{sv}_{m}"], geometry=win_T.data.geometry)
                         ps_map.data *= ps_mask.data
                         split.data += ps_map.data
 
@@ -384,7 +384,7 @@
 
                 if apply_kspace_filter and deconvolve_pixwin:
                     if k == 0:
-                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] Apply kspace filter and inv pixwin on {sv}, {m}, {snk}")
+                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] Apply kspace filter and inv pixwin on {sv}, {m}")
                     split = kspace.filter_map(split,
                                               filter,
                                               win_kspace,
@@ -393,7 +393,7 @@
                                               use_ducc_rfft=True)
                 elif apply_kspace_filter:
                     if k == 0:
-                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] WARNING: apply kspace filter but no inv pixwin on {sv}, {m}, {snk}")
+                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] WARNING: apply kspace filter but no inv pixwin on {sv}, {m}")
                     split = kspace.filter_map(split,
                                               filter,
                                               win_kspace,
@@ -403,14 +403,14 @@
 
                 elif deconvolve_pixwin:
                     if k == 0:
-                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] WARNING: inv pixwin but no kspace filter on {sv}, {m}, {snk}")
+                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] WARNING: inv pixwin but no kspace filter on {sv}, {m}")
                     split = so_map.fourier_convolution(split,
                                                        inv_pwin,
                                                        window=win_kspace,
                                                        use_ducc_rfft=True)
                 else:
                     if k == 0:
-                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] WARNING: no kspace filter and no inv pixwin on {sv}, {m}, {snk}")
+                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] WARNING: no kspace filter and no inv pixwin on {sv}, {m}")
 
 
             elif win_T.pixel == "HEALPIX":
@@ -429,10 +429,10 @@
 
                 if deconvolve_pixwin:
                     if k == 0:
-                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] WARNING: inv pixwin but no kspace filter on {sv}, {m}, {snk} (HEALPIX)")
+                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] WARNING: inv pixwin but no kspace filter on {sv}, {m} (HEALPIX)")
                 else:
                     if k == 0:
-                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] WARNING: no kspace filter and no inv pixwin on {sv}, {m}, {snk} (HEALPIX)")
+                        log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] WARNING: no kspace filter and no inv pixwin on {sv}, {m} (HEALPIX)")
 
             split = split.calibrate(cal=cal, pol_eff=pol_eff)
 
@@ -457,15 +457,20 @@
 
     # compute the power spectra
 
-    # if data, need to load alms, and so need to wait
-    # for all alms to be done
+    # if data, need to load alms, and so need to wait for all alms to be done to
+    # load safely. NOTE: but, just need to load the alms needed for this task. 
+    # this saves a lot of memory
     if which == 'data':
         t0 = time.time()
         log.info(f"[Rank {so_mpi.rank}] Loading alms")
         so_mpi.barrier()
 
         master_alms = {}
-        for sv, m in zip(sv_list, map_list): # all of them, not just this process
+        for sv, m in zip(sv_list, map_list):
+            if (sv not in sv1_iterator) and (sv not in sv2_iterator):
+                continue
+            if (m not in m1_iterator) and (m not in m2_iterator):
+                continue
             for k, snk in enumerate(splits_iterator[sv]): # k == split_idx, since data
                 master_alms[sv, m, snk] = np.load(f"{alms_dir}" + f"alms_{sv}_{m}_set{k}.npy")
         log.info(f"[Rank {so_mpi.rank}] Loaded alms: {time.time() - t0} seconds")
@@ -569,8 +574,7 @@
     master_alms = None
 
     log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] Spectra computation time: {time.time() - t0} seconds")
-    t0 = time.time()
-    
+
     if which == 'data':
         spec_name_all = f"{type}_all_sn_cross_data"
 
@@ -593,6 +597,4 @@
         # each process has separate maps in its mapset
         np.save(f"{spec_dir}" + f"{spec_name_all}.npy", ps_dict_all)
 
-    ps_dict_all = None
-    
-    log.info(f"[Rank {so_mpi.rank}, Mapset {iii}] Spectra save time: {time.time() - t0} seconds")
+    ps_dict_all = None