From a709bbcd3843ac085707945d8dd3e66b9630c26c Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Thu, 1 Oct 2026 22:14:56 -0400 Subject: [PATCH] ngmix: each chunk holds only its own rows' tile stamps Tile_cat memory-maps the tile catalogue (and the segmentation vignet catalogue) and copies the stamps of the chunk's rows only, into a ChunkStamps indexed by tile-catalogue row; rows outside the chunk raise IndexError. The per-object columns (obj_id, ra, dec, flux) stay full length. Ngmix.process passes its ID_OBJ_MIN/ID_OBJ_MAX bounds to Tile_cat. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017JQtu8oaZQEAqzbxZbPpWs --- src/shapepipe/modules/ngmix_package/ngmix.py | 61 ++++++++- tests/module/test_ngmix_tile_cat_chunk.py | 133 +++++++++++++++++++ 2 files changed, 191 insertions(+), 3 deletions(-) create mode 100644 tests/module/test_ngmix_tile_cat_chunk.py diff --git a/src/shapepipe/modules/ngmix_package/ngmix.py b/src/shapepipe/modules/ngmix_package/ngmix.py index 4adc8227b..389743a27 100644 --- a/src/shapepipe/modules/ngmix_package/ngmix.py +++ b/src/shapepipe/modules/ngmix_package/ngmix.py @@ -501,11 +501,46 @@ def position_seed(ra, dec, ccd): return cantor % (2 ** 32) +class ChunkStamps(): + """One stamp column of the tile catalogue, held for one chunk's rows. + + Indexed by tile-catalogue row like the full column it stands in for, so + ``stamps[i_tile]`` is the stamp of row ``i_tile``; rows outside the chunk + raise ``IndexError``. + + Parameters + ---------- + column : numpy.ndarray + Full stamp column, one stamp per catalogue row. A memory-mapped + column is read only over ``rows``. + rows : range + 0-based catalogue rows to hold (see :func:`chunk_rows`). + + """ + def __init__(self, column, rows): + self.rows = rows + self._stamps = np.array(column[rows.start:rows.stop]) + + def __getitem__(self, i_tile): + if not self.rows.start <= i_tile < self.rows.stop: + raise IndexError( + f"tile-catalogue row {i_tile} is outside this chunk's rows" + + f" {self.rows.start}..{self.rows.stop - 1}" + ) + return self._stamps[i_tile - self.rows.start] + + class Tile_cat(): """Tile_cat. catalog measured on a tile + The per-object columns (``obj_id``, ``ra``, ``dec``, ``flux``) are held + for every row. The stamp columns (``vign``, ``seg``) are held only for + the chunk's rows (``self.rows``) as :class:`ChunkStamps`, still indexed + by tile-catalogue row: the catalogues are memory-mapped, so a chunk reads + only its own stamps. + Parameters ---------- cat_path : str @@ -516,15 +551,23 @@ class Tile_cat(): row-aligned to ``cat_path``. When given, ``self.seg`` holds one integer seg stamp per object for the ``"uberseg"`` blend handling; ``None`` leaves ``self.seg`` unset (the noise-fill path is unaffected). + row_min, row_max : int, optional + First and last catalogue row of the chunk (1-based, inclusive; see + :func:`chunk_rows`). The default, ``-1``, is unbounded, so the whole + catalogue. """ def __init__( self, cat_path, seg_cat_path=None, + row_min=-1, + row_max=-1, ): self.cat_path = cat_path self.seg_cat_path = seg_cat_path + self.row_min = row_min + self.row_max = row_max if cat_path: self.get_data(cat_path) @@ -532,6 +575,7 @@ def get_data(self, cat_path): tile_cat = file_io.FITSCatalogue( cat_path, SEx_catalogue=True, + memmap=True, ) tile_cat.open() data = tile_cat.get_data() @@ -540,10 +584,15 @@ def get_data(self, cat_path): self.obj_id = np.copy(data['NUMBER']) self.ra = np.copy(data['XWIN_WORLD']) self.dec = np.copy(data['YWIN_WORLD']) + self.rows = chunk_rows(len(self.obj_id), self.row_min, self.row_max) # Optional columns — may be absent in external (non-SExtractor) catalogs self.flux = np.copy(data['FLUX_AUTO']) if 'FLUX_AUTO' in cols else None - self.vign = np.copy(data['VIGNET']) if 'VIGNET' in cols else None + self.vign = ( + ChunkStamps(data['VIGNET'], self.rows) + if 'VIGNET' in cols + else None + ) tile_cat.close() @@ -556,6 +605,7 @@ def get_data(self, cat_path): seg_cat = file_io.FITSCatalogue( self.seg_cat_path, SEx_catalogue=True, + memmap=True, ) seg_cat.open() seg_data = seg_cat.get_data() @@ -577,7 +627,7 @@ def get_data(self, cat_path): + " does not match the tile catalogue NUMBER; the" + " segmentation vignets are misaligned or reordered." ) - self.seg = np.copy(seg_data['VIGNET']) + self.seg = ChunkStamps(seg_data['VIGNET'], self.rows) seg_cat.close() class Postage_stamp(): @@ -1263,7 +1313,12 @@ def process(self): @sc [decision:shape_measurement.fit_initialisation,decision:shape_measurement.ngmix_seed_mode] """ - tile_cat = Tile_cat(self._tile_cat_path, self._seg_cat_path) + tile_cat = Tile_cat( + self._tile_cat_path, + self._seg_cat_path, + self._id_obj_min, + self._id_obj_max, + ) vignet_cat = self._vignet_cat check_wcs_centroid_offset( diff --git a/tests/module/test_ngmix_tile_cat_chunk.py b/tests/module/test_ngmix_tile_cat_chunk.py new file mode 100644 index 000000000..43e58b4cf --- /dev/null +++ b/tests/module/test_ngmix_tile_cat_chunk.py @@ -0,0 +1,133 @@ +"""A chunk's Tile_cat holds the stamps of its own rows, and only those. + +Each ngmix chunk reads the tile catalogue for rows ``chunk_rows(n_obj, +row_min, row_max)``. The per-object columns stay full length; the stamp +columns (``VIGNET`` and the segmentation ``VIGNET``) are held for the chunk's +rows only, indexed by tile-catalogue row exactly like the full column. +""" + +import numpy as np +import pytest +from astropy.io import fits + +from shapepipe.modules.ngmix_package.ngmix import ( + ChunkStamps, + Tile_cat, + chunk_rows, +) + +N_OBJ = 23 +STAMP = 7 + + +def _ldac(path, columns): + imhead = fits.BinTableHDU.from_columns( + [fits.Column(name="Field Header Card", format="1A", array=["x"])], + name="LDAC_IMHEAD", + ) + objects = fits.BinTableHDU.from_columns(columns, name="LDAC_OBJECTS") + fits.HDUList([fits.PrimaryHDU(), imhead, objects]).writeto(path) + return str(path) + + +@pytest.fixture +def catalogues(tmp_path): + """A tile catalogue and its row-aligned segmentation catalogue.""" + rng = np.random.default_rng(3) + number = rng.permutation(np.arange(1, 10 * N_OBJ, 10))[:N_OBJ] + vign = rng.normal(size=(N_OBJ, STAMP, STAMP)).astype(np.float32) + vign[::4, 0, :] = -1e30 + vign[1, 2, 2] = np.nan + seg = rng.integers(0, 5, size=(N_OBJ, STAMP, STAMP)).astype(np.int32) + dim = f"({STAMP}, {STAMP})" + cat = _ldac(tmp_path / "sexcat.fits", [ + fits.Column(name="NUMBER", format="J", array=number), + fits.Column(name="XWIN_WORLD", format="D", array=rng.uniform(size=N_OBJ)), + fits.Column(name="YWIN_WORLD", format="D", array=rng.uniform(size=N_OBJ)), + fits.Column(name="FLUX_AUTO", format="E", array=rng.uniform(size=N_OBJ)), + fits.Column( + name="VIGNET", format=f"{STAMP * STAMP}E", dim=dim, array=vign + ), + ]) + seg_cat = _ldac(tmp_path / "seg.fits", [ + fits.Column(name="NUMBER", format="J", array=number), + fits.Column( + name="VIGNET", format=f"{STAMP * STAMP}J", dim=dim, array=seg + ), + ]) + return cat, seg_cat + + +def _full_columns(path): + """Every column as an in-memory (not memory-mapped) read gives it.""" + with fits.open(path, memmap=False) as hdul: + data = hdul[2].data + return {name: np.copy(data[name]) for name in data.dtype.names} + + +@pytest.mark.parametrize( + "row_min, row_max", + [(-1, -1), (1, 6), (7, 15), (16, N_OBJ), (20, -1), (N_OBJ + 1, N_OBJ)], +) +def test_chunk_holds_exactly_its_rows(catalogues, row_min, row_max): + cat, seg_cat = catalogues + full = _full_columns(cat) + full_seg = _full_columns(seg_cat)["VIGNET"] + rows = chunk_rows(N_OBJ, row_min, row_max) + + tile = Tile_cat(cat, seg_cat, row_min=row_min, row_max=row_max) + + assert tile.rows == rows + # Per-object columns: full length, every row. + np.testing.assert_array_equal(tile.obj_id, full["NUMBER"]) + np.testing.assert_array_equal(tile.ra, full["XWIN_WORLD"]) + np.testing.assert_array_equal(tile.dec, full["YWIN_WORLD"]) + np.testing.assert_array_equal(tile.flux, full["FLUX_AUTO"]) + + for stamps, column in ((tile.vign, full["VIGNET"]), (tile.seg, full_seg)): + # Same stamps, bit for bit, at the same tile-catalogue rows. + for i_tile in rows: + assert stamps[i_tile].dtype == column[i_tile].dtype + assert stamps[i_tile].tobytes() == column[i_tile].tobytes() + # Only the chunk's stamps are held, in memory of their own. + assert stamps._stamps.shape == (len(rows), STAMP, STAMP) + assert stamps._stamps.flags.owndata + # The other rows are absent, and touching one says so. + for i_tile in set(range(N_OBJ)) - set(rows): + with pytest.raises(IndexError, match="outside this chunk"): + stamps[i_tile] + + +def test_stamps_outlive_the_catalogue_file(catalogues, tmp_path): + """The held stamps are copies, not views into the memory-mapped file.""" + cat, seg_cat = catalogues + expected = _full_columns(cat)["VIGNET"][4:9].tobytes() + tile = Tile_cat(cat, seg_cat, row_min=5, row_max=9) + # Overwrite the file in place with different stamps. + with fits.open(cat, mode="update") as hdul: + hdul[2].data["VIGNET"][:] = 0 + assert tile.vign._stamps.tobytes() == expected + + +def test_chunk_stamps_index_by_tile_row(): + column = np.arange(10 * 4).reshape(10, 2, 2) + stamps = ChunkStamps(column, range(3, 6)) + np.testing.assert_array_equal(stamps[np.int64(3)], column[3]) + np.testing.assert_array_equal(stamps[5], column[5]) + for i_tile in (2, 6, -1): + with pytest.raises(IndexError): + stamps[i_tile] + + +def test_misaligned_seg_catalogue_still_raises(catalogues, tmp_path): + cat, _ = catalogues + short = _ldac(tmp_path / "short_seg.fits", [ + fits.Column(name="NUMBER", format="J", array=np.arange(1, 4)), + fits.Column( + name="VIGNET", format=f"{STAMP * STAMP}J", + dim=f"({STAMP}, {STAMP})", + array=np.zeros((3, STAMP, STAMP), dtype=np.int32), + ), + ]) + with pytest.raises(ValueError, match="row-aligned"): + Tile_cat(cat, short, row_min=1, row_max=5)