From 51bf89389f9d26424620054d0c88051e6d04d55d Mon Sep 17 00:00:00 2001
From: Joshua Kim <40779225+jaejoonk@users.noreply.github.com>
Date: Tue, 29 Sep 2026 16:13:57 -0700
Subject: [PATCH 1/6] Write per-scale NILC weight maps and add cmb_weights nmap
 label

- coadd_maps_pixels now also returns the per-pixel ILC weights, and
  coadd(..., return_weights=True) returns them alongside the map.
- needlet_coadd writes the weights estimated for 'coadd' to
  {out_root}wavelet_weights_scale_{k}_{tag}.fits (removed with
  delete_intermediate like the covariance maps).
- New 'cmb_weights' nmap label applies the weights to ones maps and
  returns the per-scale weight sums (no wave2map) in
  outmaps['cmb_weights_coadd'].

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 coberus/core.py     | 37 ++++++++++++++++++++++++++--------
 coberus/pipeline.py | 49 ++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 73 insertions(+), 13 deletions(-)

diff --git a/coberus/core.py b/coberus/core.py
index 2360a59..45e5497 100644
--- a/coberus/core.py
+++ b/coberus/core.py
@@ -181,9 +181,12 @@ def coadd_maps_pixels(
     masks: np.ndarray,
     responses: np.ndarray,
     deproj_responses: np.ndarray,
-) -> np.ndarray:
+) -> tuple[np.ndarray, np.ndarray]:
     """
     Co-adds the maps in the pixel domain. Assumes all masks and maps are the same size.
+
+    Returns the coadded map (n_y, n_x) and the ILC weights (n_map, n_y, n_x)
+    applied to each map. Weights are zero for maps masked out of a pixel.
     """
 
     n_y, n_x = maps.shape[-2:]
@@ -191,6 +194,7 @@ def coadd_maps_pixels(
 
     n_deproj = len(deproj_responses)
     n_map = len(maps)
+    weights_out = np.zeros((n_map, n_y, n_x), dtype=np.float32)
 
     if n_deproj > 0:
         response_mat = (
@@ -230,6 +234,8 @@ def coadd_maps_pixels(
                 numer = np.dot(a, cinvd)
 
                 output[j, i] = numer / denom
+                # a^T C^-1 / denom, using the symmetry of C
+                weights_out[mask, j, i] = cinva / denom
 
             # Constrained ILC (see Eq. 29 & 30 of 2307.01043)
             else:
@@ -250,8 +256,9 @@ def coadd_maps_pixels(
                 a_eff = (np.dot(det_q_sub_vec, a_mix.T) / det_q).astype(np.float32)
                 weights = np.linalg.solve(cov, a_eff)
                 output[j, i] = np.dot(weights, masked_maps)
+                weights_out[mask, j, i] = weights
 
-    return output
+    return output, weights_out
 
 
 def write_to_main_array(data: np.ndarray, chunk: Chunk, main_array: da.array):
@@ -274,12 +281,13 @@ def coadded_map_wrapper(
 ) -> np.ndarray:
     """
     Wrapper for the coadd_maps_pixels function that allows you to return
-    the chunk to.
+    the chunk to. Returns (coadded map, weights, chunk).
     """
 
-    return coadd_maps_pixels(
+    output, weights = coadd_maps_pixels(
         maps, covariance_maps, masks, responses, deproj_responses
-    ), chunk
+    )
+    return output, weights, chunk
 
 
 def create_tasks_for_chunk(
@@ -308,23 +316,36 @@ def create_tasks_for_chunk(
     return coadded_map
 
 
-def coadd(client: Client, coadder: Coadder) -> da.Array:
+def coadd(client: Client, coadder: Coadder, return_weights: bool = False):
     """
     The main function for co-adding maps. Takes your Coadder object, and
     a Dask client, and returns a Dask array filled with your coadded map.
     This function uses Dask futures to coadd your maps.
+
+    If return_weights is True, also returns a Dask array of shape
+    (n_map, n_y, n_x) holding the ILC weight applied to each map in each
+    pixel (zero where a map is masked out).
     """
     chunks = coadder.chunk_task_list()
 
-    main_array = da.zeros(coadder.chunk_meta()["meta"].shape, dtype=np.float32)
+    shape = coadder.chunk_meta()["meta"].shape
+    main_array = da.zeros(shape, dtype=np.float32)
+    if return_weights:
+        weights_array = da.zeros((len(coadder.maps),) + shape, dtype=np.float32)
 
     results = [
         create_tasks_for_chunk(chunk, client, coadder, main_array) for chunk in chunks
     ]
 
     for future in dask.distributed.as_completed(results):
-        image, chunk = future.result()
+        image, weights, chunk = future.result()
 
         write_to_main_array(image, chunk, main_array)
+        if return_weights:
+            weights_array[:, chunk[0][0] : chunk[1][0], chunk[0][1] : chunk[1][1]] = (
+                weights
+            )
 
+    if return_weights:
+        return main_array, weights_array
     return main_array
diff --git a/coberus/pipeline.py b/coberus/pipeline.py
index 9221dad..d0cc486 100644
--- a/coberus/pipeline.py
+++ b/coberus/pipeline.py
@@ -362,12 +362,18 @@ def needlet_coadd(
         Whether to delete intermediate outputs
 
     nmap_labels : optional,list[str]
-        List of possible optional maps' labels
+        List of possible optional maps' labels. The special label
+        'cmb_weights' does not read any maps; instead the weights are
+        applied to wavelet maps that are 1 in all pixels, and outmaps
+        ['cmb_weights_coadd'] is a list (one per wavelet scale) of the
+        summed weights on each scale's wavelet geometry, without the
+        final wave2map. For a CMB solution (responses of 1) with no
+        deprojection this should be 1 wherever any map contributes.
 
     nmap_label_fname_func : optional,func | (nmap_label, fname) -> nmap
         Optional maps not used for covariance, but coadded with
         the same weights. Accepts the optional map's label and filename
-        and returns a path
+        and returns a path. Not called for the 'cmb_weights' label.
 
     apply_mask : optional, boolean
         If true, zero out regions of the input maps based on their masks.
@@ -385,7 +391,12 @@ def needlet_coadd(
           'mask'  : the final footprint mask (the base_tag mask, projected
                     onto the output geometry when one is provided).
         For each label in nmap_labels, also contains '{label}_coadd' with the
-        coadd of those maps using the same weights.
+        coadd of those maps using the same weights ('cmb_weights_coadd' is
+        instead a per-scale list; see nmap_labels).
+
+    The per-scale ILC weight maps estimated for 'coadd' are also written
+    to {out_root}wavelet_weights_scale_{k}_{tag}.fits for each wavelet
+    scale k and each tag included in that scale.
 
     """
     if nmap_labels is None:
@@ -441,7 +452,7 @@ def needlet_coadd(
     )
 
     # if using optional additional maps to coadd
-    do_nmaps = len(nmap_labels) > 0 and (nmap_label_fname_func is not None)
+    do_nmaps = len(nmap_labels) > 0
 
     def _get_wave(fname_func, itag, imask):
         gmap = enmap.read_map(fname_func(itag))
@@ -572,6 +583,7 @@ def _get_wave(fname_func, itag, imask):
                 lambda fname: nmap_label_fname_func(label, fname), tag, mask
             )
             for label in nmap_labels
+            if label != "cmb_weights"
         }
 
         if i == 0:
@@ -600,6 +612,16 @@ def _get_wave(fname_func, itag, imask):
 
             # optional maps
             for label in nmap_labels:
+                if label == "cmb_weights":
+                    # Maps of ones are identical for every tag, so write
+                    # one per scale and share it between tags.
+                    nwfname = f"{out_root}wavelet_{label}_scale_{j}.fits"
+                    if len(nfmaps[label][j]) == 0:
+                        filenames.append(nwfname)
+                        enmap.write_map(nwfname, enmap.ones(wmap.shape, wmap.wcs))
+                        totgibytes = totgibytes + (wmap.nbytes / 1024 / 1024.0 / 1024.0)
+                    nfmaps[label][j].append(nwfname)
+                    continue
                 nwfname = f"{out_root}wavelet_{label}_{tags[i]}_scale_{j}.fits"
                 filenames.append(nwfname)
                 nfmaps[label][j].append(nwfname)
@@ -742,11 +764,28 @@ def _get_wave(fname_func, itag, imask):
 
                 print("Number of workers: ", len(client.scheduler_info()["workers"]))
                 # Result is a dask array
-                result = coadd(client, coadder)
+                if outmaptype == "coadd":
+                    result, weights = coadd(client, coadder, return_weights=True)
+                    # Write the ILC weight map of each tag at this scale
+                    weights = weights.compute()
+                    for i, tag in enumerate(included_tags[j]):
+                        fwname = f"{out_root}wavelet_weights_scale_{j}_{tag}.fits"
+                        enmap.write_map(
+                            fwname, enmap.enmap(weights[i], owave.maps[j].wcs)
+                        )
+                        filenames.append(fwname)
+                else:
+                    result = coadd(client, coadder)
                 # This is now a numpy array
                 arr = result.compute()
                 owave.maps[j] = enmap.enmap(arr.copy(), owave.maps[j].wcs)
 
+            if outmaptype == "cmb_weights_coadd":
+                # Synthesizing constant maps is not meaningful, so keep the
+                # per-scale weight sums on the wavelet geometries.
+                outmaps[outmaptype] = [m.copy() for m in owave.maps]
+                continue
+
             coadd_map = wt_out.wave2map(owave)
             coadd_map[out_base_mask == 0] = 0
             outmaps[outmaptype] = coadd_map.copy()

From 7216d0039b75f313970c4fe06958b07e83f13175 Mon Sep 17 00:00:00 2001
From: Joshua Kim <40779225+jaejoonk@users.noreply.github.com>
Date: Tue, 29 Sep 2026 16:16:10 -0700
Subject: [PATCH 2/6] Add check script for cmb_weights and per-scale weight
 maps

Runs a block-smoothed CMB NILC on the synthetic test_pipeline dataset
with nmap_labels=["cmb_weights"] and prints, per wavelet scale, the
range of cmb_weights_coadd and the sum over tags of the written weight
maps (both should be 1 wherever any map contributes).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/check_cmb_weights.py | 66 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)
 create mode 100644 tests/check_cmb_weights.py

diff --git a/tests/check_cmb_weights.py b/tests/check_cmb_weights.py
new file mode 100644
index 0000000..b55fbec
--- /dev/null
+++ b/tests/check_cmb_weights.py
@@ -0,0 +1,66 @@
+"""
+Check the 'cmb_weights' nmap label and the per-scale weight maps written by
+coberus.pipeline.needlet_coadd, using the synthetic dataset from
+tests/test_pipeline.py.
+"""
+
+import argparse
+import glob
+
+import numpy as np
+from pixell import enmap
+from test_pipeline import BASE_TAG, TAGS, make_dataset
+
+from coberus.pipeline import gauss_beam, needlet_coadd
+
+
+def main():
+    """Run a block-smoothed CMB NILC with 'cmb_weights' and check outputs."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--out-root", default="/tmp/cob_chk_")
+    parser.add_argument("--debug", action="store_true", help="Low-res quick run.")
+    args = parser.parse_args()
+    res, lmax = (8.0, 250) if args.debug else (4.0, 500)
+    lpeaks = [lp for lp in [50, 100, 200, 350] if lp < lmax] + [lmax]
+
+    tags = [t[0] for t in TAGS]
+    info = make_dataset(f"{args.out_root}dataset", res, lmax, 1234)
+    root = f"{args.out_root}run_"
+    out = 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=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=8.0,
+        out_root=root,
+        cov_smooth_type="block",
+        cov_smooth_factor=8,
+        n_workers=2,
+        nmap_labels=["cmb_weights"],
+    )
+    print("keys:", sorted(out))
+    for k, w in enumerate(out["cmb_weights_coadd"]):
+        nz = w != 0
+        print(
+            f"cmb_weights_coadd scale {k}: shape {w.shape}, "
+            f"nonzero range {w[nz].min():.5f}..{w[nz].max():.5f}, "
+            f"zero fraction {1 - nz.mean():.3f}"
+        )
+    for k in range(len(lpeaks)):
+        files = sorted(glob.glob(f"{root}wavelet_weights_scale_{k}_*.fits"))
+        ws = np.array([enmap.read_map(f) for f in files])
+        tot = ws.sum(axis=0)
+        nz = tot != 0
+        print(
+            f"scale {k}: {len(files)} weight maps, shape {ws.shape[1:]}, "
+            f"sum over tags where nonzero: {tot[nz].min():.5f}..{tot[nz].max():.5f}"
+        )
+
+
+if __name__ == "__main__":
+    main()

From 3dc8edbeb75f8e14c2995dd6be0b18ccb398ca2b Mon Sep 17 00:00:00 2001
From: Joshua Kim <40779225+jaejoonk@users.noreply.github.com>
Date: Tue, 29 Sep 2026 16:30:06 -0700
Subject: [PATCH 3/6] Replace cmb_weights nmap label with check_cmb_weights
 flag

needlet_coadd no longer accepts 'cmb_weights' in nmap_labels. Per-scale
ILC weight maps are still written to disk; with check_cmb_weights=True,
each scale's weights are checked to sum to 1 in every covered pixel
(CMB solution), raising ValueError otherwise.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 README.md                  |  1 +
 coberus/pipeline.py        | 77 +++++++++++++++++++++++++-------------
 tests/check_cmb_weights.py | 13 ++-----
 3 files changed, 54 insertions(+), 37 deletions(-)

diff --git a/README.md b/README.md
index 881d4b9..7da7ec4 100644
--- a/README.md
+++ b/README.md
@@ -141,6 +141,7 @@ coadd_map = result['coadd']
 - **`deproj_response_funcs`** — list of response functions for constrained ILC deprojection (e.g. tSZ removal).
 - **`cov_smooth_type`** — covariance smoothing method: `'block'` (default), `'gaussian'` ([2307.01043](https://arxiv.org/abs/2307.01043)), or `'tophat'` ([2307.01258](https://arxiv.org/abs/2307.01258)).
 - **`nmap_labels` / `nmap_label_fname_func`** — coadd additional maps (e.g. simulations) using weights derived from the data maps; results returned as `result['<label>_coadd']`.
+- **`check_cmb_weights=True`** — for a CMB solution (responses of 1), check that the per-scale ILC weights (written to `{out_root}wavelet_weights_scale_{k}_{tag}.fits`) sum to 1 in every covered pixel; raises `ValueError` otherwise.
 - **`delete_intermediate=True`** — clean up wavelet/covariance files after completion.
 
 See `examples/example_notebook.ipynb` for a complete worked example with simulated CMB maps.
diff --git a/coberus/pipeline.py b/coberus/pipeline.py
index d0cc486..44c474c 100644
--- a/coberus/pipeline.py
+++ b/coberus/pipeline.py
@@ -218,6 +218,45 @@ def project_mask(imask, oshape, owcs, threshold=0.99):
         return enmap.extract(imask, oshape, owcs)
     return 1.0 * (enmap.project(imask, oshape, owcs, order=1) > threshold)
 
+
+def check_weights_sum(weights, mask_fnames, scale, atol=1e-3):
+    """Check that ILC weights sum to 1 wherever any map contributes.
+
+    Parameters
+    ----------
+    weights : ndarray
+        ILC weights of shape (n_map, n_y, n_x) for one wavelet scale.
+    mask_fnames : list[str]
+        Paths to the wavelet-geometry masks of the maps at this scale,
+        in the same order as weights.
+    scale : int
+        Index of the wavelet scale, used for reporting.
+    atol : optional, float
+        Absolute tolerance on the deviation of the weight sum from 1.
+
+    Raises
+    ------
+    ValueError
+        If the weight sum deviates from 1 by more than atol in any pixel
+        covered by at least one mask.
+    """
+    covered = np.any([enmap.read_map(f) != 0 for f in mask_fnames], axis=0)
+    wsum = weights.sum(axis=0)[covered]
+    if wsum.size == 0:
+        print(f"Scale {scale}: no covered pixels, skipping weight check")
+        return
+    dev = np.abs(wsum - 1).max()
+    print(
+        f"Scale {scale}: weight sum over {wsum.size} covered pixels in "
+        f"[{wsum.min():.6f}, {wsum.max():.6f}], max |sum-1| = {dev:.2e}"
+    )
+    if dev > atol:
+        raise ValueError(
+            f"ILC weights at scale {scale} do not sum to 1: "
+            f"max |sum-1| = {dev:.2e} > {atol:.1e}"
+        )
+
+
 def needlet_coadd(
     map_fname_func,
     mask_fname_func,
@@ -250,6 +289,7 @@ def needlet_coadd(
     nmap_labels=None,
     nmap_label_fname_func=None,
     apply_mask=False,
+    check_cmb_weights=False,
 ):
     """
     Generic function for coadding maps using an empirical
@@ -362,24 +402,23 @@ def needlet_coadd(
         Whether to delete intermediate outputs
 
     nmap_labels : optional,list[str]
-        List of possible optional maps' labels. The special label
-        'cmb_weights' does not read any maps; instead the weights are
-        applied to wavelet maps that are 1 in all pixels, and outmaps
-        ['cmb_weights_coadd'] is a list (one per wavelet scale) of the
-        summed weights on each scale's wavelet geometry, without the
-        final wave2map. For a CMB solution (responses of 1) with no
-        deprojection this should be 1 wherever any map contributes.
+        List of possible optional maps' labels.
 
     nmap_label_fname_func : optional,func | (nmap_label, fname) -> nmap
         Optional maps not used for covariance, but coadded with
         the same weights. Accepts the optional map's label and filename
-        and returns a path. Not called for the 'cmb_weights' label.
+        and returns a path.
 
     apply_mask : optional, boolean
         If true, zero out regions of the input maps based on their masks.
         This will make the maps masked sharply before wavelet transforms,
         and is hence not recommended.
 
+    check_cmb_weights : optional, boolean
+        If true, check that at each wavelet scale the ILC weights of
+        'coadd' sum over tags to 1 in every pixel where any map
+        contributes, and raise a ValueError otherwise. This only holds
+        for a CMB solution (response_func returning 1 for every tag).
 
 
     Returns
@@ -391,8 +430,7 @@ def needlet_coadd(
           'mask'  : the final footprint mask (the base_tag mask, projected
                     onto the output geometry when one is provided).
         For each label in nmap_labels, also contains '{label}_coadd' with the
-        coadd of those maps using the same weights ('cmb_weights_coadd' is
-        instead a per-scale list; see nmap_labels).
+        coadd of those maps using the same weights.
 
     The per-scale ILC weight maps estimated for 'coadd' are also written
     to {out_root}wavelet_weights_scale_{k}_{tag}.fits for each wavelet
@@ -583,7 +621,6 @@ def _get_wave(fname_func, itag, imask):
                 lambda fname: nmap_label_fname_func(label, fname), tag, mask
             )
             for label in nmap_labels
-            if label != "cmb_weights"
         }
 
         if i == 0:
@@ -612,16 +649,6 @@ def _get_wave(fname_func, itag, imask):
 
             # optional maps
             for label in nmap_labels:
-                if label == "cmb_weights":
-                    # Maps of ones are identical for every tag, so write
-                    # one per scale and share it between tags.
-                    nwfname = f"{out_root}wavelet_{label}_scale_{j}.fits"
-                    if len(nfmaps[label][j]) == 0:
-                        filenames.append(nwfname)
-                        enmap.write_map(nwfname, enmap.ones(wmap.shape, wmap.wcs))
-                        totgibytes = totgibytes + (wmap.nbytes / 1024 / 1024.0 / 1024.0)
-                    nfmaps[label][j].append(nwfname)
-                    continue
                 nwfname = f"{out_root}wavelet_{label}_{tags[i]}_scale_{j}.fits"
                 filenames.append(nwfname)
                 nfmaps[label][j].append(nwfname)
@@ -774,18 +801,14 @@ def _get_wave(fname_func, itag, imask):
                             fwname, enmap.enmap(weights[i], owave.maps[j].wcs)
                         )
                         filenames.append(fwname)
+                    if check_cmb_weights:
+                        check_weights_sum(weights, masks, j)
                 else:
                     result = coadd(client, coadder)
                 # This is now a numpy array
                 arr = result.compute()
                 owave.maps[j] = enmap.enmap(arr.copy(), owave.maps[j].wcs)
 
-            if outmaptype == "cmb_weights_coadd":
-                # Synthesizing constant maps is not meaningful, so keep the
-                # per-scale weight sums on the wavelet geometries.
-                outmaps[outmaptype] = [m.copy() for m in owave.maps]
-                continue
-
             coadd_map = wt_out.wave2map(owave)
             coadd_map[out_base_mask == 0] = 0
             outmaps[outmaptype] = coadd_map.copy()
diff --git a/tests/check_cmb_weights.py b/tests/check_cmb_weights.py
index b55fbec..0c4bfb1 100644
--- a/tests/check_cmb_weights.py
+++ b/tests/check_cmb_weights.py
@@ -1,5 +1,5 @@
 """
-Check the 'cmb_weights' nmap label and the per-scale weight maps written by
+Check the check_cmb_weights flag and the per-scale weight maps written by
 coberus.pipeline.needlet_coadd, using the synthetic dataset from
 tests/test_pipeline.py.
 """
@@ -15,7 +15,7 @@
 
 
 def main():
-    """Run a block-smoothed CMB NILC with 'cmb_weights' and check outputs."""
+    """Run a block-smoothed CMB NILC with check_cmb_weights and check outputs."""
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--out-root", default="/tmp/cob_chk_")
     parser.add_argument("--debug", action="store_true", help="Low-res quick run.")
@@ -41,16 +41,9 @@ def main():
         cov_smooth_type="block",
         cov_smooth_factor=8,
         n_workers=2,
-        nmap_labels=["cmb_weights"],
+        check_cmb_weights=True,
     )
     print("keys:", sorted(out))
-    for k, w in enumerate(out["cmb_weights_coadd"]):
-        nz = w != 0
-        print(
-            f"cmb_weights_coadd scale {k}: shape {w.shape}, "
-            f"nonzero range {w[nz].min():.5f}..{w[nz].max():.5f}, "
-            f"zero fraction {1 - nz.mean():.3f}"
-        )
     for k in range(len(lpeaks)):
         files = sorted(glob.glob(f"{root}wavelet_weights_scale_{k}_*.fits"))
         ws = np.array([enmap.read_map(f) for f in files])

From 056d57b634599e399974a39d5bfa05a2998e247d Mon Sep 17 00:00:00 2001
From: Joshua Kim <40779225+jaejoonk@users.noreply.github.com>
Date: Wed, 30 Sep 2026 16:52:21 -0700
Subject: [PATCH 4/6] Normalize smoothed covariances by the common footprint of
 covering maps

For gaussian and tophat covariance smoothing, the smoothed products and
local means averaged in zeros from outside each map's footprint, so the
covariance was biased low (by up to ~50%) near mask edges. Normalizing
each pair by its own smoothed overlap fixes the bias but makes the
per-pixel covariance indefinite when more than two footprints overlap.

Instead, at each pixel every entry is now estimated over the common mask
M_T of the set T of maps covering that pixel,
C_ij = S[M_T d_i d_j] / S[M_T], which is unbiased at edges and stays
positive semi-definite for a non-negative kernel. The local mean used
for smooth_mean_cov is normalized by each map's own mask. Block
smoothing is unchanged.

Add tests/test_cov_edges.py, which checks that the covariance is
unbiased and positive definite for three maps with overlapping edges.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 coberus/pipeline.py     | 210 +++++++++++++++++++++-----------
 tests/test_cov_edges.py | 258 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 398 insertions(+), 70 deletions(-)
 create mode 100644 tests/test_cov_edges.py

diff --git a/coberus/pipeline.py b/coberus/pipeline.py
index 44c474c..55cc7cd 100644
--- a/coberus/pipeline.py
+++ b/coberus/pipeline.py
@@ -6,6 +6,7 @@
 from contextlib import nullcontext
 import healpy as hp
 from collections import defaultdict
+from functools import partial
 import time
 import psutil
 
@@ -174,38 +175,104 @@ def cov_filter(cov_smooth_type, sigma_rad, lmax, use_annulus, annulus_fwhm_ratio
     return compute_tophat_beam(sigma_rad, lmax, w_rad_in=w_rad_in)
 
 
-def cov_smooth(
-    wmap1, wmap2, cov_smooth_type, cov_smooth_factor, sigma_rad, fft_smooth, lmax, fl
-):
+def smooth_map(imap, sigma_rad, fl, lmax):
     """
-    Smooth the product of two wavelet maps into an empirical covariance map.
-
-    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.
+    Smooth a map with the kernel used for 'gaussian' or 'tophat' covariance
+    smoothing.
 
     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.
+        imap: Input enmap.
+        sigma_rad: Gaussian smoothing scale in radians, used when fl is None.
+        fl: Harmonic filter from cov_filter, or None to smooth with FFTs.
         lmax: Maximum multipole for SHT-based smoothing.
-        fl: Harmonic filter from cov_filter; None for the block/FFT paths.
 
     Returns:
-        The smoothed covariance enmap.
+        The smoothed 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:
+    if fl is None:
         # Gaussian smoothing procedure from 2307.01043 using FFTs.
-        return enmap.smooth_gauss(prod, sigma_rad)
+        return enmap.smooth_gauss(imap, 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)
+    return band_filter(imap, fl, lmax)
+
+
+def normalized_delta(wmap, mask, smooth, ftol=1e-8):
+    """
+    Subtract the footprint-normalized local mean from a wavelet map.
+
+    The local mean is S[m w] / S[m] for smoothing operator S and binary
+    footprint m, i.e. an average over the map's own footprint only, so it
+    is not pulled towards zero near footprint edges.
+
+    Args:
+        wmap: Wavelet enmap.
+        mask: Binary footprint enmap of wmap.
+        smooth: Function applying the smoothing operator S to an enmap.
+        ftol: Floor on S[m] to avoid dividing by zero.
+
+    Returns:
+        The mean-subtracted enmap, zeroed outside the footprint.
+    """
+    return (wmap - smooth(wmap * mask) / np.maximum(smooth(mask), ftol)) * mask
+
+
+def coverage_covariances(dfnames, mfnames, cov_fname_func, smooth, ftol=1e-8):
+    """
+    Estimate footprint-normalized covariance maps for one wavelet scale.
+
+    At each pixel p the coadd combines the set T(p) of maps whose masks
+    cover p. Every entry at p is estimated over the common footprint M_T
+    of those maps,
+
+        C_ij(p) = S[M_T d_i d_j](p) / S[M_T](p),
+
+    so all entries used together at p are averages under the same weighting.
+    This removes the low bias of S[d_i d_j] near footprint edges, and keeps
+    the per-pixel covariance positive semi-definite for a non-negative
+    smoothing kernel (unlike normalizing each pair by its own overlap).
+
+    Args:
+        dfnames: Paths to the (optionally mean-subtracted) wavelet maps d_i.
+        mfnames: Paths to the binary masks of the maps, in the same order.
+        cov_fname_func: Accepts indices (i, j) with i <= j and returns the
+            path to write covariance map C_ij to.
+        smooth: Function applying the smoothing operator S to an enmap.
+        ftol: Floor on S[M_T] to avoid dividing by zero.
+
+    Returns:
+        Symmetric nested list fcovs with fcovs[i][j] the path to C_ij.
+    """
+    n = len(dfnames)
+    code = 0
+    for i, mfname in enumerate(mfnames):
+        code = code | ((enmap.read_map(mfname) != 0).astype(np.int64) << i)
+    subsets = [int(c) for c in np.unique(code) if c != 0]
+
+    # S[M_T] is only needed on the pixels covered by exactly T
+    norms = {}
+    for t in subsets:
+        mt = enmap.enmap(((code & t) == t).astype(np.float64), code.wcs)
+        norms[t] = np.maximum(smooth(mt)[code == t], ftol)
+
+    fcovs = [[""] * n for _ in range(n)]
+    for i in range(n):
+        d1 = enmap.read_map(dfnames[i])
+        for j in range(i, n):
+            d2 = d1 if i == j else enmap.read_map(dfnames[j])
+            bits = (1 << i) | (1 << j)
+            cov = enmap.zeros(d1.shape, d1.wcs)
+            for t in subsets:
+                if t & bits != bits:
+                    continue
+                region = code == t
+                prod = d1 * d2 * ((code & t) == t)
+                cov[region] = smooth(prod)[region] / norms[t]
+            fcovname = cov_fname_func(i, j)
+            enmap.write_map(fcovname, cov)
+            fcovs[i][j] = fcovname
+            fcovs[j][i] = fcovname
+    return fcovs
 
 
 def project_mask(imask, oshape, owcs, threshold=0.99):
@@ -387,7 +454,11 @@ def needlet_coadd(
     smooth_mean_cov: optional, boolean
         Only used for gaussian or top-hat smoothing.
         If True, computes the covariance from <(A-A_smooth)(B-B_smooth)>, as
-        in pyilc. Otherwise, computes the covariance from <AB>. Default is True.
+        in pyilc, with A_smooth the local mean of A over its own mask.
+        Otherwise, computes the covariance from <AB>. Default is True.
+        For gaussian or top-hat smoothing, the covariance at each pixel is
+        averaged only over the common mask of the maps covering that pixel
+        (see coverage_covariances), which avoids a low bias near mask edges.
 
     map_postprocess_func : optional,func
         A function to apply to each loaded map
@@ -687,58 +758,57 @@ def _get_wave(fname_func, itag, imask):
                 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"
+        def _cov_fname(i, j, k=k, itags=itags):
+            return f"{out_root}wavelet_cov_scale_{k}_{itags[i]}_{itags[j]}.fits"
 
-        for i in range(len(itags)):
-            wmap1 = enmap.read_map(
-                f"{out_root}wavelet_{prefix}_{itags[i]}_scale_{k}.fits"
-            )
+        if cov_smooth_type == "block":
+            for i in range(len(itags)):
+                wmap1 = enmap.read_map(fmaps[k][i])
 
-            for j in range(i, len(itags)):
-                if i == j:  # Potentially minor speedup?
-                    wmap2 = wmap1
-                else:
-                    wmap2 = enmap.read_map(
-                        f"{out_root}wavelet_{prefix}_{itags[j]}_scale_{k}.fits"
+                for j in range(i, len(itags)):
+                    if i == j:  # Potentially minor speedup?
+                        wmap2 = wmap1
+                    else:
+                        wmap2 = enmap.read_map(fmaps[k][j])
+
+                    cov = block_smooth(wmap1 * wmap2, cov_smooth_factor, slow=False)
+
+                    fcovname = _cov_fname(i, j)
+                    fcovs[k][i][j] = fcovname
+                    fcovs[k][j][i] = fcovname
+                    enmap.write_map(fcovname, cov)
+        else:
+            smooth = partial(smooth_map, sigma_rad=sigma_rad, fl=fl, lmax=lmax)
+
+            # Mean-subtract each tag's wavelet map once here rather than
+            # once per pair; smoothing the pair products of the deltas then
+            # gives the same covariance.
+            dfnames = []
+            if smooth_mean_cov:
+                for tag, wfname, mfname in zip(itags, fmaps[k], fmasks[k]):
+                    delta = normalized_delta(
+                        enmap.read_map(wfname), enmap.read_map(mfname), smooth
                     )
+                    dfname = f"{out_root}wavelet_delta_{tag}_scale_{k}.fits"
+                    enmap.write_map(dfname, delta)
+                    dfnames.append(dfname)
+
+            fcovs[k] = coverage_covariances(
+                dfnames if smooth_mean_cov else fmaps[k],
+                fmasks[k],
+                _cov_fname,
+                smooth,
+            )
 
-                cov = cov_smooth(
-                    wmap1,
-                    wmap2,
-                    cov_smooth_type,
-                    cov_smooth_factor,
-                    sigma_rad,
-                    fft_smooth,
-                    lmax,
-                    fl,
-                )
+            # 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)
 
-                fcovname = f"{out_root}wavelet_cov_scale_{k}_{itags[i]}_{itags[j]}.fits"
-                fcovs[k][i][j] = fcovname
-                fcovs[k][j][i] = fcovname
-                enmap.write_map(fcovname, cov)
-                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)
+        for i in range(len(itags)):
+            for j in range(i, len(itags)):
+                filenames.append(fcovs[k][i][j])
+                totgibytes = totgibytes + os.path.getsize(fcovs[k][i][j]) / 1024**3
 
     elapsed_time = time.time() - start_time_covariance
     print(f"Covariance finished in {elapsed_time / 60.0:.2f} minutes.")
diff --git a/tests/test_cov_edges.py b/tests/test_cov_edges.py
new file mode 100644
index 0000000..29a9cf5
--- /dev/null
+++ b/tests/test_cov_edges.py
@@ -0,0 +1,258 @@
+"""
+Test of the footprint-normalized covariance estimate in coberus.pipeline
+near overlapping mask edges.
+
+Three maps share a white "CMB" signal (variance 1) plus independent white
+noise, on masks A: dec < 4, B: dec > -4, C: |dec| < 30 (degrees), within
+|dec| < 30. In -4 < dec < 4 all three maps contribute and the edges of both
+A and B are within reach of the smoothing kernel. The covariance maps from
+normalized_delta + coverage_covariances (Gaussian SHT smoothing, as in
+needlet_coadd) are compared with the truth 1 + delta_ij * noise_i, and the
+per-pixel covariance of the maps covering each pixel is checked to be
+positive definite. The previous unnormalized estimator S[(w - S[w])^2] is
+also computed, to check that the test is sensitive to the edge bias.
+
+Example::
+
+    python tests/test_cov_edges.py --debug
+    python tests/test_cov_edges.py --output test_output
+"""
+
+import argparse
+import os
+import tempfile
+
+import numpy as np
+from pixell import enmap
+
+from coberus.pipeline import (
+    band_filter,
+    cov_filter,
+    coverage_covariances,
+    normalized_delta,
+)
+
+NOISE_VAR = np.array([0.5, 1.0, 2.0])
+DEC_LIM = 30.0
+
+
+def make_masks(shape, wcs):
+    """
+    Build the binary masks A (dec < 4), B (dec > -4) and C, all restricted
+    to |dec| < DEC_LIM degrees.
+
+    Args:
+        shape, wcs: Geometry of the maps.
+
+    Returns:
+        List of three binary enmaps.
+    """
+    dec = np.rad2deg(enmap.posmap(shape, wcs)[0])
+    inside = np.abs(dec) < DEC_LIM
+    return [
+        enmap.enmap((inside & (dec < 4)).astype(np.float64), wcs),
+        enmap.enmap((inside & (dec > -4)).astype(np.float64), wcs),
+        enmap.enmap(inside.astype(np.float64), wcs),
+    ]
+
+
+def run_edge_test(workdir, res_arcmin=30.0, sigma_deg=3.0, nsim=8, seed=0):
+    """
+    Estimate covariance maps for nsim realizations and compare with truth.
+
+    Args:
+        workdir: Directory for the intermediate FITS files.
+        res_arcmin: Pixel resolution in arcminutes.
+        sigma_deg: Gaussian covariance smoothing scale in degrees.
+        nsim: Number of realizations to average.
+        seed: RNG seed.
+
+    Returns:
+        Dict with 'dec' (row declinations in degrees), 'ratio' and
+        'naive_ratio' (dicts mapping (i, j) to the RA-averaged, sim-averaged
+        ratio of the estimated to the true C_ij per row, NaN where either
+        map is masked), and 'min_eig' (smallest per-pixel eigenvalue of the
+        covariance of the covering maps, over all pixels and sims).
+    """
+    res = np.deg2rad(res_arcmin / 60.0)
+    shape, wcs = enmap.band_geometry(np.deg2rad((-DEC_LIM - 10, DEC_LIM + 10)), res=res)
+    masks = make_masks(shape, wcs)
+    mfnames = []
+    for i, mask in enumerate(masks):
+        mfnames.append(os.path.join(workdir, f"mask_{i}.fits"))
+        enmap.write_map(mfnames[-1], mask)
+
+    sigma = np.deg2rad(sigma_deg)
+    lmax = int(180 * 60 / res_arcmin)
+    fl = cov_filter("gaussian", sigma, lmax, False, 0.5)
+
+    def smooth(imap):
+        return band_filter(imap, fl, lmax)
+
+    n = len(masks)
+    pairs = [(i, j) for i in range(n) for j in range(i, n)]
+    truth = 1.0 + np.diag(NOISE_VAR)
+    ratio = {p: 0.0 for p in pairs}
+    naive_ratio = {p: 0.0 for p in pairs}
+    min_eig = np.inf
+    covered = np.array([m != 0 for m in masks])
+    rng = np.random.default_rng(seed)
+    for _ in range(nsim):
+        sig = rng.standard_normal(shape)
+        wmaps = [
+            enmap.enmap(sig + np.sqrt(nv) * rng.standard_normal(shape), wcs) * m
+            for nv, m in zip(NOISE_VAR, masks)
+        ]
+        dfnames = []
+        for i, (w, m) in enumerate(zip(wmaps, masks)):
+            dfnames.append(os.path.join(workdir, f"delta_{i}.fits"))
+            enmap.write_map(dfnames[-1], normalized_delta(w, m, smooth))
+        fcovs = coverage_covariances(
+            dfnames,
+            mfnames,
+            lambda i, j: os.path.join(workdir, f"cov_{i}_{j}.fits"),
+            smooth,
+        )
+        cov = np.zeros((n, n) + shape)
+        for i, j in pairs:
+            cov[i, j] = cov[j, i] = enmap.read_map(fcovs[i][j])
+            both = covered[i] & covered[j]
+            ratio[(i, j)] += row_mean(cov[i, j] / truth[i, j], both) / nsim
+            dn = [w - smooth(w) for w in (wmaps[i], wmaps[j])]
+            naive = smooth(dn[0] * dn[1])
+            naive_ratio[(i, j)] += row_mean(naive / truth[i, j], both) / nsim
+        min_eig = min(min_eig, min_covering_eig(cov, covered))
+    dec = np.rad2deg(enmap.posmap(shape, wcs)[0][:, 0])
+    return {"dec": dec, "ratio": ratio, "naive_ratio": naive_ratio, "min_eig": min_eig}
+
+
+def row_mean(imap, sel):
+    """
+    Mean of imap over the selected pixels of each row.
+
+    Args:
+        imap: 2D array.
+        sel: Boolean array of the same shape.
+
+    Returns:
+        1D array of row means, NaN for rows with no selected pixels.
+    """
+    cnt = sel.sum(axis=1)
+    tot = np.where(sel, imap, 0).sum(axis=1)
+    return np.where(cnt > 0, tot / np.maximum(cnt, 1), np.nan)
+
+
+def min_covering_eig(cov, covered):
+    """
+    Smallest eigenvalue of the per-pixel covariance of the covering maps.
+
+    Args:
+        cov: Array (n, n, ny, nx) of covariance maps.
+        covered: Boolean array (n, ny, nx) of mask coverage.
+
+    Returns:
+        The smallest eigenvalue over all pixels covered by any map.
+    """
+    code = sum(covered[i].astype(np.int64) << i for i in range(len(covered)))
+    out = np.inf
+    for t in np.unique(code):
+        if t == 0:
+            continue
+        idx = [i for i in range(len(covered)) if (t >> i) & 1]
+        sub = cov[np.ix_(idx, idx)][:, :, code == t]
+        ev = np.linalg.eigvalsh(np.moveaxis(sub, -1, 0))
+        out = min(out, ev[:, 0].min())
+    return out
+
+
+def check_results(res, tol=0.05):
+    """
+    Assert that the covariance is unbiased and positive definite, and that
+    the naive estimator is visibly biased at the overlapping edges.
+
+    Args:
+        res: Output of run_edge_test.
+        tol: Allowed absolute deviation of the row-averaged ratio from 1.
+    """
+    for p, r in res["ratio"].items():
+        dev = np.nanmax(np.abs(r - 1))
+        assert dev < tol, f"C_{p} biased: max |ratio-1| = {dev:.3f}"
+    assert res["min_eig"] > 0, f"covariance not PD: min eig {res['min_eig']:.3e}"
+    strip = np.abs(res["dec"]) < 3
+    naive = res["naive_ratio"][(0, 1)][strip]
+    assert np.nanmin(naive) < 1 - 3 * tol, "naive estimator unexpectedly unbiased"
+
+
+def test_cov_edges(tmp_path):
+    """Covariance near overlapping mask edges is unbiased and PD."""
+    check_results(run_edge_test(str(tmp_path)))
+
+
+def plot_results(res, fname):
+    """
+    Plot the row-averaged estimated/true covariance ratio against dec.
+
+    Args:
+        res: Output of run_edge_test.
+        fname: Output PNG path.
+    """
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    names = "ABC"
+    _, axes = plt.subplots(1, 2, figsize=(11, 4), sharey=True)
+    for ax, key, title in zip(
+        axes,
+        ["naive_ratio", "ratio"],
+        ["previous: S[(w - S[w])^2]", "coverage-normalized"],
+    ):
+        for (i, j), r in res[key].items():
+            ax.plot(res["dec"], r, label=f"C_{names[i]}{names[j]}")
+        for edge in (-4, 4):
+            ax.axvline(edge, color="k", ls=":", lw=1)
+        ax.axhline(1, color="k", lw=1)
+        ax.set_xlabel("dec [deg]")
+        ax.set_title(title)
+    axes[0].set_ylabel("estimated / true covariance")
+    axes[1].legend(fontsize=8)
+    plt.tight_layout()
+    plt.savefig(fname, dpi=120)
+    plt.close()
+
+
+def main():
+    """Run the edge test from the command line and optionally plot it."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--res-arcmin", type=float, default=30.0)
+    parser.add_argument("--sigma-deg", type=float, default=3.0)
+    parser.add_argument("--nsim", type=int, default=8)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--output", default=None, help="Directory for a plot.")
+    parser.add_argument(
+        "--debug", action="store_true", help="Coarser pixels and fewer sims."
+    )
+    args = parser.parse_args()
+    if args.debug:
+        args.res_arcmin, args.nsim = 60.0, 2
+    with tempfile.TemporaryDirectory() as workdir:
+        res = run_edge_test(
+            workdir, args.res_arcmin, args.sigma_deg, args.nsim, args.seed
+        )
+    for p in res["ratio"]:
+        print(
+            f"C_{p}: max |ratio-1| = {np.nanmax(np.abs(res['ratio'][p] - 1)):.4f}, "
+            f"naive min ratio = {np.nanmin(res['naive_ratio'][p]):.3f}"
+        )
+    print(f"min eigenvalue = {res['min_eig']:.3e}")
+    if args.output is not None:
+        os.makedirs(args.output, exist_ok=True)
+        plot_results(res, os.path.join(args.output, "cov_edges.png"))
+    if not args.debug:
+        check_results(res)
+        print("PASSED")
+
+
+if __name__ == "__main__":
+    main()

From bb733bb4c33b415d605594b0394c272dd9f0cccf Mon Sep 17 00:00:00 2001
From: Joshua Kim <40779225+jaejoonk@users.noreply.github.com>
Date: Thu, 1 Oct 2026 17:11:02 -0700
Subject: [PATCH 5/6] Fill uncovered wavelet pixels with base map coefficients
 before wave2map

The coadder returns zero wherever no mask covers a pixel, so the
taper region and source holes of the base map were zeroed in each
wavelet scale. The sharp cut leaked into the footprint on
reconstruction, giving a ~6% rms signal error at low ell even far
from the mask edge. Those pixels now take the base map's own wavelet
coefficients; the final base mask is still applied afterwards.

Add tests/check_two_map_subset.py, a two-map test (wide noisy base map
plus a deeper map on a subset footprint) that checks per-scale ILC
weights against the analytic inverse-noise-variance solution, signal
preservation and noise of the coadd for block, gaussian and tophat
covariance smoothing.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 coberus/pipeline.py           |   9 +
 tests/check_two_map_subset.py | 457 ++++++++++++++++++++++++++++++++++
 2 files changed, 466 insertions(+)
 create mode 100644 tests/check_two_map_subset.py

diff --git a/coberus/pipeline.py b/coberus/pipeline.py
index 55cc7cd..221339f 100644
--- a/coberus/pipeline.py
+++ b/coberus/pipeline.py
@@ -877,6 +877,15 @@ def _cov_fname(i, j, k=k, itags=itags):
                     result = coadd(client, coadder)
                 # This is now a numpy array
                 arr = result.compute()
+                # The coadder returns zero where no mask covers a pixel.
+                # Fill those pixels with the base map's wavelet coefficients
+                # so that wave2map does not see a sharp cut, which leaks
+                # into the footprint at low ell. The final mask is applied
+                # after reconstruction.
+                if base_tag in included_tags[j]:
+                    ibase = included_tags[j].index(base_tag)
+                    covered = np.any([enmap.read_map(f) != 0 for f in masks], axis=0)
+                    arr = np.where(covered, arr, enmap.read_map(lmaps[ibase]))
                 owave.maps[j] = enmap.enmap(arr.copy(), owave.maps[j].wcs)
 
             coadd_map = wt_out.wave2map(owave)
diff --git a/tests/check_two_map_subset.py b/tests/check_two_map_subset.py
new file mode 100644
index 0000000..dfb9dde
--- /dev/null
+++ b/tests/check_two_map_subset.py
@@ -0,0 +1,457 @@
+"""
+Two-map test of coberus.pipeline.needlet_coadd.
+
+Map 'wide' covers a large dec band and is noisy; map 'deep' covers an
+RA/dec patch fully inside the wide footprint and has lower noise. Both
+share the same CMB realization and the same Gaussian beam (equal to the
+output beam), with independent white noise.
+
+For a common signal plus white noise with equal beams, the exact ILC
+weights at every needlet scale are inverse-noise-variance weights,
+
+    w_deep = s_wide^2 / (s_wide^2 + s_deep^2),   w_wide = 1 - w_deep,
+
+inside the deep footprint, and w_wide = 1 outside it. The script checks:
+
+1. Weight maps per scale against these values, including a profile of
+   w_deep against distance to the deep-map mask edge.
+2. Signal preservation: the signal-only maps coadded with the data
+   weights (via nmap_labels) should reproduce the input CMB.
+3. Noise: the noise-only coadd should have variance s_wide^2 outside
+   the deep footprint and s_wide^2 s_deep^2 / (s_wide^2 + s_deep^2)
+   inside it.
+
+Example::
+
+    python tests/check_two_map_subset.py --debug
+"""
+
+import argparse
+import glob
+import itertools
+import json
+import os
+import time
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+from coberus.pipeline import gauss_beam, needlet_coadd
+from pixell import curvedsky as cs
+from pixell import enmap
+
+from test_pipeline import get_camb_cl, make_apod, make_geometry, white_noise
+
+WIDE = ("wide", (-50.0, 10.0), None)
+DEEP = ("deep", (-35.0, -5.0), (40.0, 140.0))
+
+
+def make_two_maps(dset_dir, res_arcmin, lmax, fwhm, noises, taper_deg, seed):
+    """
+    Write data, signal-only and noise-only maps and binary masks for the
+    wide and deep tags.
+
+    Args:
+        dset_dir: Output directory.
+        res_arcmin: Pixel resolution in arcminutes.
+        lmax: Maximum multipole of the CMB realization.
+        fwhm: Beam FWHM in arcminutes for both maps.
+        noises: Dict tag -> white noise level in uK-arcmin.
+        taper_deg: Width of the apodization taper in degrees.
+        seed: RNG seed.
+
+    Returns:
+        (info, alm) where info maps tag -> dict of file paths for 'map',
+        'signal', 'noise' and 'mask', and alm is the CMB realization.
+    """
+    os.makedirs(dset_dir, exist_ok=True)
+    alm = cs.rand_alm(get_camb_cl(lmax), lmax=lmax, seed=seed)
+    balm = cs.almxfl(alm, gauss_beam(np.arange(lmax + 1), fwhm))
+    res = np.deg2rad(res_arcmin / 60.0)
+    info = {}
+    for i, (tag, dec_range, ra_range) in enumerate([WIDE, DEEP]):
+        shape, wcs = make_geometry(dec_range, ra_range, res)
+        apod = make_apod(shape, wcs, dec_range, taper_deg, ra_range_deg=ra_range)
+        sig = cs.alm2map(balm, enmap.empty(shape, wcs))
+        np.random.seed(seed + 1 + i)
+        noise = white_noise(shape, wcs, noises[tag])
+        info[tag] = {}
+        for kind, m in [
+            ("map", (sig + noise) * apod),
+            ("signal", sig * apod),
+            ("noise", noise * apod),
+            ("mask", (apod > 0.99).astype(np.float64)),
+        ]:
+            fname = os.path.join(dset_dir, f"{kind}_{tag}.fits")
+            enmap.write_map(fname, enmap.enmap(m, wcs))
+            info[tag][kind] = fname
+        print(f"Wrote {tag}: shape={shape}, noise={noises[tag]} uK-arcmin")
+    return info, alm
+
+
+def signed_edge_distance(mask):
+    """
+    Signed distance in degrees to the edge of a binary mask: positive
+    inside the mask, negative outside.
+
+    Args:
+        mask: Binary enmap.
+
+    Returns:
+        enmap of signed distances in degrees.
+    """
+    inside = enmap.distance_transform(mask > 0)
+    outside = enmap.distance_transform(mask == 0)
+    return enmap.enmap(np.rad2deg(np.where(mask > 0, inside, -outside)), mask.wcs)
+
+
+def analyze_run(out, root, info, alm, args, expected, run):
+    """
+    Compare one needlet_coadd run against the analytic expectations.
+
+    Args:
+        out: needlet_coadd output dict.
+        root: out_root used for the run (to find the weight maps).
+        info: Dataset info from make_two_maps.
+        alm: Input CMB alm.
+        args: Parsed CLI arguments.
+        expected: Dict of expected values.
+        run: Name of the run, used in figure names.
+
+    Returns:
+        (summary dict, list of (caption, png) figures).
+    """
+    mask_wide = enmap.read_map(info["wide"]["mask"])
+    shape, wcs = mask_wide.shape, mask_wide.wcs
+    mask_deep = enmap.extract(enmap.read_map(info["deep"]["mask"]), shape, wcs)
+    d_wide = signed_edge_distance(mask_wide)
+    d_deep = signed_edge_distance(mask_deep)
+    buf = args.edge_buffer
+    overlap = (d_deep > buf) & (d_wide > buf)
+    wide_only = (d_deep < -buf) & (d_wide > buf)
+
+    summary = {"scales": []}
+    figs = []
+    nscale = len(args.lpeaks)
+    _, axes = plt.subplots(1, 2, figsize=(11, 4))
+    bins = np.arange(-15, 15.5, 0.5)
+    cents = 0.5 * (bins[1:] + bins[:-1])
+    for k in range(nscale):
+        fw = glob.glob(f"{root}wavelet_weights_scale_{k}_wide.fits")
+        fd = glob.glob(f"{root}wavelet_weights_scale_{k}_deep.fits")
+        ww = enmap.read_map(fw[0])
+        wd = enmap.read_map(fd[0])
+        # Weight maps live on the (downgraded) per-scale wavelet geometry
+        dd = enmap.project(d_deep, ww.shape, ww.wcs, order=0)
+        dw = enmap.project(d_wide, ww.shape, ww.wcs, order=0)
+        ov = (dd > buf) & (dw > buf)
+        wo = (dd < -buf) & (dw > buf)
+        s = {
+            "scale": k,
+            "w_deep_overlap_mean": float(wd[ov].mean()),
+            "w_deep_overlap_std": float(wd[ov].std()),
+            "w_wide_wide_only_mean": float(ww[wo].mean()),
+            "w_deep_wide_only_absmax": float(np.abs(wd[wo]).max()),
+            "w_deep_near_edge_mean": float(wd[(dd > 0) & (dd < 2) & (dw > buf)].mean()),
+            "wsum_absdev_max": float(np.abs((ww + wd)[(dw > 0)] - 1).max()),
+        }
+        summary["scales"].append(s)
+        sel = dw > buf
+        prof = [
+            wd[sel & (dd >= lo) & (dd < hi)].mean()
+            if np.any(sel & (dd >= lo) & (dd < hi))
+            else np.nan
+            for lo, hi in itertools.pairwise(bins)
+        ]
+        axes[0].plot(cents, prof, label=f"scale {k} (lpeak {args.lpeaks[k]})")
+        axes[1].errorbar(
+            k, s["w_deep_overlap_mean"], s["w_deep_overlap_std"], fmt="o", color="C0"
+        )
+    axes[0].axhline(expected["w_deep"], color="k", ls="--", lw=1)
+    axes[0].axvline(0, color="gray", lw=0.5)
+    axes[0].set_xlabel("signed distance to deep-map mask edge [deg]")
+    axes[0].set_ylabel(r"mean $w_{\rm deep}$")
+    axes[0].legend(fontsize=7)
+    axes[0].set_title(f"{run}: weight profile across deep edge")
+    axes[1].axhline(expected["w_deep"], color="k", ls="--", lw=1, label="expected")
+    axes[1].set_xlabel("needlet scale")
+    axes[1].set_ylabel(r"$w_{\rm deep}$ in overlap interior (mean $\pm$ std)")
+    axes[1].legend()
+    plt.tight_layout()
+    png = f"two_map_weights_{run}.png"
+    plt.savefig(os.path.join(args.output, png), dpi=110)
+    plt.close()
+    figs.append((f"{run}: deep-map weight vs distance to its mask edge", png))
+
+    # Signal preservation: truth is the beam-convolved CMB band-limited to
+    # the needlet support (ell < lmax)
+    lmax = max(args.lpeaks)
+    fl = gauss_beam(np.arange(lmax + 1), args.fwhm)
+    fl[lmax:] = 0
+    truth = cs.alm2map(cs.almxfl(alm, fl), enmap.empty(shape, wcs))
+    sig = out["signal_coadd"]
+    resid = sig - truth
+    rms_t = truth[mask_wide > 0].std()
+    summary["signal_resid_rel_overlap"] = float(resid[overlap].std() / rms_t)
+    summary["signal_resid_rel_wide_only"] = float(resid[wide_only].std() / rms_t)
+    summary["signal_resid_rel_all"] = float(resid[mask_wide > 0].std() / rms_t)
+
+    # Noise: measured variance ratios vs expectation. Normalize by the
+    # wide-only region so the band-limit of the white noise cancels.
+    noi = out["noise_coadd"]
+    pix = noi.pixsizemap() * (180 * 60 / np.pi) ** 2  # arcmin^2
+    var_ov = np.mean(noi[overlap] ** 2 * pix[overlap])
+    var_wo = np.mean(noi[wide_only] ** 2 * pix[wide_only])
+    summary["noise_var_ratio_wide_only_over_overlap"] = float(var_wo / var_ov)
+    summary["noise_eff_overlap_rel"] = float(np.sqrt(var_ov / var_wo))
+
+    _, axes = plt.subplots(3, 1, figsize=(10, 8), sharex=True)
+    ext = None
+    for ax, m, title, vr in [
+        (axes[0], out["coadd"], "data coadd", 300),
+        (axes[1], resid, "signal coadd - input CMB", 5),
+        (
+            axes[2],
+            noi * np.sqrt(pix),
+            r"noise coadd $\times\sqrt{\Omega_{\rm pix}}$ [uK-arcmin]",
+            3 * args.noise_wide,
+        ),
+    ]:
+        im = ax.imshow(
+            m,
+            origin="lower",
+            cmap="RdBu_r",
+            vmin=-vr,
+            vmax=vr,
+            extent=ext,
+            aspect="auto",
+        )
+        ax.contour(mask_deep, levels=[0.5], colors="k", linewidths=0.6)
+        ax.set_title(f"{run}: {title}")
+        plt.colorbar(im, ax=ax)
+    plt.tight_layout()
+    png = f"two_map_maps_{run}.png"
+    plt.savefig(os.path.join(args.output, png), dpi=110)
+    plt.close()
+    figs.append(
+        (f"{run}: coadd, signal residual and noise coadd (black: deep mask edge)", png)
+    )
+    return summary, figs
+
+
+def write_report(fname, configs, expected, summaries, figs):
+    """
+    Write an HTML report with the summary tables and figures.
+
+    Args:
+        fname: Output HTML path.
+        configs: Dict config name -> argparse.Namespace of its parameters.
+        expected: Dict of expected values.
+        summaries: Dict run -> summary from analyze_run.
+        figs: List of (caption, png) tuples.
+    """
+    cfg_rows = "".join(
+        f"<tr><td>{name}</td><td>{c.noise_wide:g}</td><td>{c.noise_deep:g}</td>"
+        f"<td>{c.fwhm:.1f}</td><td>{c.lpeaks}</td><td>{c.note}</td></tr>"
+        for name, c in configs.items()
+    )
+    rows = []
+    for run, s in summaries.items():
+        for sc in s["scales"]:
+            rows.append(
+                f"<tr><td>{run}</td><td>{sc['scale']}</td>"
+                f"<td>{sc['w_deep_overlap_mean']:.4f} &plusmn; "
+                f"{sc['w_deep_overlap_std']:.4f}</td>"
+                f"<td>{sc['w_deep_near_edge_mean']:.4f}</td>"
+                f"<td>{sc['w_wide_wide_only_mean']:.4f}</td>"
+                f"<td>{sc['wsum_absdev_max']:.1e}</td></tr>"
+            )
+    glob_rows = [
+        f"<tr><td>{run}</td><td>{s['signal_resid_rel_overlap']:.2e}</td>"
+        f"<td>{s['signal_resid_rel_wide_only']:.2e}</td>"
+        f"<td>{s['signal_resid_rel_all']:.2e}</td>"
+        f"<td>{s['noise_var_ratio_wide_only_over_overlap']:.3f}</td>"
+        f"<td>{s['noise_eff_overlap_rel']:.3f}</td>"
+        f"<td>{s['time_min']:.2f}</td></tr>"
+        for run, s in summaries.items()
+    ]
+    imgs = "".join(
+        f"<figure><img src='{p}'><figcaption>{c}</figcaption></figure>" for c, p in figs
+    )
+    c0 = next(iter(configs.values()))
+    html = f"""<!doctype html><html><head><meta charset='utf-8'>
+<title>Two-map NILC test</title>
+<style>body{{font-family:sans-serif;max-width:1100px;margin:auto;padding:16px}}
+table{{border-collapse:collapse;margin-bottom:1em}}
+td,th{{border:1px solid #999;padding:3px 8px}}
+img{{max-width:100%}}</style></head><body>
+<h1>coberus needlet_coadd: two-map subset-footprint test</h1>
+<p>wide (base tag): dec {WIDE[1]}, full RA. deep: dec {DEEP[1]}, RA {DEEP[2]}.
+Same CMB and same beam (= output beam) in both maps, independent white noise,
+{c0.taper_deg} deg cosine apodization, binary masks at apod &gt; 0.99,
+res {c0.res_arcmin}'. Interior regions exclude {c0.edge_buffer} deg from mask
+edges.</p>
+<table><tr><th>config</th><th>noise wide [uK']</th><th>noise deep [uK']</th>
+<th>fwhm [']</th><th>lpeaks</th><th>purpose</th></tr>{cfg_rows}</table>
+<p>Expected in every config: w_deep = {expected["w_deep"]:.4f} in the overlap,
+w_wide = 1 outside; noise variance ratio (wide-only / overlap) =
+{expected["noise_var_ratio"]:.3f}; effective overlap noise / wide noise =
+{expected["noise_eff_rel"]:.3f}.</p>
+<h2>Per-scale weights</h2>
+<table><tr><th>run</th><th>scale</th><th>w_deep overlap interior</th>
+<th>w_deep within 2 deg inside deep edge</th><th>w_wide wide-only</th>
+<th>max |sum w - 1|</th></tr>{"".join(rows)}</table>
+<h2>Signal and noise</h2>
+<table><tr><th>run</th><th>signal resid / CMB rms (overlap)</th>
+<th>(wide-only)</th><th>(whole base mask)</th><th>noise var ratio</th>
+<th>overlap noise / wide noise</th><th>time [min]</th></tr>
+{"".join(glob_rows)}</table>
+{imgs}</body></html>"""
+    with open(fname, "w") as f:
+        f.write(html)
+
+
+def auto_fwhm(lmax, bl_min=0.01):
+    """
+    FWHM in arcminutes of the Gaussian beam with B_ell = bl_min at lmax.
+
+    Args:
+        lmax: Multipole at which the beam reaches bl_min.
+        bl_min: Beam value at lmax.
+
+    Returns:
+        FWHM in arcminutes.
+    """
+    return np.rad2deg(np.sqrt(-16 * np.log(2) * np.log(bl_min)) / lmax) * 60
+
+
+def main():
+    """Generate the two-map datasets, run needlet_coadd and check outputs."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--output", default="test_output", help="Plot/report dir.")
+    parser.add_argument(
+        "--out-root", default="/tmp/cob_two_", help="Intermediates prefix."
+    )
+    parser.add_argument("--res-arcmin", type=float, default=4.0)
+    parser.add_argument("--lmax", type=int, default=500)
+    parser.add_argument("--noise-wide", type=float, default=40.0, help="[uK-arcmin]")
+    parser.add_argument("--noise-deep", type=float, default=20.0, help="[uK-arcmin]")
+    parser.add_argument(
+        "--high-noise-factor",
+        type=float,
+        default=50.0,
+        help="Noise multiplier for the noise-dominated config.",
+    )
+    parser.add_argument("--taper-deg", type=float, default=3.0)
+    parser.add_argument("--edge-buffer", type=float, default=5.0, help="[deg]")
+    parser.add_argument("--cov-smooth-factor", type=int, default=16)
+    parser.add_argument("--runs", nargs="+", default=["block", "gaussian", "tophat"])
+    parser.add_argument(
+        "--configs",
+        nargs="+",
+        default=["baseline", "highnoise", "sharpcut"],
+        help="baseline: realistic noise, beam with B(lmax)=0.01; highnoise: "
+        "noise x high-noise-factor; sharpcut: 10' beam, signal at lmax.",
+    )
+    parser.add_argument("--workers", type=int, default=4)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--debug", action="store_true", help="Low-res quick run.")
+    args = parser.parse_args()
+    if args.debug:
+        args.res_arcmin, args.lmax = 8.0, 250
+        args.cov_smooth_factor = min(args.cov_smooth_factor, 8)
+        args.workers = min(args.workers, 2)
+    args.lpeaks = [lp for lp in [0, 50, 100, 200, 350] if lp < args.lmax] + [args.lmax]
+    os.makedirs(args.output, exist_ok=True)
+
+    nw2, nd2 = args.noise_wide**2, args.noise_deep**2
+    expected = {
+        "w_deep": nw2 / (nw2 + nd2),
+        "noise_var_ratio": (nw2 + nd2) / nd2,
+        "noise_eff_rel": np.sqrt(nd2 / (nw2 + nd2)),
+    }
+    print("expected:", expected)
+    hn = args.high_noise_factor
+    all_configs = {
+        "baseline": {
+            "noise_wide": args.noise_wide,
+            "noise_deep": args.noise_deep,
+            "fwhm": auto_fwhm(args.lmax),
+            "note": "realistic noise; beam suppresses signal near lmax",
+        },
+        "highnoise": {
+            "noise_wide": hn * args.noise_wide,
+            "noise_deep": hn * args.noise_deep,
+            "fwhm": auto_fwhm(args.lmax),
+            "note": "noise-dominated at all scales: clean check of the ILC weights",
+        },
+        "sharpcut": {
+            "noise_wide": args.noise_wide,
+            "noise_deep": args.noise_deep,
+            "fwhm": 10.0,
+            "note": "beam ~1 at lmax: signal truncated sharply by the top needlet",
+        },
+    }
+    configs = {
+        name: argparse.Namespace(**{**vars(args), **all_configs[name]})
+        for name in args.configs
+    }
+    summaries, figs = {}, []
+    for name, cfg in configs.items():
+        noises = {"wide": cfg.noise_wide, "deep": cfg.noise_deep}
+        info, alm = make_two_maps(
+            f"{args.out_root}{name}_dataset",
+            cfg.res_arcmin,
+            cfg.lmax,
+            cfg.fwhm,
+            noises,
+            cfg.taper_deg,
+            cfg.seed,
+        )
+        for run in args.runs:
+            label = f"{name}_{run}"
+            print(f"=== {label} ===")
+            root = f"{args.out_root}{label}_"
+            t0 = time.time()
+            out = needlet_coadd(
+                map_fname_func=lambda tag, info=info: info[tag]["map"],
+                mask_fname_func=lambda tag, info=info: info[tag]["mask"],
+                tags=["wide", "deep"],
+                base_tag="wide",
+                lpeaks=cfg.lpeaks,
+                lmins=[None, None],
+                lmaxs=[None, None],
+                response_func=lambda tag: 1.0,
+                beam_func=lambda tag, ells, fwhm=cfg.fwhm: gauss_beam(
+                    np.asarray(ells), fwhm
+                ),
+                out_beam_fwhm=cfg.fwhm,
+                out_root=root,
+                cov_smooth_type=run,
+                cov_smooth_factor=cfg.cov_smooth_factor,
+                n_workers=cfg.workers,
+                nmap_labels=["signal", "noise"],
+                nmap_label_fname_func=lambda lab, tag, info=info: info[tag][lab],
+                check_cmb_weights=True,
+            )
+            dt = (time.time() - t0) / 60.0
+            s, f = analyze_run(out, root, info, alm, cfg, expected, label)
+            s["time_min"] = dt
+            summaries[label] = s
+            figs += f
+            print(json.dumps(s, indent=1))
+            write_report(
+                os.path.join(args.output, "two_map_report.html"),
+                configs,
+                expected,
+                summaries,
+                figs,
+            )
+    with open(os.path.join(args.output, "two_map_summary.json"), "w") as fh:
+        json.dump({"expected": expected, "runs": summaries}, fh, indent=1)
+
+
+if __name__ == "__main__":
+    main()

From 66ba8413028c1c246fd19e7d6548fad0b96988ae Mon Sep 17 00:00:00 2001
From: Joshua Kim <40779225+jaejoonk@users.noreply.github.com>
Date: Fri, 2 Oct 2026 15:06:25 -0400
Subject: [PATCH 6/6] n_modes_eff no longer corrected by fsky (from Sam)

---
 coberus/pipeline.py | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/coberus/pipeline.py b/coberus/pipeline.py
index 221339f..84309d7 100644
--- a/coberus/pipeline.py
+++ b/coberus/pipeline.py
@@ -529,15 +529,6 @@ def needlet_coadd(
     shape, wcs = enmap.read_map_geometry(map_fname_func(base_tag))
     n_deproj = len(deproj_response_funcs) if (deproj_response_funcs is not None) else 0
 
-    # Compute fsky from base mask. This is used to determine covariance smoothing scales
-    base_mask = enmap.read_map(mask_fname_func(base_tag))
-    fsky = (
-        (np.sum(base_mask**2) / np.prod(base_mask.shape))
-        * base_mask.area()
-        / (4 * np.pi)
-    )
-    print("fsky={:.2f}".format(fsky))
-
     # Initialize Wavelets
     uht = uharm.UHT(shape, wcs, mode="curved")
     basis = wv.CosineNeedlet(lpeaks=lpeaks)
@@ -631,7 +622,6 @@ def _get_wave(fname_func, itag, imask):
                         for i in range(basis.n)
                     ]
                 )
-                * fsky
             )
             n_freq_eff = n_tag_per_scale
 
@@ -660,7 +650,6 @@ def _get_wave(fname_func, itag, imask):
                         for i in range(basis.n)
                     ]
                 )
-                * fsky
             )
             n_freq_eff = n_tag_per_scale