feat: imaging shared-state consumer — PreloadsImaging + shared source-plane mesh (phase 2/4)

[Jammy2211]

Overview

Phase 2 (of 4) of the multi-dataset shared-state epic (#599 = phase 1 design, decisions D1–D6 + the H-removal amendment). Implements the imaging shared-state consumer: the shared object is the source-plane mesh geometry only (traced mesh centres) — per-exposure PSFs and pixel offsets make the mapper, mapping matrix, curvature matrix, and regularization matrix (user amendment: adaptive regularization is data-dependent) per-dataset.

Running under --auto (Autonomy: safe, feature/medium — mirrors the shipped PyAutoArray#344 / PyAutoLens#566 pattern): plan recorded here, run proceeds through the four-leg autonomous ship gate and ends at PR-open; merge stays human.

Plan

• PyAutoArray: AbstractPreloads gains source_plane_mesh_grid; new PreloadsImaging sibling with the invariance-contract docstring (mesh geometry shareable across exposures under one lens model; mapper / F / blurred L / H are NOT — unlike the datacube whose channels share one grid, whose whole-mapper + F preload is untouched). JAX pytree registration mirroring PreloadsInterferometer.
• PyAutoLens: AnalysisImaging(shared_preloads=False) opt-in + shared_state_from (apply the lead factor's aa.DatasetModel offset → trace the lead image-mesh once → PreloadsImaging(source_plane_mesh_grid=...)); shared= threaded through log_likelihood_function → fit_from → FitImaging → TracerToInversion, mirroring PR#566. New reuse path in TracerToInversion: when preloads.source_plane_mesh_grid is set, skip image_plane_mesh_grid_pg_list + mesh tracing and build this dataset's mapper (its own offset grid) onto the preloaded mesh.
• Unit tests (numpy-only): preloads container tests in test_autoarray; in test_autolens, shared-vs-unshared figure_of_merit parity on offset datasets, opt-in default-off, single-dataset path byte-for-byte unchanged.
• Out of scope: same-λ joint stacked inversion (D2 follow-up prompt multi_joint_stack_inversion); workspace/workspace_test/profiling examples (phase 3 — includes the shared-vs-unshared assertion + runtime scripts mirroring the datacube pattern); docs (phase 4).

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary — new preloads surface)
• PyAutoLens (consumer)
• Downstream test surface at gate time: PyAutoGalaxy + PyAutoLens suites (public API touched)

Branch Survey (2026-07-10)

Repository	Current Branch	Dirty?	Note
./PyAutoArray	main	clean	claims kernel-forward-chunking, rect-adapt exist but touch only inversion/mesh/ kernel-CDF files — zero overlap with preloads/; override recorded in active.md
./PyAutoLens	main	clean

Branch: feature/multi-shared-state-core-api in worktree ~/Code/PyAutoLabs-wt/multi-shared-state-core-api/

Implementation Steps

1. PyAutoArray/autoarray/preloads/abstract.py — add source_plane_mesh_grid=None to AbstractPreloads.__init__ + docstring; keep pytree registration consistent.
2. PyAutoArray/autoarray/preloads/imaging.py (new) — PreloadsImaging(AbstractPreloads) with the invariance-contract docstring; export as aa.PreloadsImaging.
3. PyAutoLens/autolens/lens/to_inversion.py — in the mapper-construction flow, when self._preloads.source_plane_mesh_grid is set (and no whole-mapper preload), bypass image_plane_mesh_grid_pg_list + traced_mesh_grid_pg_list and build the mapper from this dataset's traced data grid onto the preloaded source-plane mesh.
4. PyAutoLens/autolens/imaging/model/analysis.py — shared_preloads flag, shared_state_from(instance) (lead DatasetModel offset honored), log_likelihood_function(..., shared=None), fit_from(..., preloads=shared) threading; JAX pytree registration for PreloadsImaging.
5. Tests as above; full suites + downstream at the gate.

Key Files

• autoarray/preloads/{abstract,interferometer}.py, autolens/interferometer/model/analysis.py — the shipped pattern to mirror
• autolens/lens/to_inversion.py:379-460 — mapper_galaxy_dict construction + existing preload short-circuit (NOT valid cross-dataset; the new path is mesh-level)
• autoarray/dataset/dataset_model.py — offset surface (no changes expected)

Original Prompt

Click to expand starting prompt

See PyAutoMind/issued/multi_shared_state_examples_phase_2_core_api.md (locked by #599 D1–D6 + amendment). Parent epic: issued/multi_shared_state_examples.md.

[Jammy2211]

Shipped to PR-open — 2026-07-10 (--auto safe, four-leg gate)

PRs (pending-release; PyAutoArray must merge first)

• PyAutoArray: feat: PreloadsImaging + shared source-plane mesh fields (multi-exposure shared-state, phase 2/4) #380
• PyAutoLens: feat: imaging shared-state consumer — AnalysisImaging.shared_state_from + mesh preload reuse (phase 2/4) PyAutoLens#600

Gate record

1. Tests ✅ — test_autoarray 896 passed; downstream test_autogalaxy 962 passed (canonical checkout against branch PyAutoArray); test_autolens 378 passed; 10 new/updated analysis tests.
2. Smoke ✅ — all six workspaces' curated lists against the worktree stack: autofit 7/7, autogalaxy 6/6, autolens 7/7, autolens_test 14/14, euclid 6/6, HowToLens 6/6 (11 initial timeouts under a 150s/6-parallel budget all pass sequentially at 600s — contention, not regression; includes the datacube shared_preloads.py JIT script).
3. Review ✅ CLEAN — verify pass drove the real FactorGraphModel end-to-end: shared-vs-unshared log-likelihood bit-identical (−58.673099306030444), shared_state_from computed once (lead only). Two CRLF-normalization diff inflations caught via the review surface and repaired byte-exactly.
4. Heart 🟡 acknowledged — user acked the exact 7-reason YELLOW set in-session (recorded in active.md heart-ack:; binds to this launch only).

What shipped

Design #599 D1–D6 + regularization amendment: aa.PreloadsImaging + mesh-geometry fields on AbstractPreloads (dataset-type-agnostic for the imaging+interferometer case); AnalysisImaging(shared_preloads=True) + shared_state_from + shared=/preloads= threading; TracerToInversion mesh-preload reuse path. Opt-in default-off; single-dataset path byte-for-byte unchanged.

Next

Merge is human (PyAutoArray #380 first). Phase 3 (workspace example sections in multi/features/{same_wavelength,wavelength_dependence,imaging_and_interferometer}, workspace_test shared-vs-unshared assertion script + autolens_profiling measurement mirroring the datacube pattern, simulator shifts default 0) is prompted and issues after this merges.

[Jammy2211]

Merged — 2026-07-10

• PyAutoArray feat: PreloadsImaging + shared source-plane mesh fields (multi-exposure shared-state, phase 2/4) #380 → f8a32d43b
• PyAutoLens #600 → 1513236da

Phase 2 complete (both pending-release for the next nightly). Phase 3 (workspace examples + workspace_test parity + profiling measurement) proceeds now under --auto. Closing.