Repository navigation
feat: stamp the small-datasets regime into FITS headers at the writer funnel #153
Description
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.fitswith 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_simulatesites 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:370is 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_simulateinautoarray/util/dataset_util.py, PyAutoArray branchclaude/jax-grad-assertions-ci-hontn3 - Sibling follow-up: bug: should_simulate rmtree's committed, gitignore-allowlisted datasets under PYAUTO_SMALL_DATASETS=1 PyAutoArray#470
Activity
Jammy2211 commentedon Aug 22, 2026
Library PRs Created
pending-release)pending-release)Merge PyAutoNerves#154 first, and release it, before PyAutoArray#474.
autoarray/pyproject.toml:30floorsautonerves>=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 viawrite_hdu_list, 14 in PyAutoGalaxy/PyAutoLens callinghdu_list.writetodirectly (neither repo importswrite_hdu_list), and 4 via PyAutoFit'spaths.save_fits(directory.py:131, confirmed towritetothe passed list as-is). All 18 build throughhdu_list_for_output_from.The stamp therefore goes in both
hdu_list_for_output_fromandwrite_hdu_list. Stamping onlyoutput_to_fitsas proposed would have missed 16 of 18. Worth noting the interferometer simulator does useoutput_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.fitsand are covered normally. Only weak lensing is FITS-free — anddataset/weak/simpleis 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_jsoncall 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(bothfrom_dictimplementations ignore it;assertionsis the precedent). Insideargumentsit would raise. The real case for deferring is that the one genuinely exposed directory,point_source/multiple_sources, is excluded from harness execution byconfig/build/no_run.yamlpending 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.csvviaoutput_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
.fitsin 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_simulateread it as "this data is capped". Those diverge —Kernel2D.from_gaussianopts out of the cap (convolver.py:729),Interferometer.from_fitsnever caps, and anyone converting real data in a shell exportingPYAUTO_SMALL_DATASETS=1stampsTon 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
Tagainst the data before acting on it, verified lossless:(16,16)and interferometer(108384,2)still delete;151×151,209×209,300×300now 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_simulatesites. 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 commentedon Aug 22, 2026
Shipped
4ae8e8f82b5d16CI 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.mdin PyAutoMind.Still to do — this is not yet live
PyAutoNerves needs a release.
autoarray/pyproject.tomlfloorsautonerves>=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 onvia_tracer_from). Namespoint_datasets.csvas a third representation no stamp reaches.should_simulate's capped branch stillrmtrees 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