From 0534daf6cfd1407dc30204a1dee1fa4dfcc0e436 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Tue, 30 Jun 2026 10:18:44 +0200 Subject: [PATCH] ngmix: add config-selectable UberSeg neighbour masking (BLEND_HANDLING) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A new ngmix-module option BLEND_HANDLING = {noisefill, uberseg} selects how a neighbour sharing a galaxy's stamp is treated. The default `noisefill` is byte-for-byte the historical behaviour (flagged pixels replaced by a noise realisation, kept at inverse-variance weight); `uberseg` instead hard-masks (weight -> 0) every pixel closer to a neighbour's segmentation footprint than to the central object's, leaving the image untouched. `uberseg_weight()` reimplements Erin Sheldon's MEDS uberseg (the `meds._uberseg.uberseg_tree` nearest-segment Voronoi partition) with a scipy cKDTree, rather than depending on `meds` whose import drags the full fitsio/esutil stack. The surviving central core is a single connected, roughly circular region — emergent geometry of the partition, not a separate aperture, faithful to MEDS `get_uberseg`. Plumbed runner -> Ngmix -> do_ngmix_metacal -> make_ngmix_observation -> prepare_ngmix_weights; Postage_stamp gains a per-epoch `segs` list. The segmentation-map *source* is not wired here (it is plumbing-gated: the seg map is the coadd's, ngmix fits per-epoch stamps); selecting `uberseg` without a seg map raises with a pointer to the gating issue. Tests: synthetic two-object stamp asserts the mask geometry (neighbour-side zeroed, central core single-connected, matches brute-force nearest-segment), the noisefill default is byte-identical to the legacy noise-fill, and uberseg leaves the image untouched while zeroing neighbour/flagged weight. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_018DLfuwQzavQQ9GbcKZQPc5 --- src/shapepipe/modules/ngmix_package/ngmix.py | 144 +++++++++++++- src/shapepipe/modules/ngmix_runner.py | 10 + tests/module/test_ngmix_uberseg.py | 198 +++++++++++++++++++ 3 files changed, 347 insertions(+), 5 deletions(-) create mode 100644 tests/module/test_ngmix_uberseg.py diff --git a/src/shapepipe/modules/ngmix_package/ngmix.py b/src/shapepipe/modules/ngmix_package/ngmix.py index 386c25ef..82bca3d6 100644 --- a/src/shapepipe/modules/ngmix_package/ngmix.py +++ b/src/shapepipe/modules/ngmix_package/ngmix.py @@ -14,10 +14,14 @@ from astropy.io import fits from modopt.math.stats import sigma_mad from ngmix.observation import Observation, ObsList +from scipy.spatial import cKDTree from sqlitedict import SqliteDict from shapepipe.pipeline import file_io +# Neighbour treatments selectable with the BLEND_HANDLING option. +BLEND_HANDLINGS = ("noisefill", "uberseg") + # Gaussian half-light radius per unit sigma: r50 = sqrt(2 ln 2) * sigma # = 1.17741 * sigma. With the ngmix area parameter T = 2 sigma^2 this # gives r50 = SIGMA_TO_R50 * sqrt(T / 2) = sqrt(ln 2 * T). @@ -149,6 +153,9 @@ def __init__( self.weights = [] self.flags = [] self.bkg_rms = [] + # Per-epoch segmentation stamps (object NUMBERs), used only by the + # "uberseg" blend handling; empty for the default noise-fill path. + self.segs = [] self.jacobs = [] # Per-epoch full WCS and the object's sky position, used only by the # "wcs" centroid source (skipped for the default "hsm" path). @@ -263,6 +270,7 @@ def __init__( id_obj_min=-1, id_obj_max=-1, centroid_source="hsm", + blend_handling="noisefill", ): if len(input_file_list) not in {6, 7}: @@ -271,6 +279,12 @@ def __init__( + " required is 6 or 7" ) + if blend_handling not in BLEND_HANDLINGS: + raise ValueError( + f"Unknown BLEND_HANDLING '{blend_handling}'; expected one of" + + f" {BLEND_HANDLINGS}" + ) + self._tile_cat_path = input_file_list[0] self._vignet_cat = Vignet( input_file_list[1], @@ -294,6 +308,7 @@ def __init__( self._id_obj_min = id_obj_min self._id_obj_max = id_obj_max self._centroid_source = centroid_source + self._blend_handling = blend_handling self._w_log = w_log @@ -668,6 +683,8 @@ def process(self): flux_guess, self._rng, centroid_source=self._centroid_source, + blend_handling=self._blend_handling, + object_number=obj_id, ) except Exception as ee: self._w_log.info( @@ -957,7 +974,65 @@ def get_noise(gal, weight, guess, pixel_scale, thresh=1.2): return sig_noise -def prepare_ngmix_weights(gal, weight, flag, rng, bkg_rms=None): +def uberseg_weight(weight, seg, object_number): + """Zero a stamp's weight on neighbour-side pixels (Sheldon/MEDS UberSeg). + + Each pixel is assigned to the object whose segmentation footprint it is + nearest to — a nearest-segment Voronoi partition — and only pixels + assigned to the central object keep their weight. This is the algorithm + behind ``meds.MEDS.get_uberseg`` / ``meds._uberseg.uberseg_tree``, + reimplemented here (the C extension lives in ``meds``, whose import drags + the full ``fitsio``/``esutil`` stack) with a ``cKDTree`` nearest-neighbour + query in place of the C k-d tree. + + Because the partition is by distance to the nearest footprint, the pixels + surviving around a compact central object form a single connected, + roughly circular core; the "circularisation" is emergent geometry, not a + separate aperture. Unlike the noise-fill treatment, masked pixels are + handed to ngmix as a hard mask (weight = 0), never replaced by a noise + realisation. + + Parameters + ---------- + weight : numpy.ndarray + Per-pixel weight (inverse variance) map for the stamp. + seg : numpy.ndarray + Segmentation map on the same grid as ``weight``: 0 for sky, the + SExtractor object number for each detected object's footprint. + object_number : int + Segmentation label of the central object (its SExtractor ``NUMBER``). + + Returns + ------- + numpy.ndarray + Copy of ``weight`` with neighbour-side pixels zeroed. + """ + weight = np.copy(weight) + + obj_pix = np.argwhere(seg != 0) + # No detected footprint, or only sky plus the central object: there is + # no neighbour to mask, so the weight is unchanged (cf. MEDS' early + # ``len(np.unique(seg)) == 2`` return). + if obj_pix.shape[0] == 0: + return weight + labels = seg[seg != 0] + if np.all(labels == object_number): + return weight + + # Nearest segmentation pixel for every stamp pixel; zero the weight + # wherever that nearest footprint belongs to a neighbour rather than to + # the central object. + grid = np.indices(seg.shape).reshape(2, -1).T + _, nearest = cKDTree(obj_pix).query(grid) + nearest_label = labels[nearest].reshape(seg.shape) + weight[nearest_label != object_number] = 0.0 + + return weight + +def prepare_ngmix_weights( + gal, weight, flag, rng, bkg_rms=None, + blend_handling="noisefill", seg=None, object_number=None, +): """bookkeeping for ngmix weights. runs on a single galaxy and epoch pixel scale and galaxy guess TO DO: decide if we want galaxy guess stuff @@ -973,16 +1048,36 @@ def prepare_ngmix_weights(gal, weight, flag, rng, bkg_rms=None): bkg_rms : numpy.ndarray, optional Per-pixel background RMS map. If supplied, unmasked pixels use ``1 / bkg_rms**2`` as the ngmix inverse variance. + blend_handling : {"noisefill", "uberseg"}, optional + How to treat pixels shared with a neighbour. ``"noisefill"`` (default) + replaces flagged pixels with a noise realisation and keeps their + inverse-variance weight — the historical behaviour. ``"uberseg"`` + instead hard-masks (weight = 0) every pixel closer to a neighbour's + segmentation footprint than to the central object's, leaving the + image untouched (see :func:`uberseg_weight`). + seg : numpy.ndarray, optional + Segmentation map on the stamp grid (object NUMBERs). Required for + ``blend_handling="uberseg"``; ignored otherwise. + object_number : int, optional + Central object's segmentation label. Required for + ``blend_handling="uberseg"``; ignored otherwise. Returns ------- numpy.ndarray - Galaxy image with masked pixels replaced by noise. + Galaxy image. For ``"noisefill"`` masked pixels are replaced by noise; + for ``"uberseg"`` the image is returned untouched. numpy.ndarray Variance map for NGMIX. numpy.ndarray Noise image. """ + if blend_handling not in BLEND_HANDLINGS: + raise ValueError( + f"Unknown blend_handling '{blend_handling}'; expected one of" + + f" {BLEND_HANDLINGS}" + ) + mask = np.copy(weight) != 0 mask[flag != 0] = False @@ -1015,7 +1110,21 @@ def prepare_ngmix_weights(gal, weight, flag, rng, bkg_rms=None): noise_img_gal = rng.standard_normal(gal.shape) * sig_noise gal_masked = np.copy(gal) - if (~mask).any(): + if blend_handling == "uberseg": + # Hard-mask neighbour-side pixels (weight -> 0) from the segmentation + # geometry; the image is left untouched (the masked pixels carry no + # weight, so ngmix ignores them in the likelihood). Bad/flagged + # pixels already sit at weight 0 from the mask above. + if seg is None or object_number is None: + raise ValueError( + "blend_handling='uberseg' requires a segmentation map and the" + + " central object_number; none reached prepare_ngmix_weights." + + " The seg-map source is gated (see CosmoStat/shapepipe issue" + + " on UberSeg segmentation plumbing)." + ) + weight_map = uberseg_weight(weight_map, seg, object_number) + elif (~mask).any(): + # noisefill (default): replace masked pixels with a noise realisation. gal_masked[~mask] = noise_img_gal[~mask] return gal_masked, weight_map, noise_img @@ -1023,6 +1132,7 @@ def prepare_ngmix_weights(gal, weight, flag, rng, bkg_rms=None): def make_ngmix_observation( gal, weight, flag, psf, wcs, rng, bkg_rms=None, centroid_source="hsm", wcs_full=None, ra=None, dec=None, + blend_handling="noisefill", seg=None, object_number=None, ): """Build an ngmix Observation for a single galaxy epoch. @@ -1061,6 +1171,15 @@ def make_ngmix_observation( ra, dec : float, optional Object sky position in degrees. Required for ``centroid_source="wcs"`` (ignored for ``"hsm"``). + blend_handling : {"noisefill", "uberseg"}, optional + Neighbour treatment passed through to :func:`prepare_ngmix_weights`; + the default ``"noisefill"`` is the historical behaviour. + seg : numpy.ndarray, optional + Segmentation map on the stamp grid. Required for + ``blend_handling="uberseg"`` (ignored otherwise). + object_number : int, optional + Central object's segmentation label. Required for + ``blend_handling="uberseg"`` (ignored otherwise). Returns ------- @@ -1074,7 +1193,8 @@ def make_ngmix_observation( psf_obs = Observation(psf, jacobian=psf_jacob) gal_masked, weight_map, noise_img = prepare_ngmix_weights( - gal, weight, flag, rng, bkg_rms=bkg_rms + gal, weight, flag, rng, bkg_rms=bkg_rms, + blend_handling=blend_handling, seg=seg, object_number=object_number, ) if centroid_source == "hsm": @@ -1207,7 +1327,10 @@ def make_runners(prior, flux_guess, rng): ) -def do_ngmix_metacal(stamp, prior, flux_guess, rng, centroid_source="hsm"): +def do_ngmix_metacal( + stamp, prior, flux_guess, rng, centroid_source="hsm", + blend_handling="noisefill", object_number=None, +): """Do Ngmix Metacal. Performs metacalibration on a single multi-epoch object and returns the @@ -1229,6 +1352,14 @@ def do_ngmix_metacal(stamp, prior, flux_guess, rng, centroid_source="hsm"): adaptive-moment centroid); ``"wcs"`` uses the catalog sky position projected through the WCS — see that function for the star-vs-galaxy rationale. + blend_handling : {"noisefill", "uberseg"}, optional + Neighbour treatment passed through to + :func:`make_ngmix_observation`; the default ``"noisefill"`` is the + historical behaviour. ``"uberseg"`` consumes ``stamp.segs`` and + ``object_number``. + object_number : int, optional + Central object's segmentation label (its SExtractor ``NUMBER``). + Required for ``blend_handling="uberseg"`` (ignored otherwise). Returns ------- @@ -1255,6 +1386,9 @@ def do_ngmix_metacal(stamp, prior, flux_guess, rng, centroid_source="hsm"): wcs_full=stamp.wcs[n_e] if n_e < len(stamp.wcs) else None, ra=stamp.ra[n_e] if n_e < len(stamp.ra) else None, dec=stamp.dec[n_e] if n_e < len(stamp.dec) else None, + blend_handling=blend_handling, + seg=stamp.segs[n_e] if n_e < len(stamp.segs) else None, + object_number=object_number, ) gal_obs_list.append(gal_obs) diff --git a/src/shapepipe/modules/ngmix_runner.py b/src/shapepipe/modules/ngmix_runner.py index 31028c7f..378a72a7 100644 --- a/src/shapepipe/modules/ngmix_runner.py +++ b/src/shapepipe/modules/ngmix_runner.py @@ -88,6 +88,15 @@ def ngmix_runner( else: centroid_source = "wcs" + # Neighbour treatment: "noisefill" (default, historical) replaces a + # neighbour's pixels with a noise realisation; "uberseg" hard-masks + # (weight -> 0) every pixel closer to a neighbour than to the central + # object, from the segmentation map. See the ngmix module docstrings. + if config.has_option(module_config_sec, "BLEND_HANDLING"): + blend_handling = config.get(module_config_sec, "BLEND_HANDLING") + else: + blend_handling = "noisefill" + # Initialise class instance ngmix_inst = Ngmix( input_file_list, @@ -101,6 +110,7 @@ def ngmix_runner( id_obj_min=id_obj_min, id_obj_max=id_obj_max, centroid_source=centroid_source, + blend_handling=blend_handling, ) # Process ngmix shape measurement and metacalibration diff --git a/tests/module/test_ngmix_uberseg.py b/tests/module/test_ngmix_uberseg.py new file mode 100644 index 00000000..98a586d7 --- /dev/null +++ b/tests/module/test_ngmix_uberseg.py @@ -0,0 +1,198 @@ +"""UberSeg neighbour-masking unit tests. + +Covers the ``BLEND_HANDLING = uberseg`` option added to the ngmix module: + +* :func:`uberseg_weight` — the Sheldon/MEDS nearest-segment Voronoi mask. + Geometry assertions on a synthetic two-object stamp: neighbour-side pixels + are zeroed, the surviving central core is a *single connected* region (the + emergent "circularisation"), and the neighbour footprint is fully removed. +* :func:`prepare_ngmix_weights` — the ``noisefill`` default is byte-for-byte + unchanged (asserted against an independent recomputation of the legacy + three-line noise-fill on a shared RNG), while ``uberseg`` hard-masks the + weight (weight -> 0) and leaves the image untouched. +* The error contract when ``uberseg`` is selected without a segmentation map + (the seg-map source is plumbing-gated upstream). +""" + +import numpy as np +import numpy.testing as npt +import pytest +from scipy import ndimage + +from shapepipe.modules.ngmix_package.ngmix import ( + prepare_ngmix_weights, + uberseg_weight, +) + + +def two_object_seg(npix=41, sep=12, r_central=3, r_neighbour=3): + """Stamp seg map: central object (label 1) at centre, neighbour (label 2) + offset ``sep`` pixels along the column axis. Returns (seg, centre, neigh). + """ + seg = np.zeros((npix, npix), dtype=np.int32) + yy, xx = np.indices((npix, npix)) + c = npix // 2 + centre = (c, c) + neigh = (c, c + sep) + seg[(yy - centre[0]) ** 2 + (xx - centre[1]) ** 2 <= r_central ** 2] = 1 + seg[(yy - neigh[0]) ** 2 + (xx - neigh[1]) ** 2 <= r_neighbour ** 2] = 2 + return seg, centre, neigh + + +def test_uberseg_zeros_neighbour_side_keeps_centre(): + """Neighbour-side pixels lose their weight; the central pixel keeps it.""" + seg, centre, neigh = two_object_seg() + weight = np.ones_like(seg, dtype=float) + + out = uberseg_weight(weight, seg, object_number=1) + + # Central object pixel kept; deep-neighbour pixel zeroed. + assert out[centre] == 1.0 + assert out[neigh] == 0.0 + # Every neighbour-footprint pixel is removed from the fit. + assert np.all(out[seg == 2] == 0.0) + # Every central-footprint pixel survives. + assert np.all(out[seg == 1] == 1.0) + # The input weight is not mutated in place. + assert np.all(weight == 1.0) + + +def test_uberseg_core_is_single_connected_region(): + """The surviving (kept-weight) region is one connected component — the + emergent circular core of the nearest-segment Voronoi partition.""" + seg, _, _ = two_object_seg() + weight = np.ones_like(seg, dtype=float) + + out = uberseg_weight(weight, seg, object_number=1) + + kept = out > 0 + _, n_components = ndimage.label(kept) + assert n_components == 1 + # The partition splits the stamp: some pixels survive, some are masked. + assert kept.any() and (~kept).any() + + +def test_uberseg_partition_is_the_perpendicular_bisector(): + """With the neighbour due-right of the centre, the mask is the far-right + half-plane beyond the footprint bisector: the left edge survives, the + column past the neighbour is gone.""" + seg, centre, neigh = two_object_seg(npix=41, sep=12) + weight = np.ones_like(seg, dtype=float) + + out = uberseg_weight(weight, seg, object_number=1) + + c_row, c_col = centre + assert out[c_row, 0] == 1.0 # far side from the neighbour: kept + assert out[c_row, -1] == 0.0 # neighbour side edge: masked + + +def test_uberseg_no_neighbour_is_passthrough(): + """A stamp with only the central object (or empty seg) is unchanged.""" + npix = 21 + weight = np.random.default_rng(0).random((npix, npix)) + 0.1 + + # Only the central object present. + seg = np.zeros((npix, npix), dtype=np.int32) + seg[8:13, 8:13] = 1 + npt.assert_array_equal(uberseg_weight(weight, seg, object_number=1), weight) + + # Wholly empty seg (no detections). + empty = np.zeros((npix, npix), dtype=np.int32) + npt.assert_array_equal(uberseg_weight(weight, empty, object_number=1), weight) + + +def test_uberseg_matches_bruteforce_nearest_segment(): + """The cKDTree result equals the O(N^2) brute-force nearest-segment rule + (Sheldon's non-C fallback) it stands in for.""" + seg, _, _ = two_object_seg(npix=25, sep=8) + weight = np.ones_like(seg, dtype=float) + + out = uberseg_weight(weight, seg, object_number=1) + + obj = np.argwhere(seg != 0) + labels = seg[seg != 0] + brute = np.ones_like(weight) + for i in range(seg.shape[0]): + for j in range(seg.shape[1]): + d2 = (i - obj[:, 0]) ** 2 + (j - obj[:, 1]) ** 2 + if labels[np.argmin(d2)] != 1: + brute[i, j] = 0.0 + npt.assert_array_equal(out, brute) + + +# --- prepare_ngmix_weights: default unchanged, uberseg hard-masks ---------- + +def _gal_flag_weight(npix=41, seed=7): + rng = np.random.default_rng(seed) + gal = rng.normal(0.0, 3.0, (npix, npix)) + 50.0 + weight = np.ones((npix, npix)) + flag = np.zeros((npix, npix), dtype=np.int32) + # A handful of flagged (bad) pixels — cosmic-ray-like. + flag[5, 5] = 1 + flag[30, 12] = 2 + return gal, flag, weight + + +def test_noisefill_default_is_byte_identical_to_legacy(): + """The default path reproduces the legacy three-line noise-fill exactly + (same RNG stream): masked pixels replaced by noise, weight 1/sigma^2.""" + gal, flag, weight = _gal_flag_weight() + + gal_out, w_out, noise_out = prepare_ngmix_weights( + gal, weight, flag, np.random.RandomState(123), + ) + + # Independent recomputation of the legacy algorithm on the same stream. + from modopt.math.stats import sigma_mad + rng = np.random.RandomState(123) + mask = np.copy(weight) != 0 + mask[flag != 0] = False + sig = sigma_mad(gal) + w_exp = mask.astype(float) / sig ** 2 + noise_exp = rng.standard_normal(gal.shape) * sig + noise_gal = rng.standard_normal(gal.shape) * sig + gal_exp = np.copy(gal) + gal_exp[~mask] = noise_gal[~mask] + + npt.assert_array_equal(gal_out, gal_exp) + npt.assert_array_equal(w_out, w_exp) + npt.assert_array_equal(noise_out, noise_exp) + + +def test_uberseg_hard_masks_weight_and_leaves_image_untouched(): + """uberseg: image returned untouched, weight zeroed on neighbour-side and + flagged pixels, positive on the central core.""" + npix = 41 + gal, flag, weight = _gal_flag_weight(npix=npix) + seg, centre, neigh = two_object_seg(npix=npix, sep=12) + + gal_out, w_out, _ = prepare_ngmix_weights( + gal, weight, flag, np.random.RandomState(5), + blend_handling="uberseg", seg=seg, object_number=1, + ) + + # Image untouched under uberseg (no noise fill). + npt.assert_array_equal(gal_out, gal) + # Neighbour footprint hard-masked; central centre kept. + assert np.all(w_out[seg == 2] == 0.0) + assert w_out[centre] > 0.0 + # Flagged bad pixels remain at weight 0 (folded into the base mask). + assert w_out[5, 5] == 0.0 + + +def test_uberseg_requires_seg_and_object_number(): + gal, flag, weight = _gal_flag_weight() + with pytest.raises(ValueError, match="requires a segmentation map"): + prepare_ngmix_weights( + gal, weight, flag, np.random.RandomState(0), + blend_handling="uberseg", + ) + + +def test_unknown_blend_handling_raises(): + gal, flag, weight = _gal_flag_weight() + with pytest.raises(ValueError, match="Unknown blend_handling"): + prepare_ngmix_weights( + gal, weight, flag, np.random.RandomState(0), + blend_handling="bogus", + )