Merge pull request #1228 from PyAutoLabs/feature/jax-visualization

feat: add use_jax_for_visualization flag and fit_for_visualization dispatch

Author: Jammy2211

autofit/non_linear/analysis/analysis.py

@@ -33,15 +33,48 @@ class Analysis(ABC):
     LATENT_KEYS = []
 
     def __init__(
-        self, use_jax : bool = False, **kwargs
+        self,
+        use_jax: bool = False,
+        use_jax_for_visualization: bool = False,
+        **kwargs,
     ):
         import os
         if os.environ.get("PYAUTO_DISABLE_JAX") == "1":
             use_jax = False
+            use_jax_for_visualization = False
+
+        if use_jax_for_visualization and not use_jax:
+            logger.warning(
+                "use_jax_for_visualization=True requires use_jax=True; "
+                "disabling use_jax_for_visualization."
+            )
+            use_jax_for_visualization = False
 
         self._use_jax = use_jax
+        self._use_jax_for_visualization = use_jax_for_visualization
         self.kwargs = kwargs
 
+    def fit_for_visualization(self, instance):
+        """
+        Build the fit used by the visualizer.
+
+        Currently a thin dispatch over ``self.fit_from``: when ``use_jax=True``
+        the fit is built on the eager JAX path (``self._xp is jnp``) and the
+        plotter materialises arrays to NumPy at the matplotlib boundary. The
+        ``use_jax_for_visualization`` flag is an explicit opt-in today — it
+        marks the intent to use JAX for visualization, and is the dispatch
+        point where full ``jax.jit``-wrapping will plug in once
+        :class:`autolens.imaging.fit_imaging.FitImaging` and its nested
+        autoarray types are registered as JAX pytrees (that work is tracked
+        separately — see ``admin_jammy/prompt/autolens/fit_imaging_pytree.md``).
+
+        ``fit_from`` is defined by Analysis subclasses (e.g. ``AnalysisImaging``),
+        not the base class — this method is only callable on subclasses that
+        provide it. Downstream visualizers should prefer this over calling
+        ``fit_from`` directly so the JIT seam stays in one place.
+        """
+        return self.fit_from(instance=instance)
+
     def __getattr__(self, item: str):
         """
         If a method starts with 'visualize_' then we assume it is associated with
@@ -306,8 +339,15 @@ def supports_background_update(self) -> bool:
 
     @property
     def supports_jax_visualization(self) -> bool:
-        """Whether the visualizer can work directly with JAX arrays."""
-        return False
+        """
+        Whether the visualizer can work directly with JAX arrays.
+
+        Derived from the ``use_jax_for_visualization`` flag passed at
+        construction time. Subclasses may override to force a specific
+        answer (e.g. an Analysis that has been audited to support JAX
+        visualization unconditionally).
+        """
+        return self._use_jax_for_visualization
 
     def perform_quick_update(self, paths, instance):
         raise NotImplementedError

test_autofit/analysis/test_use_jax_for_visualization.py

@@ -0,0 +1,73 @@
+"""Tests for the ``use_jax_for_visualization`` flag on ``Analysis``."""
+
+import pytest
+
+import autofit as af
+
+
+class _FittableAnalysis(af.Analysis):
+    """Minimal Analysis subclass with a trivial ``fit_from`` for dispatch tests."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.fit_from_calls = 0
+
+    def log_likelihood_function(self, instance):
+        return 0.0
+
+    def fit_from(self, instance):
+        self.fit_from_calls += 1
+        return ("fit", instance)
+
+
+def test_default_flag_is_false():
+    analysis = af.Analysis()
+    assert analysis._use_jax is False
+    assert analysis._use_jax_for_visualization is False
+    assert analysis.supports_jax_visualization is False
+
+
+def test_flag_requires_use_jax(caplog):
+    with caplog.at_level("WARNING"):
+        analysis = af.Analysis(use_jax=False, use_jax_for_visualization=True)
+    assert analysis._use_jax_for_visualization is False
+    assert any("requires use_jax=True" in r.message for r in caplog.records)
+
+
+def test_flag_accepted_when_use_jax_true():
+    analysis = af.Analysis(use_jax=True, use_jax_for_visualization=True)
+    assert analysis._use_jax is True
+    assert analysis._use_jax_for_visualization is True
+    assert analysis.supports_jax_visualization is True
+
+
+def test_pyauto_disable_jax_env_var_clears_both_flags(monkeypatch):
+    monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1")
+    analysis = af.Analysis(use_jax=True, use_jax_for_visualization=True)
+    assert analysis._use_jax is False
+    assert analysis._use_jax_for_visualization is False
+
+
+def test_fit_for_visualization_dispatches_to_fit_from():
+    analysis = _FittableAnalysis(use_jax=True, use_jax_for_visualization=True)
+    result = analysis.fit_for_visualization(instance="sentinel")
+    assert result == ("fit", "sentinel")
+    assert analysis.fit_from_calls == 1
+
+
+def test_fit_for_visualization_works_without_flag():
+    analysis = _FittableAnalysis()
+    result = analysis.fit_for_visualization(instance="sentinel")
+    assert result == ("fit", "sentinel")
+    assert analysis.fit_from_calls == 1
+
+
+def test_subclass_can_override_supports_jax_visualization():
+    class ForcedAnalysis(af.Analysis):
+        @property
+        def supports_jax_visualization(self):
+            return True
+
+    analysis = ForcedAnalysis()
+    assert analysis._use_jax_for_visualization is False
+    assert analysis.supports_jax_visualization is True