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

[Jammy2211]

Overview

Mechanical sweep of the Phase D pilot pattern (PR #111) across the remaining modeling_visualization_jit*.py scripts in both autolens_workspace_test and autogalaxy_workspace_test. Each script gets a __Visualization Sanity__ block with dataset-specific assertions that catch the silent-zero / collapsed-source / unconstrained-latent regression class on every JIT-cached visualization path.

Phase D.1 of z_features/fast_visualization.md. Scope deliberately limited to the JIT-cached path (highest regression risk); visualization_jax*.py and missing-dataset coverage (ellipse / weak lensing / quantity viz_jit) are Phase D.2.

Plan

• 6 scripts gain __Visualization Sanity__ blocks following the imaging pilot template from PR test(viz): __Visualization Sanity__ pilot block on imaging viz_jit (fast_viz A') #111 — same SIE sanity tracer for lensing-side assertions, same LensCalc.from_mass_obj(tracer) wrapping, same warm-call latency guard.
• Per-dataset assertion shapes per the tracker's "Per-dataset assertion shapes" section: imaging variants (delaunay / rectangular) use the imaging template; interferometer adds a model-visibilities-non-zero check; point source switches to source-plane positions; autogalaxy variants drop the lensing-side assertions (no Tracer) and assert model image / FoM finite + non-zero.
• Insertion point on every script: after the existing Part-1 caching probe block, before the Nautilus search section.
• Verify each modified script runs cleanly under PYAUTO_TEST_MODE=2 and exercises the Sanity block. Smoke list for both workspace_test repos passes (modulo pre-existing CI gaps unrelated to this PR: NSS extras on autofit_workspace, point.py JAX vmap on autolens_workspace_test).
• No library code changes, no edits to visualization_jax*.py, no new dataset coverage (those are Phase D.2).

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 to this task
./autogalaxy_workspace_test	main	pre-existing drift (dataset/*/jax_test/galaxies.json), unrelated to this task

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

Implementation Steps

1. autolens_workspace_test/scripts/imaging/modeling_visualization_jit_delaunay.py — Imaging-template SIE Sanity (CC count, finite Einstein radius, warm-call < 100 ms).
2. autolens_workspace_test/scripts/imaging/modeling_visualization_jit_rectangular.py — Same as Github actions #1; rectangular vs delaunay mesh-class difference doesn't affect lensing latents.
3. autolens_workspace_test/scripts/interferometer/modeling_visualization_jit.py — Imaging-template Sanity (SIE tracer) plus interferometer-specific fit.model_visibilities finite + non-zero assertion using the script's analysis_mge.
4. autolens_workspace_test/scripts/point_source/modeling_visualization_jit.py — Source-plane positions count match + finite + non-coincident (deflection-collapse guard), plus lensing-side Einstein-radius assertion via separate SIE tracer.
5. autogalaxy_workspace_test/scripts/imaging/modeling_visualization_jit.py — Non-lensing template: model image non-zero + finite, FoM finite.
6. autogalaxy_workspace_test/scripts/interferometer/modeling_visualization_jit.py — Non-lensing template plus fit.model_visibilities finite + non-zero.

Key Files

• 4 files in autolens_workspace_test, 2 files in autogalaxy_workspace_test (listed above).

Out of scope (Phase D.2)

• visualization_jax*.py scripts (different code path — JAX-only validation, no Nautilus search).
• Missing-dataset coverage: ellipse + weak lensing + quantity modeling_visualization_jit.py don't exist yet.
• Any library code change.

Original Prompt

Click to expand starting prompt

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

• Template: PR autolens_workspace_test#111 (imaging pilot, merged 2026-05-21) — SIE sanity tracer, LensCalc.from_mass_obj(tracer), three assertions (CC count + finite Einstein radius + < 100 ms warm).
• Per-dataset assertion shapes: documented under z_features/fast_visualization.md "Per-dataset assertion shapes" section.
• Failure modes guarded against: 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).

[Jammy2211]

Shipped — 2026-05-21

PRs

• autolens_workspace_test: test(viz): __Visualization Sanity__ rollout across modeling_visualization_jit*.py #113 (merged a2f01fe)
• autogalaxy_workspace_test: test(viz): __Visualization Sanity__ rollout across modeling_visualization_jit*.py autogalaxy_workspace_test#54 (merged 07e7ba73)

Summary

Phase D.1 of z_features/fast_visualization.md shipped. 6 scripts gain __Visualization Sanity__ blocks: 4 in autolens_workspace_test (imaging delaunay + rectangular + interferometer + point_source), 2 in autogalaxy_workspace_test (imaging + interferometer). Each block uses the SIE-sanity-tracer template from PR #111 for lensing-side assertions where applicable, with dataset-specific add-ons (interferometer: fit.model_data finite/non-zero; non-lensing autogalaxy: model_data + figure_of_merit finite/non-zero).

Session Notes

• Wrong attribute guessed first. Initial implementation used fit.model_visibilities on the interferometer variants; the correct attribute is fit.model_data (returns aa.Visibilities). Caught at local validation, fixed before pushing.
• point_source numpy import. That script imports jax.numpy as jnp but not numpy as np at module top. Added import numpy as np inside the Sanity block as a local re-import (safe + idempotent).
• point_source FoM omitted. The prior-median position in the script lies outside the image-pair basin of the prior-median lens, so figure_of_merit = -inf legitimately. The SIE lensing-side checks still catch the deflection-collapse failure class without the noisy FoM gate.
• Pixelization variants (delaunay, rectangular) not validated end-to-end locally. Their Part-1 caching-ratio assertion (cached_time < 0.5 * compile_time) fails on CPU + small datasets when speedup < 2x; this is a pre-existing fragility in those scripts unrelated to this PR. The Sanity block syntax is identical to the imaging pilot's so the run_all_scripts.sh build-server path should exercise it cleanly.
• CI red on point.py (autolens_workspace_test). Same pre-existing jax_likelihood_functions/point_source/point.py JAX vmap regression we've seen on the previous three tasks (-1e99 sentinel). Merged anyway; not introduced by this PR.

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

• visualization_jax*.py scripts (different code path — JAX-only validation, no Nautilus).
• New modeling_visualization_jit.py for: autogalaxy ellipse, autogalaxy quantity, autolens weak lensing (the missing-dataset coverage from the tracker's Coverage audit).