test(viz): quantity/modeling_visualization_jit.py — JIT-cached integration script

[Jammy2211]

Overview

Author the missing quantity/modeling_visualization_jit.py JIT-cached integration script in autogalaxy_workspace_test. The repo currently has quantity/visualization_jax.py (single-shot, no Nautilus) but lacks the JIT-cached + live-Nautilus variant that exists for every other dataset type. Mirrors the structure of the imaging variant shipped in PR #54.

Phase D.2.b.i of z_features/fast_visualization.md. Phase D.2.b.ii (ellipse) and D.2.b.iii (weak lensing) are separate follow-ups; this PR is scoped to one new file in one repo so the pattern can be validated before authoring the more divergent dataset types.

Plan

• New file autogalaxy_workspace_test/scripts/quantity/modeling_visualization_jit.py mirroring the structure of autogalaxy_workspace_test/scripts/imaging/modeling_visualization_jit.py (Part-1 caching probe + non-lensing __Visualization Sanity__ block + Part-2 live Nautilus quick-update).
• Dataset / mask / model_mge reused from quantity/visualization_jax.py. Swap ag.AnalysisImaging → ag.AnalysisQuantity; everything else stays the same shape.
• Non-lensing Sanity template from Phase D.1: fit.model_data finite + non-zero, fit.figure_of_merit finite (confirmed FitQuantity.model_data API at autogalaxy/quantity/fit_quantity.py:71-72).
• No library changes, no edits to existing visualization_jax.py, no new dataset coverage outside quantity.

Detailed implementation plan

Affected Repositories

• autogalaxy_workspace_test (primary)

Work Classification

Workspace.

Branch Survey

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

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

Implementation Steps

1. Read template autogalaxy_workspace_test/scripts/imaging/modeling_visualization_jit.py — that's the Part-1 / Sanity / Part-2 skeleton.
2. Read source for dataset/model autogalaxy_workspace_test/scripts/quantity/visualization_jax.py — copy the dataset path, mask construction, and model_mge composition. Swap ag.AnalysisImaging → ag.AnalysisQuantity (existing in viz_jax already — no API change needed).
3. Author autogalaxy_workspace_test/scripts/quantity/modeling_visualization_jit.py:
   • Module docstring matching the imaging variant's, retargeted for quantity.
   • Part 1: same caching-probe shape, asserting cached_time < compile_time * 0.5 and analysis_mge._jitted_fit_from is not None.
   • __Visualization Sanity__ block: non-lensing template from PR test(viz): __Visualization Sanity__ rollout across modeling_visualization_jit*.py #54 — fit_2.model_data finite/non-zero, fit_2.figure_of_merit finite.
   • Part 2: live Nautilus quick-update on a small model_mge with iterations_per_quick_update, assert fit.png lands in the configured output path.
4. Local validation — python scripts/quantity/modeling_visualization_jit.py (no PYAUTO_TEST_MODE). Confirm Part 1 + Sanity + Part 2 pass end-to-end and the printed |model_data|.sum() / FoM are sensible.

Key Files

• autogalaxy_workspace_test/scripts/quantity/modeling_visualization_jit.py — new file (1 script).

Out of scope (Phase D.2.b.ii / D.2.b.iii)

• autogalaxy_workspace_test/scripts/ellipse/modeling_visualization_jit.py + visualization_jax.py (both missing — D.2.b.ii).
• autolens_workspace_test/scripts/weak/ (entire directory missing — D.2.b.iii).
• Any library code change.

Original Prompt

Click to expand starting prompt

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

• Template: autogalaxy_workspace_test/scripts/imaging/modeling_visualization_jit.py (PR test(viz): __Visualization Sanity__ rollout across modeling_visualization_jit*.py #54). Same Part-1 / Sanity / Part-2 structure; only dataset / model / analysis swap.
• API references: FitQuantity.model_data confirmed at autogalaxy/quantity/fit_quantity.py:71-72 (live API as of 2026-05-21). ag.AnalysisQuantity already used in the sibling visualization_jax.py script.
• Sanity-block style: PR test(viz): __Visualization Sanity__ rollout across modeling_visualization_jit*.py #54 establishes the non-lensing autogalaxy template — np.asarray(fit.model_data), finite/non-zero + finite FoM.

[Jammy2211]

Shipped — 2026-05-22

PR

• test(viz): quantity/modeling_visualization_jit.py — JIT-cached integration script #57 (merged 314be47)

Summary

Phase D.2.b.i shipped — scripts/quantity/modeling_visualization_jit.py is the first script to exercise AnalysisQuantity + use_jax=True + JIT-cached fit_for_visualization end-to-end. Mirrors the imaging variant's structure with quantity-specific dataset/model. Part 1 caching probe shows 264× speedup on the second fit_for_visualization call (compile ~2.4 s, cached ~9 ms); Sanity block passes with |model_data|.sum() = 290.87, figure_of_merit = 1005.17.

Session Notes

• Part 2 (live Nautilus) intentionally omitted. When I mirrored the imaging variant's Part-2 block, the search died inside Nautilus's first iteration with TracerArrayConversionError at geometry_profiles.py:168 (xp.array(self.centre) on a tuple of traced scalars under jax.vmap). The imaging variant escapes because light-profile MGE Gaussians take a different decorator path; quantity + IsothermalSph is the first dataset type to surface this site under vmap. Documented in the script's docstring with the precise traceback line and the restoration instructions for when the library lands the fix.
• Library follow-up authored. PyAutoPrompt/autogalaxy/geometry_profiles_centre_jax_traceable.md proposes xp.stack([self.centre[0], self.centre[1]]) as the safe form. When that ships in a release, Part 2 can be restored by mirroring the imaging variant verbatim.
• CI green this run. No pre-existing imaging/visualization.py flake reproduced (which we hit on the previous workspace task).

Phase D.2.b remaining

• D.2.b.ii — autogalaxy_workspace_test ellipse: needs both modeling_visualization_jit.py + visualization_jax.py authored from scratch (only non-JAX visualization.py exists today).
• D.2.b.iii — autolens_workspace_test weak lensing: entire scripts/weak/ directory missing.