diff --git a/example/make_example_transfer_matrix.py b/example/make_example_transfer_matrix.py new file mode 100644 index 00000000..895be41f --- /dev/null +++ b/example/make_example_transfer_matrix.py @@ -0,0 +1,34 @@ +import numpy as np +import xfaster as xf +import os + +kerns = xf.load_and_parse("outputs_example/95x150/kernels_95x150.npz") +beams = xf.load_and_parse("outputs_example/95x150/beams_95x150.npz") +tf = np.loadtxt("maps_example/transfer_example.txt") +tf.shape +lmax = 500 +lk = slice(0, lmax + 1) +bw = beams["beam_windows"] + +os.makedirs("transfer_matrix") + +for xname in kerns["kern"].keys(): + tag1, tag2 = xname.split(":") + fname = "transfer_matrix/{}x{}_{{}}_to_{{}}_block.dat".format(tag1, tag2) + for spec in ["tt", "ee", "bb"]: + fb2 = tf[lk] * bw[spec][tag1][lk] * bw[spec][tag2][lk] + + if spec == "tt": + k = kerns["kern"][xname][..., lk] + else: + k = kerns["pkern"][xname][..., lk] + mk = kerns["mkern"][xname][..., lk] + + mat = k * fb2 + su = spec.upper() + np.savetxt(fname.format(su, su), mat.T) + + if spec != "tt": + mat = mk * fb2 + su2 = "BB" if spec == "ee" else "EE" + np.savetxt(fname.format(su, su2), mat.T) diff --git a/xfaster/xfaster_class.py b/xfaster/xfaster_class.py index 6ce451b9..bce1927e 100644 --- a/xfaster/xfaster_class.py +++ b/xfaster/xfaster_class.py @@ -115,6 +115,7 @@ class XFaster(object): "sims_transfer", "shape_transfer", "transfer", + "transfer_matrix", "sims", "beams", "data", @@ -141,6 +142,7 @@ class XFaster(object): "sims_transfer": ["transfer"], "shape_transfer": ["transfer"], "transfer": ["bandpowers"], + "transfer_matrix": ["bandpowers"], "sims": ["bandpowers"], "beams": ["transfer"], "data": ["bandpowers"], @@ -1302,7 +1304,8 @@ def get_filename( template_type, reference_type. bp_opts : bool If True, also check the following attributes (in addition to those - checked if ``data_opts`` is True): weighted_bins, return_cls. + checked if ``data_opts`` is True): weighted_bins, return_cls, + transfer_matrix. Returns ------- @@ -1342,6 +1345,8 @@ def get_filename( name += ["wbins"] if getattr(self, "return_cls", False): name += ["cl"] + if hasattr(self, "transfer_matrix"): + name += ["xfermat"] if map_tag is not None: name += ["map", map_tag] @@ -4067,7 +4072,8 @@ def get_bin_def( def kernel_precalc(self, map_tag=None, transfer=False): """ Compute the mixing kernels M_ll' = K_ll' * F_l' * B_l'^2. Called by - ``bin_cl_template`` to pre-compute kernel terms. + ``bin_cl_template`` to pre-compute kernel terms. Constructs M_ll' + directly from the ``transfer_matrix`` attribute, if present. Arguments --------- @@ -4103,6 +4109,12 @@ def kernel_precalc(self, map_tag=None, transfer=False): lk = slice(0, lmax + 1) mll = OrderedDict() + use_transfer_matrix = hasattr(self, "transfer_matrix") + if use_transfer_matrix and transfer: + raise ValueError( + "Transfer matrix cannot be used to compute transfer functions." + ) + if transfer: comps = ["cmb"] else: @@ -4118,11 +4130,21 @@ def kernel_precalc(self, map_tag=None, transfer=False): mstag = "{}_mix".format(stag) mll[mstag] = OrderedDict() - bw = self.beam_windows[spec] - if not transfer: - tf = self.transfer[stag] + if not use_transfer_matrix: + bw = self.beam_windows[spec] + if not transfer: + tf = self.transfer[stag] for xname, (m0, m1) in map_pairs.items(): + # extract transfer matrix terms + if use_transfer_matrix: + tm = self.transfer_matrix[xname][spec] + mll[stag][xname] = tm[spec][:, lk] + if spec in ["ee", "bb"]: + spec2 = "bb" if spec == "ee" else "ee" + mll[mstag][xname] = tm[spec2][:, lk] + continue + # beams fb2 = bw[m0][lk] * bw[m1][lk] @@ -5734,10 +5756,11 @@ def fisher_iterate( ) if not transfer: - out.update( - qb_transfer=self.qb_transfer, - transfer=self.transfer, - ) + if not hasattr(self, "transfer_matrix"): + out.update( + qb_transfer=self.qb_transfer, + transfer=self.transfer, + ) if self.template_type is not None: out.update(template_alpha=self.template_alpha) @@ -6110,6 +6133,99 @@ def expand_transfer(qb_transfer, bin_def, check_only=False): return self.transfer + def get_transfer_matrix(self, file_root=None): + """ + Read in an ell-by-ell transfer matrix, constructed from ratios of + `anafast` spectra using an ensemble of simulations. This matrix should + be used in place of the standard XFaster transfer function calculation, + and includes masking, filtering and beam effects. + + Arguments + --------- + file_root : str + Directory where the transfer matrix blocks are stored. Expects + column text files with naming pattern: + ``x__to__block.dat`` + Required if the transfer matrix has not already been computed. + + Returns + ------- + transfer_matrix : OrderedDict + Dictionary of matrix blocks of shape (lmax + 1, lmax + 1), + keyed by map cross (xname), output spectrum, and input spectrum. + This dictionary is also cached in the ``transfer_matrix`` attribute. + + Notes + ----- + This function is called at the ``transfer_matrix`` checkpoint, and is + meant to be an alternative to the standard ``transfer`` sequence. + """ + + tm_shape = ( + self.num_corr * (self.num_spec + 2) * (self.lmax + 1), + self.lmax + 1, + ) + + save_name = "transfer_matrix" + save_attrs = ["transfer_matrix_root", "transfer_matrix"] + ret = self.load_data( + save_name, + "transfer_matrix", + fields=save_attrs, + to_attrs=True, + shape=tm_shape, + shape_ref="transfer_matrix", + ) + + if ret is not None: + return ret["transfer_matrix"] + + transfer_matrix = OrderedDict() + for xname, (tag1, tag2) in self.map_pairs.items(): + dx = transfer_matrix.setdefault(xname, OrderedDict()) + + if tag1 == tag2: + ftag = "{}x{}".format(tag1, tag2) + else: + for t in ["{}x{}".format(tag1, tag2), "{}x{}".format(tag2, tag1)]: + fname = os.path.join(file_root, "{}_*_to_*_block.dat".format(t)) + if len(glob.glob(fname)): + ftag = t + break + else: + self.warn("Missing transfer matrix for {}".format(xname)) + continue + + # file name format string + fname = os.path.join(file_root, "{}_{{}}_to_{{}}_block.dat".format(ftag)) + + for spec_out in self.specs: + dsi = dx.setdefault(spec_out, OrderedDict()) + + specs_in = ["ee", "bb"] if spec_out in ["ee", "bb"] else [spec_out] + for spec_in in specs_in: + if spec_out in ["tt", "ee", "bb"]: + # auto- => auto- + fname1 = fname.format(spec_in.upper(), spec_out.upper()) + mat = np.loadtxt(fname1) + + # populate matrix up to lmax + dsi[spec_in] = np.zeros((self.lmax + 1, self.lmax + 1)) + nx, ny = [min([s, self.lmax + 1]) for s in mat.shape] + dsi[spec_in][:nx, :ny] = mat[:nx, :ny] + del mat + + else: + # cross1- => cross1- + s1, s2 = [s + s for s in spec_in] + dsi[spec_in] = np.sqrt(np.abs(dx[s1][s1] * dx[s2][s2])) + + self.transfer_matrix_root = file_root + self.transfer_matrix = transfer_matrix + + self.save_data(save_name, from_attrs=save_attrs) + return transfer_matrix + def get_bandpowers( self, map_tag=None, diff --git a/xfaster/xfaster_exec.py b/xfaster/xfaster_exec.py index e639e27a..fa49b1af 100644 --- a/xfaster/xfaster_exec.py +++ b/xfaster/xfaster_exec.py @@ -113,6 +113,7 @@ def xfaster_run( return_cls=False, qb_only=False, fix_bb_transfer=False, + transfer_matrix_root=None, null_first_cmb=False, delta_beta_prior=None, like_profiles=False, @@ -410,6 +411,10 @@ def xfaster_run( fix_bb_transfer : bool If True, after transfer functions have been calculated, impose that the BB transfer function is exactly equal to the EE transfer function. + transfer_matrix_root : str + If set, use the pre-computed ell-by-ell transfer matrix stored in this + directory, instead of the standard transfer function computed by the + XFaster estimator. null_first_cmb : bool If True, keep first CMB bandpowers fixed to input shape (qb=1). delta_beta_prior : float @@ -673,6 +678,7 @@ def xfaster_run( iter_max=iter_max, save_iters=save_iters, fix_bb_transfer=fix_bb_transfer, + transfer_matrix_root=transfer_matrix_root, delta_beta_prior=delta_beta_prior, cond_noise=cond_noise, cond_criteria=cond_criteria, @@ -687,6 +693,7 @@ def xfaster_run( config_vars.update(spec_opts, "Spectrum Estimation Options") config_vars.remove_option("XFaster General", "bandpower_tag") spec_opts.pop("multi_map") + spec_opts.pop("transfer_matrix_root") bandpwr_opts = spec_opts.copy() bandpwr_opts.pop("fix_bb_transfer") spec_opts.pop("file_tag") @@ -746,20 +753,25 @@ def xfaster_run( gcorr_file=gcorr_file, ) - X.log("Computing kernels...", "notice") - X.get_kernels(window_lmax=window_lmax) + if transfer_matrix_root: + X.log("Loading pre-computed transfer matrix...", "notice") + X.get_transfer_matrix(file_root=transfer_matrix_root) - X.log("Computing beam window functions...", "notice") - X.get_beams(pixwin=pixwin) + else: + X.log("Computing kernels...", "notice") + X.get_kernels(window_lmax=window_lmax) + + X.log("Computing beam window functions...", "notice") + X.get_beams(pixwin=pixwin) - X.log("Computing sim ensemble averages for transfer function...", "notice") - X.get_masked_sims(transfer=True, **sim_transfer_opts) + X.log("Computing sim ensemble averages for transfer function...", "notice") + X.get_masked_sims(transfer=True, **sim_transfer_opts) - X.log("Loading spectrum shape for transfer function...", "notice") - X.get_signal_shape(filename=signal_transfer_spec, transfer=True) + X.log("Loading spectrum shape for transfer function...", "notice") + X.get_signal_shape(filename=signal_transfer_spec, transfer=True) - X.log("Computing transfer functions...", "notice") - X.get_transfer(**transfer_opts) + X.log("Computing transfer functions...", "notice") + X.get_transfer(**transfer_opts) X.log("Computing sim ensemble averages...", "notice") X.get_masked_sims(**sim_opts) @@ -1256,6 +1268,7 @@ def add_arg( add_arg(G, "return_cls") add_arg(G, "qb_only") add_arg(G, "fix_bb_transfer") + add_arg(G, "transfer_matrix_root") add_arg(G, "null_first_cmb") add_arg(G, "delta_beta_prior", argtype=float) add_arg(G, "like_profiles")