ngmix: each chunk holds only its own rows' tile stamps

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,22 +551,31 @@ 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)
 
     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(

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)