feat: first-class lensing latent variable API in PyAutoLens

[Jammy2211]

Overview

Builds on the PyAutoGalaxy latent infrastructure shipped in PyAutoGalaxy #441 by adding the lensing-specific latents that need a Tracer. PyAutoLens's AnalysisImaging does NOT inherit the PyAutoGalaxy LATENT_KEYS / compute_latent_variables via MRO — they live on parallel imaging Analysis classes (autolens.imaging.model.AnalysisImaging vs autogalaxy.imaging.model.AnalysisImaging), not the shared dataset base — so this PR adds an explicit dispatch layer in PyAutoLens with its own registry.

After this ships, sub-prompt #3 (euclid migration) can drop the bespoke LATENT_KEYS + compute_latent_variables in euclid_strong_lens_modeling_pipeline/util.py:306-490 and inherit the library catalogue.

Plan

• Add autolens/analysis/latent.py for the one tracer-derived latent (effective_einstein_radius, delegating to PyAutoGalaxy's LensCalc.einstein_radius_jit_from so its closure cache is reused).
• Add autolens/imaging/model/latent.py for image-derived lensing latents — total_lens_flux_mujy, total_source_flux_mujy, total_lensed_source_flux_mujy, magnification. Helpers ab_mag_via_flux_from / flux_mujy_via_ab_mag_from are imported from PyAutoGalaxy (no duplication).
• Add autolens/config/latent.yaml with all 5 keys default OFF (same reasoning as Drop Binder URLs, pin Colab URLs to workspace tag 2026.4.13.6 #441: magzero requirement + compute_latent_samples runs on every fit, so default-on would crash existing users).
• Wire AnalysisImaging with a LATENT_KEYS @property reading from conf.instance["latent"] and a compute_latent_variables(parameters, model) method dispatching through a local LATENT_FUNCTIONS registry. PyAutoLens defines its own registry (no re-export of PyAutoGalaxy's total_galaxy_0_flux_mujy — the lens-light equivalent is total_lens_flux_mujy, named in the lensing vocabulary).
• Apply the same NotImplementedError-when-empty short-circuit as Drop Binder URLs, pin Colab URLs to workspace tag 2026.4.13.6 #441 so autofit's except NotImplementedError: return None skips the latent pipeline cleanly.
• Preserve LATENT_BATCH_MODE = "jit" (inherited via autogalaxy/analysis/analysis/dataset.py:28 — ZeroSolver in the Einstein-radius helper is vmap-incompatible).
• Unit tests for each latent function + AnalysisImaging end-to-end + a spy/mock test confirming effective_einstein_radius actually calls LensCalc.einstein_radius_jit_from.

Detailed implementation plan

Affected Repositories

• PyAutoLens (primary)

Work Classification

Library

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoLens	main	CLAUDE.md + README.md modified (pre-existing PyAutoPaper section adds — unrelated)

Suggested branch: feature/latent-module-autolens
Worktree root: ~/Code/PyAutoLabs-wt/latent-module-autolens/ (created later by /start_library)

Implementation Steps

1. Create autolens/analysis/latent.py with effective_einstein_radius(fit, magzero, xp=np):

   • JAX path: LensCalc.from_mass_obj(fit.tracer).einstein_radius_jit_from(init_guess=fixed_fan) where fixed_fan = jnp.array([[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0]]) (matches euclid util.py:502-509).
   • NumPy path: LensCalc.from_mass_obj(fit.tracer).einstein_radius_from(grid=fit.dataset.grids.lp).
   • Returns xp.nan on ValueError / AttributeError.
   • Cache the (closure, solver) pair properly per memory feedback_jax_closure_cache_busts. The cache lives inside LensCalc (lines 1580-1586) — just don't construct a fresh LensCalc each call within the same Analysis.

2. Create autolens/imaging/model/latent.py with the four image-derived latents. Each function takes (fit, magzero, xp=np). Use from autogalaxy.imaging.model.latent import ab_mag_via_flux_from, flux_mujy_via_ab_mag_from.

   • total_lens_flux_mujy(fit, magzero, xp=np) — fit.galaxy_image_dict[fit.galaxies[0]], sum, → AB mag → µJy. Raises ValueError if magzero is None (memory feedback_no_silent_guards).
   • total_lensed_source_flux_mujy(fit, magzero, xp=np) — fit.galaxy_image_dict[fit.tracer.galaxies[-1]] (image-plane image of source after lensing).
   • total_source_flux_mujy(fit, magzero, xp=np) — fit.tracer.galaxies[-1].image_2d_from(grid=fit.dataset.grids.lp) (source-plane intrinsic image).
   • magnification(fit, magzero, xp=np) — ratio of total_lensed_source_flux_mujy / total_source_flux_mujy. magzero param accepted but unused; dimensionless ratio.

3. Create autolens/config/latent.yaml — all 5 keys default false. Each key gets a comment explaining what it computes and that it requires magzero (except magnification which is dimensionless).

4. Wire AnalysisImaging (autolens/imaging/model/analysis.py):

   • Build a single LATENT_FUNCTIONS registry: {**lens_analysis_latents, **lens_imaging_latents} where the two source modules each export a per-module dict.
   • LATENT_KEYS @property: reads conf.instance["latent"] (lens-side latent.yaml), filters to enabled keys present in LATENT_FUNCTIONS.
   • compute_latent_variables(parameters, model): raises NotImplementedError when LATENT_KEYS is empty; otherwise dispatches through LATENT_FUNCTIONS[k](**context) for k in LATENT_KEYS, returning a tuple aligned to keys. Context dict: {"fit": fit_from(instance_from_vector(parameters)), "magzero": self.kwargs.get("magzero"), "xp": self._xp}.

5. Unit tests at test_autolens/analysis/test_latent.py and test_autolens/imaging/model/test_latent.py:

   • Per-latent against known toy models (numpy only — memory feedback_no_jax_in_unit_tests).
   • effective_einstein_radius spy test asserting it calls into LensCalc.einstein_radius_jit_from / einstein_radius_from (mock or monkeypatch).
   • On/off filtering via explicit yaml_config dict.
   • Order preservation across registry → LATENT_KEYS → return tuple.
   • End-to-end through AnalysisImaging.compute_latent_variables with at least one key enabled in the test config mirror.
   • compute_latent_variables raises NotImplementedError when LATENT_KEYS == [].
   • test_autolens/config/latent.yaml mirror with at least one latent enabled so the end-to-end test exercises the dispatch path.

Key Files

• autolens/analysis/latent.py — new module (effective_einstein_radius)
• autolens/imaging/model/latent.py — new module (4 image-derived latents)
• autolens/config/latent.yaml — new config
• autolens/imaging/model/analysis.py — wire AnalysisImaging (parallel to PyAutoGalaxy's PR Drop Binder URLs, pin Colab URLs to workspace tag 2026.4.13.6 #441 pattern)
• test_autolens/config/latent.yaml — test mirror
• test_autolens/analysis/test_latent.py — new tests
• test_autolens/imaging/model/test_latent.py — new tests

Constraints to preserve

• LATENT_BATCH_MODE = "jit" inherited via autogalaxy/analysis/analysis/dataset.py:28 — do not change.
• All yaml keys snake_case-lowercase (memory feedback_autoconf_lowercases_yaml_keys).
• Default OFF in library yaml to avoid the magzero regression that Drop Binder URLs, pin Colab URLs to workspace tag 2026.4.13.6 #441 caught.
• Reuse LensCalc.einstein_radius_jit_from rather than reimplementing the ZeroSolver loop.

Original Prompt

Click to expand starting prompt

Add first-class lensing latent-variable modules to PyAutoLens

Context

Parent epic: PyAutoPrompt/z_features/latent_refactor.md.
Depends on the PyAutoGalaxy spine in autogalaxy/latent_module.md (shipped as PyAutoGalaxy #441).

PyAutoGalaxy gets galaxy-level / image-flux latents. This task adds the lensing-specific latents that need a Tracer (mass model + sources): magnification, effective Einstein radius, lensed source flux, total source flux. Today these live in euclid_strong_lens_modeling_pipeline/util.py:306-490 as part of a bespoke AnalysisImaging subclass. After this task lands, the euclid pipeline (sub-prompt #3) can drop its custom subclass and inherit the curated set.

Task

1. Create autolens/analysis/latent.py — tracer-derived, dataset-agnostic latents. effective_einstein_radius must delegate to LensCalc.einstein_radius_jit_from() at PyAutoGalaxy/autogalaxy/operate/lens_calc.py:1520.

2. Create autolens/imaging/model/latent.py — lensing imaging-derived latents that need a FitImaging: total_lens_flux_mujy, total_lensed_source_flux_mujy, total_source_flux_mujy, magnification. Use ab_mag_via_flux_from and flux_mujy_via_ab_mag_from imported from PyAutoGalaxy (shipped in Drop Binder URLs, pin Colab URLs to workspace tag 2026.4.13.6 #441).

3. Create autolens/config/latent.yaml — flat dict of latent_key: bool. Mirror autolens/config/output.yaml. Per Drop Binder URLs, pin Colab URLs to workspace tag 2026.4.13.6 #441's regression learning, default everything OFF.

4. Wire AnalysisImaging (autolens/imaging/model/analysis.py): read config/latent.yaml, build LATENT_KEYS via @property, dispatch compute_latent_variables through a local LATENT_FUNCTIONS registry composed from the analysis-latent module and the imaging-latent module. Preserve LATENT_BATCH_MODE = "jit" (inherited via autogalaxy/analysis/analysis/dataset.py:28). Raise NotImplementedError when LATENT_KEYS == [] so autofit skips cleanly.

5. Unit tests at test_autolens/analysis/test_latent.py and test_autolens/imaging/model/test_latent.py. Spy/mock test that effective_einstein_radius calls into LensCalc.einstein_radius_jit_from. No JAX in unit tests.

Where to look

• PyAutoFit hook (do not modify): autofit/non_linear/analysis/analysis.py:34, 170, 285.
• Einstein radius JAX helper: PyAutoGalaxy/autogalaxy/operate/lens_calc.py:1520-1537 and the closure cache at 1580-1586.
• LATENT_BATCH_MODE constraint: PyAutoGalaxy/autogalaxy/analysis/analysis/dataset.py:28.
• Reference implementation: euclid_strong_lens_modeling_pipeline/util.py:306-490.
• PyAutoGalaxy mirror (now shipped): PyAutoGalaxy/autogalaxy/imaging/model/latent.py, autogalaxy/config/latent.yaml, autogalaxy/imaging/model/analysis.py LATENT_KEYS property and compute_latent_variables method.

Suggested branch

feature/latent-module-autolens

Notes

• Lens-side latent is named total_lens_flux_mujy (not total_galaxy_0_flux_mujy from PyAutoGalaxy). Per-user direction during /start_dev — PyAutoLens uses lensing-vocabulary names, doesn't re-export PyAutoGalaxy's registry.
• Default OFF in library yaml to avoid the magzero regression that PyAutoGalaxy Drop Binder URLs, pin Colab URLs to workspace tag 2026.4.13.6 #441 caught.
• Per memory feedback_jax_closure_cache_busts: when delegating to einstein_radius_jit_from, ensure the same (closure, solver) pair is reused across calls.
• Per memory feedback_no_silent_guards: if magzero is missing when an enabled latent needs it, raise loudly.

[Jammy2211]

Smoke Test Results — 2026-05-23

All 43 smoke tests passed across 6 workspaces. 1 pre-existing skip.

Worktree verified: /home/jammy/Code/PyAutoLabs-wt/latent-module-autolens/PyAutoLens/autolens/__init__.py
(sibling libraries autofit, autoarray, autogalaxy also resolved through worktree root)

Workspace	Passed	Failed	Skipped	Total
autofit_workspace	7	0	0	7
autogalaxy_workspace	6	0	0	6
autolens_workspace	7	0	0	7
autolens_workspace_test	11	0	1	12
euclid_strong_lens_modeling_pipeline	6	0	0	6
HowToLens	6	0	0	6
TOTAL	43	0	1	44

Pre-existing skip (unchanged from before this PR):

• autolens_workspace_test — database/scrape/general.py — listed in no_run.yaml as NEEDS_FIX 2026-04-27: TypeError: __hash__ method should return an integer in abstract_fit.linear_light_profile_intensity_dict during subplot rendering after search completes. Not introduced by this PR.

Wall-clock total: ~222 seconds (all 6 workspaces run in parallel, scripts within each workspace also parallel).

Note on aggregator/fit_interferometer.py and aggregator/tracer.py: logs contain WARNING - Failed to load latent variables — this is an expected autofit aggregator WARNING in test mode (no sampler internals to load), not an error. Both scripts exit cleanly with "All ... aggregator tests passed."

[Jammy2211]

Shipped — 2026-05-23

PR

• feat: first-class lensing latent variable API in PyAutoLens #534 (merged as 3b7b2ff)

Summary

Second sub-prompt of the latent_refactor epic landed. PyAutoLens now ships five lensing latents (total_lens_flux_mujy, total_lensed_source_flux_mujy, total_source_flux_mujy, magnification, effective_einstein_radius) in autolens/analysis/latent.py, plus autolens/config/latent.yaml toggle layer and AnalysisImaging dispatch (LATENT_KEYS @Property + compute_latent_variables with NotImplementedError short-circuit). All five default OFF — existing fits unchanged. Library tests 311/311; cross-workspace smoke 43/43 across 6 workspaces with the worktree active.

Session Notes

• Mid-plan correction: I initially split into autolens/analysis/latent.py (tracer-only) + autolens/imaging/model/latent.py (image-derived), but the user flagged that all five latents are dataset-agnostic — they use fit.tracer / fit.galaxy_image_dict / fit.dataset.grids.lp, APIs shared between FitImaging and FitInterferometer (both PyAutoGalaxy and PyAutoLens versions). Collapsed to a single autolens/analysis/latent.py so future AnalysisInterferometer wiring can reuse the registry without duplication.
• Lens-light latent named total_lens_flux_mujy (per user direction), NOT re-exported as total_galaxy_0_flux_mujy from PyAutoGalaxy — PyAutoLens owns its full registry with lensing-vocabulary names.
• Sub-prompt scipy.integrate Requires LowLevelCallable #3 (euclid migration) was amended mid-routing to lock aperture-flux latents as Euclid-pipeline-only (they require Euclid-specific PSF kwargs; the euclid subclass will call super().compute_latent_variables() and append aperture values).
• compute_latent_samples runs on every fit (latent_after_fit: true is autofit default), so default-on would crash users without magzero. Same regression-avoidance pattern as PyAutoGalaxy Drop Binder URLs, pin Colab URLs to workspace tag 2026.4.13.6 #441.
• feedback_autoconf_lowercases_yaml_keys and feedback_latent_modules_colocate_with_model_api memories from PyAutoGalaxy Drop Binder URLs, pin Colab URLs to workspace tag 2026.4.13.6 #441 carried over cleanly — no new memories from this task.
• Six remaining sub-prompts in PyAutoPrompt/z_features/latent_refactor.md: euclid migration (scipy.integrate Requires LowLevelCallable #3, ready), autofit_workspace tutorial expand (scipy.spatial.Voronoi does not compute list of neighbouring points of each Voronoi cell #4), autogalaxy_workspace + autolens_workspace tutorials (KMeans initialize seed points #5, Speed up computation of F matrix via clustering / Voronoi grid information? #6), smoke test (List of neighbors of each Voronoi cell #7), profiling package (Variable PSF convolution kernel size to speed up convolution step #8), prior-mapping investigation (IDEA: Perform weight KMeans clustering in the image-plane #9).