perf: modeling_visualization_jit imaging tests are ~3.5x slower on autolens than autogalaxy (>300s vs ~90s)

[Jammy2211]

Summary

Three imaging integration tests in this repo time out at the 300s per-script cap during release-prep mega-runs:

• scripts/imaging/modeling_visualization_jit.py — TIMEOUT 300s
• scripts/imaging/modeling_visualization_jit_delaunay.py — TIMEOUT 300s
• scripts/imaging/modeling_visualization_jit_rectangular.py — TIMEOUT 301s

The autogalaxy sibling (autogalaxy_workspace_test/scripts/imaging/modeling_visualization_jit.py) completes in ~88.6s. The autolens variants are ~3.5× slower than the autogalaxy equivalent — same script structure, same JIT pipeline, same visualization stack.

Background — why this surfaced now

In a prior triage round, all four scripts (autogalaxy + the three autolens variants) failed fast with:

AssertionError: expected jax.Array, got numpy.float64

This was fixed by:

• autogalaxy_workspace_test PR refactor: drop PYAUTO_DISABLE_CRITICAL_CAUSTICS env var #24 — fix(env): unblock modeling_visualization_jit tests in CI defaults
• autolens_workspace_test PR fix(env): unblock modeling_visualization_jit tests in CI defaults #70 — same title, autolens-side counterpart

Once the assertion was cleared, the autogalaxy script started passing in ~90s (visible in the slowest-25 table of the latest mega-run). The autolens variants now run far enough to hit the 300s cap — i.e. the assertion was masking a real perf issue all along.

Current parking

PR #76 (merged) parks the three scripts in config/build/no_run.yaml with the standard # SLOW <YYYY-MM-DD> convention so mega-runs surface them under the existing loud-warning banner. This issue tracks the actual perf-fix work to remove those SLOW markers.

Investigation pointers

• The autogalaxy variant uses an Imaging + Galaxy pipeline; the autolens variants use Imaging + Tracer (multi-plane ray tracing). The 3.5× gap likely lives in the ray-tracing / multi-plane visualization path or in autolens-specific decorators that don't simplify under JIT as cleanly as the autogalaxy path.
• The _delaunay and _rectangular variants add pixelization on top of the base; useful to compare runtimes between them and the base modeling_visualization_jit.py to see whether pixelization is amplifying the gap or adding a flat overhead.
• The visualization pipeline is heavy: subplot_fit_imaging invokes many sub-plots, several of which trigger ray-tracing/critical-curve/caustic computations under JAX.
• The 300s cap is enforced by PyAutoBuild's run_python.py runner; if ever pushed up to e.g. 600s the autolens variants might pass without a perf fix, but that is a band-aid — the 3.5× ratio is the real signal.

Acceptance criteria

• Identify why the autolens visualization-jit pipeline is ~3.5× slower than the autogalaxy one (profile JIT compile + visualization separately).
• Land a fix that brings autolens runtime within ~1.5× of autogalaxy (or document why >2× is structural).
• Remove the three # SLOW 2026-05-07 entries from config/build/no_run.yaml.
• Confirm by mega-run that the three scripts pass under the 300s cap.

References

• Parking PR: chore: park three modeling_visualization_jit autolens tests as SLOW #76 (merged)
• Unblocking PR (autolens): fix(env): unblock modeling_visualization_jit tests in CI defaults #70
• Unblocking PR (autogalaxy): fix(env): unblock modeling_visualization_jit tests in CI defaults autogalaxy_workspace_test#24

🤖 Filed via Claude Code from release-prep triage Cluster D.