[Jammy2211]

Summary

Adds an explicit use_jax_for_visualization flag on Analysis (closes #1227) and a new Analysis.fit_for_visualization(instance) dispatch method that the search-level visualizer now calls in place of fit_from.

Today fit_for_visualization is a thin pass-through to fit_from — this is the "Path C" pilot from the issue: with use_jax=True the fit is built on the eager JAX path (self._xp is jnp) and matplotlib plotters materialise arrays to NumPy at the boundary. No jax.jit wrapping yet. The dispatch method is the seam where the full-JIT "Path A" will plug in once FitImaging and its nested autoarray types are registered as JAX pytrees (tracked separately in admin_jammy/prompt/autolens/fit_imaging_pytree.md).

API Changes

Analysis.__init__ gains a use_jax_for_visualization: bool = False kwarg. It requires use_jax=True (otherwise it's coerced to False with a warning), and is cleared by PYAUTO_DISABLE_JAX=1 alongside use_jax. Two new public members: Analysis.fit_for_visualization(instance) (dispatch seam) and Analysis.supports_jax_visualization (property, subclass-overridable). No removals or renames — the change is additive.

See full details below.

Test Plan

• pytest test_autofit/analysis/test_use_jax_for_visualization.py — 7 new tests, all passing
• pytest test_autofit/ full suite — 1232 passed
• Downstream: PyAutoLens feature/jax-visualization routes visualizer through the new dispatch method (companion PR in PyAutoLabs/PyAutoLens)
• Downstream: autolens_workspace_test/scripts/imaging/visualization_jax.py pilot script produces fit.png + tracer.png on the eager JAX path

Full API Changes (for automation & release notes)

Added

• Analysis.__init__(..., use_jax_for_visualization: bool = False, ...) — new kwarg. Marks intent to run visualization via the JAX path. Requires use_jax=True (else coerced to False with a warning); cleared when PYAUTO_DISABLE_JAX=1.
• Analysis.fit_for_visualization(instance) — new method. Builds the fit used by the search-level visualizer. Currently a thin wrapper over self.fit_from(instance=instance); this is the seam where a jax.jit-wrapped fit will plug in once FitImaging becomes a JAX pytree.
• Analysis.supports_jax_visualization — new property. Returns self._use_jax_for_visualization; subclasses may override to force-enable regardless of the flag.

Changed Behaviour

• Analysis.__init__ now logs a warning and coerces use_jax_for_visualization to False if it is passed True without use_jax=True. Previously no such option existed.
• The PYAUTO_DISABLE_JAX=1 environment guard now clears both use_jax and use_jax_for_visualization at construction time.

Removed / Renamed

None.

Migration

No migration required — the new flag defaults to False and existing code is unaffected. To opt in:

analysis = MyAnalysis(use_jax=True, use_jax_for_visualization=True, ...)

🤖 Generated with Claude Code

[Jammy2211]

Workspace PR: PyAutoLabs/autolens_workspace_test#31