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

[cailmdaley]

Problem

Snakemake runs N ngmix chunks per tile (16 in recent campaigns), each fitting the catalogue rows chunk_rows(n_obj, ID_OBJ_MIN, ID_OBJ_MAX). On develop, Tile_cat.get_data reads the whole tile SExtractor catalogue into memory (FITSCatalogue's default memmap=False) and np.copys the full VIGNET column, and the segmentation vignets the same way when SEG_VIGNET_PATH is set. So every chunk keeps all ~36K 51×51 float32 stamps (357 MiB) for the whole run, although it only uses about 1/16 of them. While it loads, it also holds the full table read and the copy together.

Change

• ChunkStamps(column, rows) holds one stamp column for a range of rows. Like the full column, it is indexed by tile-catalogue row: stamps[i_tile] is row i_tile's stamp. A row outside the chunk raises IndexError and names the chunk's range, so code that reaches outside its chunk fails loudly rather than reading the wrong stamp.
• Tile_cat(cat_path, seg_cat_path=None, row_min=-1, row_max=-1) memory-maps the catalogue (memmap=True) and copies only the chunk's stamp rows (np.array(column[start:stop])), so the full column is never materialised. It sets self.rows = chunk_rows(...). obj_id, ra, dec and flux stay full length, so check_wcs_centroid_offset, which iterates over every obj_id, is unchanged. The seg catalogue is read the same way, and the row-alignment checks are unchanged.
• Ngmix.process passes its ID_OBJ_MIN/ID_OBJ_MAX to Tile_cat and still loops over its own chunk_rows(...). The defaults (-1, unbounded) load the whole catalogue as before.

Every consumer (tile_cat.vign[i_tile], tile_cat.seg[i_tile] in prepare_postage_stamps and the uberseg checks, and flux/ra/dec/obj_id[i_tile]) indexes with the loop's global row, so none of them changes.

Verification

• tests/module/test_ngmix_tile_cat_chunk.py uses an LDAC-layout tile catalogue with a row-aligned seg catalogue. Its NUMBER values are gapped and shuffled, and its stamps carry -1e30 markers and a NaN. The test covers six chunk shapes: unbounded, first, middle, last, open-ended, and the splitter's empty range. For each it checks that:

  • every stamp in the chunk matches an in-memory (memmap=False) read of the file in bytes and dtype, for VIGNET and the seg VIGNET;
  • the held arrays have exactly len(rows) stamps and own their memory;
  • every row outside the chunk raises IndexError;
  • the per-object columns are full length and equal to the file.

  Separate tests check that the held stamps survive the file being overwritten afterwards (copies, not views of the map), that ChunkStamps bounds work with numpy integer rows, and that a misaligned seg catalogue still raises.

• Real data, byte-identical. The input is tile 202.301 (36,022 objects) from the nibi-checks/897-sex run, whose exposure stores survive (clean: false). Its vignet and PSF stores were rebuilt with config_tile_PiViVi_psfex.ini and develop's code into a separate bench directory. ngmix ran three chunks of 150 rows (rows 1–150, 18001–18150 and 35873–36022) once with develop (615e73f) and once with this branch, in the shapepipe-develop-240b37e4 container, as Nibi job 23087595. Each pair of ngmix-202-301.fits outputs is byte-identical (md5 14764eec…, 718ba35f…, eca591e5…). The chunks fitted 107, 149 and 142 of their 150 objects (the rest have no valid epoch), so the comparison covers fitted rows and skipped rows.

• Load memory, one realistic chunk. The test loads chunk 8 of the 16-chunk partition that ngmix_range.py computes for this tile (rows 16017–18148, 2,132 objects), and reads /proc/self/status before and after Tile_cat (load_probe.py, allocation 23109257 steps 0 and 1; output saved in ngmixrows-bench/probe-23109257.txt):

  	develop	this branch
  stamps held	357.4 MiB	21.2 MiB
  anonymous RSS added by the load (RssAnon, kept for the whole run)	+360 MiB	+21 MiB
  peak RSS during load (VmHWM − baseline)	+724 MiB	+387 MiB, predominantly file-backed
  sacct MaxRSS of the load-only step	1.56 GiB	0.36 GiB

  On the branch, the load peak is mostly mapped file pages. peak_probe.py (step 2, same file) holds the map open and shows RssFile rising by 366 MiB, which is about the whole 369 MiB file: reading NUMBER/XWIN_WORLD/YWIN_WORLD/FLUX_AUTO touches every row, and readahead maps the rest. The chunk's stamp copy is the anonymous part, RssAnon +21 MiB. The mapped pages are clean page cache: the kernel can reclaim them, and the chunks on a node share them because they all map the same file. The mappings go when get_data() returns and its local data/seg_data references are dropped; close() alone does not unmap while those views are alive. After that, RssFile is back at its baseline. On develop, the corresponding +724 MiB is anonymous memory in each chunk. Load time is dominated by I/O and varies with the page cache: in this run it was 1.8 s for develop and 0.07 s for the branch.

• Full ngmix run of that chunk. Steps 7 and 8 of allocation 23095876 ran develop and the branch concurrently. Both fitted 2,121 of 2,132 objects, and their outputs are byte-identical (md5 6edc824d…). Wall time was 7,032 s for the branch and 7,084 s for develop. Most of that is store I/O from NFS /scratch: production stages the stores node-local, and the bench does not. sacct MaxRSS is 2.83 GiB for the branch and 2.86 GiB for develop, so it does not resolve the difference. Nibi uses jobacct_gather/cgroup, and its MaxRSS is the step cgroup's memory, which includes the page cache of the ~3 GB of sqlite stores each chunk reads (inferred). Tile_cat lives for the whole of process(), so the +360 MiB versus +21 MiB of anonymous memory measured above is what each chunk carries throughout the fit. On this tile, with 16 chunks, that is about 5.3 GiB less anonymous memory per tile.

• Full suite in the develop SIF (allocation 23095876, pytest tests -m "not slow"): 851 passed, 4 skipped, 2 deselected.

Composition with #925

#925 moves the seg stamps into the tile catalogue as a SEG_VIGNET column and holds VIGNET/SEG_VIGNET as views of the in-memory table. That avoids the copy, but still keeps the whole table in every chunk. The two PRs conflict in Tile_cat and at the Tile_cat(...) call in process. The resolution is mechanical: drop seg_cat_path, keep row_min/row_max, and read

self.vign = ChunkStamps(data['VIGNET'], self.rows) if 'VIGNET' in cols else None
self.seg = ChunkStamps(data['SEG_VIGNET'], self.rows) if 'SEG_VIGNET' in cols else None

with Tile_cat(self._tile_cat_path, self._id_obj_min, self._id_obj_max). #925's comment above self.flux ("views into it, not copies, so it is held once") auto-merges, and has to go. Two of #925's tests assert the whole-column view design and need adapting: test_tile_cat_reads_seg_vignet compares with_seg.seg against the full array, and should compare row by row instead; test_tile_cat_holds_the_stamp_columns_once asserts that vign and seg share one buffer, and test_ngmix_tile_cat_chunk.py replaces it. With that resolution, plus this PR's chunk test switched to a SEG_VIGNET column, the module tests matching ngmix or tile_cat pass on the merge (134 passed, allocation 23095876).

#925 is now retargeted onto #933, and the resolution recipe (the diff saved as ngmixrows-bench/compose925-resolution.diff) applies cleanly there; that last point comes from an independent review, not a rerun on my part.

Merge notes

git merge-tree against #933 (fix/dr6-windowed-positions) and #886 (feat/wire-external-masks) is clean. Against #925, the conflict is in ngmix.py (above). Against #887 (feat/instrument-defect-map), the only conflict is in workflow/README.md, which #887 already has with develop itself; this branch doesn't touch that file.

Claude Opus 5.5 on behalf of Cail

🤖 Generated with Claude Code