fix: guard point-source parity scripts against the SMALL_DATASETS cap

[Jammy2211]

Overview

point.py's JAX-vmap parity assert failing with a different value on every run is not a
numerical bug in the point-source likelihood. It is on-disk dataset poisoning.

dataset/point_source/simple/ is deliberately committed and gitignore-allowlisted
(.gitignore:13), and it is JSON-only — three tracked files, no data.fits. Any script
that touches it while PYAUTO_SMALL_DATASETS=1 is in force calls
al.util.dataset.should_simulate(), which rmtrees the directory and re-runs the simulator
under the cap. Because there is no data.fits, there is no SMALLDAT header card for
PyAutoArray#471's guard to read, so _is_capped_at_the_current_cap() returns False, the
delete fires, and a later full-regime reader gets stamp None, finds no shape to check, and
keeps the degenerate data. CI clones fresh so git restores it; locally, one run poisons the
next — which is exactly the reported "run A (parallel batch) → -1e+99, runs B/C/D (serial)
→ 16.131221" signature.

The replacement data is degenerate, not merely coarse: under the cap PointSolver.solve
short-circuits at autolens/point/solver/point_solver.py:119 and returns the fixed pair
[(1.0, 0.0), (0.0, 1.0)], unrelated to the lens model. So the parity assert is structurally
meaningless in that regime — the prompt's own leading hypothesis, confirmed.

autolens_workspace_test#264 (2026-08-22) root-caused this and fixed the one at-risk site its
sweep found (point_source/visualization/visualization.py, which declared ENV: real_plots
and so never released the cap). That sweep missed one site, and it left the parity scripts
themselves undefended. This task closes both, and makes the library stop failing silently.

Note on the prompt's repro command: it sets PYAUTO_SMALL_DATASETS=1 by hand, which overrides
the ENV: jax full_datasets declaration the script carries at line 21. That invocation cannot
pass and never could — it is operator error in the repro, not evidence of a live bug. It is
corrected as part of this task.

Plan

• Close the remaining at-risk site. scripts/point_source/visualization/modeling_visualization_jit.py
  declares only ENV: real_output — a marker that does not release the dataset cap — while its
  own __Env__ prose claims "full-resolution data". It calls should_simulate() on the
  committed dataset. Same defect Feature/iterate refactor #264 fixed next door; add full_datasets and correct the prose.
• Guard the parity scripts. Five point_source/ consumers of the committed dataset assert
  hardcoded likelihood literals with no check that the cap is off. Add a loud pre-should_simulate
  guard following the scripts/interferometer/nufft.py:250-274 precedent, so a misconfigured or
  hand-run invocation fails naming the cause instead of producing a meaningless number.
• Make the library short-circuit non-silent. PointSolver.solve returns a fixed,
  model-independent position pair under the cap with no signal whatsoever. Emit a one-time
  logger.warning, consistent with the organism's no-silent-guards doctrine.
• Correct and reconcile the Mind record. The prompt's paths are stale (the script moved to
  scripts/point_source/jax_likelihood/point.py) and its repro is invalid. Three prompts describe
  this same surface; fold them into one line of truth.

Detailed implementation plan

Work Classification

Both (library + workspace). Library PR ships first so the workspace can consume its
API-change summary.

Affected Repositories

• PyAutoLens (primary)
• autolens_workspace_test

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoLens	main	clean
./autolens_workspace_test	main	clean

worktree_check_conflict → exit 0 (no other active.md task claims either repo).

Suggested branch: feature/point-source-dataset-cap-guard

Worktree root: ~/Code/PyAutoLabs-wt/point-source-dataset-cap-guard/ (created later by /start_library)

Implementation Steps

1. PyAutoLens — autolens/point/solver/point_solver.py. In the
   PYAUTO_SMALL_DATASETS == "1" branch at line 119, emit a module-level-deduplicated
   logger.warning before returning the fixed pair, naming the flag, the returned coordinates,
   and the fact that the result is independent of the lens model. Keep it one-shot (a module
   _WARNED flag) so a vmap batch does not print thousands of lines. Update the Notes docstring
   at lines 95-102 to record that the short-circuit now announces itself.
2. PyAutoLens — test_autolens/weak/test_simulator_small_datasets.py. Add a test asserting
   the warning fires once and that the returned grid is the documented fixed pair. Unit tests are
   NumPy-only here, which suits this branch (it is reached before any xp dispatch).
3. autolens_workspace_test — scripts/point_source/visualization/modeling_visualization_jit.py.
   Change ENV: real_output → ENV: full_datasets real_output (line 35) and rewrite the __Env__
   paragraph to state why, mirroring the wording Feature/iterate refactor #264 landed in visualization.py.
4. autolens_workspace_test — cap guard. Add the guard to the five consumers of
   dataset/point_source/simple that assert pinned values:
   jax_likelihood/point.py, jax_likelihood/image_plane.py, jax_likelihood/source_plane.py,
   jax_likelihood/fluxes_time_delays.py, jax_grad/gradient.py. It must run before the
   should_simulate() call (that call is what deletes the dataset), and read the env var directly
   rather than inferring from a grid shape — this dataset is JSON, so there is no capped geometry
   to detect. Message names PYAUTO_SMALL_DATASETS, the ENV: jax full_datasets declaration that
   should have released it, and the degenerate-solver consequence.
5. Sweep for the same shape. Re-run the Feature/iterate refactor #264 audit properly: every !dataset/** allowlisted
   directory in every workspace, cross-referenced against every should_simulate() call site whose
   ENV: line lacks full_datasets. Feature/iterate refactor #264 claimed this was the only at-risk site in the organism
   and it was not; report what the corrected sweep finds.
6. Verification. One run of point.py under the real smoke profile (not the prompt's
   hand-rolled env), from a clean git checkout dataset/, plus a second run immediately after to
   prove the first no longer poisons it.
7. Mind reconciliation. Rewrite draft/bug/autolens/point_jax_vmap_parity_nondeterministic.md
   with the correct path and a valid repro; retire the jax-point-source-point-smoke-sentinel
planned.md entry and its prompt (superseded — its own 2026-08-09 note already says the file
   moved and the -1e99 mechanism was fixed by PyAutoLens#662); and strike the "absorb
   point_jax_vmap_parity_nondeterministic.md" bullet from
   draft/bug/autolens/point_solver_error_bisect_health.md, which no longer needs to carry it.

Key Files

• PyAutoLens/autolens/point/solver/point_solver.py:95-120 — the smoke short-circuit and its docstring
• PyAutoLens/test_autolens/weak/test_simulator_small_datasets.py — existing home for cap-flag tests
• autolens_workspace_test/scripts/point_source/visualization/modeling_visualization_jit.py:35 — the missed ENV: declaration
• autolens_workspace_test/scripts/point_source/jax_likelihood/point.py:56 — unguarded should_simulate() ahead of the pinned assert at line 231
• autolens_workspace_test/scripts/interferometer/nufft.py:250-274 — the guard precedent to mirror
• autolens_workspace_test/config/build/profile_smoke.yaml:16,48 — where the cap is set and how full_datasets releases it
• autolens_workspace_test/.gitignore:13 — the !dataset/point_source/simple/** allowlist that makes the delete destructive

Key Trade-offs

• Guard on the env var, not on geometry. nufft.py guards by asserting its mask shape, which
  works because it owns a FITS dataset with detectable capped geometry. This dataset is JSON with
  no shape to inspect, so the guard has to read PYAUTO_SMALL_DATASETS directly. Less elegant,
  but it is the only signal that exists here.
• Warn rather than raise in PointSolver. Raising under the cap would break every smoke script
  that legitimately relies on the short-circuit for speed. A one-shot warning keeps the speedup and
  removes the silence.
• Not folded into the bisect epic. point_solver_error_bisect_health.md (phase 1 of the
  cluster arc, sized too-large) declares it absorbs this prompt on the grounds of a shared
  investigation surface. That is no longer true: this is a workspace env-declaration defect, not a
  triangle-solve accuracy question, and it should not wait on a 12-phase epic.

Testing Approach

• python -m pytest test_autolens/weak/test_simulator_small_datasets.py for the library change.
• Full python -m pytest test_autolens/ before the library PR (CI gate is Python 3.12 + 3.13).
• Smoke-profile run of point.py per step 6, twice, from a clean dataset/ — the only check that
  actually settles the original report.
• Deliberate negative test: run one guarded script with PYAUTO_SMALL_DATASETS=1 and confirm it
  fails with the new message and leaves git status on dataset/point_source/simple/ clean.

Original Prompt

Click to expand starting prompt

point.py JAX-vmap parity assert is non-deterministic under the smoke env

Type: bug
Target: autolens
Repos:

• autolens_workspace_test
• PyAutoLens
  Difficulty: small
  Autonomy: supervised
  Priority: normal
  Status: formalised

autolens_workspace_test/scripts/jax_likelihood_functions/point_source/point.py fails its
JAX-vs-numpy parity assert under the smoke env, and fails differently between runs:

AssertionError: point: JAX vmap likelihood mismatch
  run A (parallel batch):  ACTUAL [-1.e+99]     DESIRED -83.380498
  run B/C/D (serial):      ACTUAL [16.131221]   DESIRED -83.380498

-1e+99 is the failed-fit sentinel; 16.131221 is a finite but wrong value. Same script,
same env, same commit — so the assert is not measuring what it intends to.

Observed while smoke-gating PyAutoArray#398 (convolver-gaussian-small-datasets-cap, merged
2026-07-22). Confirmed unrelated to that change: A/B'd by checking out the pre-fix
autoarray/operators/convolver.py and re-running — identical failure and identical value.
The script also contains no Convolver / from_gaussian / PSF usage at all, so a
convolution change cannot reach it. It is one of the ~10 already-failing workspace scripts
Heart reported on 2026-07-20, i.e. it pre-dates that work.

Lead worth checking first: PointSolver has its own PYAUTO_SMALL_DATASETS short-circuit
at PyAutoLens/autolens/point/solver/point_solver.py:111 that skips the triangle-tiling
solve entirely under the smoke flag. If the parity assert runs against that short-circuited
solve, the comparison may be structurally meaningless in smoke mode — in which case the fix
is either to unset the flag for this script (config/build/env_vars.yaml override) or to
skip the assert when the short-circuit is active, rather than to chase a numerical bug.

Second possibility to rule out: genuine non-determinism in the triangle solve (ordering /
tie-breaking) that the parity tolerance rtol=1e-4 cannot absorb.

Repro (from autolens_workspace_test/):

PYAUTO_TEST_MODE=2 PYAUTO_SMALL_DATASETS=1 PYAUTO_SKIP_FIT_OUTPUT=1 \
PYAUTO_SKIP_VISUALIZATION=1 PYAUTO_SKIP_CHECKS=1 PYAUTO_FAST_PLOTS=1 JAX_ENABLE_X64=True \
NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib \
python scripts/jax_likelihood_functions/point_source/point.py

Run it several times, and once inside a parallel batch — the failure value changes. Per
feedback_flaky_test_sample_size, a few passing runs will not settle this; decide on the
mechanism, not on a run tally.