diff --git a/pipeline/misc/validate_mask_isohits.py b/pipeline/misc/validate_mask_isohits.py new file mode 100644 index 0000000..7685ce4 --- /dev/null +++ b/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 .. Tag may be empty (''). + """ + ext = ext.lstrip('.') + pattern = re.compile(rf'^{base}(?P(?:_[^.]*)?)\.{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}.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) diff --git a/soopercool/map_utils.py b/soopercool/map_utils.py index 1e7cfe9..ee19642 100644 --- a/soopercool/map_utils.py +++ b/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 \ No newline at end of file diff --git a/soopercool/utils.py b/soopercool/utils.py index 13dfa13..598d7f8 100644 --- a/soopercool/utils.py +++ b/soopercool/utils.py @@ -500,38 +500,49 @@ def bin_validation_power_spectra(cls_dict, nmt_binning, return cls_binned_dict -def plot_transfer_function(lb, tf_dict, lmin, lmax, field_pairs, file_name): +def plot_transfer_function(lb, tf_dict, lmin, lmax, field_pairs, + file_name=None): """ Plot the transfer function given an input dictionary. """ - plt.figure(figsize=(25, 25)) - grid = plt.GridSpec(9, 9, hspace=0.3, wspace=0.3) - - for id1, f1 in enumerate(field_pairs): - for id2, f2 in enumerate(field_pairs): - ax = plt.subplot(grid[id1, id2]) - expected = 1. if f1 == f2 else 0. - ylims = [0, 1.05] if f1 == f2 else [-0.01, 0.01] - - ax.axhline(expected, color="k", ls="--", zorder=6) - # We need to understand the offdigonal TF panels in the presence of - # NaMaster purification - we don't have a clear interpretation. - ax.set_title(f"{f1} $\\rightarrow$ {f2}", fontsize=14) - ax.plot(lb, tf_dict[f"{f1}_to_{f2}"], color="navy") - - if id1 == 8: - ax.set_xlabel(r"$\ell$", fontsize=14) - else: - ax.set_xticks([]) - - if f1 != f2: - ax.ticklabel_format(axis="y", style="scientific", - scilimits=(0, 0), useMathText=True) - - ax.set_xlim(lmin, lmax) - ax.set_ylim(ylims[0], ylims[1]) - - plt.savefig(file_name, bbox_inches="tight") + nfield = len(field_pairs) + plt.figure(figsize=(25*nfield/9., 25*nfield/9.)) + grid = plt.GridSpec(nfield, nfield, hspace=0.3, wspace=0.3) + + for label, tf in tf_dict.items(): + for id1, f1 in enumerate(field_pairs): + for id2, f2 in enumerate(field_pairs): + ax = plt.subplot(grid[id1, id2]) + expected = 1. if f1 == f2 else 0. + ylims = [0, 1.05] if f1 == f2 else [-0.01, 0.01] + + ax.axhline(expected, color="k", ls="--", zorder=6) + # We need to understand the offdigonal TF panels in the presence of + # NaMaster purification - we don't have a clear interpretation. + ax.set_title(f"{f1} $\\rightarrow$ {f2}", fontsize=14) + ax.plot(lb, tf[f"{f1}_to_{f2}"], label=label) + + if id2 == 0: + ax.set_ylabel(r"$T_\ell$", fontsize=14) + if id1 == nfield-1: + ax.set_xlabel(r"$\ell$", fontsize=14) + else: + ax.set_xticks([]) + + if f1 != f2: + ax.ticklabel_format(axis="y", style="scientific", + scilimits=(0, 0), useMathText=True) + + ax.set_xlim(lmin, lmax) + ax.set_ylim(ylims[0], ylims[1]) + if label is not None and id1 == 0 and id2 == nfield-1: + ax.legend(fontsize=10, bbox_to_anchor=[1,1], + loc="upper left") + + if file_name is not None: + plt.savefig(file_name, bbox_inches="tight") + else: + return def get_binary_mask_from_nhits(nhits_map, nside, zero_threshold=1e-3):