feat: stamp the small-datasets regime into FITS headers at the writer funnel

[Jammy2211]

Overview

PYAUTO_SMALL_DATASETS=1 caps simulated datasets to a reduced resolution, but nothing on disk records which regime a dataset was written under. That gap let a capped dataset survive into a later full-resolution run and be loaded silently, producing deterministic assertion failures that pass in CI on the same commit — diagnosed in autolens_workspace_test#260.

The imaging manifestation is fixed on PyAutoArray branch claude/jax-grad-assertions-ci-hontn3 (commit 5ef6eea), which infers the regime from data.fits's shape. That inference is deliberately narrow and, by construction, cannot cover two dataset families:

• Point-source and weak-lensing datasets are JSON with no FITS at all, so there is nothing to infer from.
• Interferometer datasets keep their shape under the cap. The visibility count is fixed by the committed uv file while the real-space grid behind it is capped 256x256 -> 16x16, so a capped run writes a data.fits with identical NAXIS and different values. This one fails silently — there is no shape mismatch and no assertion to trip — which is strictly worse than the loud imaging failure that started this.

Proposal

Record the regime at write time instead of inferring it at read time.

Every FITS write in the entire stack funnels through a single function — autonerves/fitsable.py:89 output_to_fits (verified: it is the only output_to_fits definition across PyAutoNerves, PyAutoArray, PyAutoGalaxy and PyAutoLens, and is re-exported as aa.output_to_fits). Adding a header card there when PYAUTO_SMALL_DATASETS=1 is active would be:

• truthful by construction — the stamp is written by the same call that writes the data, so it cannot disagree with it, and there is no stamped-but-empty-directory failure mode;
• zero-call-site — no changes across the ~420 should_simulate sites in the workspaces;
• the only discriminant that can catch the interferometer case, since it does not depend on shape.

PyAutoArray.should_simulate would then prefer the stamp and keep the shape heuristic as the legacy fallback for datasets already on disk that carry no stamp.

Scope and risks worth weighing

• This changes a header card on every FITS the stack writes. Round-trip tests, file-hash regression pins, and any golden-file comparison could be disturbed. That is the main reason this was deliberately kept out of the PyAutoArray fix rather than riding along with it.
• JSON datasets (autonerves/dictable.py:370 is the equivalent funnel) would need their own treatment, and stamping there risks round-trip pollution of the dictable schema. Point-source and weak lensing stay exposed until that is settled — worth deciding explicitly rather than by omission.
• The stamp only helps datasets written after it lands, so the shape fallback is not throwaway work.

Context

• Root-cause investigation and full reproduction: fix: jax_grad assertions fail locally but pass in CI autolens_workspace_test#260
• The narrow imaging fix this follows up: should_simulate in autoarray/util/dataset_util.py, PyAutoArray branch claude/jax-grad-assertions-ci-hontn3
• Sibling follow-up: bug: should_simulate rmtree's committed, gitignore-allowlisted datasets under PYAUTO_SMALL_DATASETS=1 PyAutoArray#470

[Jammy2211]

Library PRs Created

• PyAutoNerves — fix: stamp the small-datasets regime into FITS headers at the writer funnel #154 (write side, pending-release)
• PyAutoArray — fix: prefer the FITS regime stamp over the shape heuristic in should_simulate PyAutoArray#474 (read side, pending-release)

Merge PyAutoNerves#154 first, and release it, before PyAutoArray#474. autoarray/pyproject.toml:30 floors autonerves>=2026.8.22.1, which is currently the newest release on PyPI — so the newest legitimately-resolvable autonerves predates the stamp. Not a blocker (absent card → unknown → existing shape fallback, the designed degradation), but until a stamped autonerves ships, the read side has nothing to read.

Two corrections to this issue's text

1. The funnel is not output_to_fits. It is the only definition of that name, but not the only write path. Verified inventory: 18 library FITS writes — 2 via write_hdu_list, 14 in PyAutoGalaxy/PyAutoLens calling hdu_list.writeto directly (neither repo imports write_hdu_list), and 4 via PyAutoFit's paths.save_fits (directory.py:131, confirmed to writeto the passed list as-is). All 18 build through hdu_list_for_output_from.

The stamp therefore goes in both hdu_list_for_output_from and write_hdu_list. Stamping only output_to_fits as proposed would have missed 16 of 18. Worth noting the interferometer simulator does use output_to_fits, so the original proposal would still have caught the motivating case — the hole is elsewhere and larger.

2. Point-source datasets are not JSON-only. They write a top-level data.fits and are covered normally. Only weak lensing is FITS-free — and dataset/weak/simple is regime-invariant (nothing in its write path reads the env var), so there is no bug there to fix either.

The JSON decision (scope item 4), taken explicitly

Deferred, with evidence rather than by omission. The entire JSON exposure is 5 of 268 call sites across two directories, and of 146 non-test output_to_json call sites only 18 serialise a dataset — ~88% are tracers, galaxies and position lists with no regime to record, mostly written into autofit search-output directories where a regime marker is meaningless.

The stated worry — that in-payload stamping would pollute the dictable round-trip — turned out to be false: an unknown key survives as a top-level sibling of type/class_path/arguments (both from_dict implementations ignore it; assertions is the precedent). Inside arguments it would raise. The real case for deferring is that the one genuinely exposed directory, point_source/multiple_sources, is excluded from harness execution by config/build/no_run.yaml pending PyAutoLens#480.

Filed as a follow-up scoped to point-source only, with both expiring justifications recorded as explicit re-check triggers — plus a third representation nobody had flagged: point_datasets.csv via output_to_csv, which neither a FITS nor a JSON stamp reaches.

Scope item 3 (the round-trip surface) — the risk was smaller than feared

No hash, golden-file or checksum pin over any .fits in the stack. Every header consumer indexes named keys; nothing iterates or **kwargs-expands, so mask/geometry reconstruction cannot be perturbed. The card is byte-size neutral in practice (5760 → 5760 B; a header block holds 36 cards, real headers carry ~10).

Adversarial review found a data-loss regression, and it is fixed

The first cut deleted genuinely full-resolution datasets. The stamp records "the env var was set at write time"; should_simulate read it as "this data is capped". Those diverge — Kernel2D.from_gaussian opts out of the cap (convolver.py:729), Interferometer.from_fits never caps, and anyone converting real data in a shell exporting PYAUTO_SMALL_DATASETS=1 stamps T on full-resolution data. Reproduced: a 300×300 image was destroyed, where the #471 shape heuristic had explicitly refused to delete it — a strict weakening of the safety property #471 established.

Fixed by corroborating a destructive T against the data before acting on it, verified lossless: (16,16) and interferometer (108384,2) still delete; 151×151, 209×209, 300×300 now kept. Details in PyAutoArray#474.

Review ran 6 diverse lenses plus independent refuters: 18 findings, 13 refuted, every survivor actioned.

Status

Tests green — PyAutoNerves 165 passed, PyAutoArray 1101 passed, 0 failed. Workspace impact: none — no symbol removed, renamed or resignatured; the zero-call-site property held across all 253 should_simulate sites. Heart gate not consulted (no PyAutoHeart checkout in this environment); the documented per-repo suite fallback was used instead.

Stopping at PR-open. Merge is a human decision.

---

Generated by Claude Code

[Jammy2211]

Shipped

• PyAutoNerves#154 merged — 4ae8e8f
• PyAutoArray#474 merged — 82b5d16

CI green on both across the full matrix (Python 3.12, 3.13, nojax — 3/3 each). Merged library-first.

Completion record: complete/2026/08/small-datasets-regime-stamp.md in PyAutoMind.

Still to do — this is not yet live

PyAutoNerves needs a release. autoarray/pyproject.toml floors autonerves>=2026.8.22.1, which was the newest release on PyPI when this shipped — so it predates the stamp. Until a stamped autonerves is published, an installed-from-PyPI autoarray sees no card, falls back to the shape heuristic, and the interferometer case stays open. That degradation is by design and nothing breaks, but the fix does not take effect until the release lands.

The workspace smoke suite was never run against this. It needs the full PyAutoLens stack and datasets, unavailable in the session that shipped it. This changes the bytes of every FITS the stack writes, so it is worth a run before the release goes out.

The Heart gate was not consulted — no PyAutoHeart checkout was available; the documented per-repo suite fallback was used instead (PyAutoNerves 165 passed; PyAutoArray 1101 passed, 0 failed).

Follow-ups

• draft/bug/pyautolens/point_source_json_datasets_record_no_regime.md — the deferred JSON decision, scoped to point-source only, with both expiring justifications recorded as re-check triggers (PyAutoLens#480; weak/simple's dependence on via_tracer_from). Names point_datasets.csv as a third representation no stamp reaches.
• Unfiled: should_simulate's capped branch still rmtrees unconditionally and ignores the stamp it now has, so every smoke run re-simulates every dataset. Pre-existing; if stamp is not True: is now a cheap fix.

Closing — both halves are merged.

---

Generated by Claude Code