test(viz): __Visualization Sanity__ rollout across visualization_jax*.py

[Jammy2211]

Overview

Phase D.2.a of z_features/fast_visualization.md. Companion to Phase D.1 (PR autolens_workspace_test#113, PR autogalaxy_workspace_test#54). Rolls the __Visualization Sanity__ block pattern out to the remaining visualization_jax*.py scripts — the single-shot JAX-backed visualization path (no Nautilus search) that's the sibling code path to the JIT-cached modeling_visualization_jit*.py already covered.

Phase D.2.b — authoring new scripts for missing dataset coverage (ellipse modeling_visualization_jit + viz_jax in autogalaxy_workspace_test, quantity modeling_visualization_jit in autogalaxy_workspace_test, new weak-lensing scripts in autolens_workspace_test) — is a separate follow-up after this lands.

Plan

• 6 existing visualization_jax*.py scripts gain a __Visualization Sanity__ block following the same pattern as Phase D.1. Insertion point: after the existing PILOT SUCCEEDED end-of-script print (these scripts have no Part-1/Part-2 split — just a single-shot validation).
• Per-dataset assertion shapes — lensing variants use the SIE sanity tracer (1+ tangential CC, finite positive Einstein radius, warm-call < 100 ms); interferometer adds fit.model_data (complex Visibilities) finite + non-zero; non-lensing autogalaxy variants use fit.model_data + figure_of_merit finite/non-zero (no Tracer); point_source omits a script-specific FoM assertion (prior-median position can legitimately give chi² = -inf).
• Build the fit via analysis.fit_from(instance=instance) since these scripts don't keep a cached fit_2 reference (unlike the modeling_visualization_jit scripts).
• Verify the FitQuantity attribute name (model_data or similar) at implementation time — the prompt notes this as a known check point.
• No library changes, no Nautilus search added, no edits outside the new Sanity blocks.

Detailed implementation plan

Affected Repositories

• autolens_workspace_test (primary)
• autogalaxy_workspace_test

Work Classification

Workspace.

Branch Survey

Repository	Current Branch	Dirty?
./autolens_workspace_test	main	pre-existing drift (README + dataset/build/*), unrelated
./autogalaxy_workspace_test	main	pre-existing drift (dataset/*/jax_test/galaxies.json), unrelated

Suggested branch: feature/viz-sanity-rollout-jax-scripts
Worktree root: ~/Code/PyAutoLabs-wt/viz-sanity-rollout-jax-scripts/ (created later by /start_workspace)

Implementation Steps

1. autolens_workspace_test/scripts/imaging/visualization_jax.py — SIE Sanity (1 CC + finite Einstein r + warm < 100 ms).
2. autolens_workspace_test/scripts/interferometer/visualization_jax.py — SIE Sanity + fit.model_data (complex Visibilities) finite + non-zero via analysis.fit_from(instance=instance).
3. autolens_workspace_test/scripts/point_source/visualization_jax.py — SIE Sanity only (no FoM assertion).
4. autogalaxy_workspace_test/scripts/imaging/visualization_jax.py — Non-lensing template: fit.model_data finite + non-zero, fit.figure_of_merit finite.
5. autogalaxy_workspace_test/scripts/interferometer/visualization_jax.py — Non-lensing + fit.model_data (Visibilities) finite + non-zero.
6. autogalaxy_workspace_test/scripts/quantity/visualization_jax.py — Non-lensing applied to FitQuantity. Verify attribute name (fit.model_data vs fit.model vs other) at implementation time.

Key Files

• 3 files in autolens_workspace_test, 3 files in autogalaxy_workspace_test (listed above).

Out of scope (Phase D.2.b)

• New modeling_visualization_jit.py + visualization_jax.py for autogalaxy ellipse (no JIT script yet), autogalaxy quantity (no JIT script yet), autolens weak lensing (entire dir missing).
• Any library code change.

Original Prompt

Click to expand starting prompt

See PyAutoPrompt/issued/visualization_sanity_rollout_jax_scripts.md for the verbatim authored prompt. Key reference points kept inline:

• Template: PR autolens_workspace_test#111 (original imaging pilot), PR autolens_workspace_test#113 + PR autogalaxy_workspace_test#54 (Phase D.1 rollout). Same SIE sanity tracer + LensCalc.from_tracer(...) + 3-assertion shape.
• Insertion point: end of script (after PILOT SUCCEEDED print), since visualization_jax*.py has no Part-1/Part-2 split.
• Failure modes guarded: silent-zero source plane (PyAutoGalaxy abd7b717 2026-04-19, PyAutoFit #1280 2026-05-17); closure cache-busting (PyAutoGalaxy #433/#434 2026-05-21); deflection collapse (point source); model-image collapse (non-lensing autogalaxy); NUFFT collapse (interferometer).

[Jammy2211]

Shipped — 2026-05-21

PRs

• autolens_workspace_test: test(viz): __Visualization Sanity__ rollout across visualization_jax*.py #115 (merged 6b19026)
• autogalaxy_workspace_test: test(viz): __Visualization Sanity__ rollout across visualization_jax*.py autogalaxy_workspace_test#55 (merged 27faf2a7)

Summary

Phase D.2.a of z_features/fast_visualization.md. 6 existing visualization_jax*.py scripts gain __Visualization Sanity__ blocks (3 in each workspace_test repo). All 6 validated end-to-end locally with the expected SIE einstein_radius=1.1938 + warm-call latencies in the 88-98 ms range on autolens variants, and finite model_data + figure_of_merit on the non-lensing autogalaxy variants (FitQuantity.model_data confirmed at autogalaxy/quantity/fit_quantity.py:71-72).

Session Notes

• Insertion point is the end of the script for visualization_jax variants (no Part-1/Part-2 split — these are single-shot JAX-backed visualization scripts, no Nautilus search). Different from the modeling_visualization_jit variants where the Sanity block sits between Part 1 and Part 2.
• Used analysis.fit_from(instance=instance) to materialise the fit for the model-data assertions since these scripts don't keep a cached fit_2 reference (unlike modeling_visualization_jit).
• Workspace CI red on both repos:
  • autolens_workspace_test: pre-existing point.py JAX vmap regression (-1e99 sentinel) — same one we've seen on every recent task.
  • autogalaxy_workspace_test: pre-existing imaging/visualization.py dataset.png assertion failure (sub-agent verified it fails on canonical main too).
• Both pre-existing failures merged anyway; identical decision to the previous task.

Phase D.2.b follow-ups (not yet authored)

• autolens_workspace_test: entire scripts/weak/ directory missing — needs modeling_visualization_jit.py + visualization_jax.py authored from scratch for weak-lensing dataset.
• autogalaxy_workspace_test: scripts/ellipse/modeling_visualization_jit.py missing (visualization_jax.py already exists).
• autogalaxy_workspace_test: scripts/quantity/modeling_visualization_jit.py missing (visualization_jax.py already exists).