Merge pull request #534 from PyAutoLabs/feature/latent-module-autolens

feat: first-class lensing latent variable API in PyAutoLens

Author: Jammy2211

autolens/analysis/latent.py

@@ -0,0 +1,176 @@
+"""
+Latent variables for PyAutoLens analyses.
+
+All latents take a generic ``fit`` argument and access ``fit.tracer``,
+``fit.galaxy_image_dict`` and ``fit.dataset.grids.lp`` — APIs that exist
+identically on both ``FitImaging`` (``autolens/imaging/fit_imaging.py:176``)
+and ``FitInterferometer`` (``autolens/interferometer/fit_interferometer.py:176``).
+The registry is dataset-agnostic; a future ``AnalysisInterferometer``
+wiring can reuse it without code duplication.
+
+User-level enable/disable: each key in ``autolens/config/latent.yaml`` maps
+to a bool. All five default ``false`` because ``compute_latent_samples``
+runs on every fit (``latent_after_fit: true`` in autofit's default
+``output.yaml``) and the latents that require ``magzero`` would otherwise
+crash existing fits where ``magzero`` is not passed.
+"""
+import logging
+from typing import Callable, Dict, List, Optional
+
+import numpy as np
+
+from autoconf import conf
+from autogalaxy.imaging.model.latent import (
+    ab_mag_via_flux_from,
+    flux_mujy_via_ab_mag_from,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _require_magzero(magzero, name):
+    if magzero is None:
+        raise ValueError(
+            f"magzero must be passed to the Analysis via kwargs to compute "
+            f"the '{name}' latent. Disable it in config/latent.yaml or "
+            f"pass magzero=<value>."
+        )
+
+
+def total_lens_flux_mujy(fit, magzero, xp=np):
+    """
+    Total integrated flux of the lens galaxy (``fit.tracer.galaxies[0]``),
+    magzero-converted to microjanskies.
+
+    Returns NaN when galaxy 0 has no light profile (raises ``KeyError`` /
+    ``AttributeError`` inside ``fit.galaxy_image_dict``).
+    """
+    _require_magzero(magzero, "total_lens_flux_mujy")
+    try:
+        image = fit.galaxy_image_dict[fit.tracer.galaxies[0]]
+    except (AttributeError, KeyError, IndexError):
+        return xp.nan
+    total_flux = xp.sum(image.array)
+    return flux_mujy_via_ab_mag_from(
+        ab_mag=ab_mag_via_flux_from(flux=total_flux, magzero=magzero, xp=xp),
+        xp=xp,
+    )
+
+
+def total_lensed_source_flux_mujy(fit, magzero, xp=np):
+    """
+    Image-plane integrated flux of the source galaxy after lensing
+    (``fit.galaxy_image_dict[fit.tracer.galaxies[-1]]``).
+    """
+    _require_magzero(magzero, "total_lensed_source_flux_mujy")
+    try:
+        image = fit.galaxy_image_dict[fit.tracer.galaxies[-1]]
+    except (AttributeError, KeyError, IndexError):
+        return xp.nan
+    total_flux = xp.sum(image.array)
+    return flux_mujy_via_ab_mag_from(
+        ab_mag=ab_mag_via_flux_from(flux=total_flux, magzero=magzero, xp=xp),
+        xp=xp,
+    )
+
+
+def total_source_flux_mujy(fit, magzero, xp=np):
+    """
+    Source-plane intrinsic flux of the source galaxy, via
+    ``fit.tracer.galaxies[-1].image_2d_from(grid=fit.dataset.grids.lp)``.
+    """
+    _require_magzero(magzero, "total_source_flux_mujy")
+    try:
+        source_image = fit.tracer.galaxies[-1].image_2d_from(
+            grid=fit.dataset.grids.lp, xp=xp
+        )
+    except (AttributeError, IndexError):
+        return xp.nan
+    total_flux = xp.sum(source_image.array)
+    return flux_mujy_via_ab_mag_from(
+        ab_mag=ab_mag_via_flux_from(flux=total_flux, magzero=magzero, xp=xp),
+        xp=xp,
+    )
+
+
+def magnification(fit, magzero, xp=np):
+    """
+    Ratio of image-plane to source-plane source flux — the integrated
+    magnification implied by the lens model and source light profile.
+
+    ``magzero`` is accepted but unused (the µJy conversions cancel in the
+    ratio). It's still required in the signature so the dispatcher can
+    pass a uniform context dict to every latent function.
+    """
+    lensed = total_lensed_source_flux_mujy(fit=fit, magzero=magzero, xp=xp)
+    intrinsic = total_source_flux_mujy(fit=fit, magzero=magzero, xp=xp)
+    return lensed / intrinsic
+
+
+def effective_einstein_radius(fit, magzero, xp=np):
+    """
+    Effective Einstein radius via the tangential critical curve.
+
+    JAX path: ``LensCalc.einstein_radius_jit_from(init_guess=fan)``, where
+    ``fan`` is a fixed 4-seed fan at ±1 arcsec from the lens centre — the
+    JIT-compatible variant required because ``ZeroSolver`` (line 1520 of
+    ``autogalaxy/operate/lens_calc.py``) uses ``lax.cond`` /
+    ``lax.while_loop`` early termination that is incompatible with
+    ``jax.vmap`` but fine under ``jax.jit``.
+
+    NumPy path: ``LensCalc.einstein_radius_from(grid=fit.dataset.grids.lp)``.
+    """
+    from autogalaxy.operate.lens_calc import LensCalc
+
+    try:
+        lens_calc = LensCalc.from_mass_obj(fit.tracer)
+        if xp is not np:
+            import jax.numpy as jnp
+            init_guess = jnp.array(
+                [[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0]]
+            )
+            return lens_calc.einstein_radius_jit_from(init_guess=init_guess)
+        return lens_calc.einstein_radius_from(grid=fit.dataset.grids.lp)
+    except (ValueError, AttributeError):
+        return xp.nan
+
+
+LATENT_FUNCTIONS: Dict[str, Callable] = {
+    "total_lens_flux_mujy": total_lens_flux_mujy,
+    "total_lensed_source_flux_mujy": total_lensed_source_flux_mujy,
+    "total_source_flux_mujy": total_source_flux_mujy,
+    "magnification": magnification,
+    "effective_einstein_radius": effective_einstein_radius,
+}
+
+
+def latent_keys_enabled(yaml_config: Optional[Dict[str, bool]] = None) -> List[str]:
+    """
+    Return the ordered list of enabled latent keys.
+
+    Reads ``conf.instance["latent"]`` (a flat ``key: bool`` dict from
+    ``autolens/config/latent.yaml``) unless ``yaml_config`` is passed
+    explicitly — tests pass a literal dict to avoid pushing a temporary
+    config directory.
+
+    Unknown keys (present in the yaml but not in :data:`LATENT_FUNCTIONS`)
+    are dropped with a logger warning rather than raising — yaml carries
+    forward-compat entries for latents that ship in later releases.
+    """
+    if yaml_config is None:
+        yaml_config = dict(conf.instance["latent"])
+
+    enabled: List[str] = []
+    for key, on in yaml_config.items():
+        if not on:
+            continue
+        if key not in LATENT_FUNCTIONS:
+            logger.warning(
+                "latent.yaml lists '%s' but no such latent is registered; "
+                "dropping. Known latents: %s",
+                key,
+                sorted(LATENT_FUNCTIONS),
+            )
+            continue
+        enabled.append(key)
+    return enabled

autolens/config/latent.yaml

@@ -0,0 +1,43 @@
+# Toggles for the catalogue of latent variables computed by `AnalysisImaging`
+# (and, when wired in a follow-up, `AnalysisInterferometer`).
+#
+# Each entry maps a registered latent name (see
+# `autolens/analysis/latent.py::LATENT_FUNCTIONS`) to a bool. Setting `false`
+# excludes that latent from `LATENT_KEYS` so it is neither computed nor
+# written to `latent/samples.csv` / `latent/latent_summary.json`.
+#
+# Workspaces should mirror this file in their own `config/latent.yaml` to
+# override defaults locally (workspace values shadow library values).
+#
+# All keys ship `false` because:
+#  - `compute_latent_samples` runs on every fit (`latent_after_fit: true`
+#    in autofit's default output.yaml), so on-by-default would crash any
+#    existing fit that doesn't pass `magzero` to the Analysis.
+#  - autoconf lowercases yaml keys at read time, so the registry/yaml
+#    names must be snake_case-lowercase (this leaks into the `latent.csv`
+#    column header — e.g. `total_lens_flux_mujy`, not `_muJy`).
+
+# `total_lens_flux_mujy` — total integrated flux of the lens galaxy
+# (`fit.tracer.galaxies[0]`) in microjanskies. Requires `magzero` via
+# Analysis kwargs. Returns NaN when the lens has no light profile.
+total_lens_flux_mujy: false
+
+# `total_lensed_source_flux_mujy` — image-plane integrated flux of the
+# source galaxy after lensing (`fit.galaxy_image_dict[tracer.galaxies[-1]]`)
+# in microjanskies. Requires `magzero`.
+total_lensed_source_flux_mujy: false
+
+# `total_source_flux_mujy` — source-plane intrinsic flux of the source
+# galaxy (computed via the source's light profile on `fit.dataset.grids.lp`)
+# in microjanskies. Requires `magzero`.
+total_source_flux_mujy: false
+
+# `magnification` — ratio of image-plane lensed source flux to source-plane
+# intrinsic source flux. Dimensionless; `magzero` is accepted but unused.
+magnification: false
+
+# `effective_einstein_radius` — effective Einstein radius in arcseconds,
+# from the tangential critical curve via
+# `LensCalc.einstein_radius_jit_from` (JAX) or `einstein_radius_from`
+# (numpy). Does NOT require `magzero`.
+effective_einstein_radius: false

autolens/imaging/model/analysis.py

@@ -32,6 +32,38 @@ class AnalysisImaging(AnalysisDataset):
     Result = ResultImaging
     Visualizer = VisualizerImaging
 
+    @property
+    def LATENT_KEYS(self):
+        from autolens.analysis.latent import latent_keys_enabled
+        return latent_keys_enabled()
+
+    def compute_latent_variables(self, parameters, model):
+        """
+        Compute the catalogue of lensing 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 autolens.analysis.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)
+
     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