Mask validation

pipeline/misc/validate_mask_isohits.py

@@ -0,0 +1,220 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from pixell import enmap
+from soopercool import BBmeta, utils as su, map_utils as mu
+import pymaster as nmt
+import os
+import re
+import argparse
+
+
+X_OFFSET_UNIT = 2  # Plots: Shift in x to avoid overlapping lines
+
+
+def get_filename_tags(dir, base, ext):
+    """
+    Extract all tags from filenames in a directory matching the
+    pattern <base><tag>.<extension>. Tag may be empty ('').
+    """
+    ext = ext.lstrip('.')
+    pattern = re.compile(rf'^{base}(?P<tag>(?:_[^.]*)?)\.{ext}$')
+
+    tags = []
+    for filename in os.listdir(dir):
+        match = pattern.match(filename)
+        if match:
+            tags.append(match.group('tag'))
+
+    return tags
+
+def main(args):
+    """
+    Validate masks generated by get_analysis_mask.py.
+    Use the same param file to ensure paths are correct, with the
+    following requirements:
+
+    Parameters
+    ----------
+    output_directory: Output directory used for get_analysis_mask.py
+
+    general_pars:
+      pix_type: car or hp
+      car_template: Path to car template (if pix_type == car)
+      binning_file: Path to binning file
+
+    mask_validation:
+      num_sims: Number of sims to use
+      sim_id_start: Sim ID to start from
+      unfiltered_map_dir: Directory containing unfiltered maps (sims)
+      unfiltered_map_template: Filename template
+    """
+
+    meta = BBmeta(args.globals)
+    purify_b = args.no_purify
+    fwhm_amin = args.fwhm
+    pix_type = meta.pix_type
+    sim_type = args.sim_type
+    conv = args.no_unit_conversion
+    norm_std = args.no_norm_std
+    simdir_unfiltered = meta.mask_validation["unfiltered_map_dir"]
+    unfiltered_map_tmpl = meta.mask_validation["unfiltered_map_template"]
+
+    if os.path.exists(meta.masks["analysis_mask"]):
+        # Mask directory specified in config file
+        mask_dir = os.path.dirname(meta.masks["analysis_mask"])
+        plot_dir = f"{mask_dir}/plots"
+        BBmeta.make_dir(plot_dir)
+
+        # Mask filename
+        mask_basename = os.path.basename(meta.masks["analysis_mask"]).split('.')[0]
+        mask_tags = ['']
+    else:
+        # Directory tree as structured in get_analysis_mask.py
+        out_dir = meta.output_directory
+        mask_dir = f"{out_dir}/masks"
+        plot_dir = f"{out_dir}/plots/masks"
+        BBmeta.make_dir(plot_dir)
+
+        # Check mask directory for masks
+        mask_basename = "analysis_mask"
+        mask_tags = get_filename_tags(mask_dir, mask_basename, ".fits")
+        if len(mask_tags) == 0:
+            raise FileNotFoundError(f"No masks found in {mask_dir} with names matching {mask_basename}<tag>.fits.")
+
+    # Binning
+    nmt_bins = meta.read_nmt_binning()
+    leff = nmt_bins.get_effective_ells()
+    lmax = nmt_bins.lmax
+
+    # Load theory Cls
+    _, clth = su.get_theory_cls(lmax=lmax, fwhm_amin=fwhm_amin)
+    cbtt = nmt_bins.bin_cell(clth['TT'][:lmax+1])
+    cbee = nmt_bins.bin_cell(clth['EE'][:lmax+1])
+    cbbb = nmt_bins.bin_cell(clth['BB'][:lmax+1])
+
+    # Car template geometry
+    shape, wcs = enmap.read_map_geometry(meta.car_template) if pix_type == 'car' else (None, None)
+
+
+    ##########
+    # Validate
+    ##########
+    sim_id_start = meta.mask_validation["sim_id_start"]
+    nsims = meta.mask_validation["num_sims"]
+
+    norm_factor = np.sqrt(nsims) if norm_std else 1
+
+    cl00 = {}; cl22 = {}; cl22b = {}
+    means_tt = {}; means_ee = {}; means_bb = {}
+    std_tt = {}; std_ee = {}; std_bb = {}
+    std_bb_bonly = {}
+
+    for tag in mask_tags:
+        mask_path = os.path.join(mask_dir, f'{mask_basename}{tag}.fits')
+        mask = mu.read_map(mask_path, pix_type=pix_type, geometry=(shape, wcs))
+
+        cl00[tag] = []; cl22[tag] = []; cl22b[tag] = []
+        for i in range(sim_id_start, sim_id_start+nsims):
+            cmb_map_file = os.path.join(simdir_unfiltered, unfiltered_map_tmpl.format(sim_type=sim_type, id_sim=i))
+            mpt, mpq, mpu = mu.read_map(cmb_map_file, convert_K_to_muK=conv, pix_type=pix_type, geometry=(shape, wcs))
+
+            f0 = nmt.NmtField(mask, [mpt], spin=0, wcs=wcs, lmax=lmax, lmax_mask=lmax)
+            f2 = nmt.NmtField(mask, [mpq, mpu], purify_b=purify_b, wcs=wcs, lmax=lmax, lmax_mask=lmax)
+
+            cl00[tag].append(nmt.compute_full_master(f0, f0, nmt_bins))
+            cl22[tag].append(nmt.compute_full_master(f2, f2, nmt_bins))
+
+            # B-only sims: no purification
+            bonly_map_file = os.path.join(simdir_unfiltered, unfiltered_map_tmpl.format(sim_type="cmbB", id_sim=i))
+            _, mpqb, mpub = mu.read_map(bonly_map_file, convert_K_to_muK=conv, pix_type=pix_type, geometry=(shape, wcs))
+
+            f2b = nmt.NmtField(mask, [mpqb, mpub], purify_b=False, wcs=wcs, lmax=lmax, lmax_mask=lmax)
+            cl22b[tag].append(nmt.compute_full_master(f2b, f2b, nmt_bins))
+
+        means_tt[tag] = np.mean(cl00[tag], axis=0)[0]
+        means_ee[tag] = np.mean(cl22[tag], axis=0)[0]
+        means_bb[tag] = np.mean(cl22[tag], axis=0)[3]
+        std_tt[tag] = np.std(cl00[tag], axis=0)[0]/norm_factor
+        std_ee[tag] = np.std(cl22[tag], axis=0)[0]/norm_factor
+        std_bb[tag] = np.std(cl22[tag], axis=0)[3]/norm_factor
+
+        std_bb_bonly[tag] = np.std(cl22b[tag], axis=0)[3]/norm_factor
+
+
+    ##########
+    # Plots
+    ##########
+    formatting = {'fmt': 'o', 'markersize': 0, 'linewidth': 2, 'capsize': 3, 'capthick': 2}
+    offsets = (np.arange(len(mask_tags)) - len(mask_tags)/2 + 0.5) * X_OFFSET_UNIT
+    ells = {tag: leff+offsets[i] for i, tag in enumerate(mask_tags)}
+    labels = {tag: tag[1:] if tag.startswith('_') else tag for tag in mask_tags}
+
+    plt.figure()
+    plt.semilogy(leff, cbtt, label='TT theory', color='k', linestyle='-')
+    for tag in mask_tags:
+        plt.errorbar(ells[tag], means_tt[tag], yerr=std_tt[tag], label=labels[tag], **formatting)
+    plt.legend()
+    plt.xlabel(r"Multipole $\ell$")
+    plt.ylabel(r"$C_b^{TT} \; [\mu\mathrm{K}^2]$")
+    plt.savefig(os.path.join(plot_dir, "mask_validation_TT.png"), bbox_inches='tight')
+    plt.close()
+
+    plt.figure()
+    plt.semilogy(leff, cbee, label='EE theory', color='k', linestyle='-')
+    for tag in mask_tags:
+        plt.errorbar(ells[tag], means_ee[tag], yerr=std_ee[tag], label=labels[tag], **formatting)
+    plt.legend()
+    plt.xlabel(r"Multipole $\ell$")
+    plt.ylabel(r"$C_b^{EE} \; [\mu\mathrm{K}^2]$")
+    plt.savefig(os.path.join(plot_dir, "mask_validation_EE.png"), bbox_inches='tight')
+    plt.close()
+
+    plt.figure()
+    plt.semilogy(leff, cbbb, label='BB theory', color='k', linestyle='-')
+    for tag in mask_tags:
+        plt.errorbar(ells[tag], means_bb[tag], yerr=std_bb[tag], label=labels[tag], **formatting)
+    plt.legend()
+    plt.xlabel(r"Multipole $\ell$")
+    plt.ylabel(r"$C_b^{BB} \; [\mu\mathrm{K}^2]$")
+    plt.savefig(os.path.join(plot_dir, "mask_validation_BB.png"), bbox_inches='tight')
+    plt.close()
+
+    plt.figure()
+    for tag in mask_tags:
+        plt.plot(ells[tag], (means_ee[tag]-cbee)/std_ee[tag], label=labels[tag])
+    plt.xlabel(r"Multipole $\ell$")
+    plt.ylabel(r"$(\^{C}_b - C_b^{\mathrm{th}})/\sigma(\^{C}_b)$")
+    plt.savefig(os.path.join(plot_dir, "mask_validation_EE_sigmas.png"), bbox_inches='tight')
+    plt.close()
+
+    plt.figure()
+    for tag in mask_tags:
+        plt.plot(ells[tag], (means_bb[tag]-cbbb)/std_bb[tag], label=labels[tag])
+    plt.xlabel(r"Multipole $\ell$")
+    plt.ylabel(r"$(\^{C}_b - C_b^{\mathrm{th}})/\sigma(\^{C}_b)$")
+    plt.savefig(os.path.join(plot_dir, "mask_validation_BB_sigmas.png"), bbox_inches='tight')
+    plt.close()
+
+    plt.figure()
+    for tag in mask_tags:
+        plt.semilogy(ells[tag], std_bb[tag]/std_bb_bonly[tag], label=labels[tag])
+    plt.xlabel(r"Multipole $\ell$")
+    plt.ylabel(r"$\sigma(C_b^{\mathrm{CMB}})/\sigma(C_b^{\mathrm{B\ only}})$")
+    plt.savefig(os.path.join(plot_dir, "mask_validation_compare_bonly.png"), bbox_inches='tight')
+    plt.close()
+
+    print(f"Mask validation plots saved to {plot_dir}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--globals", required=True, help="Path to paramfile")
+    parser.add_argument("--fwhm", default=30, help="Beam size (arcmin). Default: 30")
+    parser.add_argument("--sim-type", default='cmbTEB', help="Input sim type (T/E/B). Default: 'cmbTEB'")
+    parser.add_argument("--no-purify", action="store_false", help="Do not purify B-modes with NaMaster")
+    parser.add_argument("--no-unit-conversion", action="store_false", help="Do not convert map from K to muK")
+    parser.add_argument("--no-norm-std", default="store_false", help="Do not normalize errorbars by sqrt(nsims)")
+
+    args = parser.parse_args()
+
+    main(args)

soopercool/map_utils.py

@@ -1,6 +1,6 @@
 import numpy as np
 import healpy as hp
-from pixell import enmap, enplot, curvedsky
+from pixell import enmap, enplot, curvedsky, utils
 import matplotlib.pyplot as plt
 from pixell import uharm
 import pymaster as nmt
@@ -85,7 +85,8 @@ def map2alm(map, pix_type="hp"):
         return curvedsky.map2alm(map, lmax=lmax)
 
 
-def alm2map(alm, pix_type="hp", nside=None, car_map_template=None):
+def alm2map(alm, pix_type="hp", nside=None, car_template=None,
+            geometry=None):
     """
     """
     _check_pix_type(pix_type)
@@ -96,11 +97,16 @@ def alm2map(alm, pix_type="hp", nside=None, car_map_template=None):
         assert nside is not None, "nside is required"
         return hp.alm2map(alm, nside=nside)
     else:
-        if isinstance(car_map_template, str):
-            shape, wcs = enmap.read_map_geometry(car_map_template)
+        assert geometry is not None or car_template is not None, "geometry info required"
+        if isinstance(car_template, str):
+            shape, wcs = enmap.read_map_geometry(car_template)
+        elif car_template is None:
+            shape, wcs = geometry
         else:
-            shape, wcs = car_map_template.geometry
-        map = enmap.zeros((3,) + shape, wcs)
+            shape, wcs = car_template.geometry
+        # shape = (dim,) + shape if dim > 1 else shape
+        map = enmap.zeros(shape, wcs)
+
         return curvedsky.alm2map(alm, map)
 
 
@@ -205,7 +211,10 @@ def read_map(map_file,
         m = hp.read_map(map_file, **kwargs)
     else:
         if geometry is None:
-            geometry = enmap.read_map_geometry(car_template)
+            try:
+                geometry = enmap.read_map_geometry(car_template)
+            except AttributeError:
+                pass
         m = enmap.read_map(map_file, geometry=geometry)
 
     return conv*m
@@ -304,7 +313,6 @@ def _plot_map_hp(map, lims=None, file_name=None, title=None):
     for i in range(ncomp):
         if ncomp != 1:
             f = "TQU"[i]
-        print("np.shape(map)", np.shape(np.atleast_2d(map)))
         hp.mollview(
             np.atleast_2d(map)[i],
             cmap=cmap,
@@ -431,22 +439,29 @@ def apodize_mask(mask, apod_radius_deg, apod_type, pix_type="hp"):
             apod_type
         )
     else:
-        distance = enmap.distance_transform(mask)
-        distance = np.rad2deg(distance)
-
-        mask_apo = mask.copy()
-        idx = np.where(distance > apod_radius_deg)
-
+        ## OLD - not identical to enmap.apod_mask - TODO: check
+        # distance = enmap.distance_transform(mask)
+        # distance = np.rad2deg(distance)
+
+        # mask_apo = mask.copy()
+        # idx = np.where(distance > apod_radius_deg)
+
+        # if apod_type == "C1":
+        #     mask_apo = 0.5 - 0.5 * np.cos(-np.pi * distance / apod_radius_deg)
+        # elif apod_type == "C2":
+        #     mask_apo = (
+        #         distance / apod_radius_deg -
+        #         np.sin(2 * np.pi * distance / apod_radius_deg) / (2 * np.pi)
+        #     )
+        # else:
+        #     raise ValueError(f"Unknown apodization type {apod_type}")
+        # mask_apo[idx] = 1
         if apod_type == "C1":
-            mask_apo = 0.5 - 0.5 * np.cos(-np.pi * distance / apod_radius_deg)
-        elif apod_type == "C2":
-            mask_apo = (
-                distance / apod_radius_deg -
-                np.sin(2 * np.pi * distance / apod_radius_deg) / (2 * np.pi)
+            mask_apo = enmap.apod_mask(
+                mask, width=apod_radius_deg*utils.degree
             )
         else:
-            raise ValueError(f"Unknown apodization type {apod_type}")
-        mask_apo[idx] = 1
+            raise NotImplementedError(f"Unknown apodization type {apod_type}")
 
     return mask_apo
 
@@ -526,3 +541,34 @@ def binary_mask_from_map(map, pix_type="hp", geometry=None):
     binary[hits_proxy > 0.1] = 1.
 
     return binary
+
+
+def get_spin_derivatives(map_file):
+    """
+    First and second spin derivatives of a given spin-0 map.
+    """
+    from matplotlib import cm
+
+    pix_type = _get_pix_type(map_file)
+    map = read_map(map_file, pix_type=pix_type)
+    nside, geometry = (None, None)
+    ell = np.arange(lmax_from_map(map, pix_type=pix_type) + 1)
+    alpha1i = np.sqrt(ell*(ell + 1.))
+    alpha2i = np.sqrt((ell - 1.)*ell*(ell + 1.)*(ell + 2.))
+
+    if pix_type == "car":
+        geometry = (map.shape, map.wcs)
+    else:
+        nside = hp.npix2nside(np.shape(map)[-1])
+
+    def almtomap(alm):
+        return alm2map(alm, pix_type=pix_type, nside=nside, geometry=geometry)
+    def maptoalm(map):
+        return map2alm(map, pix_type=pix_type)
+
+    first = almtomap(hp.almxfl(maptoalm(map), alpha1i))
+    second = almtomap(hp.almxfl(maptoalm(map), alpha2i))
+    cmap = cm.YlOrRd
+    cmap.set_under("w")
+
+    return first, second