diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..c580cb4 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,12 @@ +name: Tests +on: pull_request +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - run: pip install -e ".[test]" + - run: pytest diff --git a/coberus/pipeline.py b/coberus/pipeline.py index 0985549..16357df 100644 --- a/coberus/pipeline.py +++ b/coberus/pipeline.py @@ -11,7 +11,7 @@ def gauss_beam(ell, fwhm): - """ + r""" Compute Gaussian beam window function $B_\ell$ for FWHM in arcminutes. Args: @@ -46,6 +46,35 @@ def block_smooth(imap, factor, slow=False): return omap +def band_filter(imap, fl, lmax, tol=1e-8): + """ + Apply an isotropic harmonic filter with SHTs truncated to the filter's + effective band-limit. + + Equivalent to pixell.curvedsky.filter(imap, fl, lmax=lmax) up to ~tol, + but much faster for low-pass filters: both SHTs run at the largest + multipole where the filter amplitude still exceeds tol times its peak, + rather than at lmax, and modes where the filter is negligible are never + analyzed. + + Args: + imap: Input enmap. + fl: 1D harmonic filter array. + lmax: Maximum multipole of the untruncated filter operation. + tol: Relative filter amplitude below which multipoles are dropped. + + Returns: + The filtered enmap on the input geometry. + """ + fl = np.asarray(fl[: lmax + 1]) + nz = np.where(np.abs(fl) > tol * np.abs(fl).max())[0] + lcut = max(int(nz[-1]), 1) if nz.size else 1 + if lcut >= lmax: + return cs.filter(imap, fl, lmax=lmax) + alm = cs.almxfl(cs.map2alm(imap, lmax=lcut), fl[: lcut + 1]) + return cs.alm2map(alm, enmap.empty(imap.shape, imap.wcs, imap.dtype)) + + def free_mem(): return f"{psutil.virtual_memory()[1] / 1024 / 1024 / 1024:.1f} GiB" @@ -122,94 +151,61 @@ def compute_tophat_beam(w_rad, lmax, w_rad_in=None, n_theta=20000): return filt_harm -# Helper function for covariance smoothing -def cov_smooth( - wmap1, - wmap2, - i, - j, - cov_smooth_type, - cov_smooth_factor, - sigma_rad, - smooth_mean_cov, - fft_smooth, - lmax, - ells, - use_annulus, - annulus_fwhm_ratio, -): - if cov_smooth_type == "block": - cov = block_smooth(wmap1 * wmap2, cov_smooth_factor, slow=False) - - elif cov_smooth_type == "gaussian": - # Applies Gaussian smoothing procedure from 2307.01043. - - if smooth_mean_cov: - if fft_smooth: - wmap1_smooth = enmap.smooth_gauss(wmap1, sigma_rad) - - if i == j: - wmap2_smooth = wmap1_smooth - else: - wmap2_smooth = enmap.smooth_gauss(wmap2, sigma_rad) - - cov = enmap.smooth_gauss( - (wmap1 - wmap1_smooth) * (wmap2 - wmap2_smooth), - sigma_rad, - ) - - else: - gauss_beam = np.exp(-0.5 * ells * (ells + 1) * sigma_rad**2) - wmap1_smooth = cs.filter(wmap1, gauss_beam, lmax=lmax) - - if i == j: - wmap2_smooth = wmap1_smooth - else: - wmap2_smooth = cs.filter(wmap2, gauss_beam, lmax=lmax) +def cov_filter(cov_smooth_type, sigma_rad, lmax, use_annulus, annulus_fwhm_ratio): + """ + Build the harmonic filter used for SHT-based covariance smoothing. - cov = cs.filter( - (wmap1 - wmap1_smooth) * (wmap2 - wmap2_smooth), - gauss_beam, - lmax=lmax, - ) + Args: + cov_smooth_type: 'gaussian' (procedure from 2307.01043) or 'tophat' + (procedure from 2307.01258). + sigma_rad: Smoothing scale in radians. + lmax: Maximum multipole. + use_annulus: For 'tophat', whether to exclude modes below a second + inner filter scale. + annulus_fwhm_ratio: Ratio of the inner annulus scale to sigma_rad. - else: - # Compute covariance without smoothing the maps. This - # uses only one SHT/FFT, but is less stable than the - # smooth_mean_cov_approach - if fft_smooth: - cov = enmap.smooth_gauss(wmap1 * wmap2, sigma_rad) - else: - cov = cs.filter((wmap1) * (wmap2), gauss_beam, lmax=lmax) - - elif cov_smooth_type == "tophat": - # Compute beam for top-hat smoothing procedure from 2307.01258 - w_rad = sigma_rad - - if use_annulus: - w_rad_in = annulus_fwhm_ratio * w_rad - else: - w_rad_in = None + Returns: + 1D harmonic filter array up to lmax. + """ + if cov_smooth_type == "gaussian": + ells = np.arange(lmax + 1) + return np.exp(-0.5 * ells * (ells + 1) * sigma_rad**2) + w_rad_in = annulus_fwhm_ratio * sigma_rad if use_annulus else None + return compute_tophat_beam(sigma_rad, lmax, w_rad_in=w_rad_in) - tophat_beam = compute_tophat_beam(w_rad, lmax, w_rad_in=w_rad_in) - if smooth_mean_cov: - wmap1_smooth = cs.filter(wmap1, tophat_beam, lmax=lmax) +def cov_smooth( + wmap1, wmap2, cov_smooth_type, cov_smooth_factor, sigma_rad, fft_smooth, lmax, fl +): + """ + Smooth the product of two wavelet maps into an empirical covariance map. - if i == j: - wmap2_smooth = wmap1_smooth - else: - wmap2_smooth = cs.filter(wmap2, tophat_beam, lmax=lmax) + Mean subtraction (smooth_mean_cov) is handled by the caller, which passes + mean-subtracted delta maps in place of the raw wavelet maps; this function + only smooths the product. - cov = cs.filter( - (wmap1 - wmap1_smooth) * (wmap2 - wmap2_smooth), - tophat_beam, - lmax=lmax, - ) + Args: + wmap1: First wavelet enmap (or its mean-subtracted delta). + wmap2: Second wavelet enmap (or its mean-subtracted delta). + cov_smooth_type: 'block', 'gaussian' or 'tophat'. + cov_smooth_factor: Block downgrade factor for 'block' smoothing. + sigma_rad: Smoothing scale in radians (unused for 'block'). + fft_smooth: For 'gaussian', smooth with FFTs instead of SHTs. + lmax: Maximum multipole for SHT-based smoothing. + fl: Harmonic filter from cov_filter; None for the block/FFT paths. - else: - cov = cs.filter((wmap1) * (wmap2), tophat_beam, lmax=lmax) # Eq. 11 - return cov + Returns: + The smoothed covariance enmap. + """ + prod = wmap1 * wmap2 + if cov_smooth_type == "block": + return block_smooth(prod, cov_smooth_factor, slow=False) + if cov_smooth_type == "gaussian" and fft_smooth: + # Gaussian smoothing procedure from 2307.01043 using FFTs. + return enmap.smooth_gauss(prod, sigma_rad) + # SHT-based gaussian (2307.01043) or tophat (2307.01258, Eq. 11) + # smoothing, with the SHTs truncated to the filter's band-limit. + return band_filter(prod, fl, lmax) def needlet_coadd( @@ -392,6 +388,10 @@ def needlet_coadd( f"Duplicate keys would be produced in output dictionary: {_output_keys}. " "Check nmap_labels for collisions with 'coadd' or 'mask'." ) + # These would collide with internal wavelet_{map,mask,delta}_* file names + _reserved = {"map", "mask", "delta"} & set(nmap_labels) + if _reserved: + raise ValueError(f"nmap_labels uses reserved label(s): {sorted(_reserved)}") start_time = time.time() lmax = max(lpeaks) # Cosine needlets have zero support beyond lpeak @@ -613,36 +613,60 @@ def _get_wave(fname_func, itag, imask): itags.append(tag) included_tags[k] = list(itags) + if cov_smooth_type == "block": + sigma_rad = None + else: + sigma_rad = cov_smooth_scales[k] + + # Build the harmonic smoothing filter once per scale (the tophat + # beam in particular is expensive to compute). fl is None for the + # paths that do not use SHT filtering. + if cov_smooth_type == "block" or (cov_smooth_type == "gaussian" and fft_smooth): + fl = None + else: + fl = cov_filter( + cov_smooth_type, sigma_rad, lmax, use_annulus, annulus_fwhm_ratio + ) + + # For mean-subtracted covariances, smooth each tag's wavelet map once + # here rather than once per pair inside cov_smooth; smoothing the pair + # products of the deltas then gives the same covariance. + hoist = cov_smooth_type != "block" and smooth_mean_cov + dfnames = [] + if hoist: + for tag in itags: + wmap = enmap.read_map(f"{out_root}wavelet_map_{tag}_scale_{k}.fits") + if fl is None: + smap = enmap.smooth_gauss(wmap, sigma_rad) + else: + smap = band_filter(wmap, fl, lmax) + dfname = f"{out_root}wavelet_delta_{tag}_scale_{k}.fits" + enmap.write_map(dfname, wmap - smap) + dfnames.append(dfname) + prefix = "delta" if hoist else "map" + for i in range(len(itags)): - wmap1 = enmap.read_map(f"{out_root}wavelet_map_{itags[i]}_scale_{k}.fits") + wmap1 = enmap.read_map( + f"{out_root}wavelet_{prefix}_{itags[i]}_scale_{k}.fits" + ) for j in range(i, len(itags)): if i == j: # Potentially minor speedup? wmap2 = wmap1 else: wmap2 = enmap.read_map( - f"{out_root}wavelet_map_{itags[j]}_scale_{k}.fits" + f"{out_root}wavelet_{prefix}_{itags[j]}_scale_{k}.fits" ) - if cov_smooth_type == "block": - sigma_rad = None - else: - sigma_rad = cov_smooth_scales[k] - cov = cov_smooth( wmap1, wmap2, - i, - j, cov_smooth_type, cov_smooth_factor, sigma_rad, - smooth_mean_cov, fft_smooth, lmax, - ells, - use_annulus, - annulus_fwhm_ratio, + fl, ) fcovname = f"{out_root}wavelet_cov_scale_{k}_{itags[i]}_{itags[j]}.fits" @@ -652,6 +676,11 @@ def _get_wave(fname_func, itag, imask): filenames.append(fcovname) totgibytes = totgibytes + (cov.nbytes / 1024 / 1024.0 / 1024.0) + # The delta maps are only needed within this scale, so free the + # (RAM)disk space immediately rather than at the end of the run. + for dfname in dfnames: + os.remove(dfname) + elapsed_time = time.time() - start_time_covariance print(f"Covariance finished in {elapsed_time / 60.0:.2f} minutes.") diff --git a/pyproject.toml b/pyproject.toml index 1a6cabf..1a931e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,10 @@ dev = [ "ruff", "uv" ] +test = [ + "pytest", + "camb", +] [project.urls] Homepage = "https://github.com/simonsobs/map-coaddition/" diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..193666f --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,498 @@ +""" +End-to-end test of coberus.pipeline.needlet_coadd on a small synthetic +ACT+Planck-like dataset. + +Generates apodized maps on overlapping but unequal footprints (wide +Planck-like band, narrower ACT-like band with apodized source holes, and a +small day-time sub-band) with a realistic lensed CMB from a low-accuracy +CAMB call, per-tag Gaussian beams and white noise. The maps are written to +a RAMdisk, coadded with 'block', 'gaussian' (SHT) and 'tophat' (SHT) +covariance smoothing, and compared. + +Example:: + + python scripts/test_needlet_coadd.py --debug + python scripts/test_needlet_coadd.py # full run +""" + +import argparse +import os +import shutil +import time + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +from pixell import curvedsky as cs, enmap, enplot +from scipy.stats import binned_statistic as binnedstat + +from coberus.pipeline import gauss_beam, needlet_coadd + +TAGS = [ + # tag, fwhm_arcmin, noise_uK_arcmin, footprint + ("planck_143", 7.3, 33.0, "wide"), + ("planck_217", 5.0, 47.0, "wide"), + ("act_f090", 2.0, 18.0, "act"), + ("act_f150", 1.4, 17.0, "act_narrow"), + ("act_f150_day", 1.4, 25.0, "patch"), +] +BASE_TAG = "planck_143" +# Footprints with very different sky areas: (dec range, optional RA range) +FOOTPRINTS = { + "wide": ((-75.0, 30.0), None), + "act": ((-60.0, 20.0), None), + "act_narrow": ((-40.0, 8.0), None), + "patch": ((-25.0, 8.0), (20.0, 100.0)), +} +# (dec, ra) source cuts +HOLES_DEG = [(-45.0, 50.0), (-20.0, 150.0), (0.0, 280.0), (-10.0, 60.0)] +CACHE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "cache") + + +# Copied from msyriac/orphics +class bin1D: + ''' + * Takes data defined on x0 and produces values binned on x. + * Assumes x0 is linearly spaced and continuous in a domain? + * Assumes x is continuous in a subdomain of x0. + * Should handle NaNs correctly. + ''' + + + def __init__(self, bin_edges): + + self.update_bin_edges(bin_edges) + + + def update_bin_edges(self,bin_edges): + + self.bin_edges = bin_edges + self.numbins = len(bin_edges)-1 + self.cents = (self.bin_edges[:-1]+self.bin_edges[1:])/2. + + self.bin_edges_min = self.bin_edges.min() + self.bin_edges_max = self.bin_edges.max() + + def bin(self,ix,iy,stat=np.nanmean): + x = ix.copy() + y = iy.copy() + # this just prevents an annoying warning (which is otherwise informative) everytime + # all the values outside the bin_edges are nans + y[xself.bin_edges_max] = 0 + + bin_means = binnedstat(x,y,bins=self.bin_edges,statistic=stat)[0] + + return self.cents,bin_means + +# Copied from msyriac/orphics +def white_noise(shape=None,wcs=None,noise_muK_arcmin=None,seed=None): + """ + Generate a non-band-limited white noise map. + """ + if seed is not None: np.random.seed(seed) + ipsizemap = enmap.pixsizemap(shape,wcs) + pmap = ipsizemap*((180.*60./np.pi)**2.) + div = pmap/noise_muK_arcmin**2. + return np.random.standard_normal(shape) / np.sqrt(div) + + +def get_camb_cl(lmax, cache_dir=CACHE_DIR): + """ + Return the lensed CMB TT power spectrum C_ell in uK^2 up to lmax. + + Uses a low-accuracy CAMB call with fiducial LCDM parameters. The result + has no free parameters besides lmax, so it is cached to disk. + + Args: + lmax: Maximum multipole. + cache_dir: Directory for the cached spectrum. + + Returns: + 1D array of C_ell from ell=0 to lmax. + """ + os.makedirs(cache_dir, exist_ok=True) + fname = os.path.join(cache_dir, f"camb_cltt_lmax_{lmax}.npy") + if os.path.exists(fname): + return np.load(fname) + import camb + + pars = camb.set_params( + H0=67.5, + ombh2=0.022, + omch2=0.122, + As=2.1e-9, + ns=0.965, + lmax=lmax + 500, + lens_potential_accuracy=0, + ) + pars.set_accuracy(AccuracyBoost=0.5, lAccuracyBoost=0.5) + results = camb.get_results(pars) + cl = results.get_cmb_power_spectra(pars, CMB_unit="muK", raw_cl=True)["total"][ + : lmax + 1, 0 + ] + np.save(fname, cl) + return cl + + +def cos_taper(x): + """Cosine taper rising from 0 at x<=0 to 1 at x>=1 for array x.""" + return 0.5 - 0.5 * np.cos(np.pi * np.clip(x, 0, 1)) + + +def make_geometry(dec_range_deg, ra_range_deg, res): + """ + Build a CAR geometry for a dec band, optionally cropped in RA. + + The band is sliced from the global full-sky pixelization (and the RA + crop is pixel-snapped from it), so all footprints are pixel-compatible + with each other. + + Args: + dec_range_deg: (dec_lo, dec_hi) of the footprint in degrees. + ra_range_deg: Optional (ra_lo, ra_hi) in degrees; None keeps the + full RA circle. + res: Pixel resolution in radians. + + Returns: + (shape, wcs) tuple. + """ + shape, wcs = enmap.band_geometry(np.deg2rad(dec_range_deg), res=res) + if ra_range_deg is not None: + # RA ordered high -> low to match the native decreasing-RA axis + box = np.deg2rad( + [ + [dec_range_deg[0], ra_range_deg[1]], + [dec_range_deg[1], ra_range_deg[0]], + ] + ) + shape, wcs = enmap.subgeo(shape, wcs, box=box) + return shape, wcs + + +def make_apod( + shape, + wcs, + dec_range_deg, + taper_deg, + ra_range_deg=None, + holes_deg=None, + hole_deg=2.0, +): + """ + Build an apodization map for a footprint with optional holes. + + Args: + shape, wcs: Geometry of the map. + dec_range_deg: (dec_lo, dec_hi) of the footprint in degrees. + taper_deg: Width of the cosine taper at the footprint edges in + degrees. + ra_range_deg: Optional (ra_lo, ra_hi) in degrees for RA-cropped + patches; adds a taper at the RA edges. + holes_deg: Optional list of (dec, ra) hole centers in degrees. + hole_deg: Hole radius in degrees; the taper extends over another + hole_deg beyond the radius. + + Returns: + Apodization enmap in [0, 1]. + """ + dec, ra = enmap.posmap(shape, wcs) + lo, hi = np.deg2rad(dec_range_deg) + w = np.deg2rad(taper_deg) + apod = cos_taper(np.minimum(dec - lo, hi - dec) / w) + if ra_range_deg is not None: + rlo, rhi = np.deg2rad(ra_range_deg) + apod = apod * cos_taper(np.minimum(ra - rlo, rhi - ra) / w) + if holes_deg: + pts = np.deg2rad(np.asarray(holes_deg)).T + r = enmap.distance_from(shape, wcs, pts) + apod = apod * cos_taper((r - np.deg2rad(hole_deg)) / np.deg2rad(hole_deg)) + return enmap.enmap(apod, wcs) + + +def make_dataset(dset_dir, res_arcmin, lmax, seed): + """ + Generate the synthetic dataset and write maps and binary masks to disk. + + Each tag's map is a common CMB realization convolved with the tag's + beam, plus independent white noise, multiplied by an apodization over + the tag's footprint. The binary mask keeps apod > 0.99. + + Args: + dset_dir: Output directory (ideally on a RAMdisk). + res_arcmin: Pixel resolution in arcminutes. + lmax: Maximum multipole for the CMB realization. + seed: Master RNG seed. + + Returns: + Dict mapping tag to {'map': path, 'mask': path, 'fwhm': fwhm}. + """ + os.makedirs(dset_dir, exist_ok=True) + cl = get_camb_cl(lmax) + alm = cs.rand_alm(cl, lmax=lmax, seed=seed) + res = np.deg2rad(res_arcmin / 60.0) + ells = np.arange(lmax + 1) + info = {} + for i, (tag, fwhm, noise, foot) in enumerate(TAGS): + dec_range, ra_range = FOOTPRINTS[foot] + shape, wcs = make_geometry(dec_range, ra_range, res) + omap = cs.alm2map( + cs.almxfl(alm.copy(), gauss_beam(ells, fwhm)), enmap.empty(shape, wcs) + ) + rng = np.random.default_rng(seed + 100 + i) + omap = omap + white_noise(shape,wcs,noise) + holes = HOLES_DEG if foot != "wide" else None + apod = make_apod( + shape, + wcs, + dec_range, + taper_deg=3.0, + ra_range_deg=ra_range, + holes_deg=holes, + ) + mfile = os.path.join(dset_dir, f"map_{tag}.fits") + kfile = os.path.join(dset_dir, f"mask_{tag}.fits") + enmap.write_map(mfile, omap * apod) + enmap.write_map(kfile, (apod > 0.99).astype(np.float64)) + info[tag] = {"map": mfile, "mask": kfile, "fwhm": fwhm} + print(f"Wrote {tag}: shape={shape}, fwhm={fwhm}', noise={noise} uK-arcmin") + return info + + +def eplot(imap, fname, downgrade, **kwargs): + """Render an enmap with pixell.enplot and write it to fname (PNG).""" + p = enplot.plot(imap, downgrade=downgrade, colorbar=True, ticks=20, **kwargs) + enplot.write(fname, p) + + +def binned_coadd_cl(imap, mask, lmax, bin_edges, out_beam_fwhm): + """ + Compute the binned, beam-deconvolved power spectrum of a masked coadd. + + Uses w2 = int(mask^2)dA / 4pi normalization for the binary + footprint mask. + + Args: + imap: Coadded enmap (already zeroed outside mask). + mask: Binary footprint enmap. + lmax: Maximum multipole. + bin_edges: Multipole bin edges. + out_beam_fwhm: Output beam FWHM in arcminutes to deconvolve. + + Returns: + (bin centers, binned C_ell) tuple. + """ + w2 = np.sum(mask**2 * mask.pixsizemap()) / (4 * np.pi) + cl = cs.alm2cl(cs.map2alm(imap, lmax=lmax)) / w2 + cl /= gauss_beam(np.arange(cl.size), out_beam_fwhm) ** 2 + return bin1D(bin_edges).bin(np.arange(cl.size), cl) + + +def plot_spectra(outmaps, cl_theory, lmax, out_beam_fwhm, fname): + """ + Plot theory vs measured coadd D_ell for each smoothing run. + + The top panel shows the theory curve, binned theory and binned coadd + bandpowers; the bottom panel shows the ratio of each coadd to the + binned theory, where the deficit from unity reflects the empirical-ILC + bias of each covariance smoothing scheme. + + Args: + outmaps: Dict mapping run name to needlet_coadd output dict. + cl_theory: Theory C_ell used to generate the CMB. + lmax: Maximum multipole. + out_beam_fwhm: Output beam FWHM in arcminutes. + fname: Output PNG path. + """ + ells = np.arange(lmax + 1) + dfact = ells * (ells + 1) / (2 * np.pi) + bin_edges = np.arange(40, lmax - 10, 40) + cents, clt = bin1D(bin_edges).bin(ells, cl_theory) + bfact = cents * (cents + 1) / (2 * np.pi) + fig, (ax1, ax2) = plt.subplots( + 2, 1, figsize=(8, 7), sharex=True, height_ratios=[2, 1] + ) + ax1.plot(ells, dfact * cl_theory, "k-", lw=1, label="CAMB theory") + ax1.plot(cents, bfact * clt, "k_", ms=10, label="binned theory") + for run, marker in zip(outmaps, ["o", "s", "^"]): + _, cl = binned_coadd_cl( + outmaps[run]["coadd"], + outmaps[run]["mask"], + lmax, + bin_edges, + out_beam_fwhm, + ) + ax1.plot(cents, bfact * cl, marker, ms=4, label=f"coadd ({run})") + ax2.plot(cents, cl / clt, marker, ms=4) + ax1.set_ylabel(r"$D_\ell$ [$\mu$K$^2$]") + ax1.legend() + ax2.axhline(1, color="k", lw=1) + ax2.set_xlabel(r"$\ell$") + ax2.set_ylabel(r"$C_\ell^{\rm coadd} / C_\ell^{\rm theory}$") + plt.tight_layout() + plt.savefig(fname, dpi=120) + plt.close() + + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output", + default="test_output", + help="Directory for plots.", + ) + parser.add_argument( + "--out-root", + default=None, + help="Prefix for the dataset and intermediates. Defaults to a " + "RAMdisk (/dev/shm/) when it has enough free space, else /tmp/.", + ) + parser.add_argument( + "--res-arcmin", type=float, default=4.0, help="Pixel resolution in arcminutes." + ) + parser.add_argument("--lmax", type=int, default=500, help="Maximum multipole.") + parser.add_argument( + "--lpeaks", + type=int, + nargs="+", + default=None, + help="Cosine needlet lpeaks (default: a few scales).", + ) + parser.add_argument( + "--out-beam-fwhm", + type=float, + default=8.0, + help="Output beam FWHM in arcminutes.", + ) + parser.add_argument( + "--ilc-bias-tol", + type=float, + default=0.01, + help="ILC bias tolerance for gaussian smoothing.", + ) + parser.add_argument( + "--cov-smooth-factor", + type=int, + default=32, + help="Block downgrade factor for block smoothing.", + ) + parser.add_argument( + "--workers", type=int, default=4, help="Number of dask workers." + ) + parser.add_argument("--seed", type=int, default=1234, help="RNG seed.") + parser.add_argument( + "--debug", + action="store_true", + help="Coarser resolution and lower lmax for quick tests.", + ) + args = parser.parse_args(argv) + + if args.out_root is None: + # Prefer a RAMdisk, but fall back to /tmp on small /dev/shm mounts + free_gib = shutil.disk_usage("/dev/shm").free / 1024**3 + args.out_root = "/dev/shm/cob_test_" if free_gib > 4 else "/tmp/cob_test_" + print(f"out_root = {args.out_root} (/dev/shm has {free_gib:.1f} GiB free)") + + if args.debug: + args.res_arcmin = max(args.res_arcmin, 8.0) + args.lmax = min(args.lmax, 250) + args.cov_smooth_factor = min(args.cov_smooth_factor, 8) + args.workers = min(args.workers, 2) + if args.lpeaks is None: + args.lpeaks = [lp for lp in [50, 100, 200, 350] if lp < args.lmax] + args.lpeaks.append(args.lmax) + print(f"lpeaks = {args.lpeaks}") + + os.makedirs(args.output, exist_ok=True) + tags = [t[0] for t in TAGS] + info = make_dataset( + f"{args.out_root}dataset", args.res_arcmin, args.lmax, args.seed + ) + + figs = [] + for tag in [BASE_TAG, "act_f090", "act_f150_day"]: + png = f"input_{tag}" + eplot( + enmap.read_map(info[tag]["map"]), + os.path.join(args.output, png), + downgrade=2, + range=400, + ) + figs.append((f"Input map: {tag}", f"{png}.png")) + + timings, outmaps = {}, {} + for run in ["block", "gaussian", "tophat"]: + print(f"=== Running needlet_coadd with {run} covariance smoothing ===") + t0 = time.time() + outmaps[run] = needlet_coadd( + map_fname_func=lambda tag: info[tag]["map"], + mask_fname_func=lambda tag: info[tag]["mask"], + tags=tags, + base_tag=BASE_TAG, + lpeaks=args.lpeaks, + lmins=[None] * len(tags), + lmaxs=[None] * len(tags), + response_func=lambda tag: 1.0, + beam_func=lambda tag, ells: gauss_beam(np.asarray(ells), info[tag]["fwhm"]), + out_beam_fwhm=args.out_beam_fwhm, + out_root=f"{args.out_root}{run}_", + cov_smooth_type=run, + cov_smooth_factor=args.cov_smooth_factor, + ilc_bias_tol=args.ilc_bias_tol, + n_workers=args.workers, + delete_intermediate=True, + ) + timings[run] = time.time() - t0 + print(f"{run} run took {timings[run] / 60.0:.2f} min") + png = f"coadd_{run}" + eplot( + outmaps[run]["coadd"], + os.path.join(args.output, png), + downgrade=2, + range=400, + ) + figs.append((f"Coadd: {run} smoothing", f"{png}.png")) + + eplot( + outmaps["block"]["mask"], + os.path.join(args.output, "mask"), + downgrade=2, + min=0, + max=1, + ) + figs.append(("Final footprint mask", "mask.png")) + for run in ["gaussian", "tophat"]: + png = f"diff_{run}_block" + eplot( + outmaps[run]["coadd"] - outmaps["block"]["coadd"], + os.path.join(args.output, png), + downgrade=2, + range=50, + ) + figs.append((f"Difference: {run} - block", f"{png}.png")) + + plot_spectra( + outmaps, + get_camb_cl(args.lmax), + args.lmax, + args.out_beam_fwhm, + os.path.join(args.output, "spectra.png"), + ) + figs.append(("Coadd power spectra vs theory", "spectra.png")) + return outmaps[run]["coadd"] + + + +def test_main(): + + x = main([ + "--debug" + ]) + assert x is not None + assert(x.wcs is not None) + +if __name__ == "__main__": + main()