Skip to content

feat: ag visualizer dispatch via fit_for_visualization #389

@Jammy2211

Description

@Jammy2211

Overview

PyAutoLens visualizers were updated in PR #443 (2026-04-19) to dispatch through analysis.fit_for_visualization(instance=instance) so the fit reused for plotting goes through autofit's cached jax.jit wrapper when use_jax_for_visualization=True. PyAutoGalaxy received matching pytree registration for FitImaging (PR #364, 2026-04-22) and FitInterferometer (PR #376), but the visualizer dispatch was never switched over for either. As a result use_jax_for_visualization=True on ag.AnalysisImaging / AnalysisInterferometer is currently a silent no-op for visualization. This task swaps the three call sites.

Plan

  • Three fit_from(instance=...) calls in PyAutoGalaxy visualizer files become fit_for_visualization(instance=...).
  • The change is safe for the NumPy default path because fit_for_visualization falls back to fit_from whenever use_jax_for_visualization=False.
  • Existing imaging JAX visualization smoke test now actually exercises the jit-cached path instead of silently no-opping.
  • Existing NumPy interferometer / imaging visualization smoke tests stay green.
  • No new pytree registration in this task — both imaging and interferometer pytrees are already shipped.
Detailed implementation plan

Affected Repositories

  • PyAutoGalaxy (primary)
  • autogalaxy_workspace_test (verification only — no edits)

Work Classification

Library

Branch Survey

Repository Current Branch Dirty?
./PyAutoGalaxy main clean

Suggested branch: feature/autogalaxy-viz-dispatch-swap
Worktree root: ~/Code/PyAutoLabs-wt/autogalaxy-viz-dispatch-swap/ (created later by /start_library)

Implementation Steps

  1. autogalaxy/imaging/model/visualizer.py:79 — replace fit = analysis.fit_from(instance=instance) with fit = analysis.fit_for_visualization(instance=instance) (single-analysis visualize path).
  2. autogalaxy/imaging/model/visualizer.py:176 — same replacement inside the visualize_combined list-comprehension that builds per-analysis fits for multi-analysis fits.
  3. autogalaxy/interferometer/model/visualizer.py:81 — same replacement (single-analysis visualize path; no visualize_combined exists in the interferometer visualizer).
  4. Run pytest test_autogalaxy/imaging and pytest test_autogalaxy/interferometer to confirm no regressions.
  5. Run /smoke_test autogalaxy_workspace_test imaging/visualization.py imaging/visualization_jax.py interferometer/visualization.py — all three must pass. The visualization_jax.py script will now actually exercise the jit-cached path rather than silently no-opping.

Key Files

  • autogalaxy/imaging/model/visualizer.py — two-line change (lines 79, 176)
  • autogalaxy/interferometer/model/visualizer.py — one-line change (line 81)

Reference (PyAutoLens equivalents — pattern to mirror)

  • PyAutoLens/autolens/imaging/model/visualizer.py:97 — single-analysis dispatch
  • PyAutoLens/autolens/imaging/model/visualizer.py:239 — multi-analysis combined dispatch
  • PyAutoLens/autolens/interferometer/model/visualizer.py:96 + :209

Original Prompt

Click to expand starting prompt

PyAutoLens visualizers were updated to dispatch through analysis.fit_for_visualization(instance=instance)
in PR #443 (commit 761c1778e, 2026-04-19) so that — when use_jax_for_visualization=True is set on
the analysis — the fit reused for plotting goes through the cached jax.jit wrapper added to the
autofit base Analysis in PR #1228.

PyAutoGalaxy was later given matching JAX pytree registration for FitImaging + DatasetModel +
Galaxies (PR #364, 2026-04-22) but the visualizer dispatch was never switched over. As a result
use_jax_for_visualization=True on ag.AnalysisImaging is currently a no-op for visualization —
the pytree machinery is in place but the call sites still go through the eager
analysis.fit_from(instance=instance) path.

What to change

  1. @PyAutoGalaxy/autogalaxy/imaging/model/visualizer.py:79 — swap
    fit = analysis.fit_from(instance=instance) for
    fit = analysis.fit_for_visualization(instance=instance).
  2. @PyAutoGalaxy/autogalaxy/imaging/model/visualizer.py:176 — same swap inside the
    visualize_combined path that builds per-analysis fits in a multi-analysis scenario.
  3. @PyAutoGalaxy/autogalaxy/interferometer/model/visualizer.py:81 — same swap.
    Note: AnalysisInterferometer pytree registration shipped in PyAutoGalaxy PR feat: register pytrees for AnalysisInterferometer #376
    (see complete.md 2026-04-? entry — _register_fit_interferometer_pytrees is now
    active in autogalaxy/interferometer/model/analysis.py:165-184), so dispatching
    through fit_for_visualization is now safe for interferometer. Originally this
    bullet was deferred — it was added to scope on 2026-05-08 once interferometer
    pytree registration landed.

fit_for_visualization is defined on the autofit base Analysis. It dispatches to a
jax.jit-cached wrapper when use_jax_for_visualization=True and falls back to plain fit_from
otherwise — so this change is safe for the NumPy default path.

Reference (PyAutoLens equivalent)

  • @PyAutoLens/autolens/imaging/model/visualizer.py:97 — single-analysis dispatch
  • @PyAutoLens/autolens/imaging/model/visualizer.py:239 — multi-analysis visualize_combined dispatch

Verification

  • autogalaxy_workspace_test/scripts/imaging/visualization_jax.py should still pass (uses
    use_jax_for_visualization=True) and will now actually exercise the jit-cached path rather
    than silently no-op.
  • autogalaxy_workspace_test/scripts/imaging/visualization.py should still pass (NumPy path —
    fit_for_visualization falls back to fit_from when the flag is off).
  • Run both via /smoke_test.
  • For interferometer the equivalent workspace_test scripts don't exist yet — coverage is added
    by autogalaxy_workspace_test/jax_viz_dataset_coverage.md (Phase 1C of the JAX
    visualization roadmap). Verify the dispatch swap doesn't break existing
    autogalaxy_workspace_test/scripts/interferometer/ (NumPy path — should be unchanged).

Out of scope

  • Production workspace adoption (autogalaxy_workspace scripts opting into
    use_jax_for_visualization=True) — defer until Path A from issued/fit_imaging_pytree.md lands.
  • AnalysisEllipse / AnalysisQuantity — these have no pytree registration today;
    covered by autogalaxy/fit_pytree_registration_other_datasets.md (which is now
    scoped to ellipse + quantity only — interferometer pytree registration shipped in
    PR feat: register pytrees for AnalysisInterferometer #376).

Background

Original feature: complete.md entries jax-visualization and mge-jit-visualization
(both 2026-04-19). Autogalaxy imaging pytree registration: complete.md entry for the
imaging port (2026-04-22).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions