design: multi-dataset shared-state — shifted shared Delaunay mesh (phase 1/4)

[Jammy2211]

Overview

Phase 1 (of 4) of generalising the shipped datacube shared-state likelihood machinery (#565, PyAutoFit#1307 — both closed; epic z_features/complete/analysis_shared_state.md) to imaging multi-dataset fitting: one Delaunay mesh shared and shifted across exposures of the same lens. This phase is design-only — its deliverable is a design note + locked API sketch posted on this issue; no source edits. Phases 2 (core library API), 3 (workspace example sections + parity test) and 4 (docs) have prompts in PyAutoMind but are deliberately un-issued until their predecessor nears shipping.

Plan

• Decide what is genuinely shareable for imaging multi-exposure fits: per-dataset PSFs + pixel offsets rule out sharing the curvature matrix F and mapping matrix L; the candidate shared object is the source-plane Delaunay mesh geometry. Quantify the saving honestly (speed win vs consistency win).
• Lock the same-wavelength "one source reconstruction" semantics against the FactorGraph per-factor contract — a joint solve (F = Σᵢ LᵢᵀWᵢLᵢ, D = Σᵢ Lᵢᵀdᵢ, one s) does not decompose into the per-factor likelihood sum.
• Confirm the different-wavelength case is "shared shifted mesh geometry only, independent per-dataset reconstructions" (colour differences).
• Design shift composition with aa.DatasetModel: shifts known/fixed by default (PyAutoReduce target_pixel differences), optionally free per-frame (dy, dx) with Gaussian priors of width = the recorded registration residuals (floor ~0.1–0.3 px); simulator shifts default 0.
• Cover the imaging_and_interferometer variant (imaging preloads meeting the existing PreloadsInterferometer).
• Deliverable: design note (options weighed, decision + rationale per question) + locked PreloadsImaging / AnalysisImaging.shared_state_from API sketch for Phase 2 + example-section outline for Phase 3.

Detailed implementation plan

Affected Repositories

• PyAutoLens (primary — design decisions land here in Phase 2)
• PyAutoArray, autolens_workspace, autolens_workspace_test, autolens_profiling (downstream phases)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean (worktree claims: kernel-forward-chunking, rect-adapt)
./PyAutoLens	main	clean
./autolens_workspace	main	1 modified file (worktree claim: markdown-example-renderings)
./autolens_workspace_test	main	clean

Suggested branch: none — design-only phase, no worktree. (Existing claims above only matter for Phases 2–3.)

Design questions to resolve (the work)

1. What is shareable for imaging? Unlike the datacube (channel-invariant uv/noise → share F = LᵀW̃L + mapper), imaging exposures have per-dataset PSFs and pixel offsets → neither F nor L is invariant. Candidate shared object: the source-plane Delaunay mesh geometry (image-mesh sparse grid traced once from the lead dataset), with each dataset building its own L_i by tracing its offset grid onto that shared mesh. Quantify what this actually saves before committing — if the shareable fraction is small, the honest answer may be "share the mesh for consistency, not speed" and the design should say so explicitly.
2. Same-wavelength "one reconstruction" semantics. A genuinely joint solve does not decompose into FactorGraphModel's per-factor likelihood sum — decide between: (a) a joint-inversion Analysis that owns all N same-λ datasets (one factor), (b) per-factor reconstructions with the mesh shared (N reconstructions on the same mesh — not literally one source), or (c) shared-state carrying the joint solve with per-factor likelihood contributions derived from it. Justify against the sub-task A contract ("shared object computed once, forwarded to factors").
3. Different-wavelength semantics: confirm "share mesh geometry only" falls out of whichever mechanism Q2 picks.
4. Shift mechanism. aa.DatasetModel already provides per-dataset (y,x) offset (+ rotation) as optional free parameters (multi/features/dataset_offsets/). Design how DatasetModel offsets compose with the shared mesh (mesh built once in the source plane; each dataset's offset applies to its image-plane grid before tracing). Default = known/fixed from PyAutoReduce target_pixel differences; optional free (dy,dx) with Gaussian priors of width = registration residuals. Simulator support: per-dataset shifts, default 0.
5. imaging_and_interferometer: how the shared mesh crosses the two dataset types (imaging preloads vs the existing aa.PreloadsInterferometer).
6. Preloads API shape: PreloadsImaging fields (mesh/mapper geometry vs mapper_galaxy_dict reuse), what AnalysisImaging.shared_state_from returns, and JAX pytree registration (recompute inside the jitted region each eval — no instance memoisation).

Grounding (verified 2026-07-10)

• Only the interferometer consumer exists (aa.PreloadsInterferometer, AnalysisInterferometer.shared_preloads/shared_state_from); no PreloadsImaging, nothing on AnalysisImaging.
• The multi/ examples already use af.FactorGraphModel(*factors, use_jax=True) — the PyAutoFit#1308 mechanism applies directly.
• aa.DatasetModel (autoarray/dataset/dataset_model.py) is the existing offset/rotation surface.
• Datacube consumer measured 2026-07-10 at 4.06× per-likelihood (sma × 4 ch) — the pattern being generalised works.

Key Files

• PyAutoArray/autoarray/preloads/{abstract,interferometer}.py — the Preloads pattern to extend
• PyAutoLens/autolens/interferometer/model/analysis.py — shared_state_from consumer pattern (PR#566)
• PyAutoArray/autoarray/dataset/dataset_model.py — existing offset/rotation surface
• autolens_workspace/scripts/multi/features/{same_wavelength,wavelength_dependence,imaging_and_interferometer,dataset_offsets}/ — Phase 3 targets
• autolens_workspace_test/scripts/jax_likelihood_functions/datacube/shared_preloads.py — parity-test pattern

Process

• Consult autolens_assistant/ (AGENTS.md → relevant skills//wiki/core/ pages) before locking the source-reconstruction design; PyAutoMemory lensing_wiki/methods_wiki for source-reconstruction + likelihood context.
• Output lands as a design-note comment on this issue; Phases 2–4 prompts get updated with the locked decisions before they are issued.

Original Prompt

Click to expand starting prompt (phase 1; parent prompt with the user's verbatim request + PyAutoReduce registration semantics is at PyAutoMind feature/autolens/multi_shared_state_examples.md → issued/)

See PyAutoMind/issued/multi_shared_state_examples_phase_1_design.md — the six design questions above are its verbatim core; the parent prompt issued/multi_shared_state_examples.md carries the original user request verbatim, the PyAutoReduce frame-products registration semantics (HST/JWST vs Keck, target_pixel, residual_reliable, data_units, PSF-tier provenance), and the use-case → example-script table.

[Jammy2211]

Design note — multi-dataset shared-state (phase 1/4)

Decisions D1–D6 answer the issue's design questions Q1–Q6. Grounded in: TracerToInversion (autolens/lens/to_inversion.py), the shipped interferometer consumer (PyAutoArray#344 / PyAutoLens#566), aa.DatasetModel + multi/features/dataset_offsets/, and autolens_assistant/skills/al_multi_dataset.md (factor-graph composition contract).

D1 — The shared object is the source-plane mesh geometry, not the mapper

The datacube consumer could preload the whole mapper_galaxy_dict + curvature F because every channel shares one real-space mask/grid and one uv/noise. Imaging exposures break both: per-exposure PSFs make F (and the blurred mapping matrix) per-dataset, and per-exposure pixel offsets make even the unblurred mapping matrix per-dataset. Crucially, a Mapper object embeds the dataset's own traced data-grid mappings, so the existing whole-mapper preload short-circuit (to_inversion.py:403) is not valid across datasets with different grids — it must not be reused as-is.

What IS channel-... exposure-invariant (given one lens model per evaluation):

Quantity	Invariant across exposures?	Notes
image-plane mesh grid (image-mesh e.g. Overlay/Hilbert on lead dataset)	✅ (computed once from lead)	adapt-image-dependent → lead dataset's adapt image is the reference
traced source-plane mesh grid + Delaunay triangulation/neighbours	✅	pure function of lens model + lead mesh
regularization matrix H	✅	mesh-geometry-only (Constant/ConstantSplit)
traced data grid, mapping matrix L_i, blurred L_i, F_i, D_i, solve	❌ per-dataset	offsets + PSFs

So: shared_state_from (lead factor) ray-traces the lead image-mesh once and returns the source-plane mesh geometry (+ H); every factor builds its own mapper by tracing its own (offset) data grid onto that shared mesh.

Honest saving: for imaging, per-dataset costs (PSF-blurred L_i, F_i, solve) dominate the inversion budget, so the speed win is modest — the primary win is consistency: every exposure reconstructs on the identical source-pixel grid (pixel-by-pixel comparable reconstructions, prerequisite for any joint/co-added source analysis and for colour maps across λ). The design commits to this framing; Phase 3's measurement script (sibling of likelihood_runtime/datacube/shared_preloads.py) records the actual ratio rather than promising one.

D2 — Same-wavelength "one reconstruction": shared mesh per-factor now; joint stacked inversion as an explicit follow-up

A literal joint solve (F = Σᵢ LᵢᵀWᵢLᵢ, D = Σᵢ LᵢᵀWᵢdᵢ, one s, one log-det pair) does not decompose into FactorGraphModel's per-factor likelihood sum — the log-det and regularization terms are global. Three options were weighed:

• (a) Joint-inversion Analysis owning all N same-λ exposures as ONE factor (a stacked inversion: block-stacked L_i, summed F/D). Statistically the true "one source". Requires a new stacked imaging inversion variant + a group-of-datasets Analysis; bypasses the shared-state mechanism for the group (nothing cross-factor left to share) and loses per-exposure Result granularity.
• (b) Shared mesh, per-factor solves — N reconstructions that live on the identical mesh (consistent, comparable, averageable) but are not literally one solve.
• (c) Shared-state carrying the joint solve — rejected: shared_state_from sees only the lead analysis + instance; feeding it all datasets breaks the sub-task A contract and the per-factor evidence split is arbitrary.

Locked: (b) for phases 2–3 — it is incremental on the shipped mechanism, delivers the docstring'd example sections requested, and gives mesh-consistent per-exposure reconstructions whose differences are themselves diagnostic (registration errors, PSF-tier quality). (a) is filed as an explicitly-scoped follow-up prompt (feature/autolens/multi_joint_stack_inversion.md, to be written when phase 3 nears shipping) for the precision "one source" case; it composes with this design — the same-λ group becomes one factor whose internal inversion consumes the same shared mesh when combined with other-λ/interferometer factors.

D3 — Different wavelengths: falls out of D1/D2(b) unchanged

Shared shifted mesh geometry, independent per-dataset reconstructions — this is exactly what (b) does with zero extra machinery. Colour freedom is already the factor-graph idiom (per-band prior overrides). No further design needed; phase 3 documents it in wavelength_dependence/.

D4 — Shifts: aa.DatasetModel composition, known-by-default

• DatasetModel (offset + rotation, multi/features/dataset_offsets/) is already threaded per-instance through fit_from; the shared mesh composes with it naturally: the lead factor's own offset applies when tracing the mesh (the mesh is defined in the lead exposure's frame), and each factor's offset applies to its own data grid before mapping onto the shared mesh. No new shift API.
• Default = shifts known/fixed: instantiate DatasetModel with fixed (y,x) from PyAutoReduce target_pixel differences (HST/JWST; WCS-pair equivalent) or offset_dy/dx_px (Keck). Optional free: per-frame (dy,dx) GaussianPriors with σ = the recorded registration residuals (residual_dy/dx_px, floor ~0.1–0.3 px; honor residual_reliable, never header RMS_RA/RMS_DEC).
• Simulator: per-dataset shifts added to the relevant simulator.py scripts, default 0 (so existing outputs are unchanged); phase 3 wires them.

D5 — imaging_and_interferometer: mesh field lives on AbstractPreloads

The shared mesh is source-plane geometry — dataset-type-agnostic. Put the new field(s) on aa.AbstractPreloads so PreloadsImaging and the existing PreloadsInterferometer both carry them; the lead factor may be either type. Each side keeps its own per-dataset machinery (imaging: blurred L; interferometer: F = LᵀW̃L, which for a single interferometer dataset in the pair remains its own). The datacube fast path (whole mapper + F preload) is untouched — it remains valid precisely because its channels share one grid.

D6 — Locked API sketch (Phase 2 scope)

# PyAutoArray
class AbstractPreloads:
    def __init__(self, curvature_matrix=None, mapper_galaxy_dict=None,
                 source_plane_mesh_grid=None, regularization_matrix=None): ...
    # source_plane_mesh_grid: traced mesh centres (+ triangulation via the existing
    # mesh-geometry objects); ONLY valid to share across datasets under one lens model.

class PreloadsImaging(AbstractPreloads): ...   # imaging sibling; docstring records the
    # invariance contract: mesh/H shareable across exposures; mapper/F are NOT (unlike datacube)

# PyAutoLens
class AnalysisImaging:
    def __init__(..., shared_preloads: bool = False): ...
    def shared_state_from(self, instance):        # lead factor, once per eval, in-jit
        # apply lead DatasetModel offset -> trace lead image-mesh -> source-plane mesh (+ H)
        return aa.PreloadsImaging(source_plane_mesh_grid=..., regularization_matrix=...)
    def log_likelihood_function(self, instance, shared=None): ...
    # fit_from -> FitImaging -> TracerToInversion(preloads=shared): NEW reuse path —
    # when preloads.source_plane_mesh_grid is set, skip image_plane_mesh_grid_pg_list +
    # mesh tracing; build this dataset's mapper onto the preloaded mesh.

JAX: preloads registered as pytrees (interferometer precedent); recomputed inside the jitted region every eval — never memoised on the instance. Opt-in default-off; shared is None path byte-for-byte unchanged (single-dataset fits untouched).

Phase 3 outline (unchanged from prompt, anchored by D1–D5)

same_wavelength/ (headline: offsets + shared mesh + fixed-vs-free shift variants), wavelength_dependence/ (D3), imaging_and_interferometer/ (D5), simulators (shifts default 0), dataset_offsets/ cross-references; workspace_test parity script mirroring datacube/shared_preloads.py; profiling measurement recording the honest ratio. Real-data prose: manifest target_pixel / residual_reliable / data_units / psf.method provenance notes (docstring-level, not a pipeline).

Follow-ups filed by this design

• multi_joint_stack_inversion prompt (D2 option (a)) — written when phase 3 nears shipping (no-bulk-issue rule).

[Jammy2211]

Design amendment (user review, 2026-07-10)

D1/D6 amended: the regularization matrix H is removed from the shared object. Regularization can be data-dependent (e.g. adaptive/luminosity-weighted regularization uses per-dataset adapt images / reconstructions), so H is not exposure-invariant in general. The shared object is the source-plane mesh geometry only; each factor builds its own H as today (for mesh-geometry-only schemes like Constant/ConstantSplit this recomputes identical matrices — a potential future preload behind an explicit invariance guard, but NOT in scope).

Phase 2 proceeds under --auto per the user's instruction, with this amendment locked.

[Jammy2211]

Epic complete — 2026-07-10

All four phases + the scoped-preloads follow-up shipped and merged in one day:

Phase	PRs	Merged
1 — design (this issue)	D1–D6 + regularization amendment (comments above)	locked
2 — core API	PyAutoArray#380 + PyAutoLens#600	✅
— scoped-preloads follow-up	PyAutoArray#381 + PyAutoLens#601 (cross-type consumption hazard found in phase 3 review)	✅
3 — workspace	autolens_workspace#261 + autolens_workspace_test#162 + autolens_profiling#60	✅
4 — docs	autolens_workspace#262	✅

Delivered: AnalysisImaging(shared_preloads=True) — one shared/shifted source-plane Delaunay mesh across exposures (same-λ, multi-λ, imaging+interferometer), mesh-geometry-only sharing (no F/L/mapper/H — per-exposure PSFs/offsets + the adaptive-regularization amendment), dataset-scoped cross-type consumption, bit-identical parity proven (numpy+JAX, in smoke), measured 1.14× at hst × 4 (consistency is the win, as D1 predicted).

Remaining follow-up (separate, un-issued by design): multi_joint_stack_inversion — the literal single joint reconstruction for same-λ exposure groups (D2 option (a)). Closing.