Transfer matrix integration

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)

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:
+            ``<tag1>x<tag2>_<spec_in>_to_<spec_out>_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,

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")