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)