feat: first-class latent variable API in PyAutoGalaxy

[Jammy2211]

Overview

PyAutoGalaxy currently has no library-level latent variable API. Users who want latents from a fit must subclass AnalysisImaging and hand-write LATENT_KEYS + compute_latent_variables (see autogalaxy_workspace/scripts/guides/results/workflow/csv_make.py, and the heavier reference in euclid_strong_lens_modeling_pipeline/util.py:306-490). This issue introduces first-class library modules — autogalaxy/analysis/latent.py (tracer/galaxy-agnostic) and autogalaxy/imaging/model/latent.py (image-derived) — plus a config/latent.yaml for user-level enable/disable, so users get a curated default catalogue without writing Python and can extend it cleanly.

This is the dependency root of a broader cross-repo latent refactor tracked at PyAutoPrompt/z_features/latent_refactor.md. PyAutoLens, the Euclid pipeline, workspace tutorials, smoke tests, and profiling all build on what's defined here.

Plan

• Add a autogalaxy/analysis/latent.py module for galaxy-level / tracer-agnostic latents (derived structural quantities that don't need an image array).
• Add a autogalaxy/imaging/model/latent.py module for image-derived latents (total galaxy flux from image, magzero → muJy conversions).
• Add autogalaxy/config/latent.yaml — a flat dict of latent_key: bool toggles, default-true, mirroring config/output.yaml.
• Wire AnalysisImaging so its LATENT_KEYS and compute_latent_variables are built from the yaml at init: only enabled keys are computed, in the order PyAutoFit's vmap path expects.
• Unit tests covering one latent per category, the on/off toggle behaviour, and ordering.
• No PyAutoFit changes. LATENT_KEYS is purely class-level in PyAutoFit; the autoconf-driven filtering lives entirely in PyAutoGalaxy's Analysis subclass.

Detailed implementation plan

Affected Repositories

• PyAutoGalaxy (primary)

Work Classification

Library

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoGalaxy	main	CLAUDE.md modified (pre-existing PyAutoPaper section add — unrelated)

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

Implementation Steps

1. Create autogalaxy/analysis/latent.py with galaxy-level / tracer-agnostic latents. Each latent is a function (instance, xp) -> value with a one-paragraph docstring explaining what it is and when it's useful.

2. Create autogalaxy/imaging/model/latent.py with image-derived latents. Port ab_mag_via_flux_from / flux_mujy_via_ab_mag_from from euclid_strong_lens_modeling_pipeline/util.py into PyAutoGalaxy. Galaxy total-flux-from-image latent lives here.

3. Create autogalaxy/config/latent.yaml — flat dict, default-true, mirroring the structure of config/output.yaml. Wired into conf.instance[\"latent\"] per the existing conf.instance[\"...\"] pattern (autogalaxy/analysis/plotter.py:22, analysis/adapt_images/adapt_images.py:49).

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

   • At init (or as a cached @property), read conf.instance[\"latent\"].
   • Build LATENT_KEYS = [k for k, on in latent_yaml.items() if on] — filtered, deterministically ordered.
   • compute_latent_variables returns a dict of only enabled keys, in the same order as LATENT_KEYS (PyAutoFit zips positionally with vmap output at autofit/non_linear/analysis/analysis.py:285).

5. Unit tests at test_autogalaxy/analysis/test_latent.py + an imaging counterpart:

   • One test per latent category against a known toy model.
   • Toggle behaviour: disabling a key removes it from both LATENT_KEYS and the dict.
   • Ordering: enabled-key order in LATENT_KEYS matches dict iteration order.
   • No JAX in unit tests (memory rule feedback_no_jax_in_unit_tests).

Key Files

• autogalaxy/analysis/latent.py — new module.
• autogalaxy/imaging/model/latent.py — new module.
• autogalaxy/config/latent.yaml — new config.
• autogalaxy/imaging/model/analysis.py — wire AnalysisImaging to consume the yaml.
• test_autogalaxy/analysis/test_latent.py — new tests.
• test_autogalaxy/imaging/test_latent.py — new tests.

Constraints to preserve

• AnalysisDataset.LATENT_BATCH_MODE = \"jit\" at autogalaxy/analysis/analysis/dataset.py:28 — do not change.
• Existing JAX-aware helper LensCalc.einstein_radius_jit_from() at autogalaxy/operate/lens_calc.py:1520-1537 and its closure cache at 1580-1586 — do not break (although Einstein-radius latent itself ships in the PyAutoLens follow-up, not here).
• Module starts as a single .py file. Convert to a latent/ package if it grows past ~150 lines.
• Per memory `feedback_workspace_config_default_true`, new library yaml keys need mirroring into workspace `config/latent.yaml` later (sub-prompt Refactor/complex #5 for autogalaxy_workspace).

Original Prompt

Click to expand starting prompt

Add first-class latent-variable modules to PyAutoGalaxy

Context

Parent epic: `PyAutoPrompt/z_features/latent_refactor.md`.

Currently, anyone who wants latent variables out of a PyAutoGalaxy fit must subclass `AnalysisImaging`, hand-write a `LATENT_KEYS` list, and hand-write a `compute_latent_variables` method (see `autogalaxy_workspace/scripts/guides/results/workflow/csv_make.py:140-160` and the heavier reference in `euclid_strong_lens_modeling_pipeline/util.py:306-490`). This task promotes that into first-class library API so users can opt into a curated catalogue of galaxy latents via config, and so each latent has a single home with docstrings + unit tests.

This is the library dependency root of the broader latent-refactor epic. Sub-prompts #2, #3, #5, #7, #8 all build on what's defined here.

Task

1. Create `autogalaxy/analysis/latent.py` — galaxy-level / tracer-agnostic latents (anything that derives only from a `Plane` / galaxies list, not from imaging arrays). Each latent is a single function taking `(instance, xp)` (or similar — finalize during implementation), with a one-paragraph docstring explaining what it is and when it's useful.

2. Create `autogalaxy/imaging/model/latent.py` — imaging-derived latents (need the image array from a `FitImaging`). Examples from the euclid reference: total galaxy flux from image, magzero-converted muJy fluxes (see `ab_mag_via_flux_from` / `flux_mujy_via_ab_mag_from` helpers in `euclid_strong_lens_modeling_pipeline/util.py` — port these into PyAutoGalaxy where they belong).

3. Create `autogalaxy/config/latent.yaml` — flat dict of `latent_key: bool` toggles, all defaulting `True`. Mirror the structure of `autogalaxy/config/output.yaml`.

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

   • At init (or as a cached `@property`), read `config/latent.yaml`.
   • Build `LATENT_KEYS = [k for k, on in latent_yaml.items() if on]` — filtered, ordered.
   • `compute_latent_variables` must return a dict containing only the enabled keys, in the same order as `LATENT_KEYS` (critical — PyAutoFit zips positionally with vmap output at `autofit/non_linear/analysis/analysis.py:285`).

5. Unit tests at `test_autogalaxy/analysis/test_latent.py` (and an imaging-flavour counterpart). Minimum coverage:

   • One test per latent category that values are sensible against a known toy model.
   • Toggle behaviour: disabling a key removes it from both `LATENT_KEYS` and the dict.
   • No JAX in unit tests — per memory `feedback_no_jax_in_unit_tests`. Cross-xp checks belong in workspace_test (sub-prompt Feature/release process #7).

Where to look

• PyAutoFit hook (do not modify): `autofit/non_linear/analysis/analysis.py` lines 34 (`LATENT_KEYS`), 170 (`compute_latent_samples`), 285 (vmap positional zip).
• Existing JAX-aware latent helper to delegate to: `PyAutoGalaxy/autogalaxy/operate/lens_calc.py:1520` (`einstein_radius_jit_from`). Reuse, do not reimplement.
• `LATENT_BATCH_MODE` constraint: `PyAutoGalaxy/autogalaxy/analysis/analysis/dataset.py:28` is `"jit"`. Do not change; lensing latents downstream depend on it (also for sub-prompt Coordinates #2).
• Reference implementation to mine: `euclid_strong_lens_modeling_pipeline/util.py:306-490` — the full LATENT_KEYS list + `compute_latent_variables` body. Pick the galaxy-only / image-derived bits here; the lensing-only bits (magnification, effective Einstein radius, lensed source flux) belong in sub-prompt Coordinates #2.
• Config dir layout to mirror: `PyAutoGalaxy/autogalaxy/config/output.yaml`.

Verification

Tests + a quick end-to-end manual check using an autogalaxy_workspace results-workflow script — confirm that with the new yaml, the user no longer needs to subclass `AnalysisImaging` to get a sensible default latent set; the same csv is produced.

Notes

• Start as single `.py` files (~5-6 latents on this side, per the euclid reference). If either grows past ~150 lines, convert to a `latent/` package with one file per latent category.
• No PyAutoFit changes. The config-driven on/off lives in `AnalysisImaging` subclass logic; PyAutoFit core stays dataset-agnostic.
• Order matters in `LATENT_KEYS` because of the vmap positional zip. Sort keys deterministically (yaml insertion order or alphabetical — pick one and document it).

[Jammy2211]

Smoke Test Results — 2026-05-23

All 42 smoke tests passed across 6 workspaces (PR #441 feature/latent-module-autogalaxy).

Worktree active: /home/jammy/Code/PyAutoLabs-wt/latent-module-autogalaxy/PyAutoGalaxy/

Workspace	Passed	Failed	Skipped	Total
autofit_workspace	6	0	1	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: 42 passed / 0 failed / 2 skipped / 44 total — wall clock 222.8s

Skips (pre-existing, not related to this PR):

• autofit_workspace / searches/mcmc.py — mcmc is in no_run.yaml (Zeus section fails Test Model Initialization)
• autolens_workspace_test / database/scrape/general.py — NEEDS_FIX 2026-04-27 (__hash__ returning non-int in subplot_fit_imaging)

[Jammy2211]

Shipped — 2026-05-23

PR

• feat: first-class latent variable API in PyAutoGalaxy #441 (merged as a5cc42a)

Summary

First-class latent variable API landed in PyAutoGalaxy. New module autogalaxy.imaging.model.latent (helpers, registry, one concrete latent), new autogalaxy/config/latent.yaml toggle layer, and AnalysisImaging overrides for LATENT_KEYS (config-driven @property) and compute_latent_variables (tuple return aligned to keys, NotImplementedError short-circuit when empty). Library tests 917/917 pass; cross-workspace smoke 42/42 pass with the worktree active. Latent defaults OFF so existing fits without magzero keep working unchanged — users opt in by editing the yaml and passing magzero via Analysis kwargs.

Session Notes

• Caught autoconf lowercases yaml keys (saved as feedback_autoconf_lowercases_yaml_keys memory) — muJy became mujy in conf.instance["latent"], so the registry key and yaml key both use lowercase. The user-facing csv column header is total_galaxy_0_flux_mujy.
• Caught regression risk mid-PR: latent_after_fit: true is the autofit default, so leaving the latent on-by-default would have crashed every existing fit that doesn'''t pass magzero. Flipped library default to false and added an empty-LATENT_KEYS short-circuit; only the test config keeps the latent on so the new tests still exercise the dispatch path.
• Architectural decision (saved as feedback_latent_modules_colocate_with_model_api memory): latent modules colocate with the model API, i.e. <dataset>/model/latent.py for dataset-specific, not top-level <lib>/latent.py.
• PyAutoLens follow-up (sub-prompt Coordinates #2 of the latent_refactor epic at PyAutoPrompt/z_features/latent_refactor.md) will add lensing latents (magnification, effective Einstein radius, lensed source flux). Workspace tutorials, the smoke test, and the profiling package are the other open sub-prompts.