diff --git a/bbpower/bandpasses.py b/bbpower/bandpasses.py index 6163548..5a809c5 100644 --- a/bbpower/bandpasses.py +++ b/bbpower/bandpasses.py @@ -70,8 +70,8 @@ def convolve_sed(self, sed, params): if self.do_dphi1: dphi1 = params[self.name_dphi1] - normed_dphi1 = dphi1 * np.pi / 180. * (self.nu - self.nu_mean) / self.nu_mean - dphi1_phase = np.cos(2.*normed_dphi1) + 1j * np.sin(2.*normed_dphi1) + normed_dphi1 = dphi1 * np.pi / 180. * (self.nu - self.nu_mean) / self.nu_mean # noqa + dphi1_phase = np.cos(2.*normed_dphi1) + 1j * np.sin(2.*normed_dphi1) # noqa nu_prime = self.nu + dnu # CMB sed diff --git a/bbpower/compsep.py b/bbpower/compsep.py index a69e8d3..7f6cbc9 100644 --- a/bbpower/compsep.py +++ b/bbpower/compsep.py @@ -3,7 +3,7 @@ from scipy.linalg import sqrtm from bbpipe import PipelineStage -from .types import NpzFile, FitsFile, YamlFile, DirFile +from .types import FitsFile, YamlFile, DirFile from .fg_model import FGModel from .param_manager import ParameterManager from .bandpasses import (Bandpass, rotate_cells, rotate_cells_mat, @@ -46,7 +46,7 @@ def get_moments_lmax(self): return self.config['fg_model'].get('moments_lmax', 384) def precompute_w3j(self): - from pyshtools.utils import Wigner3j + from pyshtools.utils import Wigner3j # noqa lmax = self.get_moments_lmax() ells_w3j = np.arange(0, lmax) @@ -189,7 +189,7 @@ def parse_sacc_file(self): # Get power spectra and covariances if self.config['bands'] == 'all': - if not (self.s_cov.covariance.covmat.shape[-1] == len(self.s.mean) == self.n_bpws * self.ncross): + if not (self.s_cov.covariance.covmat.shape[-1] == len(self.s.mean) == self.n_bpws * self.ncross): # noqa: E501 raise ValueError("C_ell vector's size is wrong") v2d = np.zeros([self.n_bpws, self.ncross]) @@ -198,7 +198,7 @@ def parse_sacc_file(self): v2d_fid = np.zeros([self.n_bpws, self.ncross]) cv2d = np.zeros([self.n_bpws, self.ncross, self.n_bpws, self.ncross]) - self.vector_indices = self.vector_to_matrix(np.arange(self.ncross, dtype=int)).astype(int) + self.vector_indices = self.vector_to_matrix(np.arange(self.ncross, dtype=int)).astype(int) # noqa: E501 self.indx = [] # Parse into the right ordering @@ -227,7 +227,7 @@ def parse_sacc_file(self): pol2b = self.pols[p2b].lower() cl_typb = f'cl_{pol1b}{pol2b}' ind_b = self.s.indices(cl_typb, (t1b, t2b)) - cv2d[:, ind_vec, :, ind_vecb] = self.s_cov.covariance.covmat[ind_a][:, ind_b] + cv2d[:, ind_vec, :, ind_vecb] = self.s_cov.covariance.covmat[ind_a][:, ind_b] # noqa: E501 # Store data self.bbdata = self.vector_to_matrix(v2d) @@ -244,7 +244,8 @@ def load_cmb(self): """ Loads the CMB BB spectrum as defined in the config file. """ - cmb_lensingfile = np.loadtxt(self.config['cmb_model']['cmb_templates'][0]) + cmb_lensingfile = np.loadtxt( + self.config['cmb_model']['cmb_templates'][0]) cmb_bbfile = np.loadtxt(self.config['cmb_model']['cmb_templates'][1]) self.cmb_ells = cmb_bbfile[:, 0] @@ -310,7 +311,7 @@ def sed(nu): for i_c1, c_name1 in enumerate(self.fg_model.component_names): fg_scaling[i_c1, i_c1] = comp_scaling[i_c1] - for c_name2, epsname in self.fg_model.components[c_name1]['names_x_dict'].items(): + for c_name2, epsname in self.fg_model.components[c_name1]['names_x_dict'].items(): # noqa: E501 i_c2 = self.fg_model.component_order[c_name2] eps = params[epsname] fg_scaling[i_c1, i_c2] = eps * np.outer(single_sed[i_c1], @@ -396,8 +397,9 @@ def model(self, params): cl_cross = np.zeros((self.n_ell, self.npol, self.npol)) for i in range(self.npol): - cl_cross[:, i, i] = np.sqrt(fg_cell[c1, :, i, i] * - fg_cell[c2, :, i, i]) + cl_cross[:, i, i] = np.sqrt( + fg_cell[c1, :, i, i] * fg_cell[c2, :, i, i] + ) clrot = rotate_cells_mat(mat2, mat1, cl_cross) cls += clrot * fg_scaling[c1, c2, f1, f2] cls_array_fg[f1, f2] = cls @@ -452,7 +454,7 @@ def model(self, params): cls += 0.5 * (fg_scaling_d2[f1, c1] * (fg_scaling[c1, c1, f2, f2])**0.5 + fg_scaling_d2[f2, c1] * - (fg_scaling[c1, c1, f1, f1])**0.5) * cls_02[c1] + (fg_scaling[c1, c1, f1, f1])**0.5) * cls_02[c1] # noqa: E501 cls_array_fg[f1, f2] += cls # Window convolution @@ -478,7 +480,7 @@ def model(self, params): for f2 in range(self.nfreqs): cls_array_list[:, f1, :, f2, :] = rotate_cells(self.bpss[f2], self.bpss[f1], - cls_array_list[:, f1, :, f2, :], + cls_array_list[:, f1, :, f2, :], # noqa: E501 params) return cls_array_list.reshape([self.n_bpws, self.nmaps, self.nmaps]) @@ -555,7 +557,7 @@ def h_and_l_dx(self, params): """ Hamimeche and Lewis likelihood. Taken from Cobaya written by H, L and Torrado - See: https://github.com/CobayaSampler/cobaya/blob/master/cobaya/likelihoods/_cmblikes_prototype/_cmblikes_prototype.py + See: https://github.com/CobayaSampler/cobaya/blob/master/cobaya/likelihoods/_cmblikes_prototype/_cmblikes_prototype.py # noqa: E501 """ model_cls = self.model(params) dx_vec = [] @@ -571,7 +573,7 @@ def h_and_l_dx(self, params): def h_and_l(self, C, Chat, Cfl_sqrt): try: diag, U = np.linalg.eigh(C) - except: + except: # noqa return [np.inf] rot = U.T.dot(Chat).dot(U) roots = np.sqrt(diag) @@ -581,7 +583,7 @@ def h_and_l(self, C, Chat, Cfl_sqrt): U.dot(rot.dot(U.T), rot) try: diag, rot = np.linalg.eigh(rot) - except: + except: # noqa return [np.inf] diag = (np.sign(diag - 1) * np.sqrt(2 * np.maximum(0, diag - np.log(diag) - 1))) @@ -645,7 +647,7 @@ def emcee_sampler(self): pos = None nsteps_use = max(n_iters-nchain, 0) - with Pool() as pool: + with Pool() as pool: # noqa import time start = time.time() sampler = emcee.EnsembleSampler(nwalkers, ndim, @@ -653,7 +655,7 @@ def emcee_sampler(self): backend=backend) if nsteps_use > 0: sampler.run_mcmc(pos, nsteps_use, store=True, progress=False) - end = time.time() + end = time.time() return sampler, end-start @@ -673,14 +675,14 @@ def prior(hypercube): prior = [] for h, pr in zip(hypercube, self.params.p_free_priors): if pr[1] == 'Gaussian': - prior.append(GaussianPrior(float(pr[2][0]), float(pr[2][1]))(h)) + prior.append(GaussianPrior(float(pr[2][0]), float(pr[2][1]))(h)) # noqa: E501 else: - prior.append(UniformPrior(float(pr[2][0]), float(pr[2][2]))(h)) + prior.append(UniformPrior(float(pr[2][0]), float(pr[2][2]))(h)) # noqa: E501 return prior # Optional dumper function giving run-time read access to # the live points, dead points, weights and evidences - def dumper(live, dead, logweights, logZ, logZerr): + def dumper(live, dead, logweights, logZ, logZerr): # noqa print("Last dead point:", dead[-1]) settings = PolyChordSettings(ndim, nder) @@ -688,12 +690,12 @@ def dumper(live, dead, logweights, logZ, logZerr): settings.file_root = 'pch' settings.nlive = self.config['nlive'] settings.num_repeats = self.config['nrepeat'] - settings.do_clustering = False # Assume unimodal posterior - settings.boost_posterior = 10 # Increase number of posterior samples - settings.nprior = 200 # Draw nprior initial prior samples - settings.maximise = True # Maximize posterior at the end - settings.read_resume = False # Read from resume file of earlier run - settings.feedback = 2 # Verbosity {0,1,2,3} + settings.do_clustering = False # Assume unimodal posterior + settings.boost_posterior = 10 # Increase number of posterior samples + settings.nprior = 200 # Draw nprior initial prior samples + settings.maximise = True # Maximize posterior at the end + settings.read_resume = False # Read from resume file of earlier run + settings.feedback = 2 # Verbosity {0,1,2,3} output = pypolychord.run_polychord(likelihood, ndim, nder, settings, prior, dumper) @@ -729,9 +731,9 @@ def chi2(par): method="Powell") def lnprobd(p): - l = self.lnprob(p) + l = self.lnprob(p) # noqa if l == -np.inf: - l = -1E100 + l = -1E100 # noqa return l fisher = - nd.Hessian(lnprobd)(res.x) @@ -758,8 +760,8 @@ def timing(self, n_eval=300): def predicted_spectra(self, at_min=True, save_npz=True): """ - Evaluates model at a the maximum likelihood and - writes predicted spectra into a numpy array + Evaluates model at a the maximum likelihood and + writes predicted spectra into a numpy array with shape (nbpws, nmaps, nmaps). """ if at_min: @@ -768,7 +770,6 @@ def predicted_spectra(self, at_min=True, save_npz=True): else: p = self.params.p0 pars = self.params.build_params(p) - print(pars) model_cls = self.model(pars) if self.config['bands'] == 'all': tr_names = sorted(list(self.s.tracers.keys())) @@ -799,7 +800,6 @@ def predicted_spectra(self, at_min=True, save_npz=True): s.add_covariance(self.bbcovar) s.save_fits(self.get_output('output_dir')+'/cells_model.fits', overwrite=True) - return def run(self): @@ -808,12 +808,15 @@ def run(self): self.setup_compsep() if self.config.get('sampler') == 'emcee': sampler, timing = self.emcee_sampler() + chi2 = -2*self.lnprob(self.minimizer()) np.savez(self.get_output('output_dir')+'/emcee.npz', chain=sampler.chain, names=self.params.p_free_names, - time=timing) + time=timing, + chi2=chi2, + ndof=len(self.bbcovar)) print("Finished sampling", timing) - elif self.config.get('sampler')=='polychord': + elif self.config.get('sampler') == 'polychord': sampler = self.polychord_sampler() print("Finished sampling") elif self.config.get('sampler') == 'fisher': @@ -848,7 +851,7 @@ def run(self): names=self.params.p_free_names) print("Total time:", sampler[0]) print("Time per eval:", sampler[1]) - elif self.config.get('sampler')=='predicted_spectra': + elif self.config.get('sampler') == 'predicted_spectra': at_min = self.config.get('predict_at_minimum', True) save_npz = not self.config.get('predict_to_sacc', False) sampler = self.predicted_spectra(at_min=at_min, save_npz=save_npz) diff --git a/bbpower/compsep_nopipe.py b/bbpower/compsep_nopipe.py new file mode 100644 index 0000000..8d5e606 --- /dev/null +++ b/bbpower/compsep_nopipe.py @@ -0,0 +1,993 @@ +import numpy as np +import os +import argparse +import yaml +import time +import sacc # noqa +import sys +sys.path.append( + os.path.abspath(os.path.join(os.path.dirname(__file__), '.')) +) + +import mpi_utils as mpi # noqa +from fg_model import FGModel # noqa +from param_manager import ParameterManager # noqa +from bandpasses import (Bandpass, rotate_cells, rotate_cells_mat, # noqa + decorrelated_bpass) + + +def _yaml_loader(config): + """ + Custom yaml loader to load the configuration file. + """ + def path_constructor(loader, node): + return "/".join(loader.construct_sequence(node)) + yaml.SafeLoader.add_constructor("!path", path_constructor) + with open(config, "r") as f: + return yaml.load(f, Loader=yaml.SafeLoader) + + +class BBCompSep(object): + """ + Component separation stage + This stage does harmonic domain foreground cleaning (e.g. BICEP). + The foreground model parameters are defined in the config.yml file. + """ + def __init__(self, args): + """ + Initialize from the command line arguments. + + Parameters + ---------- + args : str + Command line arguments. + """ + + # Load the configuration file + config = _yaml_loader(args.config) + self.config = config["global"] | config["BBCompSep"] + setattr(self, "config_fname", getattr(args, "config")) + setattr(self, "data", self.config["data"]) + + def setup_compsep(self): + """ + Pre-load the data, CMB BB power spectrum, and foreground models. + """ + self.parse_sacc_file() + if self.config['fg_model'].get('use_moments'): + self.precompute_w3j() + self.load_cmb() + self.fg_model = FGModel(self.config) + self.params = ParameterManager(self.config) + if self.use_handl: + self.prepare_h_and_l() + return + + def get_moments_lmax(self): + return self.config['fg_model'].get('moments_lmax', 384) + + def precompute_w3j(self): + from pyshtools.utils import Wigner3j + + lmax = self.get_moments_lmax() + ells_w3j = np.arange(0, lmax) + w3j = np.zeros_like(ells_w3j, dtype=float) + self.big_w3j = np.zeros((lmax, lmax, lmax)) + for ell1 in ells_w3j[1:]: + for ell2 in ells_w3j[1:]: + w3j_array, ellmin, ellmax = Wigner3j(ell1, ell2, 0, 0, 0) + w3j_array = w3j_array[:ellmax - ellmin + 1] + # make the w3j_array the same shape as the w3j + if len(w3j_array) < len(ells_w3j): + reference = np.zeros(len(w3j)) + reference[:w3j_array.shape[0]] = w3j_array + w3j_array = reference + + w3j_array = np.concatenate([w3j_array[-ellmin:], + w3j_array[:-ellmin]]) + w3j_array = w3j_array[:len(ells_w3j)] + w3j_array[:ellmin] = 0 + + self.big_w3j[:, ell1, ell2] = w3j_array + + self.big_w3j = self.big_w3j**2 + + def matrix_to_vector(self, mat): + return mat[..., self.index_ut[0], self.index_ut[1]] + + def vector_to_matrix(self, vec): + if vec.ndim == 1: + mat = np.zeros([self.nmaps, self.nmaps]) + mat[self.index_ut] = vec + mat = mat + mat.T - np.diag(mat.diagonal()) + elif vec.ndim == 2: + mat = np.zeros([len(vec), self.nmaps, self.nmaps]) + mat[..., self.index_ut[0], self.index_ut[1]] = vec[..., :] + for i, m in enumerate(mat): + mat[i] = m + m.T - np.diag(m.diagonal()) + else: + raise ValueError("Input vector can only be 1- or 2-D") + return mat + + def _freq_pol_iterator(self): + icl = -1 + map_sets = list(self.config["map_sets"]) + for b1 in range(len(map_sets)): + for p1 in range(self.npol): + m1 = p1 + self.npol * b1 + for b2 in range(b1, len(map_sets)): + if b1 == b2: + p2_r = range(p1, self.npol) + else: + p2_r = range(self.npol) + for p2 in p2_r: + m2 = p2 + self.npol * b2 + icl += 1 + yield map_sets[b1], map_sets[b2], p1, p2, m1, m2, icl + + def parse_sacc_file(self): + """ + Reads the data in the sacc file included the power spectra, + bandpasses, and window functions. + """ + from copy import deepcopy + + # Decide if you're using H&L + self.use_handl = self.config['likelihood_type'] == 'h&l' + + # Read data + cells_coadded = self.data["cells_coadded"].format(sim_id=self.sim_id) + self.s = sacc.Sacc.load_fits(cells_coadded) + self.s_cov = sacc.Sacc.load_fits(self.data["cells_coadded_cov"]) + if self.use_handl: + s_fid = sacc.Sacc.load_fits(self.data["cells_fiducial"]) + s_noi = sacc.Sacc.load_fits(self.data["cells_noise"]) + + # Keep only desired tracers + tr_names = list(self.config['map_sets'].keys()) + tr_names_before = deepcopy(self.s.tracers) + for tr in tr_names_before: + if tr not in tr_names: + self.s.remove_tracers([tr]) + self.s_cov.remove_tracers([tr]) + tr_comb = self.s.get_tracer_combinations() + + # Keep only desired correlations + self.pols = self.config['pol_channels'] + corr_all = ['cl_00', 'cl_0e', 'cl_0b', 'cl_e0', 'cl_b0', + 'cl_ee', 'cl_eb', 'cl_be', 'cl_bb'] + corr_keep = [] + for m1 in self.pols: + for m2 in self.pols: + clname = 'cl_' + m1.lower() + m2.lower() + if "0" in clname: + continue + corr_keep.append(clname) + for c in corr_all: + if c not in corr_keep: + self.s.remove_selection(c) + self.s_cov.remove_selection(c) + if self.use_handl: + s_fid.remove_selection(c) + s_noi.remove_selection(c) + + # Scale cuts + self.s.remove_selection(ell__gt=self.config['l_max']) + self.s.remove_selection(ell__lt=self.config['l_min']) + self.s_cov.remove_selection(ell__gt=self.config['l_max']) + self.s_cov.remove_selection(ell__lt=self.config['l_min']) + if self.use_handl: + s_fid.remove_selection(ell__gt=self.config['l_max']) + s_fid.remove_selection(ell__lt=self.config['l_min']) + s_noi.remove_selection(ell__gt=self.config['l_max']) + s_noi.remove_selection(ell__lt=self.config['l_min']) + + for tr1, tr2 in tr_comb: + ind1 = self.s.indices(data_type='cl_bb', tracers=(tr1, tr2)) + ind2 = self.s_cov.indices(data_type='cl_bb', tracers=(tr1, tr2)) + assert np.all(ind1 == ind2), "Covariance sacc ordering is wrong" + if self.use_handl: + ind3 = self.s_fid.indices(data_type='cl_bb', + tracers=(tr1, tr2)) + ind4 = self.s_noi.indices(data_type='cl_bb', + tracers=(tr1, tr2)) + assert np.all(ind1 == ind3), "Fiducial sacc ordering is wrong" + assert np.all(ind1 == ind4), "Noise sacc ordering is wrong" + + self.nfreqs = len(tr_names) + self.npol = len(self.pols) + self.nmaps = self.nfreqs * self.npol + self.index_ut = np.triu_indices(self.nmaps) + self.ncross = (self.nmaps * (self.nmaps + 1)) // 2 + self.pol_order = dict(zip(self.pols, range(self.npol))) + + # Collect bandpasses + self.bpss = [] + for i_t, tn in enumerate(tr_names): + t = self.s.tracers[tn] + nu = t.nu + dnu = np.zeros_like(nu) + dnu[1:-1] = 0.5 * (nu[2:] - nu[:-2]) + dnu[0] = nu[1] - nu[0] + dnu[-1] = nu[-1] - nu[-2] + bnu = t.bandpass + self.bpss.append(Bandpass(nu, dnu, bnu, i_t+1, self.config)) + + # Get ell sampling + # Example power spectrum + self.ell_b, _ = self.s.get_ell_cl('cl_' + 2 * self.pols[0].lower(), + tr_names[0], tr_names[0]) + # Avoid l<2 + win0 = self.s.data[0]['window'] + mask_w = win0.values > 1 + self.bpw_l = win0.values[mask_w] + self.n_ell = len(self.bpw_l) + self.n_bpws = len(self.ell_b) + # D_ell factor + self.dl2cl = 2 * np.pi / (self.bpw_l * (self.bpw_l + 1)) + self.windows = np.zeros([self.ncross, self.n_bpws, self.n_ell]) + + # Get power spectra and covariances + if not (self.s_cov.covariance.covmat.shape[-1] == len(self.s.mean) + == self.n_bpws * (self.nfreqs * (self.nfreqs + 1)) // 2 * self.npol**2): # noqa + raise ValueError("C_ell vector's size is wrong") + + v2d = np.zeros([self.n_bpws, self.ncross]) + if self.use_handl: + v2d_noi = np.zeros([self.n_bpws, self.ncross]) + v2d_fid = np.zeros([self.n_bpws, self.ncross]) + cv2d = np.zeros([self.n_bpws, self.ncross, self.n_bpws, self.ncross]) + + self.vector_indices = self.vector_to_matrix( + np.arange(self.ncross, dtype=int) + ).astype(int) + self.indx = [] + + # Parse into the right ordering + itr1 = self._freq_pol_iterator() + for b1, b2, p1, p2, m1, m2, ind_vec in itr1: + pol1 = self.pols[p1].lower() + pol2 = self.pols[p2].lower() + cl_typ = f'cl_{pol1}{pol2}' + ind_a = self.s.indices(cl_typ, (b1, b2)) + if len(ind_a) != self.n_bpws: + raise ValueError("All power spectra need to be " + "sampled at the same ells") + w = self.s.get_bandpower_windows(ind_a) + self.windows[ind_vec, :, :] = w.weight[mask_w, :].T + v2d[:, ind_vec] = np.array(self.s.mean[ind_a]) + if self.use_handl: + _, v2d_noi[:, ind_vec] = s_noi.get_ell_cl(cl_typ, b1, b2) + _, v2d_fid[:, ind_vec] = s_fid.get_ell_cl(cl_typ, b1, b2) + itr2 = self._freq_pol_iterator() + for b1b, b2b, p1b, p2b, _, _, ind_vecb in itr2: + pol1b = self.pols[p1b].lower() + pol2b = self.pols[p2b].lower() + cl_typb = f'cl_{pol1b}{pol2b}' + ind_b = self.s.indices(cl_typb, (b1b, b2b)) + cv2d[:, ind_vec, :, ind_vecb] = self.s_cov.covariance.covmat[ind_a][:, ind_b] # noqa + + # Store data + self.bbdata = self.vector_to_matrix(v2d) + if self.use_handl: + self.bbnoise = self.vector_to_matrix(v2d_noi) + self.bbfiducial = self.vector_to_matrix(v2d_fid) + self.bbcovar = cv2d.reshape([self.n_bpws * self.ncross, + self.n_bpws * self.ncross]) + self.invcov = np.linalg.solve(self.bbcovar, + np.identity(len(self.bbcovar))) + np.savez(self.output_dir + '/data_ell_cl_invcov.npz', + ell=self.ell_b, cl=self.bbdata, invcov=self.invcov) + return + + def load_cmb(self): + """ + Loads the CMB BB spectrum as defined in the config file. + """ + cmb_lensingfile = np.loadtxt( + self.config['cmb_model']['cmb_templates'][0] + ) + cmb_bbfile = np.loadtxt(self.config['cmb_model']['cmb_templates'][1]) + + self.cmb_ells = cmb_bbfile[:, 0] + mask = (self.cmb_ells <= self.bpw_l.max()) & (self.cmb_ells > 1) + self.cmb_ells = self.cmb_ells[mask] + + # TODO: this is a patch + nell = len(self.cmb_ells) + self.cmb_tens = np.zeros([self.npol, self.npol, nell]) + self.cmb_lens = np.zeros([self.npol, self.npol, nell]) + self.cmb_scal = np.zeros([self.npol, self.npol, nell]) + if 'B' in self.config['pol_channels']: + ind = self.pol_order['B'] + self.cmb_tens[ind, ind] = (cmb_bbfile[:, 3][mask] - + cmb_lensingfile[:, 3][mask]) + self.cmb_lens[ind, ind] = cmb_lensingfile[:, 3][mask] + if 'E' in self.config['pol_channels']: + ind = self.pol_order['E'] + self.cmb_tens[ind, ind] = (cmb_bbfile[:, 2][mask] - + cmb_lensingfile[:, 2][mask]) + self.cmb_scal[ind, ind] = cmb_lensingfile[:, 2][mask] + return + + def integrate_seds(self, params): + """ + Integrated the SEDs of sky components and integrate them + over the given bandpasses. Output the relative scaling amplitudes + for cross correlations of different frequency channels, and the + corresponding polarization angle rotation matrices. + """ + single_sed = np.zeros([self.fg_model.n_components, + self.nfreqs]) + comp_scaling = np.zeros([self.fg_model.n_components, + self.nfreqs, self.nfreqs]) + fg_scaling = np.zeros([self.fg_model.n_components, + self.fg_model.n_components, + self.nfreqs, self.nfreqs]) + rot_matrices = [] + + for i_c, c_name in enumerate(self.fg_model.component_names): + comp = self.fg_model.components[c_name] + units = comp['cmb_n0_norm'] + sed_params = [params[comp['names_sed_dict'][k]] + for k in comp['sed'].params] + rot_matrices.append([]) + + def sed(nu): + return comp['sed'].eval(nu, *sed_params) + + for tn in range(self.nfreqs): + sed_b, rot = self.bpss[tn].convolve_sed(sed, params) + single_sed[i_c, tn] = sed_b * units + rot_matrices[i_c].append(rot) + + if comp['decorr']: + d_amp = params[comp['decorr_param_names']['decorr_amp']] + d_nu01 = params[comp['decorr_param_names']['decorr_nu01']] + d_nu02 = params[comp['decorr_param_names']['decorr_nu02']] + decorr_delta = d_amp**(1./np.log(d_nu01/d_nu02)**2) + for f1 in range(self.nfreqs): + for f2 in range(f1, self.nfreqs): + sed_12 = decorrelated_bpass(self.bpss[f1], + self.bpss[f2], + sed, params, + decorr_delta) + comp_scaling[i_c, f1, f2] = sed_12 * units * units + else: + comp_scaling[i_c] = np.outer(single_sed[i_c], single_sed[i_c]) + + for i_c1, c_name1 in enumerate(self.fg_model.component_names): + fg_scaling[i_c1, i_c1] = comp_scaling[i_c1] + for c_name2, epsname in self.fg_model.components[c_name1]['names_x_dict'].items(): # noqa + i_c2 = self.fg_model.component_order[c_name2] + eps = params[epsname] + fg_scaling[i_c1, i_c2] = eps * np.outer(single_sed[i_c1], + single_sed[i_c2]) + fg_scaling[i_c2, i_c1] = eps * np.outer(single_sed[i_c2], + single_sed[i_c1]) + return fg_scaling, np.array(rot_matrices) + + def evaluate_power_spectra(self, params): + """ + Evaluates the power spectra for each component and each ell + """ + fg_pspectra = np.zeros([self.fg_model.n_components, self.npol, + self.npol, self.n_ell]) + + # Fill diagonal + for i_c, c_name in enumerate(self.fg_model.component_names): + comp = self.fg_model.components[c_name] + for cl_comb, clfunc in comp['cl'].items(): + m1, m2 = cl_comb + ip1 = self.pol_order[m1] + ip2 = self.pol_order[m2] + pspec_params = [params[comp['names_cl_dict'][cl_comb][k]] + for k in clfunc.params] + p_spec = clfunc.eval(self.bpw_l, *pspec_params) * self.dl2cl + fg_pspectra[i_c, ip1, ip2] = p_spec + if m1 != m2: + fg_pspectra[i_c, ip2, ip1] = p_spec + + return fg_pspectra + + def model(self, params): + """ + Defines the total model and integrates over bandpasses and windows. + """ + # [npol,npol,nell] + cmb_cell = (params['r_tensor'] * self.cmb_tens + + params['A_lens'] * self.cmb_lens + + self.cmb_scal) * self.dl2cl + # [nell,npol,npol] + cmb_cell = np.transpose(cmb_cell, axes=[2, 0, 1]) + + if self.config['cmb_model'].get('use_birefringence'): + bi_angle = np.radians(params['birefringence']) + c = np.cos(2*bi_angle) + s = np.sin(2*bi_angle) + bmat = np.array([[c, s], + [-s, c]]) + cmb_cell = rotate_cells_mat(bmat, bmat, cmb_cell) + # [ncomp, ncomp, nfreq, nfreq], [ncomp, nfreq,[matrix]] + fg_scaling, rot_m = self.integrate_seds(params) + # [ncomp,npol,npol,nell] + fg_cell = self.evaluate_power_spectra(params) + # [nfreq, nfreq, nell, npol, npol] + cls_array_fg = np.zeros([self.nfreqs, self.nfreqs, + self.n_ell, self.npol, self.npol]) + # [ncomp,nell,npol,npol] + fg_cell = np.transpose(fg_cell, axes=[0, 3, 1, 2]) + + # SED scaling + cmb_scaling = np.ones(self.nfreqs) + cmb_rot = [] + for f1 in range(self.nfreqs): + cs, crot = self.bpss[f1].convolve_sed(None, params) + cmb_scaling[f1] = cs + cmb_rot.append(crot) + + for f1 in range(self.nfreqs): + # Note that we only need to fill in half of the frequencies + for f2 in range(f1, self.nfreqs): + cls = (rotate_cells_mat(cmb_rot[f2], cmb_rot[f1], cmb_cell) * + cmb_scaling[f1] * cmb_scaling[f2]) + + # Loop over component pairs + for c1 in range(self.fg_model.n_components): + for c2 in range(self.fg_model.n_components): + mat1 = rot_m[c1, f1] + mat2 = rot_m[c2, f2] + if c1 == c2: + clrot = rotate_cells_mat(mat2, mat1, fg_cell[c1]) + else: + # For cross component, enforcing EB term is zero. + cl_cross = np.zeros((self.n_ell, + self.npol, self.npol)) + for i in range(self.npol): + cl_cross[:, i, i] = np.sqrt( + fg_cell[c1, :, i, i] * fg_cell[c2, :, i, i] + ) + clrot = rotate_cells_mat(mat2, mat1, cl_cross) + cls += clrot * fg_scaling[c1, c2, f1, f2] + cls_array_fg[f1, f2] = cls + + # Add moment terms if needed + if self.config['fg_model'].get('use_moments'): + # TODO: moments work with: + # - B-only + # - No polarization angle business + # - Only power-law beta power spectra at l_pivot=80 + + # Evaluate 1st/2nd order SED derivatives. + # [nfreq, ncomp] + fg_scaling_d1 = self.integrate_seds_der(params, order=1) + fg_scaling_d2 = self.integrate_seds_der(params, order=2) + + # Compute 1x1 for each component + # Compute 0x2 for each component (essentially this is sigma_beta) + # Evaluate beta power spectra. + lmax_mom = self.get_moments_lmax() + # [ncomp, nell, npol, npol] + cls_11 = np.zeros([self.fg_model.n_components, self.n_ell, + self.npol, self.npol]) + # [ncomp, nell, npol, npol] + cls_02 = np.zeros([self.fg_model.n_components, self.n_ell, + self.npol, self.npol]) + for i_c, c_name in enumerate(self.fg_model.component_names): + comp = self.fg_model.components[c_name] + gamma = params[comp['names_moments_dict']['gamma_beta']] + amp = params[comp['names_moments_dict']['amp_beta']] * 1E-6 + cl_betas = self.bcls(lmax=lmax_mom, gamma=gamma, amp=amp) + cl_cc = fg_cell[i_c, :] + # cls_1x1 = 0 + cls_1x1 = self.evaluate_1x1(params, lmax=lmax_mom, + cls_cc=cl_cc, + cls_bb=cl_betas) + cls_11[i_c, :lmax_mom, :, :] = cls_1x1 + # cls_0x2 = 0 + cls_0x2 = self.evaluate_0x2(params, lmax=lmax_mom, + cls_cc=cl_cc, + cls_bb=cl_betas) + cls_02[i_c, :lmax_mom, :, :] = cls_0x2 + + # Add components scaled in frequency + for f1 in range(self.nfreqs): + # Note that we only need to fill in half of the frequencies + for f2 in range(f1, self.nfreqs): + cls = np.zeros([self.n_ell, self.npol, self.npol]) + for c1 in range(self.fg_model.n_components): + cls += (fg_scaling_d1[f1, c1] * fg_scaling_d1[f2, c1] * + cls_11[c1]) + cls += 0.5 * (fg_scaling_d2[f1, c1] * + (fg_scaling[c1, c1, f2, f2])**0.5 + + fg_scaling_d2[f2, c1] * + (fg_scaling[c1, c1, f1, f1])**0.5) * cls_02[c1] # noqa + cls_array_fg[f1, f2] += cls + + # Bandpower window convolution + cls_array_list = np.zeros([self.n_bpws, self.nfreqs, + self.npol, self.nfreqs, + self.npol]) + for f1 in range(self.nfreqs): + for p1 in range(self.npol): + m1 = f1*self.npol+p1 + for f2 in range(f1, self.nfreqs): + p0 = p1 if f1 == f2 else 0 + for p2 in range(p0, self.npol): + m2 = f2*self.npol+p2 + windows = self.windows[self.vector_indices[m1, m2]] + clband = np.dot(windows, cls_array_fg[f1, f2, :, + p1, p2]) + cls_array_list[:, f1, p1, f2, p2] = clband + if m1 != m2: + cls_array_list[:, f2, p2, f1, p1] = clband + + # Polarization angle rotation + for f1 in range(self.nfreqs): + for f2 in range(self.nfreqs): + cls_array_list[:, f1, :, f2, :] = rotate_cells( + self.bpss[f2], self.bpss[f1], + cls_array_list[:, f1, :, f2, :], params + ) + + return cls_array_list.reshape([self.n_bpws, self.nmaps, self.nmaps]) + + def bcls(self, lmax, gamma, amp): + ls = np.arange(lmax) + bcls = np.zeros(len(ls)) + bcls[2:] = (ls[2:] / 80.)**gamma + return bcls*amp + + def integrate_seds_der(self, params, order=1): + """ + Define the first order derivative of the SED + """ + fg_scaling_der = np.zeros([self.fg_model.n_components, + self.nfreqs]) + + for i_c, c_name in enumerate(self.fg_model.component_names): + comp = self.fg_model.components[c_name] + units = comp['cmb_n0_norm'] + sed_params = [params[comp['names_sed_dict'][k]] + for k in comp['sed'].params] + + # Set SED function with scaling beta + def sed_der(nu): + nu0 = params[comp['names_sed_dict']['nu0']] + x = np.log(nu / nu0) + # This is only valid for spectral indices + return x**order * comp['sed'].eval(nu, *sed_params) + + for tn in range(self.nfreqs): + sed_b = self.bpss[tn].convolve_sed(sed_der, params)[0] + fg_scaling_der[i_c, tn] = sed_b * units + + return fg_scaling_der.T + + def evaluate_1x1(self, params, lmax, cls_cc, cls_bb): + """ + Evaluate the 1x1 moment for auto-spectra + """ + + ls = np.arange(lmax) + v_left = (2*ls+1)[:, None, None] * cls_cc[:lmax, :, :] + v_right = (2*ls+1) * cls_bb[:lmax] + + mat = self.big_w3j + v_left = np.transpose(v_left, axes=[1, 0, 2]) + moment1x1 = np.dot(np.dot(mat, v_right), v_left) / (4*np.pi) + return moment1x1 + + def evaluate_0x2(self, params, lmax, cls_cc, cls_bb): + """ + Evaluate the 0x2 moment for auto-spectra + Assume power law for beta + """ + ls = np.arange(lmax) + prefac = np.sum((2 * ls + 1) * cls_bb) / (4*np.pi) + return cls_cc[:lmax] * prefac + + def chi_sq_dx(self, params): + """ + Chi^2 likelihood. + """ + model_cls = self.model(params) + return self.matrix_to_vector(self.bbdata - model_cls).flatten() + + def prepare_h_and_l(self): + """ + Prepare the HL likelihood. + """ + from scipy.linalg import sqrtm # noqa + fiducial_noise = self.bbfiducial + self.bbnoise + self.Cfl_sqrt = np.array([sqrtm(f) for f in fiducial_noise]) + self.observed_cls = self.bbdata + self.bbnoise + return + + def h_and_l_dx(self, params): + """ + Hamimeche and Lewis likelihood. + Taken from Cobaya written by H, L and Torrado + See: https://github.com/CobayaSampler/cobaya/blob/master/cobaya/likelihoods/_cmblikes_prototype/_cmblikes_prototype.py # noqa + """ + model_cls = self.model(params) + dx_vec = [] + for k in range(model_cls.shape[0]): + C = model_cls[k] + self.bbnoise[k] + X = self.h_and_l(C, self.observed_cls[k], self.Cfl_sqrt[k]) + if np.any(np.isinf(X)): + return [np.inf] + dx = self.matrix_to_vector(X).flatten() + dx_vec = np.concatenate([dx_vec, dx]) + return dx_vec + + def h_and_l(self, C, Chat, Cfl_sqrt): + """ + Evaluate the Hamimeche and Lewis likelihood. + """ + try: + diag, U = np.linalg.eigh(C) + except: # noqa + return [np.inf] + rot = U.T.dot(Chat).dot(U) + roots = np.sqrt(diag) + for i, root in enumerate(roots): + rot[i, :] /= root + rot[:, i] /= root + U.dot(rot.dot(U.T), rot) + try: + diag, rot = np.linalg.eigh(rot) + except: # noqa + return [np.inf] + diag = (np.sign(diag - 1) * + np.sqrt(2 * np.maximum(0, diag - np.log(diag) - 1))) + Cfl_sqrt.dot(rot, U) + for i, d in enumerate(diag): + rot[:, i] = U[:, i] * d + return rot.dot(U.T) + + def lnprob(self, par): + """ + Likelihood with priors. + """ + prior = self.params.lnprior(par) + if not np.isfinite(prior): + return -np.inf + + return prior + self.lnlike(par) + + def lnlike(self, par): + """ + Likelihood without priors. + """ + params = self.params.build_params(par) + if self.use_handl: + dx = self.h_and_l_dx(params) + if np.any(np.isinf(dx)): + return -np.inf + else: + dx = self.chi_sq_dx(params) + like = -0.5 * np.dot(dx, np.dot(self.invcov, dx)) + + return like + + def emcee_sampler(self): + """ + Sample the model with MCMC. + """ + import emcee # noqa + from multiprocessing import Pool + + fname_temp = self.output_dir + '/emcee.npz.h5' + backend = emcee.backends.HDFBackend(fname_temp) + + nwalkers = self.config['nwalkers'] + n_iters = self.config['n_iters'] + ndim = len(self.params.p0) + found_file = os.path.isfile(fname_temp) + + try: + nchain = len(backend.get_chain()) + except AttributeError: + found_file = False + + if found_file and self.config['resume']: + print("Restarting from previous run") + pos = None + nsteps_use = max(n_iters-nchain, 0) + else: + backend.reset(nwalkers, ndim) + pos = [self.params.p0 + 1.e-3*np.random.randn(ndim) + for i in range(nwalkers)] + nsteps_use = n_iters + + with Pool() as pool: # noqa + import time + start = time.time() + sampler = emcee.EnsembleSampler(nwalkers, ndim, + self.lnprob, + backend=backend) + if nsteps_use > 0: + sampler.run_mcmc(pos, nsteps_use, store=True, progress=True) + end = time.time() + + return sampler, end-start + + def polychord_sampler(self): + """ + Sample the model with PolyChord nested sampler. + """ + import pypolychord # noqa + from pypolychord.settings import PolyChordSettings # noqa + from pypolychord.priors import UniformPrior, GaussianPrior # noqa + ndim = len(self.params.p0) + nder = 0 + + # Log-likelihood compliant with PolyChord's input + def likelihood(theta): + return self.lnlike(theta), [0] + + def prior(hypercube): + prior = [] + for h, pr in zip(hypercube, self.params.p_free_priors): + if pr[1] == 'Gaussian': + prior.append( + GaussianPrior(float(pr[2][0]), float(pr[2][1]))(h) + ) + else: + prior.append( + UniformPrior(float(pr[2][0]), float(pr[2][2]))(h) + ) + return prior + + # Optional dumper function giving run-time read access to + # the live points, dead points, weights and evidences + def dumper(live, dead, logweights, logZ, logZerr): + print("Last dead point:", dead[-1]) + + settings = PolyChordSettings(ndim, nder) + settings.base_dir = self.output_dir + '/polychord' + os.makedirs(settings.base_dir, exist_ok=True) + settings.file_root = 'pch' + settings.nlive = self.config['nlive'] + settings.num_repeats = self.config['nrepeat'] + settings.do_clustering = False # Assume unimodal posterior + settings.boost_posterior = 10 # Increase number of posterior samples + settings.nprior = 200 # Draw nprior initial prior samples + settings.maximise = True # Maximize posterior at the end + settings.read_resume = self.config['resume'] # Resume for earlier run + settings.feedback = 0 # Verbosity {0,1,2,3} + + output = pypolychord.run_polychord(likelihood, ndim, nder, settings, + prior, dumper) + + return output + + def minimizer(self): + """ + Find maximum posterior value + """ + from scipy.optimize import minimize # noqa + + def chi2(par): + c2 = -2*self.lnprob(par) + return c2 + + res = minimize(chi2, self.params.p0, method="Powell") + return res.x + + def fisher(self): + """ + Evaluate Fisher matrix (corresponding to the posterior distribution) + """ + import numdifftools as nd # noqa + from scipy.optimize import minimize # noqa + + def chi2(par): + c2 = -2*self.lnprob(par) + return c2 + + res = minimize(chi2, self.params.p0, method="Powell") + + def lnprobd(p): + ll = self.lnprob(p) + if ll == -np.inf: + ll = -1E100 + return ll + + fisher = - nd.Hessian(lnprobd)(res.x) + return res.x, fisher + + def singlepoint(self): + """ + Evaluate at a single point + """ + chi2 = -2 * self.lnprob(self.params.p0) + return chi2 + + def timing(self, n_eval=300): + """ + Evaluate n times and benchmark + """ + import time + start = time.time() + for i in range(n_eval): + self.lnprob(self.params.p0) + end = time.time() + + return end-start, (end-start)/n_eval + + def predicted_spectra(self, at_min=True, save_npz=True): + """ + Evaluates model at a the maximum likelihood and + writes predicted spectra into a numpy array + with shape (nbpws, nmaps, nmaps). + """ + if at_min: + sampler = self.minimizer() + p = np.array(sampler) + else: + p = self.params.p0 + pars = self.params.build_params(p) + model_cls = self.model(pars) + print("model_cls", model_cls.shape) + tr_names = list(self.config['map_sets'].keys()) + + if save_npz: + np.savez(self.output_dir+'/cells_model.npz', + tracers=tr_names, + ls=self.ell_b, + dls=model_cls) + return + s = sacc.Sacc() + for tn in tr_names: + t = self.s.tracers[tn] + s.add_tracer('NuMap', tn, quantity='cmb_polarization', + spin=2, nu=t.nu, bandpass=t.bandpass, + ell=t.ell, beam=t.beam, nu_unit='GHz', + map_unit='uK_CMB') + for b1, b2, p1, p2, m1, m2, ind in self._freq_pol_iterator(): + cl = model_cls[:, m1, m2] + pol1 = self.pols[p1].lower() + pol2 = self.pols[p2].lower() + cltyp = f'cl_{pol1}{pol2}' + win = sacc.BandpowerWindow(self.bpw_l, self.windows[ind].T) + s.add_ell_cl(cltyp, b1, b2, self.ell_b, cl, window=win) + + s.add_covariance(self.bbcovar) + s.save_fits(self.output_dir+'/cells_model.fits', + overwrite=True) + return + + def run(self): + """ + Run the BB component separation stage. + """ + # Make output directory + self.output_dir = self.config["output_dir"].format(sim_id=self.sim_id) + os.makedirs(self.output_dir, exist_ok=True) + + # Make duplicate of config file + from shutil import copyfile + copyfile(self.config_fname, + self.config["config_copy"].format(sim_id=self.sim_id)) + self.setup_compsep() + + # Get the chi2 and best-fit estimates + from scipy.stats import chi2 as scipy_chi2 + sampler = self.minimizer() + chi2 = -2*self.lnprob(sampler) + ndof = len(self.bbcovar) + kwargs = { + "params": sampler, + "names": self.params.p_free_names, + "chi2": chi2, + "ndof": len(self.bbcovar), + "pte": scipy_chi2.sf(chi2, ndof, loc=0, scale=1) + } + np.savez(self.output_dir+'/chi2.npz', **kwargs) + with open(self.output_dir+'/chi2.txt', 'w') as f: + for key, value in kwargs.items(): + f.write('%s: %s\n' % (key, value)) + print("Saved best-fit parameters") + + # Get best-fit power spectra + at_min = self.config.get('predict_at_minimum', True) + save_npz = not self.config.get('predict_to_sacc', False) + sampler = self.predicted_spectra(at_min=at_min, save_npz=save_npz) + print("Predicted spectra saved") + + if self.config.get('sampler') == 'emcee': + sampler, timing = self.emcee_sampler() + np.savez(self.output_dir+'/emcee.npz', + chain=sampler.chain, + names=self.params.p_free_names, + time=timing) + print("Finished sampling", timing) + + elif self.config.get('sampler') == 'polychord': + os.makedirs(self.output_dir + '/polychord/clusters', exist_ok=True) + sampler = self.polychord_sampler() + print("Finished sampling") + + elif self.config.get('sampler') == 'fisher': + p0, fisher = self.fisher() + cov = np.linalg.inv(fisher) + for i, (n, p) in enumerate(zip(self.params.p_free_names, p0)): + print(n+" = %.3lE +- %.3lE" % (p, np.sqrt(cov[i, i]))) + np.savez(self.output_dir+'/fisher.npz', + params=p0, fisher=fisher, + names=self.params.p_free_names) + + elif self.config.get('sampler') == 'single_point': + sampler = self.singlepoint() + np.savez(self.output_dir+'/single_point.npz', + chi2=sampler, ndof=len(self.bbcovar), + names=self.params.p_free_names) + print("Chi2:", sampler, len(self.bbcovar)) + + elif self.config.get('sampler') == 'timing': + sampler = self.timing() + np.savez(self.output_dir+'/timing.npz', + timing=sampler[1], + names=self.params.p_free_names) + print("Total time:", sampler[0]) + print("Time per eval:", sampler[1]) + else: + raise ValueError("Unknown sampler") + + return + + +def main(args): + """ + Execute the BBCompSep stage with arguments args: + * config: string. + Path to configuration file with input parameters. + """ + config = _yaml_loader(args.config) + + # Creating the simulation indices range to loop over + sim_ids = config["global"]["sim_ids"] + if isinstance(sim_ids, list): + sim_ids = np.array(sim_ids, dtype=int) + elif isinstance(sim_ids, str): + if "," in sim_ids: + id_min, id_max = sim_ids.split(",") + sim_ids = np.arange(int(id_min), int(id_max)+1) + else: + sim_ids = np.array([int(sim_ids)]) + else: + sim_ids = [None] + + # MPI related initialization + rank, size, comm = mpi.init(True) + + # Initialize tasks for MPI sharing + mpi_shared_list = sim_ids + + # Every rank must have the same shared list + mpi_shared_list = comm.bcast(mpi_shared_list, root=0) + task_ids = mpi.distribute_tasks(size, rank, len(mpi_shared_list), + logger=None) + local_mpi_list = [mpi_shared_list[i] for i in task_ids] + + for sim_id in local_mpi_list: + start = time.time() + compsep = BBCompSep(args) + setattr(compsep, "sim_id", sim_id) + compsep.run() + mpi.print_rnk0(f"Processed {len(sim_ids)} simulations " + f"in {time.time() - start:.1f} seconds.", rank) + comm.Barrier() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description="SO SAT BB likelihood") + parser.add_argument( + "--config", type=str, + help="Path to yaml file with pipeline configuration" + ) + + args = parser.parse_args() + main(args) diff --git a/bbpower/fg_model.py b/bbpower/fg_model.py index 3fc7ef2..5101714 100644 --- a/bbpower/fg_model.py +++ b/bbpower/fg_model.py @@ -1,5 +1,11 @@ import fgbuster.component_model as fgc -import bbpower.fgcls as fgl +import sys +import os +sys.path.append( + os.path.abspath(os.path.join(os.path.dirname(__file__), '.')) +) + +import fgcls as fgl # noqa class FGModel: @@ -46,7 +52,7 @@ def load_foregrounds(self, config): "is not a valid component" + "to correlate %s with" % key) if par[0] == key: - raise KeyError("%s is cross correlated with itself." % par[0]) + raise KeyError("%s is cross correlated with itself." % par[0]) # noqa comp['names_x_dict'][par[0]] = pn # Loop through SED parameters. diff --git a/bbpower/mpi_utils.py b/bbpower/mpi_utils.py new file mode 100644 index 0000000..4560d23 --- /dev/null +++ b/bbpower/mpi_utils.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python +import sys +import numpy as np +import math + +# set the default value +_initialized = False +_switch = False +rank = 0 +size = 1 +comm = None + + +def print_rnk0(text, rank): + if rank == 0: + print(text) + + +def init(switch=False): + ''' initialize MPI set-up ''' + global _initialized, _switch + global rank, size, comm + + exit_code = 0 + + if not _initialized: + _initialized = True + else: + print("MPI is already intialized") + return exit_code + + if not switch: + print("WARNING: MPI is turned off by default. " + "Use mpi.init(switch=True) to initialize MPI") + print("MPI is turned off") + return exit_code + else: + _switch = True + + try: + from mpi4py import MPI + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + print("MPI: rank %d is initalized" % rank) + + except ImportError as exc: + sys.stderr.write("IMPORT ERROR: " + __file__ + " (" + str(exc) + "). " + "Could not load mpi4py. MPI will not be used.\n") + return rank, size, comm + + +def is_initialized(): + global _initialized + return _initialized + + +def is_mpion(): + global _switch + return _switch + + +def taskrange(imax, imin=0, shift=0): + """ + """ + global rank, size + + if (not isinstance(imin, int) or not isinstance(imax, int) + or not isinstance(shift, int)): + raise TypeError("imin, imax and shift must be integers") + elif not is_initialized(): + print("MPI is not yet properly initialized. " + "Are you sure this is what you want to do?") + + if not is_mpion(): + return np.arange(imin, imax + 1) + + ntask = math.ceil((imax - imin + 1)/size)*size + + subrange = None + if ntask <= 0: + print_rnk0("number of task can't be zero", rank) + subrange = np.arange(0, 0) # return zero range + else: + if ntask != imax - imin + 1: + print_rnk0(f"WARNING: setting ntask={ntask}", rank) + perrank = ntask // size + print_rnk0(f"Running {ntask} simulations on {size} nodes", rank) + subrange = np.arange(rank*perrank, (rank + 1)*perrank) + + return subrange + + +def distribute_tasks(size, rank, ntasks, logger=None): + """ + Distributes [ntasks] tasks among [size] workers, and outputs + the list of tasks assigned to a given rank. + + Parameters + ---------- + size: int + The number of workers available. + rank: int + The number (ID) of the current worker. + ntasks: int + The number of tasks. + logger: logging.Logger + Logging instance to write output. If None, ignore. + Returns + ------- + local_task_ids: list + List with indices corresponding to the tasks assigned + to worker [rank] + """ + if size > ntasks: + local_start = rank + local_stop = rank + 1 + else: + local_start = rank * (ntasks // size) + local_stop = local_start + (ntasks // size) + + local_task_ids = list(range(ntasks))[local_start:local_stop] + + if rank >= ntasks: + local_task_ids = [] + + # If ntasks is not divisible by size, there will be a set of + # ntasks_left < size leftover tasks. Distribute one of them each to the + # first ntasks_left workers. + if ntasks % size != 0: + leftover = np.arange(ntasks)[-(ntasks % size):] + if rank < len(leftover): + local_task_ids.append(leftover[rank]) + + if logger is not None: + logger.info(f"Rank {rank} has {len(local_task_ids)} tasks.") + logger.info(f"Total number of tasks is {ntasks}") + logger.info(f"local_task_ids: {np.array(local_task_ids, dtype=int)}") + + return local_task_ids diff --git a/bbpower/param_manager.py b/bbpower/param_manager.py index 67da208..c3306d9 100644 --- a/bbpower/param_manager.py +++ b/bbpower/param_manager.py @@ -79,7 +79,7 @@ def __init__(self, config): i_bps = 1 while 'bandpass_%d' % i_bps in cnf_bps: if cnf_bps['bandpass_%d' % i_bps].get('parameters'): - self._add_parameters(cnf_bps['bandpass_%d' % i_bps]['parameters']) + self._add_parameters(cnf_bps['bandpass_%d' % i_bps]['parameters']) # noqa i_bps += 1 self.p0 = np.array(self.p0) @@ -95,6 +95,6 @@ def lnprior(self, par): if np.char.lower(pr[1]) == 'gaussian': # Gaussian prior lnp += -0.5 * ((p - pr[2][0])/pr[2][1])**2 else: # Only other option is top-hat - if not(float(pr[2][0]) <= p <= float(pr[2][2])): + if not (float(pr[2][0]) <= p <= float(pr[2][2])): return -np.inf return lnp diff --git a/bbpower/plotter_nopipe.py b/bbpower/plotter_nopipe.py new file mode 100644 index 0000000..1b652ee --- /dev/null +++ b/bbpower/plotter_nopipe.py @@ -0,0 +1,540 @@ +import healpy as hp +import numpy as np +import argparse +import matplotlib +import matplotlib.pyplot as plt +import dominate as dom +import dominate.tags as dtg +import sacc +import os +import yaml +import time +from getdist import MCSamples +from getdist import plots as gplots +from itertools import combinations_with_replacement as cwr + +import sys +sys.path.append( + os.path.abspath(os.path.join(os.path.dirname(__file__), '.')) +) + +import mpi_utils as mpi # noqa + +matplotlib.use('Agg') + + +def _yaml_loader(config): + """ + Custom yaml loader to load the configuration file. + """ + def path_constructor(loader, node): + return "/".join(loader.construct_sequence(node)) + yaml.SafeLoader.add_constructor("!path", path_constructor) + with open(config, "r") as f: + return yaml.load(f, Loader=yaml.SafeLoader) + + +labels_dict = { + 'A_lens': '$A_{\\rm lens}$', + 'r_tensor': '$r$', + 'beta_d': '$\\beta_d$', + 'epsilon_ds': '$\\epsilon_{ds}$', + 'alpha_d_bb': '$\\alpha_d^{BB}$', + 'amp_d_bb': '$A_d^{BB}$', + 'alpha_d_ee': '$\\alpha_d^{EE}$', + 'amp_d_ee': '$A_d^{EE}$', + 'beta_s': '$\\beta_s$', + 'alpha_s_bb': '$\\alpha_s^{BB}$', + 'amp_s_bb': '$A_s^{BB}$', + 'alpha_s_ee': '$\\alpha_s^{EE}$', + 'amp_s_ee': '$A_s^{EE}$', + 'amp_s_beta': '$B_s$', + 'amp_d_beta': '$B_d$', + 'gamma_d_beta': '$\\gamma_d', + 'gamma_s_beta': '$\\gamma_s$' +} + + +class BBPlotter(object): + """ + Plotting stage for BBPower. Plots power spectra, best-fit models, and + posterior chains. + """ + def __init__(self, args): + """ + Initialize from the command line arguments. + + Parameters + ---------- + args : str + Command line arguments. + """ + # Load the configuration file + config = _yaml_loader(args.config) + self.config = config["global"] | config["BBPlotter"] + setattr(self, "data", self.config["data"]) + setattr(self, "config_fname", getattr(args, "config")) + + # Load global parameters + self.nside = self.config['nside'] + self.npix = hp.nside2npix(self.nside) + + def create_page(self): + """ + """ + # Open plots directory + self.plot_dir = self.config["plot_dir"].format(sim_id=self.sim_id) + if not os.path.isdir(self.plot_dir): + os.mkdir(self.plot_dir) + + # Create HTML page + self.doc = dom.document(title='BBPower plots page') + with self.doc.head: + dtg.link(rel='stylesheet', href='style.css') + dtg.script(type='text/javascript', src='script.js') + with self.doc: + dtg.h1("Pipeline outputs") + dtg.h2("Contents:", id='contents') + lst = dtg.ul() + lst += dtg.li(dtg.a('Bandpasses', href='#bandpasses')) + lst += dtg.li(dtg.a('Coadded power spectra', href='#coadded')) + if self.config['plot_cells_null']: + lst += dtg.li(dtg.a('Null tests', href='#nulls')) + if self.config['plot_likelihood']: + lst += dtg.li(dtg.a('Likelihood', href='#like')) + + def add_bandpasses(self): + """ + """ + with self.doc: + dtg.h2("Bandpasses", id='bandpasses') + lst = dtg.ul() + + # Overall plot + title = 'Bandpasses summary' + fname = self.plot_dir + '/bpass_summary.png' + plt.figure() + plt.title(title, fontsize=14) + for n, t in self.s_cd_x.tracers.items(): + nu_mean = np.sum(t.bandpass*t.nu**3)/np.sum(t.bandpass*t.nu**2) + plt.plot(t.nu, t.bandpass/np.amax(t.bandpass), + label=n+', $\\langle\\nu\\rangle=%.1lf\\,{\\rm GHz}$' % nu_mean) # noqa + plt.xlabel('$\\nu\\,[{\\rm GHz}]$', fontsize=14) + plt.ylabel('Transmission', fontsize=14) + plt.ylim([0., 1.3]) + plt.legend(frameon=0, ncol=2, labelspacing=0.1, loc='upper left') + plt.xscale('log') + plt.savefig(fname, bbox_inches='tight') + plt.close() + lst += dtg.li(dtg.a(title, href=fname)) + + for n, t in self.s_cd_x.tracers.items(): + fname = self.plot_dir + '/bpass_' + n + '.png' + plt.figure() + plt.title(title, fontsize=14) + plt.plot(t.nu, t.bandpass/np.amax(t.bandpass)) + plt.xlabel('$\\nu\\,[{\\rm GHz}]$', fontsize=14) + plt.ylabel('Transmission', fontsize=14) + plt.ylim([0., 1.05]) + plt.savefig(fname, bbox_inches='tight') + plt.close() + lst += dtg.li(dtg.a(title, href=fname)) + dtg.div(dtg.a('Back to TOC', href='#contents')) + + def add_coadded(self): + (do_best, do_fid, do_cross, + do_tot, do_noise) = ( + self.best_fit is not None, + self.s_fid is not None, + self.s_cd_x is not None, + self.s_cd_t is not None, + self.s_cd_n is not None + ) + print("do_best", do_best) + if do_best: + best_fit_label = "\n".join( + [fr"{ll}: {b}"for ll, b in self.best_fit.items()] + ) + else: + best_fit_label = "" + if not do_cross: + print("No cross-coadded spectra found. Do not plot any spectra.") + return + with self.doc: + dtg.h2("Coadded power spectra", id='coadded') + lst = dtg.ul() + offset = 2 + for t1, t2 in self.s_cd_x.get_tracer_combinations(): + for p1, p2 in cwr(self.pols, 2): + x = p1 + p2 + typ = 'cl_' + x + l_cro, cl_cro = self.s_cd_x.get_ell_cl( + typ, t1, t2, return_cov=False + ) + l_cov, _, cov = self.s_cov.get_ell_cl( + typ, t1, t2, return_cov=True + ) + msk_cov = np.logical_and(l_cov <= self.lmax, + l_cov >= self.lmin) + if len(l_cro) == 0: + print(f"No data for {typ} found. Skip.") + continue + + # Plot title + title = f"{t1} x {t2}, {typ}" + # Plot file + fname = self.plot_dir + '/cls_' + fname += f"{t1}_x_{t2}_{typ}.png" + print(fname) + f, (main, sub) = plt.subplots( + 2, 1, sharex=True, figsize=(6, 4), + gridspec_kw={'height_ratios': [3, 1]} + ) + f.suptitle(title, fontsize=14) + if do_fid: + l_fid, cl_fid = self.s_fid.get_ell_cl(typ, t1, t2) + msk_fid = np.logical_and(l_fid <= self.lmax, + l_fid >= self.lmin) + main.plot(l_fid[msk_fid], cl_fid[msk_fid], 'k--', + label='Fiducial') + if do_best: + l_best, cl_best = self.s_best.get_ell_cl( + typ, t1, t2 + ) + msk_best = np.logical_and(l_best <= self.lmax, + l_best >= self.lmin) + main.plot(l_best[msk_best], cl_best[msk_best], + 'k-', label="Best fit") + main.plot([], [], "w-", label=best_fit_label) + sub.axhspan(-3, 3, facecolor="k", alpha=0.1) + sub.axhspan(-2, 2, facecolor="k", alpha=0.2) + sub.axhspan(-1, 1, facecolor="k", alpha=0.3) + sub.axhline(0, color="k") + + if do_tot: + l_tot, cl_tot = self.s_cd_t.get_ell_cl( + typ, t1, t2 + ) + msk_tot = np.logical_and(l_tot <= self.lmax, + l_tot >= self.lmin) + x, y = (l_tot[msk_tot], cl_tot[msk_tot]) + main.plot( + x - offset, y, label='Total coadd', + color="darkred", marker=".", + markerfacecolor='darkred', linestyle="", + ) + eb_tot = main.plot( + x + offset, -y, color="darkred", + marker=".", markerfacecolor='white', + linestyle="", + ) + eb_tot[-1][0].set_linestyle('--') + + if do_noise: + l_noi, cl_noi = self.s_cd_n.get_ell_cl( + typ, t1, t2 + ) + msk_noi = np.logical_and(l_noi <= self.lmax, + l_noi >= self.lmin) + x, y = (l_noi[msk_noi], cl_noi[msk_noi]) + main.plot( + x, y, label='Noise', + color="darkorange", marker=".", + markerfacecolor='darkorange', linestyle="", + ) + eb_noi = main.plot( + x + offset, -y, color="darkorange", + marker=".", markerfacecolor="white", + linestyle="", + ) + eb_noi[-1][0].set_linestyle('--') + + msk_cro = np.logical_and(l_cro <= self.lmax, + l_cro >= self.lmin) + x, y, yerr = ( + l_cro[msk_cro], cl_cro[msk_cro], + np.sqrt(np.fabs(np.diag(cov)))[msk_cov] + ) + main.errorbar( + x - offset, y, yerr, label='Cross coadd', + color="navy", marker=".", + markerfacecolor="navy", linestyle="", + ) + eb_cro = main.errorbar( + x + offset, -y, yerr, color="navy", marker=".", + markerfacecolor="white", linestyle="", + ) + eb_cro[-1][0].set_linestyle('--') + + if do_best: + if do_fid: + sub.plot( + l_fid[msk_fid], + (cl_fid[msk_fid] - cl_best[msk_best])/yerr, + "k--" + ) + sub.plot( + x, (y - cl_best[msk_best])/yerr, + color="navy", marker=".", markerfacecolor="navy", + linestyle="" + ) + sub.set_ylabel( + r"$(C_\ell-C_\ell^{\rm best})/\sigma(C_\ell)$" + ) + sub.set_xlabel('$\\ell$', fontsize=15) + sub.set_ylim((-5, 5)) + main.set_yscale('log') + if self.config['compute_dell']: + main.set_ylabel('$D_\\ell$', fontsize=15) + else: + main.set_ylabel('$C_\\ell$', fontsize=15) + main.legend(loc="upper left", bbox_to_anchor=(1, 1), + fontsize=12) + plt.savefig(fname, bbox_inches='tight') + plt.close() + lst += dtg.li(dtg.a(title, href=fname)) + + dtg.div(dtg.a('Back to TOC', href='#contents')) + + def add_nulls(self): + """ + """ + do_nulls = self.s_null is not None + if not do_nulls: + return + with self.doc: + dtg.h2("Null tests", id='nulls') + lst = dtg.ul() + + for t1, t2 in self.s_null.get_tracer_combinations(): + title = f"{t1} x {t2}" + fname = self.plot_dir + '/cls_null_' + fname += f"{t1}_x_{t2}.png" + print(fname) + plt.figure() + plt.title(title, fontsize=15) + for p1 in range(2): + for p2 in range(2): + x = self.pols[p1] + self.pols[p2] + typ = 'cl_' + x + l, cl, cv = self.s_null.get_ell_cl(typ, t1, t2, + return_cov=True) + msk = l < self.lmax + el = np.sqrt(np.fabs(np.diag(cv)))[msk] + plt.errorbar(l[msk], cl[msk]/el, + yerr=np.ones_like(el), + fmt=self.cols_typ[x]+'-', label=x) + plt.xlabel('$\\ell$', fontsize=15) + plt.ylabel('$C_\\ell/\\sigma_\\ell$', fontsize=15) + plt.legend() + plt.savefig(fname, bbox_index='tight') + plt.close() + lst += dtg.li(dtg.a(title, href=fname)) + + dtg.div(dtg.a('Back to TOC', href='#contents')) + + def add_contours(self): + """ + """ + with self.doc: + dtg.h2("Likelihood", id='like') + lst = dtg.ul() + + # Select only selected parameters for which we have labels + names_common = list(set(list(self.chain['names'])) + & self.config['truth'].keys() + & set(self.config['params_plot'])) + msk_common = np.array([n in names_common + for n in self.chain['names']]) + _, nsamp, npar_chain = self.chain['chain'].shape + chain = self.chain['chain'][:, nsamp//4:, :].reshape([-1, npar_chain])[:, msk_common] # noqa + names_common = np.array(self.chain['names'])[msk_common] + labels = [labels_dict[n].replace("$", "") for n in names_common] + + # Getdist + samples = MCSamples( + samples=chain, + names=names_common, + labels=labels + ) + + # Plot posteriors + g = gplots.getSubplotPlotter() + g.settings.title_limit_fontsize = 12 + g.settings.axes_fontsize = 14 + g.settings.axes_labelsize = 18 + g.triangle_plot([samples], + filled=True, + contour_colors=['navy'], + title_limit=1) + + # Plot truth values + for i, n in enumerate(names_common): + v = self.config["truth"][n] + if v is None: + continue + g.subplots[i, i].plot([v, v], [0, 1], 'r-') + for j in range(i + 1, len(names_common)): + u = self.config["truth"][names_common[j]] + g.subplots[j, i].plot([v], [u], 'r*') + + # Save + fname = self.plot_dir + '/triangle.pdf' + g.export(fname) + print(fname) + lst += dtg.li(dtg.a("Likelihood contours", href=fname)) + + dtg.div(dtg.a('Back to TOC', href='#contents')) + + def write_page(self): + """ + """ + plots_page = self.config['plots_page'].format(sim_id=self.sim_id) + with open(plots_page, 'w') as f: + f.write(self.doc.render()) + + def read_inputs(self): + """ + """ + print("Reading inputs") + + # Power spectra + saccs_dict = { + "cells_fiducial": "s_fid", + "cells_best_fit": "s_best", + "cells_coadded": "s_cd_x", + "cells_coadded_total": "s_cd_t", + "cells_noise": "s_cd_n", + "cells_null": "s_null", + "cells_coadded_cov": "s_cov", + } + + # Load relevant sacc files for plotting + for cl, sc in saccs_dict.items(): + setattr(self, sc, None) + if self.data[cl] is None: + print(cl, "is None") + continue + data = self.data[cl].format(sim_id=self.sim_id) + + if not os.path.isfile(data): + print(data, "is not file") + continue + if hasattr(self, f"plot_{cl}"): + if not self.config[f"plot_{cl}"]: + print(cl, "not plotted") + continue + setattr(self, sc, sacc.Sacc.load_fits(data)) + + # Keep only desired tracers + print("Selecting tracers") + for sn, s in saccs_dict.items(): + print("Reading", sn) + sc = getattr(self, s) + if sc is None: + continue + tr_list = list(sc.tracers.keys()) + for t in tr_list: + if t not in self.config["map_sets"]: + sc.remove_tracers([t]) + setattr(self, s, sc) + + # Chains + print("Reading chains") + chain_file = self.config["chain_file"].format(sim_id=self.sim_id) + + if not os.path.isfile(chain_file): + self.config['plot_likelihood'] = False + print("No posterior chain file found. Do not plot it.") + if self.config['plot_likelihood']: + self.chain = np.load(chain_file) + + # Best-fit parameters + chi2_file = self.config["chi2_file"].format(sim_id=self.sim_id) + self.best_fit, self.chisq, self.pte = (None, None, None) + if os.path.isfile(chi2_file): + chisq = np.load(chi2_file) + self.best_fit = { + labels_dict[n]: f"{float(p):.3f}" + for n, p in zip(chisq["names"], chisq['params']) + } + self.best_fit["chi2"] = str(int(chisq['chi2'])) + self.best_fit["ndof"] = str(int(chisq['ndof'])) + self.best_fit["pte"] = f"{chisq['pte']:.1e}" + + self.chisq = chisq["chi2"] + self.ndof = chisq["ndof"] + self.pte = chisq["pte"] + + self.cols_typ = {'ee': 'r', 'eb': 'g', 'be': 'y', 'bb': 'b'} + self.lmin = self.config['lmin_plot'] + self.lmax = self.config['lmax_plot'] + self.pols = [p.lower() for p in self.config['pol_channels']] + + def run(self): + """ + Run the BB pipeline plotting stage. + """ + self.read_inputs() + self.create_page() + self.add_bandpasses() + self.add_coadded() + if self.config['plot_cells_null']: + self.add_nulls() + if self.config['plot_likelihood']: + self.add_contours() + self.write_page() + + +def main(args): + """ + Execute the BBPlotter stage with arguments args: + * config: string. + Path to configuration file with input parameters. + """ + config = _yaml_loader(args.config) + + # Creating the simulation indices range to loop over + sim_ids = config["global"]["sim_ids"] + if isinstance(sim_ids, list): + sim_ids = np.array(sim_ids, dtype=int) + elif isinstance(sim_ids, str): + if "," in sim_ids: + id_min, id_max = sim_ids.split(",") + sim_ids = np.arange(int(id_min), int(id_max)+1) + else: + sim_ids = np.array([int(sim_ids)]) + else: + sim_ids = [None] + + # MPI related initialization + rank, size, comm = mpi.init(True) + mpi_shared_list = sim_ids + + # Every rank must have the same shared list + mpi_shared_list = comm.bcast(mpi_shared_list, root=0) + task_ids = mpi.distribute_tasks(size, rank, len(mpi_shared_list), + logger=None) + local_mpi_list = [mpi_shared_list[i] for i in task_ids] + + for sim_id in local_mpi_list: + start = time.time() + plotter = BBPlotter(args) + setattr(plotter, "sim_id", sim_id) + plotter.run() + comm.Barrier() + mpi.print_rnk0(f"Processed {len(sim_ids)} simulations " + f"in {time.time() - start:.1f} seconds.", rank) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser( + description="Plotting stage for BB Cross-CL pipeline" + ) + parser.add_argument( + "--config", type=str, + help="Path to yaml file with pipeline configuration" + ) + + args = parser.parse_args() + main(args) diff --git a/bbpower/power_specter.py b/bbpower/power_specter.py index 8c767e1..033a0e5 100644 --- a/bbpower/power_specter.py +++ b/bbpower/power_specter.py @@ -40,7 +40,7 @@ def read_beams(self, nbeams): # Check that there are enough beams if len(beam_fnames) != nbeams: raise ValueError("Couldn't find enough beams: " - "%d != %d" (len(beam_fnames), nbeams)) + f"{len(beam_fnames)} != {nbeams}") self.larr_all = np.arange(3*self.nside) self.beams = {} @@ -65,8 +65,7 @@ def compute_cells_from_splits(self, splits_list): fname = fname + '.gz' if not os.path.isfile(fname): raise ValueError("Can't find file ", splits_list[s]) - mp_q, mp_u = hp.read_map(fname, field=[2*b, 2*b+1], - verbose=False) + mp_q, mp_u = hp.read_map(fname, field=[2*b, 2*b+1]) fields[name] = self.get_field(b, [mp_q, mp_u]) # Iterate over field pairs @@ -104,8 +103,7 @@ def read_bandpasses(self): def read_masks(self, nbands): self.masks = [] for i in range(nbands): - m = hp.read_map(self.get_input('masks_apodized'), - verbose=False) + m = hp.read_map(self.get_input('masks_apodized')) self.masks.append(hp.ud_grade(m, nside_out=self.nside)) def get_bandpowers(self): @@ -114,7 +112,6 @@ def get_bandpowers(self): # Custom spacing edges = np.loadtxt(self.config['bpw_edges']).astype(int) bpws = np.zeros(3*self.nside, dtype=int)-1 - weights = np.ones(3*self.nside) for ibpw, (l0, lf) in enumerate(zip(edges[:-1], edges[1:])): if lf < 3*self.nside: bpws[l0:lf] = ibpw @@ -127,17 +124,15 @@ def get_bandpowers(self): bpws[l0:l0+dell] = ibpw l0 += dell - is_dell = False + f_ell = np.ones_like(self.larr_all, dtype=np.float64) if self.config.get('compute_dell'): - is_dell = True - self.bins = nmt.NmtBin(self.nside, - bpws=bpws, + f_ell = self.larr_all*(self.larr_all+1)/(2*np.pi) + self.bins = nmt.NmtBin(bpws=bpws, ells=self.larr_all, - weights=weights, - is_Dell=is_dell) + lmax=3*self.nside-1, + f_ell=f_ell) else: # otherwise it could be a constant integer interval - self.bins = nmt.NmtBin(self.nside, - nlb=int(self.config['bpw_edges'])) + raise NotImplementedError("Constant-width ell bins not supported.") def get_fname_workspace(self, band1, band2): b1 = min(band1, band2) @@ -167,8 +162,7 @@ def compute_workspace(self, band1, band2): mdum = np.zeros([2, self.npix]) f1 = self.get_field(b1, mdum) f2 = self.get_field(b2, mdum) - w.compute_coupling_matrix(f1, f2, self.bins, - n_iter=self.config['n_iter']) + w.compute_coupling_matrix(f1, f2, self.bins) w.write_to(fname) return w @@ -203,7 +197,7 @@ def get_cell_iterator(self): splits_range = range(self.nsplits) for s2 in splits_range: l2 = self.get_map_label(b2, s2) - yield(b1, b2, s1, s2, l1, l2) + yield (b1, b2, s1, s2, l1, l2) def get_sacc_tracers(self): sacc_t = [] @@ -228,13 +222,13 @@ def get_sacc_windows(self): wsp = self.workspaces[name] bpw_win = wsp.get_bandpower_windows() windows_wsp[name]['EE'] = sacc.BandpowerWindow(self.larr_all, - bpw_win[0, :, 0, :].T) + bpw_win[0, :, 0, :].T) # noqa: E501 windows_wsp[name]['EB'] = sacc.BandpowerWindow(self.larr_all, - bpw_win[1, :, 1, :].T) + bpw_win[1, :, 1, :].T) # noqa: E501 windows_wsp[name]['BE'] = sacc.BandpowerWindow(self.larr_all, - bpw_win[2, :, 2, :].T) + bpw_win[2, :, 2, :].T) # noqa: E501 windows_wsp[name]['BB'] = sacc.BandpowerWindow(self.larr_all, - bpw_win[3, :, 3, :].T) + bpw_win[3, :, 3, :].T) # noqa: E501 return windows_wsp def save_cell_to_file(self, cell, tracers, fname, with_windows=False): diff --git a/examples/config_nopipe.yml b/examples/config_nopipe.yml new file mode 100644 index 0000000..ae4a3ba --- /dev/null +++ b/examples/config_nopipe.yml @@ -0,0 +1,301 @@ +# These parameters are accessible to all stages +chains_dir: &chains_dir /scratch/gpfs/SIMONSOBS/users/kw6905/sat-iso/v1/bbpower/sat_rotated_wmap_planck_20260217/fiducial_model_Alens ## YOUR OUTPUT DIRECTORY +covar_sacc: &covar_sacc /scratch/gpfs/SIMONSOBS/external_data_e2e/v1/cl_and_cov_sacc_rotated.fits + +global: + nside: 256 # only important for plotting + delta_ell: 10 # needed to generate fiducial spectra + compute_dell: True + # Number of simulations to run (ignored if empty) + sim_ids: null + # Path where copy of config file will be stored + # config_copy: !path [*chains_dir, config_copy.yaml] + config_copy: /scratch/gpfs/SIMONSOBS/users/kw6905/sat-iso/v1/bbpower/sat_rotated_wmap_planck_20260217/fiducial_model_Alens/config_copy.yaml + + # Input data or sims + data: + # Path to sacc file with cross-split power spectra + cells_coadded: *covar_sacc # In this case, same as covar_sacc + # Path to sacc file with noise power spectra + cells_noise: null + # Path to sacc file with fiducial power spectra + cells_fiducial: null + # Path to sacc file with cross-split power spectrum covariance + # cells_coadded_cov: *covar_sacc + cells_coadded_cov: *covar_sacc + # Path to sacc file with best-fit power spectra + cells_best_fit: !path [*chains_dir, cells_model.fits] + # Path to sacc file with auto- and cross-bundle power spectra + cells_coadded_total: null + # Path to sacc file with null-test power spectra + cells_null: null + + map_sets: + SATp3_f150_south_science: + beam_file: null + SATp3_f090_south_science: + beam_file: null + wmap_f023_filtered_SATp3_f023: + beam_file: null + planck_f030_filtered_SATp3_f090_south_science: + beam_file: null + planck_f100_filtered_SATp3_f090_south_science: + beam_file: null + planck_f143_filtered_SATp3_f090_south_science: + beam_file: null + planck_f217_filtered_SATp3_f090_south_science: + beam_file: null + planck_f353_filtered_SATp3_f090_south_science: + beam_file: null + +theory_params: + CMB: + cmb_templates: + - /scratch/gpfs/SIMONSOBS/external_data_e2e/camb_lens_nobb.dat + - /scratch/gpfs/SIMONSOBS/external_data_e2e/camb_lens_r1.dat + A_lens: 1. + r_tensor: 0. + dust: + beta_d: 1.57 + alpha_d_EE: -0.12 + alpha_d_BB: -0.15 + A_d_EE: 45. + A_d_BB: 23. + nu0_d: 353. + T_d: 20. + synch: + beta_s: -3.3 + alpha_s_EE: 1. # -0.79 + alpha_s_BB: 1.7 # -0.82 + A_s_EE: 8.3 + A_s_BB: 1.9 + nu0_s: 23. + +BBCompSep: + output_dir: *chains_dir + # Sampler type. Options are: + # - 'emcee': a full MCMC will be run using emcee. + # - 'fisher': only the Fisher matrix (i.e. the likelihood + # Hessian matrix) will be calculated around the fiducial + # parameters chosen as prior centers. + # - 'single_point': only the chi^2 at the value used as the + # center for all parameter priors will be calculated. + # - 'timing': will compute the average time taken by one + # likelihood computation. + sampler: emcee + # If you chose polychord: + nlive: 50 + nrepeat: 50 + # If you chose emcee: + nwalkers: 48 + n_iters: 5000 + # Resume sampling from earlier runs + resume: False + # Whether to write model best fits into sacc (otherwise into npz archive) + predict_to_sacc: True + # Likelihood type. Options are: + # - 'chi2': a standard chi-squared Gaussian likelihood. + # - 'h&l': Hamimeche & Lewis likelihood. + likelihood_type: chi2 + # What is the starting point? + r_init: 1.e-3 + # Which polarization channels do you want to include? + # Can be ['E'], ['B'] or ['E','B']. + pol_channels: ['B'] + # Scale cuts (will apply to all frequencies) + l_min: 30 + l_max: 300 + + # CMB model + cmb_model: + # Template power spectrum. Should contained the lensed power spectra + # with r=0 and r=1 respectively. + cmb_templates: + - /scratch/gpfs/SIMONSOBS/external_data_e2e/camb_lens_nobb.dat + - /scratch/gpfs/SIMONSOBS/external_data_e2e/camb_lens_r1.dat + + # Free parameters + params: + # tensor-to-scalar ratio + # See below for the meaning of the different elements in the list. + r_tensor: ['r_tensor', 'tophat', [-2., 0.00, 2.]] + #r_tensor: ['r_tensor', 'fixed', [0.00]] + # Lensing amplitude + A_lens: ['A_lens', 'tophat', [0., 1., 2.]] + #A_lens: ['A_lens', 'fixed', [1.0]] + + # Foreground model + fg_model: + # Include moment parameters? + use_moments: False + moments_lmax: 192 + + # Add one section per component. They should be called `component_X`, + # starting with X=1 + component_1: + # Name for this component + name: Dust + # Type of SED. Should be one of the classes stored in fgbuster.components + # https://github.com/fgbuster/fgbuster/blob/master/fgbuster/component_model.py + sed: Dust + # Type of power spectra for all possible polarization channel combinations. + # Any combinations not added here will be assumed to be zero. + # The names should be one of the classes in bbpower/fgcls.py + cl: + EE: ClPowerLaw + BB: ClPowerLaw + # Parameters of the SED + sed_parameters: + # The key can be anything you want, but each parameter in the model + # must have a different name. + # The first item in the list is the name of the parameter used by fgbuster + # The second item is the type of prior. The last item are the numbers + # necessary to define the prior. They should be: + # - Gaussian: [mean,sigma] + # - tophat: [lower edge, start value, upper edge] + # - fixed: [parameter value] + # nu0-type parameters can only be fixed. + beta_d: ['beta_d', 'Gaussian', [1.5, 1.0]] + temp_d: ['temp', 'fixed', [20.]] + nu0_d: ['nu0', 'fixed', [353.]] + cl_parameters: + # Same for power spectrum parameters + # (broken down by polarization channel combinations) + EE: + amp_d_ee: ['amp', 'tophat', [0., 10., "inf"]] + alpha_d_ee: ['alpha', 'tophat', [-1., -0.42, 0.]] + l0_d_ee: ['ell0', 'fixed', [80.]] + BB: + amp_d_bb: ['amp', 'tophat', [0., 5., "inf"]] + alpha_d_bb: ['alpha', 'tophat', [-2., -0.2, 1.]] + l0_d_bb: ['ell0', 'fixed', [80.]] + # If this component should be correlated with any other, list them here + cross: + # In this case the list should contain: + # [component name, prior type, prior parameters] + #epsilon_ds: ['component_2', 'tophat', [-1., 0., 1.]] + epsilon_ds: ['component_2', 'fixed', [0.]] + # Moment parameters. Only relevant if set "use_moments" to "True" + moments: + # Define gammas for varying spectral indices of components + gamma_d_beta : ['gamma_beta', 'tophat', [-6., -3.5, -2.]] + amp_d_beta : ['amp_beta', 'tophat', [0., 0., 1]] + + component_2: + name: Synchrotron + sed: Synchrotron + cl: + EE: ClPowerLaw + BB: ClPowerLaw + sed_parameters: + #beta_s: ['beta_pl', 'fixed', [-3.0]] + beta_s: ['beta_pl', 'Gaussian', [-3.0, 0.5]] + nu0_s: ['nu0', 'fixed', [23.]] + cl_parameters: + EE: + #amp_s_ee: ['amp', 'fixed', [3.4]] + amp_s_ee: ['amp', 'tophat', [0., 4., "inf"]] + #alpha_s_ee: ['alpha', 'fixed', [-0.79]] + alpha_s_ee: ['alpha', 'tophat', [-2., -0.6, 1.]] + l0_s_ee: ['ell0', 'fixed', [80.]] + BB: + #amp_s_bb: ['amp', 'fixed', [1.8] ] + amp_s_bb: ['amp', 'tophat', [0., 2., "inf"]] + #alpha_s_bb: ['alpha', 'fixed', [-0.82]] + alpha_s_bb: ['alpha', 'tophat', [-2., -0.4, 3.]] + l0_s_bb: ['ell0', 'fixed', [80.]] + # Moment parameters. Only relevant if set "use_moments" to "True" + moments: + # Define gammas for varying spectral indices of components + gamma_s_beta : ['gamma_beta', 'tophat', [-6., -3.5, -2.]] + amp_s_beta : ['amp_beta', 'tophat', [0., 0., 1]] + + # Add extra parameters for different systematics + systematics: + # Bandpass-related systematics + bandpasses: + # Add a section called bandpass_X for any bandpass you + # want to add extra information for. They will be matched + # with the tracers stored in the input SACC files. + planck_f030_filtered_SATp3_f090_south_science: + # Add a file for a frequency-dependent polarization angle + # 2 columns: freq (GHz), phi (deg.). + # If not present, this effect will be ignored. + #phase_nu: "./examples/data/phase_3layer_lf_0.txt" + # Parametrize different bandpass systematics + parameters: + # Relative shift in mean bandpass frequency. + #shift_1: ['shift', 'Gaussian', [0., 0.01]] + # Relative gain variation. + #gain_1: ['gain', 'Gaussian', [1., 0.01]] + # Polarization angle (prior in units of radians). + # angle_1: ['angle', 'Gaussian', [0., 5.]] + # angle_1: ['angle', 'fixed', [5.]] + planck_f100_filtered_SATp3_f090_south_science: + #phase_nu: "./examples/data/phase_3layer_lf_0.txt" + parameters: + #shift_2: ['shift', 'Gaussian', [0., 0.01]] + #gain_2: ['gain', 'Gaussian', [1., 0.01]] + #angle_2: ['angle', 'Gaussian', [0., 5.]] + # angle_2: ['angle', 'fixed', [5.]] + planck_f143_filtered_SATp3_f090_south_science: + #phase_nu: "./examples/data/phase_3layer_mf_0.txt" + parameters: + #shift_4: ['shift', 'Gaussian', [0., 0.01]] + #gain_4: ['gain', 'Gaussian', [1., 0.01]] + #angle_3: ['angle', 'Gaussian', [0., 5.]] + planck_f217_filtered_SATp3_f090_south_science: + #phase_nu: "./examples/data/phase_3layer_uhf_0.txt" + parameters: + #shift_5: ['shift', 'Gaussian', [0., 0.01]] + #gain_5: ['gain', 'Gaussian', [1., 0.01]] + #angle_4: ['angle', 'Gaussian', [0., 5.]] + planck_f353_filtered_SATp3_f090_south_science: + #phase_nu: "./examples/data/phase_3layer_uhf_0.txt" + parameters: + #shift_6: ['shift', 'Gaussian', [0., 0.01]] + #gain_6: ['gain', 'Gaussian', [1., 0.01]] + #angle_5: ['angle', 'Gaussian', [0., 5.]] + +BBPlotter: + # MCMC chain file + chain_file: !path [*chains_dir, emcee.npz] + # chi2 npz file + chi2_file: !path [*chains_dir, chi2.npz] + + # Output directory for plots + plot_dir: !path [*chains_dir, plots] + # HTML page with plots + plots_page: !path [*chains_dir, plots_page.html] + + # Which polarization channels to plot + pol_channels: ["B"] + # Maximum and minimum ell between which things will be plotted + lmin_plot: 30 + lmax_plot: 300 + # Add total coadds to C_ell plots? + plot_cells_coadded_total: False + # Add noise spectra to C_ell plots? + plot_cells_noise: False + # Plot null tests? + plot_cells_null: False + # Plot likelihood? + plot_likelihood: True + # List of parameters to plot + params_plot: [r_tensor, A_lens, beta_d, amp_d_bb, alpha_d_bb, beta_s, amp_s_bb, alpha_s_bb] + + # "Truth" parameter values for plotting purposes + truth: + A_lens: 1 + r_tensor: 0 + beta_d: 1.54 + epsilon_ds: 0. + alpha_d_bb: -0.16 + amp_d_bb: 28 + beta_s: -3. + alpha_s_bb: -0.93 + amp_s_bb: 1.6 + gamma_d_beta: null + amp_d_beta: 0 + gamma_s_beta: null + amp_s_beta: 0 diff --git a/examples/generate_SO_maps.py b/examples/generate_SO_maps.py index e535eaf..f77b248 100644 --- a/examples/generate_SO_maps.py +++ b/examples/generate_SO_maps.py @@ -1,6 +1,5 @@ -from utils import * -import os import numpy as np +import utils as ut import noise_calc as nc from optparse import OptionParser import healpy as hp @@ -17,13 +16,15 @@ np.random.seed(o.seed) npix = hp.nside2npix(o.nside) +band_names = ['LF1', 'LF2', 'MF1', 'MF2', 'UHF1', 'UHF2'] + # Signal maps lmax = 3*o.nside - 1 larr = np.arange(lmax+1) -cl2dl=larr*(larr+1)/(2*np.pi) +cl2dl = larr*(larr+1)/(2*np.pi) dl2cl = np.zeros(lmax+1) -dl2cl[1:]=2*np.pi/(larr[1:]*(larr[1:]+1)) -clsee, clsbb, cldee, cldbb, clcee, clcbb = get_component_spectra(lmax) +dl2cl[1:] = 2*np.pi/(larr[1:]*(larr[1:]+1)) +clsee, clsbb, cldee, cldbb, clcee, clcbb = ut.get_component_spectra(lmax) clsee *= dl2cl clsbb *= dl2cl cldee *= dl2cl @@ -31,35 +32,33 @@ clcee *= dl2cl clcbb *= dl2cl cl0 = 0*clsee -_, Qs, Us = hp.synfast([cl0, clsee, clsbb, cl0, cl0, cl0], - o.nside, verbose=False, new=True) -_, Qd, Ud = hp.synfast([cl0, cldee, cldbb, cl0, cl0, cl0], - o.nside, verbose=False, new=True) -_, Qc, Uc = hp.synfast([cl0, clcee, clcbb, cl0, cl0, cl0], - o.nside, verbose=False, new=True) +_, Qs, Us = hp.synfast([cl0, clsee, clsbb, cl0, cl0, cl0], o.nside, new=True) +_, Qd, Ud = hp.synfast([cl0, cldee, cldbb, cl0, cl0, cl0], o.nside, new=True) +_, Qc, Uc = hp.synfast([cl0, clcee, clcbb, cl0, cl0, cl0], o.nside, new=True) map_comp = np.array([[Qc, Uc], [Qs, Us], [Qd, Ud]]) -bpss = {n: Bpass(n, f'examples/data/bandpasses/{n}.txt') +bpss = {n: ut.Bpass(n, f'examples/data/bandpasses/{n}.txt') for n in band_names} -seds = get_convolved_seds(band_names, bpss) +seds = ut.get_convolved_seds(band_names, bpss) _, nfreqs = seds.shape map_freq = np.einsum('ij,ikl', seds, map_comp) # Noise maps nsplits = 4 -sens=1 -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, - include_beam=False) +sens = 1 +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, include_beam=False +) map_noise = np.zeros([nsplits, nfreqs, 2, npix]) for i_s in range(nsplits): for i_f in range(nfreqs): _, mpq, mpu = hp.synfast([cl0, nell[i_f], nell[i_f], cl0, cl0, cl0], - o.nside, verbose=False, new=True) + o.nside, new=True) map_noise[i_s, i_f, 0, :] = mpq * np.sqrt(nsplits) map_noise[i_s, i_f, 1, :] = mpu * np.sqrt(nsplits) @@ -68,8 +67,7 @@ for i_f, s in enumerate(s_fwhm): fwhm = s * np.pi/180./60. for i_p in [0, 1]: - map_freq[i_f, i_p, :] = hp.smoothing(map_freq[i_f, i_p, :], - fwhm=fwhm, verbose=False) + map_freq[i_f, i_p, :] = hp.smoothing(map_freq[i_f, i_p, :], fwhm=fwhm) # Write output for s in range(nsplits): diff --git a/examples/generate_SO_spectra.py b/examples/generate_SO_spectra.py index aba63ba..8da5cbc 100644 --- a/examples/generate_SO_spectra.py +++ b/examples/generate_SO_spectra.py @@ -1,73 +1,82 @@ import numpy as np -from utils import * +import utils as ut import noise_calc as nc import sacc import sys prefix_out = sys.argv[1] +# Use forecast params from Wolz et al. (2302.04276) +so_forecast = "--so_forecast" in sys.argv +print("so_forecast:", so_forecast) # Bandpasses -bpss = {n: Bpass(n,f'examples/data/bandpasses/{n}.txt') for n in band_names} +band_names = ['LF1', 'LF2', 'MF1', 'MF2', 'UHF1', 'UHF2'] +bpss = {n: ut.Bpass(n, f'examples/data/bandpasses/{n}.txt') + for n in band_names} # Bandpowers dell = 10 nbands = 100 lmax = 2+nbands*dell larr_all = np.arange(lmax+1) -lbands = np.linspace(2,lmax,nbands+1,dtype=int) +lbands = np.linspace(2, lmax, nbands+1, dtype=int) leff = 0.5*(lbands[1:]+lbands[:-1]) -windows = np.zeros([nbands,lmax+1]) -cl2dl=larr_all*(larr_all+1)/(2*np.pi) -dl2cl=np.zeros_like(cl2dl) +windows = np.zeros([nbands, lmax+1]) +cl2dl = larr_all*(larr_all+1)/(2*np.pi) +dl2cl = np.zeros_like(cl2dl) dl2cl[1:] = 1/(cl2dl[1:]) -for b,(l0,lf) in enumerate(zip(lbands[:-1],lbands[1:])): - windows[b,l0:lf] = (larr_all * (larr_all + 1)/(2*np.pi))[l0:lf] - windows[b,:] /= dell +for b, (l0, lf) in enumerate(zip(lbands[:-1], lbands[1:])): + windows[b, l0:lf] = (larr_all * (larr_all + 1)/(2*np.pi))[l0:lf] + windows[b, :] /= dell s_wins = sacc.BandpowerWindow(larr_all, windows.T) # Beams -beams = {band_names[i]: b for i, b in enumerate(nc.Simons_Observatory_V3_SA_beams(larr_all))} +beams = {band_names[i]: b + for i, b in enumerate(nc.Simons_Observatory_V3_SA_beams(larr_all))} 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) +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, :]) = ut.get_component_spectra(lmax, + so_forecast) dls_comp *= dl2cl[None, None, None, None, :] # Convolve with windows -bpw_comp=np.sum(dls_comp[:,:,:,:,None,:]*windows[None,None,None,None,:,:],axis=5) +bpw_comp = np.sum(dls_comp[:, :, :, :, None, :] * windows[None, None, None, None, :, :], axis=5) # noqa # Convolve with bandpasses -seds = get_convolved_seds(band_names, bpss) +seds = ut. get_convolved_seds(band_names, bpss, so_forecast) _, nfreqs = seds.shape # Component -> frequencies -bpw_freq_sig=np.einsum('ik,jm,iljno',seds,seds,bpw_comp) +bpw_freq_sig = np.einsum('ik,jm,iljno', seds, seds, bpw_comp) # N_ell -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,:]*windows[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,:] +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, :]*windows[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 to signal -bpw_freq_tot=bpw_freq_sig+bpw_freq_noi -bpw_freq_tot=bpw_freq_tot.reshape([nfreqs*2,nfreqs*2,nbands]) -bpw_freq_sig=bpw_freq_sig.reshape([nfreqs*2,nfreqs*2,nbands]) -bpw_freq_noi=bpw_freq_noi.reshape([nfreqs*2,nfreqs*2,nbands]) +bpw_freq_tot = bpw_freq_sig+bpw_freq_noi +bpw_freq_tot = bpw_freq_tot.reshape([nfreqs*2, nfreqs*2, nbands]) +bpw_freq_sig = bpw_freq_sig.reshape([nfreqs*2, nfreqs*2, nbands]) +bpw_freq_noi = bpw_freq_noi.reshape([nfreqs*2, nfreqs*2, nbands]) # Creating Sacc files s_d = sacc.Sacc() @@ -92,10 +101,10 @@ # Adding power spectra print("Adding spectra") -nmaps=2*nfreqs -ncross=(nmaps*(nmaps+1))//2 -indices_tr=np.triu_indices(nmaps) -map_names=[] +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('band%d' % (ib+1) + '_E') map_names.append('band%d' % (ib+1) + '_B') @@ -105,9 +114,9 @@ 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, leff, bpw_freq_sig[i1, i2, :], window=s_wins) - s_f.add_ell_cl(cl_type, n1, n2, leff, bpw_freq_sig[i1, i2, :], window=s_wins) - s_n.add_ell_cl(cl_type, n1, n2, leff, bpw_freq_noi[i1, i2, :], window=s_wins) + s_d.add_ell_cl(cl_type, n1, n2, leff, bpw_freq_sig[i1, i2, :], window=s_wins) # noqa + s_f.add_ell_cl(cl_type, n1, n2, leff, bpw_freq_sig[i1, i2, :], window=s_wins) # noqa + s_n.add_ell_cl(cl_type, n1, n2, leff, bpw_freq_noi[i1, i2, :], window=s_wins) # noqa # Add covariance print("Adding covariance") @@ -115,8 +124,8 @@ factor_modecount = 1./((2*leff+1)*dell*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])): - covar = (bpw_freq_tot[i1, j1, :]*bpw_freq_tot[i2, j2, :]+ - bpw_freq_tot[i1, j2, :]*bpw_freq_tot[i2, j1, :]) * factor_modecount + 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) cov_bpw = cov_bpw.reshape([ncross * nbands, ncross * nbands]) s_d.add_covariance(cov_bpw) diff --git a/examples/run_compsep_nopipe.sh b/examples/run_compsep_nopipe.sh new file mode 100644 index 0000000..251fbbd --- /dev/null +++ b/examples/run_compsep_nopipe.sh @@ -0,0 +1,13 @@ +#!/bin/bash -l + +## Software environment +export OMP_NUM_THREADS=4 +export PYTHONPATH="${PYTHONPATH}:/home/kw6905/bbdev/BBPower" # YOUR LOCAL BBPower INSTALL PATH + +basedir=/home/kw6905/bbdev/BBPower/examples ## YOUR RUNNING DIRECTORY +bbpower_dir=/home/kw6905/bbdev/BBPower ## PATH TO YOUR LOCAL BBPOWER +cd $basedir + +bbpower_config=${basedir}/config_nopipe.yml +python -u ${bbpower_dir}/bbpower/compsep_nopipe.py --config $bbpower_config +python -u ${bbpower_dir}/bbpower/plotter_nopipe.py --config $bbpower_config diff --git a/examples/utils.py b/examples/utils.py index 6c7d1ab..314a898 100644 --- a/examples/utils.py +++ b/examples/utils.py @@ -1,19 +1,10 @@ import numpy as np -#Foreground model -A_sync_BB = 2.0 +# Foreground model (1) EB_sync = 2. -alpha_sync_EE = -0.6 -alpha_sync_BB = -0.4 +EB_dust = 2. beta_sync = -3. nu0_sync = 23. - -A_dust_BB = 5.0 -EB_dust = 2. -alpha_dust_EE = -0.42 -alpha_dust_BB = -0.2 -beta_dust = 1.59 -temp_dust = 19.6 nu0_dust = 353. Alens = 1.0 @@ -21,53 +12,53 @@ band_names = ['LF1', 'LF2', 'MF1', 'MF2', 'UHF1', 'UHF2'] -#CMB spectrum +# CMB spectrum def fcmb(nu): x = 0.017608676067552197*nu ex = np.exp(x) return ex*(x/(ex-1))**2 -#All spectra -def comp_sed(nu,nu0,beta,temp,typ): +# All spectra +def comp_sed(nu, nu0, beta, temp, typ): if typ == 'cmb': return fcmb(nu) elif typ == 'dust': - x_to=0.04799244662211351*nu/temp - x_from=0.04799244662211351*nu0/temp + x_to = 0.04799244662211351*nu/temp + x_from = 0.04799244662211351*nu0/temp return (nu/nu0)**(1+beta)*(np.exp(x_from)-1)/(np.exp(x_to)-1)*fcmb(nu0) elif typ == 'sync': return (nu/nu0)**beta*fcmb(nu0) return None -#Component power spectra -def dl_plaw(A,alpha,ls): +# Component power spectra +def dl_plaw(A, alpha, ls): return A*((ls+0.001)/80.)**alpha def read_camb(fname, lmax): larr_all = np.arange(lmax+1) - l,dtt,dee,dbb,dte = np.loadtxt(fname,unpack=True) - l = l.astype(int) - msk = l <= lmax - l = l[msk] + ell, dtt, dee, dbb, dte = np.loadtxt(fname, unpack=True) + ell = ell.astype(int) + msk = ell <= lmax + ell = ell[msk] dltt = np.zeros(len(larr_all)) - dltt[l] = dtt[msk] + dltt[ell] = dtt[msk] dlee = np.zeros(len(larr_all)) - dlee[l] = dee[msk] + dlee[ell] = dee[msk] dlbb = np.zeros(len(larr_all)) - dlbb[l] = dbb[msk] + dlbb[ell] = dbb[msk] dlte = np.zeros(len(larr_all)) - dlte[l] = dte[msk] - return dltt,dlee,dlbb,dlte + dlte[ell] = dte[msk] + return dltt, dlee, dlbb, dlte -#Bandpasses +# Bandpasses class Bpass(object): - def __init__(self,name,fname): + def __init__(self, name, fname): self.name = name - self.nu,self.bnu = np.loadtxt(fname,unpack=True) + self.nu, self.bnu = np.loadtxt(fname, unpack=True) self.dnu = np.zeros_like(self.nu) self.dnu[1:] = np.diff(self.nu) self.dnu[0] = self.dnu[1] @@ -75,28 +66,52 @@ def __init__(self,name,fname): norm = np.sum(self.dnu*self.bnu*self.nu**2*fcmb(self.nu)) self.bnu /= norm - def convolve_sed(self,f): + def convolve_sed(self, f): sed = np.sum(self.dnu*self.bnu*self.nu**2*f(self.nu)) return sed -def get_component_spectra(lmax): +def get_component_spectra(lmax, so_forecast=False): + if so_forecast: # foreground parameters from 2302.04276 + A_sync_BB = 1.6 + alpha_sync_EE = -0.7 + alpha_sync_BB = -0.93 + A_dust_BB = 28. + alpha_dust_EE = -0.32 + alpha_dust_BB = -0.16 + else: # foreground parameters from 2011.02449 + A_sync_BB = 2.0 + alpha_sync_EE = -0.6 + alpha_sync_BB = -0.4 + A_dust_BB = 5. + alpha_dust_EE = -0.42 + alpha_dust_BB = -0.2 larr_all = np.arange(lmax+1) - dls_sync_ee=dl_plaw(A_sync_BB*EB_sync,alpha_sync_EE,larr_all) - dls_sync_bb=dl_plaw(A_sync_BB,alpha_sync_BB,larr_all) - dls_dust_ee=dl_plaw(A_dust_BB*EB_dust,alpha_dust_EE,larr_all) - dls_dust_bb=dl_plaw(A_dust_BB,alpha_dust_BB,larr_all) - _,dls_cmb_ee,dls_cmb_bb,_=read_camb("./examples/data/camb_lens_nobb.dat", lmax) + dls_sync_ee = dl_plaw(A_sync_BB*EB_sync, alpha_sync_EE, larr_all) + dls_sync_bb = dl_plaw(A_sync_BB, alpha_sync_BB, larr_all) + dls_dust_ee = dl_plaw(A_dust_BB*EB_dust, alpha_dust_EE, larr_all) + dls_dust_bb = dl_plaw(A_dust_BB, alpha_dust_BB, larr_all) + _, dls_cmb_ee, dls_cmb_bb, _ = read_camb( + "./examples/data/camb_lens_nobb.dat", lmax + ) return (dls_sync_ee, dls_sync_bb, dls_dust_ee, dls_dust_bb, dls_cmb_ee, Alens*dls_cmb_bb) -def get_convolved_seds(names, bpss): + +def get_convolved_seds(names, bpss, so_forecast=False): + if so_forecast: # foreground parameters from 2302.04276 + beta_dust = 1.54 + temp_dust = 20. + else: # foreground parameters from 2011.02449 + beta_dust = 1.59 + temp_dust = 19.6 + nfreqs = len(names) - seds = np.zeros([3,nfreqs]) + seds = np.zeros([3, nfreqs]) for ib, n in enumerate(names): b = bpss[n] - seds[0,ib] = b.convolve_sed(lambda nu : comp_sed(nu,None,None,None,'cmb')) - seds[1,ib] = b.convolve_sed(lambda nu : comp_sed(nu,nu0_sync,beta_sync,None,'sync')) - seds[2,ib] = b.convolve_sed(lambda nu : comp_sed(nu,nu0_dust,beta_dust,temp_dust,'dust')) + seds[0, ib] = b.convolve_sed(lambda nu: comp_sed(nu, None, None, None, 'cmb')) # noqa + seds[1, ib] = b.convolve_sed(lambda nu: comp_sed(nu, nu0_sync, beta_sync, None, 'sync')) # noqa + seds[2, ib] = b.convolve_sed(lambda nu: comp_sed(nu, nu0_dust, beta_dust, temp_dust, 'dust')) # noqa return seds diff --git a/test/run_compsep_test.sh b/test/run_compsep_test.sh index 52c51ba..b6071d7 100755 --- a/test/run_compsep_test.sh +++ b/test/run_compsep_test.sh @@ -6,11 +6,17 @@ mkdir -p test/test_out python ./examples/generate_SO_spectra.py test/test_out # Run pipeline -python -m bbpower BBCompSep --cells_coadded=./test/test_out/cls_coadd.fits --cells_noise=./test/test_out/cls_noise.fits --cells_fiducial=./test/test_out/cls_fid.fits --cells_coadded_cov=./test/test_out/cls_coadd.fits ---output_dir=./test/test_out --config_copy=./test/test_out/config_copy.yml --config=./test/test_config.yml +python -m bbpower BBCompSep \ + --cells_coadded=./test/test_out/cls_coadd.fits \ + --cells_noise=./test/test_out/cls_noise.fits \ + --cells_fiducial=./test/test_out/cls_fid.fits \ + --cells_coadded_cov=./test/test_out/cls_coadd.fits \ + --output_dir=./test/test_out \ + --config_copy=./test/test_out/config_copy.yml \ + --config=./test/test_config_emcee.yml # Check chi2 -if python -c "import numpy as np; chi2=np.load('test/test_out/param_chains.npz')['chi2']; assert chi2<1E-5"; then +if python -c "import numpy as np; chi2=np.load('test/test_out/emcee.npz')['chi2']; assert chi2<1E-5"; then echo "Test passed" else echo "Test did not pass" diff --git a/test/run_polychord_test.sh b/test/run_polychord_test.sh index fcbb6e3..ff68cf9 100755 --- a/test/run_polychord_test.sh +++ b/test/run_polychord_test.sh @@ -5,16 +5,16 @@ mkdir -p test/test_out # Generate fiducial cls python ./examples/generate_SO_spectra.py test/test_out -# Generate simulations -for seed in {1001..1100} -do - mkdir -p test/test_out/s${seed} - echo ${seed} - python examples/generate_SO_maps.py --output-dir test/test_out/s${seed} --seed ${seed} --nside 64 -done +# # Generate simulations +# for seed in {1001..1005} +# do +# mkdir -p test/test_out/s${seed} +# echo ${seed} +# python examples/generate_SO_maps.py --output-dir test/test_out/s${seed} --seed ${seed} --nside 64 +# done # Run pipeline -python -m bbpower BBPowerSpecter --splits_list=./examples/test_data/splits_list.txt --masks_apodized=./examples/data/maps/norm_nHits_SA_35FOV.fits --bandpasses_list=./examples/data/bpass_list.txt --sims_list=./examples/test_data/sims_list.txt --beams_list=./examples/data/beams_list.txt --cells_all_splits=./test/test_out/cells_all_splits.fits --cells_all_sims=./test/test_out/cells_all_sims.txt --mcm=./test/test_out/mcm.dum --config=./test/test_config_polychord.yml +python -m bbpower BBPowerSpecter --splits_list=./examples/test_data/splits_list.txt --masks_apodized=./examples/data/maps/norm_nHits_SA_35FOV.fits --bandpasses_list=./examples/data/bpass_list.txt --sims_list=./examples/test_data/sims_list_10.txt --beams_list=./examples/data/beams_list.txt --cells_all_splits=./test/test_out/cells_all_splits.fits --cells_all_sims=./test/test_out/cells_all_sims.txt --mcm=./test/test_out/mcm.dum --config=./test/test_config_polychord.yml python -m bbpower BBPowerSummarizer --splits_list=./examples/test_data/splits_list.txt --bandpasses_list=./examples/data/bpass_list.txt --cells_all_splits=./test/test_out/cells_all_splits.fits --cells_all_sims=./test/test_out/cells_all_sims.txt --cells_coadded_total=./test/test_out/cells_coadded_total.fits --cells_coadded=./test/test_out/cells_coadded.fits --cells_noise=./test/test_out/cells_noise.fits --cells_null=./test/test_out/cells_null.fits --config=./test/test_config_polychord.yml @@ -29,5 +29,5 @@ else echo "Test passed" fi -# Cleanup -rm -r test/test_out +# # Cleanup +# rm -r test/test_out diff --git a/test/run_power_specter_test.sh b/test/run_power_specter_test.sh index 506ab55..686d740 100755 --- a/test/run_power_specter_test.sh +++ b/test/run_power_specter_test.sh @@ -22,12 +22,11 @@ python -m bbpower BBCompSep --cells_coadded=./test/test_out/cells_coadded.fits python -m bbpower BBPlotter --cells_coadded_total=./test/test_out/cells_coadded_total.fits --cells_coadded=./test/test_out/cells_coadded.fits --cells_noise=./test/test_out/cells_noise.fits --cells_null=./test/test_out/cells_null.fits --cells_fiducial=./test/test_out/cls_fid.fits --param_chains=./test/test_out/emcee.npz --plots=./test/test_out/plots.dir --plots_page=./test/test_out/plots_page.html --config=./test/test_config_emcee.yml -fchain="test/test_out/param_chains.npz" if python -c "import numpy as np; a=np.load('${fchain}'); rchi2 = a['chi2'] / a['ndof']; print('chi2/dof = ', rchi2) ; assert rchi2 < 2"; then echo "Test passed" else echo "Test did not pass" fi -# Cleanup -rm -r test/test_out +# # Cleanup +# rm -r test/test_out diff --git a/test/run_sampling_test.sh b/test/run_sampling_test.sh deleted file mode 100755 index 74fc665..0000000 --- a/test/run_sampling_test.sh +++ /dev/null @@ -1,24 +0,0 @@ -#!/bin/bash - -# Create the output directory -mkdir -p test/test_out - -# Generate some fake data -python ./examples/generate_SO_spectra.py test/test_out - -# Run pipeline -# This runs the likelihood sampler -python -m bbpower BBCompSep --cells_coadded=./test/test_out/cls_coadd.fits --cells_noise=./test/test_out/cls_coadd.fits --cells_fiducial=./test/test_out/cls_fid.fits --cells_coadded_cov=./test/test_out/cls_coadd.fits --param_chains=./test/test_out/param_chains.npz --config_copy=./test/test_out/config_copy.yml --config=./test/test_config_sampling.yml - -# This plots the results of the pipeline -python -m bbpower BBPlotter --cells_coadded_total=./test/test_out/cls_coadd.fits --cells_coadded=./test/test_out/cls_coadd.fits --cells_noise=./test/test_out/cls_coadd.fits --cells_null=./test/test_out/cls_coadd.fits --cells_fiducial=./test/test_out/cls_fid.fits --param_chains=./test/test_out/param_chains.npz --plots=./test/test_out/plots.dir --plots_page=./test/test_out/plots_page.html --config=./test/test_config_sampling.yml - -# Check the final plots exist -if [ ! -f ./test/test_out/plots.dir/triangle.png ]; then - echo "Test did not pass" -else - echo "Test passed" -fi - -# Cleanup -rm -r test/test_out diff --git a/test/test_config_sampling.yml b/test/test_config_predicted_spectra.yml similarity index 67% rename from test/test_config_sampling.yml rename to test/test_config_predicted_spectra.yml index 5435113..fc52090 100644 --- a/test/test_config_sampling.yml +++ b/test/test_config_predicted_spectra.yml @@ -1,38 +1,40 @@ # These parameters are accessible to all stages global: - # HEALPix resolution parameter nside: 64 - # Use D_l = l*(l+1)*C_l/(2*pi) instead of C_l? compute_dell: True +BBPowerSpecter: + # Bandpower definition + bpw_edges: "./examples/test_data/bpw_edges.txt" + purify_B: True + n_iter : 3 + +BBPowerSummarizer: + # Covariance types + nulls_covar_type: "diagonal" + nulls_covar_diag_order: 0 + data_covar_type: "block_diagonal" + data_covar_diag_order: 0 + BBCompSep: - # Sampler type. Options are: - # - 'emcee': a full MCMC will be run using emcee. - # - 'fisher': only the Fisher matrix (i.e. the likelihood - # Hessian matrix) will be calculated around the fiducial - # parameters chosen as prior centers. - # - 'maximum_likelihood': only the best-fit parameters will - # be searched for. - # - 'single_point': only the chi^2 at the value used as the - # center for all parameter priors will be calculated. - # - 'timing': will compute the average time taken by one - # likelihood computation. - sampler: 'maximum_likelihood' + # Sampler type (choose 'emcee', 'polychord', 'maximum_likelihood', + # 'single_point' or 'timing') + sampler: 'predicted_spectra' + # If you chose polychord: + nlive: 50 + nrepeat: 50 # If you chose emcee: - # Number of walkers nwalkers: 24 - # Number of iterations per walker n_iters: 1000 - # Likelihood type. Options are: - # - 'chi2': a standard chi-squared Gaussian likelihood. - # - 'h&l': Hamimeche & Lewis likelihood. - likelihood_type: 'h&l' + # Likelihood type (choose 'chi2' or 'h&l') + likelihood_type: 'chi2' + # What is the starting point? + r_init: 1.e-3 # Which polarization channels do you want to include? - # Can be ['E'], ['B'] or ['E','B']. - pol_channels: ['E','B'] + pol_channels: ['B'] # Scale cuts (will apply to all frequencies) l_min: 30 - l_max: 120 + l_max: 300 # CMB model cmb_model: @@ -46,7 +48,7 @@ BBCompSep: # See below for the meaning of the different elements in the list. r_tensor: ['r_tensor', 'tophat', [-0.1, 0.00, 0.1]] # Lensing amplitude - A_lens: ['A_lens', 'tophat', [0.00,1.0,2.00]] + A_lens: ['A_lens', 'tophat', [0.00, 1.0, 2.00]] # Foreground model fg_model: @@ -60,8 +62,7 @@ BBCompSep: sed: Dust # Type of power spectra for all possible polarization channel combinations. # Any combinations not added here will be assumed to be zero. - # The names should be one of the classes in bbpower/fgcls.py. This is quite - # limiter for now, so consider adding to it if you want something fancier. + # The names should be one of the classes in bbpower/fgcls.py cl: EE: ClPowerLaw BB: ClPowerLaw @@ -83,20 +84,17 @@ BBCompSep: # Same for power spectrum parameters # (broken down by polarization channel combinations) EE: - amp_d_ee: ['amp', 'tophat', [0., 10., "inf"]] + amp_d_ee: ['amp', 'tophat', [0., 10., 20.]] alpha_d_ee: ['alpha', 'tophat', [-1., -0.42, 0.]] l0_d_ee: ['ell0', 'fixed', [80.]] BB: - amp_d_bb: ['amp', 'tophat', [0., 5., "inf"]] + amp_d_bb: ['amp', 'tophat', [0., 5., 10.]] alpha_d_bb: ['alpha', 'tophat', [-1., -0.2, 0.]] l0_d_bb: ['ell0', 'fixed', [80.]] # If this component should be correlated with any other, list them here cross: # In this case the list should contain: # [component name, prior type, prior parameters] - # Each of this will create a new parameter, corresponding to a constant - # scale- and frequency-independend correlation coefficient between - # the two components. epsilon_ds: ['component_2', 'tophat', [-1., 0., 1.]] component_2: @@ -110,22 +108,17 @@ BBCompSep: nu0_s: ['nu0', 'fixed', [23.]] cl_parameters: EE: - amp_s_ee: ['amp', 'tophat', [0., 4., "inf"]] + amp_s_ee: ['amp', 'tophat', [0., 4., 8.]] alpha_s_ee: ['alpha', 'tophat', [-1., -0.6, 0.]] l0_s_ee: ['ell0', 'fixed', [80.]] BB: - amp_s_bb: ['amp', 'tophat', [0., 2., "inf"]] + amp_s_bb: ['amp', 'tophat', [0., 2., 4.]] alpha_s_bb: ['alpha', 'tophat', [-1., -0.4, 0.]] l0_s_bb: ['ell0', 'fixed', [80.]] BBPlotter: - # Maximum ell up to which things will be plotted - lmax_plot: 128 - # Add total coadds to C_ell plots? + lmax_plot: 300 plot_coadded_total: False - # Add noise spectra to C_ell plots? plot_noise: False - # Plot null tests? plot_nulls: False - # Plot likelihood? plot_likelihood: True diff --git a/test/test_powerspecter.yml b/test/test_powerspecter.yml index 6469cdf..6c73d3e 100644 --- a/test/test_powerspecter.yml +++ b/test/test_powerspecter.yml @@ -33,7 +33,7 @@ inputs: cells_fiducial: ./test/test_out/cls_fid.fits # Overall configuration file -config: ./test/test_config.yml +config: ./test/test_config_emcee.yml # If all the outputs for a stage already exist then do not re-run that stage resume: False diff --git a/test/test_sampling.yml b/test/test_sampling.yml deleted file mode 100644 index 351c7b8..0000000 --- a/test/test_sampling.yml +++ /dev/null @@ -1,61 +0,0 @@ -# Python modules that are imported to find -# stage classes. Any stages imported in these -# modules are automatically detected and their names can -# be used below -modules: bbpower - - -# The launcher to use -# These are defined in bbpipe/sites -launcher: local - - -# The list of stages to run and the number of processors -# to use for each. -stages: - - name: BBCompSep - nprocess: 1 - - name: BBPlotter - nprocess: 1 - -# Definitions of where to find inputs for the overall pipeline. -# Any input required by a pipeline stage that is not generated by -# a previous stage must be defined here. They are listed by tag. -inputs: - # This is the file containing the data power spectra, - # covariance and other metadata (window functions, bandpasses, - # beams etc.). - cells_coadded: ./test/test_out/cls_coadd.fits - # This contains a fiducial power spectrum for a model that - # roughly resembles the data. This is needed for the - # Hamimeche and Lewis (H&L) likelihood. If you're not using that, - # you can put garbage here. - cells_fiducial: ./test/test_out/cls_fid.fits - # This contains the noise power spectrum. Only needed for - # H&L. - cells_noise: ./test/test_out/cls_noise.fits - # Null power spectra. Not really needed if you only care about - # the likelihood part of the pipeline. Otherwise it should - # contain a set of power spectrum combinations that should - # be compatible with zero. - cells_null: ./test/test_out/cls_coadd.fits - # Power spectra computed by coadding all the different data - # splits (including auto-correlations). As above, ignore this - # if you only care about the likelihood part of the pipeline. - cells_coadded_total: ./test/test_out/cls_coadd.fits - -# Overall configuration file describing the options for -# the different pipeline stages. -config: ./test/test_config_sampling.yml - -# If all the outputs for a stage already exist then do not re-run that stage -resume: False - -# Put all the output files in this directory: -output_dir: ./test/test_out - -# Put the logs from the individual stages in this directory: -log_dir: ./test/test_out - -# Put the log for the overall pipeline infrastructure in this file: -pipeline_log: ./test/test_out/log.txt