Merge pull request #472 from PyAutoLabs/feature/latent-class-phase2

refactor(latent): LatentGalaxy class + declare Analysis.Latent (Phase 2)

Author: Jammy2211

autogalaxy/__init__.py

@@ -74,6 +74,8 @@
 from .gui.scribbler import Scribbler
 from .imaging.fit_imaging import FitImaging
 from .imaging.model.analysis import AnalysisImaging
+from autofit import Latent
+from .imaging.model.latent import LatentGalaxy
 from .imaging.simulator import SimulatorImaging
 from .interferometer.simulator import SimulatorInterferometer
 from .interferometer.fit_interferometer import FitInterferometer

autogalaxy/imaging/model/analysis.py

@@ -22,6 +22,7 @@
 from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages
 from autogalaxy.analysis.analysis.dataset import AnalysisDataset
 from autogalaxy.cosmology.model import LensingCosmology
+from autogalaxy.imaging.model.latent import LatentGalaxy
 from autogalaxy.imaging.model.result import ResultImaging
 from autogalaxy.imaging.model.visualizer import VisualizerImaging
 from autogalaxy.imaging.fit_imaging import FitImaging
@@ -33,6 +34,7 @@
 class AnalysisImaging(AnalysisDataset):
     Result = ResultImaging
     Visualizer = VisualizerImaging
+    Latent = LatentGalaxy
 
     def __init__(
         self,
@@ -88,11 +90,6 @@ def __init__(
     def imaging(self):
         return self.dataset
 
-    @property
-    def LATENT_KEYS(self):
-        from autogalaxy.imaging.model.latent import latent_keys_enabled
-        return latent_keys_enabled()
-
     def log_likelihood_function(self, instance: af.ModelInstance) -> float:
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
@@ -176,33 +173,6 @@ def fit_from(self, instance: af.ModelInstance) -> FitImaging:
             xp=self._xp,
         )
 
-    def compute_latent_variables(self, parameters, model):
-        """
-        Compute the catalogue of latent variables enabled in
-        ``config/latent.yaml`` for the given parameter vector.
-
-        Returns a tuple positionally aligned with :attr:`LATENT_KEYS` —
-        PyAutoFit zips it with the keys at
-        ``autofit/non_linear/analysis/analysis.py:285`` and stacks per
-        sample for the JIT batch path at lines 223-234.
-
-        Raises ``NotImplementedError`` when no latents are enabled so
-        PyAutoFit's outer ``except NotImplementedError`` short-circuits
-        the latent pipeline cleanly (no empty ``latent.csv`` written).
-        """
-        from autogalaxy.imaging.model.latent import LATENT_FUNCTIONS
-
-        keys = self.LATENT_KEYS
-        if not keys:
-            raise NotImplementedError
-
-        xp = self._xp
-        instance = model.instance_from_vector(vector=parameters)
-        fit = self.fit_from(instance=instance)
-        magzero = self.kwargs.get("magzero", None)
-        context = {"fit": fit, "magzero": magzero, "xp": xp}
-        return tuple(LATENT_FUNCTIONS[k](**context) for k in keys)
-
     @staticmethod
     def _register_fit_imaging_pytrees() -> None:
         """Register every type reachable from a ``FitImaging`` return value

autogalaxy/imaging/model/latent.py

@@ -16,6 +16,7 @@
 
 import numpy as np
 
+import autofit as af
 from autoconf import conf
 
 logger = logging.getLogger(__name__)
@@ -160,3 +161,38 @@ def latent_keys_enabled(yaml_config: Optional[Dict[str, bool]] = None) -> List[s
             continue
         enabled.append(key)
     return enabled
+
+
+class LatentGalaxy(af.Latent):
+    """
+    Latent-variable catalogue for galaxy imaging fits, declared on
+    ``AnalysisImaging`` as ``Latent = LatentGalaxy`` (mirrors ``Visualizer`` /
+    ``Result``).
+
+    :meth:`keys` returns the config-enabled latent names; :meth:`variables`
+    dispatches the :data:`LATENT_FUNCTIONS` registry on a per-sample fit.
+    Subclass to add project-specific latents, composing the library values via
+    ``LatentGalaxy.variables(analysis, parameters, model)``.
+    """
+
+    # Preserves the previous ``AnalysisDataset.LATENT_BATCH_MODE = "jit"``: the
+    # per-sample jit path is used because the lensing Einstein-radius latent
+    # (shared dataset base) routes through ``ZeroSolver``, which is
+    # vmap-incompatible.
+    BATCH_MODE = "jit"
+
+    @staticmethod
+    def keys(analysis) -> List[str]:
+        return latent_keys_enabled()
+
+    @staticmethod
+    def variables(analysis, parameters, model):
+        keys = latent_keys_enabled()
+        if not keys:
+            raise NotImplementedError
+        xp = analysis._xp
+        instance = model.instance_from_vector(vector=parameters)
+        fit = analysis.fit_from(instance=instance)
+        magzero = analysis.kwargs.get("magzero", None)
+        context = {"fit": fit, "magzero": magzero, "xp": xp}
+        return tuple(LATENT_FUNCTIONS[k](**context) for k in keys)

test_autogalaxy/imaging/model/test_latent.py

@@ -10,6 +10,7 @@
 from autogalaxy.imaging.model import latent as _latent_module
 from autogalaxy.imaging.model.latent import (
     LATENT_FUNCTIONS,
+    LatentGalaxy,
     ab_mag_via_flux_from,
     flux_mujy_via_ab_mag_from,
     latent_keys_enabled,
@@ -139,7 +140,7 @@ def test_latent_keys_enabled_drops_unknown_with_warning(caplog):
     assert any("never_registered_latent" in rec.message for rec in caplog.records)
 
 
-def test_analysis_imaging_compute_latent_variables_aligns_with_latent_keys(
+def test_latent_galaxy_variables_aligns_with_keys(
     masked_imaging_7x7,
 ):
     galaxy = ag.Galaxy(redshift=0.5, light=ag.lp.Sersic(intensity=0.1))
@@ -150,35 +151,34 @@ def test_analysis_imaging_compute_latent_variables_aligns_with_latent_keys(
     )
 
     parameters = np.array(model.physical_values_from_prior_medians)
-    values = analysis.compute_latent_variables(parameters=parameters, model=model)
+    values = LatentGalaxy.variables(analysis, parameters=parameters, model=model)
+    keys = LatentGalaxy.keys(analysis)
 
     assert isinstance(values, tuple)
-    assert len(values) == len(analysis.LATENT_KEYS)
+    assert len(values) == len(keys)
     # test_autogalaxy/config/latent.yaml enables both keys, raw flux first.
-    assert analysis.LATENT_KEYS == ["total_galaxy_0_flux", "total_galaxy_0_flux_mujy"]
+    assert keys == ["total_galaxy_0_flux", "total_galaxy_0_flux_mujy"]
     assert all(np.isfinite(v) for v in values)
 
 
-def test_analysis_imaging_compute_latent_variables_raises_when_empty(monkeypatch):
+def test_latent_galaxy_variables_raises_when_empty(monkeypatch):
     # When no latents are enabled, autofit's `except NotImplementedError`
-    # at autofit/non_linear/analysis/analysis.py:304 short-circuits the
-    # latent pipeline. We match that contract by raising explicitly.
-    monkeypatch.setattr(
-        ag.AnalysisImaging,
-        "LATENT_KEYS",
-        property(lambda self: []),
-    )
+    # short-circuits the latent pipeline. LatentGalaxy.variables matches that
+    # contract by raising explicitly.
+    monkeypatch.setattr(_latent_module, "latent_keys_enabled", lambda *a, **k: [])
     analysis = ag.AnalysisImaging(dataset=MagicMock(), use_jax=False)
 
     with pytest.raises(NotImplementedError):
-        analysis.compute_latent_variables(parameters=np.array([]), model=MagicMock())
+        LatentGalaxy.variables(analysis, parameters=np.array([]), model=MagicMock())
 
 
-def test_analysis_imaging_latent_keys_property_reads_config():
+def test_analysis_imaging_declares_latent_galaxy_and_keys_read_config():
     # The autouse fixture in test_autogalaxy/conftest.py pushes the test
     # config dir whose latent.yaml enables both keys.
     dataset = MagicMock()
     analysis = ag.AnalysisImaging(dataset=dataset, use_jax=False)
 
-    assert analysis.LATENT_KEYS == ["total_galaxy_0_flux", "total_galaxy_0_flux_mujy"]
-    assert set(analysis.LATENT_KEYS).issubset(LATENT_FUNCTIONS.keys())
+    assert ag.AnalysisImaging.Latent is LatentGalaxy
+    keys = ag.AnalysisImaging.Latent.keys(analysis)
+    assert keys == ["total_galaxy_0_flux", "total_galaxy_0_flux_mujy"]
+    assert set(keys).issubset(LATENT_FUNCTIONS.keys())