perf: cache zero_contour solver + viz sanity net (fast_viz A')

[Jammy2211]

Overview

Two earlier attempts to put zero_contour-based critical curves on the default visualization path were silently broken (2026-04-19 yaml flip revert, 2026-05-16 Euclid pipeline regression). A 2026-05-21 perf benchmark identified a third, independent issue: LensCalc._critical_curve_list_via_zero_contour rebuilds its closure + ZeroSolver on every call, busting the JIT cache and making every call cost ~10 s. With the closure cached, warm calls drop to ~66 ms — fast enough to be a default for any JIT'd likelihood function.

This task lands all three coupled fixes that the future Phase A config dispatch depends on: the perf bug, the broad-except that silently swallowed both prior failures, and the first __Visualization Sanity__ regression-net block.

Plan

• Cache (f, solver) inside LensCalc keyed on (kind, pixel_scales, tol, max_newton) so subsequent _via_zero_contour_from() calls reuse JAX's compile cache.
• Tighten the broad except Exception: in PyAutoLens/autolens/imaging/plot/fit_imaging_plots.py:52 so future regressions of this class fail loudly with a logged warning, not silently.
• Add the first __Visualization Sanity__ block on autolens_workspace_test/scripts/imaging/modeling_visualization_jit.py — asserts non-empty critical curves, finite/positive Einstein radius, AND warm-call latency under 100 ms (regression net for the closure-cache-busting bug).
• Verify via /smoke_test against euclid_strong_lens_modeling_pipeline that the pipeline still runs cleanly.
• NO config flip — that's the next sub-prompt (revised Phase A: context-aware dispatch).
• NO z_projects/euclid edits — live science work, separate Phase B prompt targets the pipeline workspace.

Detailed implementation plan

Affected Repositories

• PyAutoGalaxy (primary — perf fix in LensCalc)
• PyAutoLens (broad-except tighten)
• autolens_workspace_test (Visualization Sanity pilot block)

Work Classification

Library (then workspace follow-up bundled into the same PR pair)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoGalaxy	main	clean
./PyAutoLens	main	clean
./autolens_workspace_test	main	dirty — pre-existing drift (README + dataset/build/*), unrelated to this task

Suggested branch: feature/fast-viz-zero-contour-perf
Worktree root: ~/Code/PyAutoLabs-wt/fast-viz-zero-contour-perf/ (created later by /start_library)

Implementation Steps

1. PyAutoGalaxy autogalaxy/operate/lens_calc.py — in _critical_curve_list_via_zero_contour (around line 1121), introduce self._zero_contour_cache: dict (lazy init) keyed on (kind, pixel_scales, tol, max_newton). On cache hit, reuse the stored (f, solver). Unit test: two calls with same params reuse same (f, solver) identity; different params produce distinct entries.
2. PyAutoLens autolens/imaging/plot/fit_imaging_plots.py:52 — replace bare except Exception: with except (ModuleNotFoundError, ValueError): (silent — known recoverable) + except Exception as exc: logger.warning(..., exc_info=True) (loud — anything else). Unit test re-broadens to confirm the WARNING log fires.
3. autolens_workspace_test scripts/imaging/modeling_visualization_jit.py — append a __Visualization Sanity__ prose block following the same style as the existing __Likelihood Sanity__ at line ~170. Three assertions: len(tangential_critical_curve_list) > 0, np.isfinite(er) and er > 0, and warm-call latency < 100 ms on the second call.

Key Files

• PyAutoGalaxy/autogalaxy/operate/lens_calc.py:1121-1178 — perf cache fix
• PyAutoLens/autolens/imaging/plot/fit_imaging_plots.py:19-53 — broad-except tighten
• autolens_workspace_test/scripts/imaging/modeling_visualization_jit.py — Sanity block

Out of scope (deferred to follow-up sub-prompts)

• Config flip / context-aware dispatch (revised Phase A)
• Latent migration in euclid_strong_lens_modeling_pipeline/util.py:491 (Phase B)
• IPython update_display(fig, display_id=...) wiring (Phase C)
• Rollout of __Visualization Sanity__ across other dataset types (Phase D)

Original Prompt

Click to expand starting prompt

See PyAutoPrompt/issued/fast_viz_zero_contour_perf_fix.md for the verbatim authored prompt — the implementation steps above already mirror it. Key reference points kept inline for issue-time visibility:

• April 2026 revert (PyAutoGalaxy commit abd7b717): ZeroSolver raised inside model-fits and the exception was swallowed by the broad except Exception: at fit_imaging_plots.py:52 — critical curves silently vanished on HPC runs.

• May 2026 revert (PyAutoFit PR #1280): same failure shape on Euclid pipeline; source-plane FITS files wrote all-zero, Einstein-radius posteriors collapsed to the full prior, nothing raised.

• Perf benchmark (2026-05-21) on SIE + circular source:

  Method	First call	Warm call
  marching_squares	32 ms	32 ms
  zero_contour (current code)	10300 ms	10300 ms
  zero_contour (reused f / solver)	10679 ms	66 ms

  The closure cache-busting bug means every call rebuilds JAX's compile cache. Fix: cache (f, solver) on LensCalc keyed on (kind, pixel_scales, tol, max_newton). Warm calls drop to ~66 ms.

• Out of scope (deferred sub-prompts): config flip, Euclid latent migration, IPython display wiring, Phase D rollout. The full out-of-scope list and reference memory entries live in the authored prompt file.

[Jammy2211]

Smoke Test Results — 2026-05-21

All 6 smoke tests passed across 1 workspace (euclid_strong_lens_modeling_pipeline).

Workspace	Passed	Failed	Total
euclid_strong_lens_modeling_pipeline	6	0	6

Scripts: start_here.py, scripts/initial_lens_model.py, scripts/full_model.py, scripts/lens_model_waveband.py, scripts/mge_lens_only.py, scripts/sersic_lens_model.py.

Run against the feature/fast-viz-zero-contour-perf worktree (PyAutoGalaxy LensCalc cache fix + PyAutoLens broad-except tighten). PYAUTO_TEST_MODE=2 with PYAUTO_DISABLE_JAX=1 per the workspace's config/build/env_vars.yaml. PYAUTO_SKIP_WORKSPACE_VERSION_CHECK=1 set to bypass a pre-existing workspace/library version mismatch (2026.5.14.2 vs 2026.5.21.1) unrelated to this task.

Confirms no pipeline regression from the cache fix and broad-except tightening.

[Jammy2211]

Library PRs Created — 2026-05-21

PRs

• PyAutoGalaxy: perf(lens_calc): cache (f, ZeroSolver) for zero_contour critical curves #434
• PyAutoLens: fix(viz): make _compute_critical_curve_lines failures loud, not silent PyAutoLens#527

Status

Library work shipped; workspace work pending. Next: /ship_workspace to add the __Visualization Sanity__ block as a workspace PR on autolens_workspace_test. The Sanity block's perf assertion (< 100 ms warm call) depends on the PyAutoGalaxy cache fix being released, so the workspace PR will hold until #434 ships in a release.

API Changes

• LensCalc._zero_contour_cache: dict — new private attribute, transparently caches (f, ZeroSolver) across calls. No public surface.
• _compute_critical_curve_lines — broad except Exception tightened: ModuleNotFoundError/ValueError still silent, anything else logs WARNING with exc_info=True. Same return tuple.

Verification

• PyAutoGalaxy pytest test_autogalaxy/: 924 passed.
• PyAutoLens pytest test_autolens/: 297 passed.
• CPU benchmark: zero_contour tangential warm 10300 ms → 67 ms; Einstein radius warm 10300 ms → 68 ms.
• euclid_strong_lens_modeling_pipeline smoke (6 scripts): all PASS (see perf: cache zero_contour solver + viz sanity net (fast_viz A') #433 (comment)).

[Jammy2211]

Shipped — 2026-05-21

PRs

• PyAutoGalaxy: perf(lens_calc): cache (f, ZeroSolver) for zero_contour critical curves #434 (merged a9a887a)
• PyAutoLens: fix(viz): make _compute_critical_curve_lines failures loud, not silent PyAutoLens#527 (merged a91febcb)
• autolens_workspace_test: test(viz): __Visualization Sanity__ pilot block on imaging viz_jit (fast_viz A') autolens_workspace_test#111 (merged 099df899)

Summary

Phase A′ of z_features/fast_visualization.md done. LensCalc._critical_curve_list_via_zero_contour now caches (f, ZeroSolver) keyed on (kind, pixel_scales, tol, max_newton) — warm CPU calls drop from ~10300 ms to ~67 ms (tangential) and ~68 ms (Einstein radius). The broad except in _compute_critical_curve_lines now logs WARNING with exc_info=True for unexpected exceptions while keeping ModuleNotFoundError/ValueError silent (both are pre-handled upstream). First __Visualization Sanity__ block lands on modeling_visualization_jit.py with curve-count + Einstein-radius + warm-call < 100 ms assertions to catch any future regression of either failure mode.

Session Notes

• The 2026-05-21 perf benchmark turned up the cache-busting bug as a third independent root cause behind the prior viz reverts (alongside the broad except and the JAX-trace issue). The library-PR perf table covers all three regimes (JIT'd, plain NumPy outer, JAX_DISABLE_JIT).
• Pre-existing point-source JAX vmap regression in autolens_workspace_test/scripts/jax_likelihood_functions/point_source/point.py (likelihood returns -1e99 sentinel instead of -83.38) reproduces on canonical main and is out of scope for this task — separate prompt to follow.
• Tracker scope revisions worth knowing for the next sub-prompts: Phase A is now "make plotter dispatch context-aware" (not "flip YAML default"); Phase B retargeted from z_projects/euclid (live science, off-limits) to euclid_strong_lens_modeling_pipeline/util.py:491 where the effective_einstein_radius latent is currently commented out and ready to re-enable; Phase C reduces to wiring IPython.display.update_display(fig, display_id=...) since BackgroundQuickUpdate already exists in PyAutoFit (commit 1fee93174).
• The workspace PR merged with red smoke CI (the point.py failure above). The __Visualization Sanity__ block runs in the build-server's run_all_scripts.sh rather than the smoke list, so its < 100 ms perf assertion will start firing against the next released library.