From 7e235bad9051884406ab0ed09f6ea42525ac3872 Mon Sep 17 00:00:00 2001 From: Bai-Chiang Date: Wed, 12 Jun 2024 16:17:47 -0400 Subject: [PATCH] Add script with correct theoretical covariance matrix calculation. If estimated power spectra is average over all cross spectra across all splits, the covariance matrix calculation in the simple script was not correct. Add a separate script with correct one. Also add a comment to old script to avoid confusion. --- examples/generate_fiducial_spectra.py | 3 +- examples/generate_fiducial_spectra_xsplit.py | 266 +++++++++++++++++++ 2 files changed, 268 insertions(+), 1 deletion(-) create mode 100644 examples/generate_fiducial_spectra_xsplit.py diff --git a/examples/generate_fiducial_spectra.py b/examples/generate_fiducial_spectra.py index ddf59c9..c84aae0 100644 --- a/examples/generate_fiducial_spectra.py +++ b/examples/generate_fiducial_spectra.py @@ -175,6 +175,7 @@ def generate_SO_spectra(args): factor_modecount = 1./((2*lb + 1)*delta_ell*fsky) for ii, (i1, i2) in enumerate(zip(indices_tr[0], indices_tr[1])): for jj, (j1, j2) in enumerate(zip(indices_tr[0], indices_tr[1])): + # simple covariance matrix for quick tests covar = (bpw_freq_tot[i1, j1, :]*bpw_freq_tot[i2, j2, :] + bpw_freq_tot[i1, j2, :]*bpw_freq_tot[i2, j1, :]) * factor_modecount # noqa cov_bpw[ii, :, jj, :] = np.diag(covar) @@ -198,4 +199,4 @@ def generate_SO_spectra(args): args = parser.parse_args() - generate_SO_spectra(args) \ No newline at end of file + generate_SO_spectra(args) diff --git a/examples/generate_fiducial_spectra_xsplit.py b/examples/generate_fiducial_spectra_xsplit.py new file mode 100644 index 0000000..bcfd18a --- /dev/null +++ b/examples/generate_fiducial_spectra_xsplit.py @@ -0,0 +1,266 @@ +import numpy as np +from utils import (Bpass, get_nmt_binning, get_component_spectra, + get_convolved_seds, read_beam_window) +import noise_calc as nc +import sacc +import sys +import os +import yaml +import argparse + + +def generate_SO_spectra(args): + """ + Compute theory CMB, synchrotron and dust power spectra from fiducial + template, beam- and bandpass-convolves them, and saves the multifrequency + power spectra inside a SACC file. + """ + # Load the configuration file + print(" Loading configuration file") + fname_config = args.globals + + with open(fname_config) as f: + config = yaml.load(f, Loader=yaml.FullLoader) + + nside = config["global"]["nside"] + delta_ell = config["global"]["delta_ell"] + band_names = list(config["global"]["map_sets"].keys()) + beam_files = list(config["global"]["map_sets"][b]["beam_file"] + for b in band_names) + bandpass_files = list(config["global"]["map_sets"][b]["bandpass_file"] + for b in band_names) + + theory_params = config["theory_params"] + cmb_templates = config["theory_params"]["CMB"]["cmb_templates"] + + output_dir = args.outdir + + # Bandpasses + print(" Loading bandpasses") + bpss = {n: Bpass(n, b) for n, b in zip(band_names, bandpass_files)} + + # Bandpowers + print(" Computing bandpowers") + nmt_binning = get_nmt_binning(nside, delta_ell) + lmax = nmt_binning.lmax + ls = np.arange(lmax + 1) + cl2dl = ls*(ls + 1)/(2*np.pi) + dl2cl = np.zeros_like(cl2dl) + dl2cl[1:] = 1/(cl2dl[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] + + s_wins = sacc.BandpowerWindow(ls, lwin) + + # lbands = np.linspace(2, lmax, n_bins+1, dtype=int) + # windows = np.zeros([n_bins, lmax+1]) + # for b, (l0, lf) in enumerate(zip(lbands[:-1], lbands[1:])): + # windows[b, l0:lf] = (ls * (ls + 1)/(2*np.pi))[l0:lf] + # windows[b, :] /= delta_ell + + # Beams + print(" Loading beams") + beams = {n: read_beam_window(b, lmax) + for n, b in zip(band_names, beam_files)} + + print(" Calculating power spectra") + # Component spectra + dls_comp = np.zeros([3, 2, 3, 2, lmax + 1]) #[ncomp,np,ncomp,np,nl] + (dls_comp[1, 0, 1, 0, :], + dls_comp[1, 1, 1, 1, :], + dls_comp[2, 0, 2, 0, :], + dls_comp[2, 1, 2, 1, :], + dls_comp[0, 0, 0, 0, :], + dls_comp[0, 1, 0, 1, :]) = get_component_spectra( + lmax, config["theory_params"], fname_camb_lens_nobb=cmb_templates[0] + ) + dls_comp *= dl2cl[None, None, None, None, :] + + # Convolve with bandpower windows + #bpw_comp=np.sum(dls_comp[:,:,:,:,None,:]*windows[None,None,None,None,:,:],axis=5) + bpw_comp = np.sum( + dls_comp[:, :, :, :, None, :] * lwin.T[None, None, None, None, :, :], + axis=5 + ) + + # Convolve with bandpasses + seds = get_convolved_seds(band_names, bpss, theory_params) + _, nfreqs = seds.shape + + # Components -> frequencies + bpw_freq_sig = np.einsum('ik,jm,iljno', seds, seds, bpw_comp) + + # N_ell + # TODO: accept general theory power spectra + # sens = 2 + # knee = 1 + # ylf = 1 + fsky = 0.1 + nell = np.zeros([nfreqs, lmax+1]) + # _, nell[:,2:], _ = nc.Simons_Observatory_V3_SA_noise(sens, knee, + # ylf, fsky, + # lmax + 1, 1) + n_bpw = np.sum(nell[:, None, :]*lwin.T[None,:,:], axis=2) + bpw_freq_noi = np.zeros_like(bpw_freq_sig) + for ib, n in enumerate(n_bpw): + bpw_freq_noi[ib, 0, ib, 0, :] = n_bpw[ib, :] + bpw_freq_noi[ib, 1, ib, 1, :] = n_bpw[ib, :] + + # Add noise to signal + bpw_freq_tot = bpw_freq_sig + bpw_freq_noi + bpw_freq_tot = bpw_freq_tot.reshape([nfreqs*2, nfreqs*2, n_bins]) + bpw_freq_sig = bpw_freq_sig.reshape([nfreqs*2, nfreqs*2, n_bins]) + bpw_freq_noi = bpw_freq_noi.reshape([nfreqs*2, nfreqs*2, n_bins]) + + # Creating Sacc files + s_d = sacc.Sacc() + s_f = sacc.Sacc() + s_n = sacc.Sacc() + + # Adding tracers + print(" Adding sacc tracers") + for ib, n in enumerate(band_names): + bandpass = bpss[n] + beam = beams[n] + for s in [s_d, s_f, s_n]: + # print(n) + # print(bandpass.nu) + # print(bandpass.bnu) + # print(ls) + # print(beam) + s.add_tracer('NuMap', n, + quantity='cmb_polarization', + spin=2, + nu=bandpass.nu, + bandpass=bandpass.bnu, + ell=ls, + beam=beam, + nu_unit='GHz', + map_unit='uK_CMB') + + # Adding power spectra + print(" Adding power spectra to sacc") + nmaps = 2*nfreqs + ncross = (nmaps*(nmaps + 1))//2 + indices_tr = np.triu_indices(nmaps) + map_names=[] + + for ib, n in enumerate(band_names): + map_names.append(n + '_E') + map_names.append(n + '_B') + + for ii, (i1, i2) in enumerate(zip(indices_tr[0], indices_tr[1])): + n1 = map_names[i1][:-2] + n2 = map_names[i2][:-2] + p1 = map_names[i1][-1].lower() + p2 = map_names[i2][-1].lower() + cl_type = f'cl_{p1}{p2}' + s_d.add_ell_cl(cl_type, n1, n2, lb, bpw_freq_sig[i1, i2, :], + window=s_wins) + s_f.add_ell_cl(cl_type, n1, n2, lb, bpw_freq_sig[i1, i2, :], + window=s_wins) + s_n.add_ell_cl(cl_type, n1, n2, lb, bpw_freq_noi[i1, i2, :], + window=s_wins) + + # Add covariance + n_split = 4 + def is_cross(s1,f1,s2,f2): + """ + Given two maps with split s1 frequence f1 and split s2 and frequency f2 + Determine whether the power spectrum is cross spectrum. + """ + if s1 != s2: + return True + else: + if f1 != f2: + return True + else: + return False + def get_cross_splits(f1,f2): + """ + Given two frequency f1 and f2, get all map splits pairs that the power + spectrum only contain cross spectrum. + Return split pairs and total number of cross spectrum. + """ + if f1 == f2: + # If freq are the same, then all cross pairs would have different + # split, so the number of total cross pairs is: + num_cross = n_split*(n_split-1)/2 + split_pairs = np.triu_indices(n_split,1) + assert num_cross == split_pairs[0].size + return split_pairs, num_cross + else: + # If freq are different, then any split combination will be cross pair, + # so the number of total cross pairs is: + num_cross = n_split**2 + split_pairs = np.triu_indices(n_split,-n_split) + assert num_cross == split_pairs[0].size + return split_pairs, num_cross + print(" Adding covariance to sacc") + cov_bpw = np.zeros([ncross, n_bins, ncross, n_bins]) + factor_modecount = 1./((2*lb + 1)*delta_ell*fsky) + for ii, (m1, m2) in enumerate(zip(indices_tr[0], indices_tr[1])): + for jj, (m3, m4) in enumerate(zip(indices_tr[0], indices_tr[1])): + # frequency index of each map + f1 = m1//2 + f2 = m2//2 + f3 = m3//2 + f4 = m4//2 + split_pairs_12, num_cross_12 = get_cross_splits(f1,f2) + split_pairs_34, num_cross_34 = get_cross_splits(f3,f4) + for s_alpha, s_beta in zip(*split_pairs_12): + for s_mu, s_nu in zip(*split_pairs_34): + if is_cross(s_alpha, f1, s_mu, f3): + Cell_13 = bpw_freq_sig[m1,m3,:] + else: + Cell_13 = bpw_freq_tot[m1,m3,:] + + if is_cross(s_beta, f2, s_nu, f4): + Cell_24 = bpw_freq_sig[m2,m4,:] + else: + Cell_24 = bpw_freq_tot[m2,m4,:] + + if is_cross(s_alpha, f1, s_nu, f4): + Cell_14 = bpw_freq_sig[m1,m4,:] + else: + Cell_14 = bpw_freq_tot[m1,m4,:] + + if is_cross(s_beta, f2, s_mu, f3): + Cell_23 = bpw_freq_sig[m2,m3,:] + else: + Cell_23 = bpw_freq_tot[m2,m3,:] + + cov_bpw[ii,:,jj,:] += np.diag( + factor_modecount + * (1/num_cross_12) + * (1/num_cross_34) + * ( Cell_13 * Cell_24 + Cell_14 * Cell_23) + ) + cov_bpw = cov_bpw.reshape([ncross * n_bins, ncross * n_bins]) + s_d.add_covariance(cov_bpw) + + # Write output + print(" Writing sacc files") + s_d.save_fits(f"{output_dir}/cls_coadd.fits", overwrite=True) + s_f.save_fits(f"{output_dir}/cls_fid.fits", overwrite=True) + s_n.save_fits(f"{output_dir}/cls_noise.fits", overwrite=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Pre-processing external data for the pipeline") + parser.add_argument("--globals", type=str, + help="Path to the yaml with global parameters") + parser.add_argument("--outdir", type=str, + help="Path to the output directory") + + args = parser.parse_args() + + generate_SO_spectra(args)