feat: array-free fits with ordinary light profiles via the data-term identity (streaming P4)

[Jammy2211]

Overview

Phase 4 of the streaming-visibilities epic (Discussion https://github.com/orgs/PyAutoLabs/discussions/13). Phases 1–3 made an array-free Interferometer.from_stream dataset (no data/noise_map/uv_wavelengths/transformer, just SparseTerms + the sparse operator with cached data_term/noise_normalization) fit, save/reload and plot end to end — but only for pixelization and linear-light-profile models. Fits with an ordinary (non-linear) light profile still raise, because profile_visibilities = F i_p is formed over N_vis and subtracted from the data. This task applies the identity Σ|d−p|²/σ² = data_term − 2 i_pᵀd̃ + i_pᵀW̃i_p (one operated_matrix_slim_from product, already computed by sparse_dirty_image_from) so such fits run array-free, with and without an inversion, in numpy and under jax.jit.

Trap: fast_chi_squared reads term 3 from the raw sparse_operator.data_term whenever dataset.data is None; removing the current raise alone would make the inversion silently use the unsubtracted data term. The inversion DatasetInterface therefore gains a per-fit data_term override, and the guard stays until it is wired.

Plan

• Give the inversion DatasetInterface an optional precomputed data_term; fast_chi_squared prefers it over the operator scalar when data is None.
• In autogalaxy compute the corrected data term next to the sparse dirty image from one W̃ i_p product and pass both (with data=None) into the inversion for sparse fits with ordinary light.
• Make light-profile-only (no inversion) fits work via a chi_squared hook on the same identity, so log_likelihood/figure_of_merit work on array-free datasets.
• Replace the typed raise in profile_visibilities with None on array-free datasets; model_data and the visibility dicts raise a clear typed DatasetException instead of None + array.
• PyAutoLens mirrors autogalaxy (it already imports the ag helpers).
• Tests: parity at rel 1e-8 vs the in-memory dense fit for LP-only, LP+pixelization, LP+linear-light in numpy and under jax.jit; a transformer spy proves the N_vis path is never formed; autoarray unit tests for the override.

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary)
• PyAutoGalaxy
• PyAutoLens

Branch Survey

Repository	Current Branch	Dirty?
array/PyAutoArray	main	clean
galaxy/PyAutoGalaxy	main	clean
lens/PyAutoLens	main	clean

Suggested branch: feature/streaming-p4-light-profile-identity (same name in all three repos; worktree ~/Code/PyAutoLabs-wt/streaming-p4-light-profile-identity/)

Implementation Steps

A. PyAutoArray (array/PyAutoArray)

1. autoarray/inversion/inversion/dataset_interface.py (DatasetInterface.__init__, L2-76): add
   data_term: Optional[float] = None (store as self.data_term). Docstring: "χ² data term
   Σ|d|²/σ² of this interface's (possibly profile-subtracted) data; used when data is None."
2. autoarray/inversion/inversion/interferometer/abstract.py fast_chi_squared (L222-241): when
   self.dataset.data is None, term 3 = getattr(self.dataset, "data_term", None), falling back to
   sparse_operator.data_term, else the existing InversionException. Pure scalar → jit-safe.
3. autoarray/fit/fit_interferometer.py (aa.FitInterferometer): add an overridable hook used by the
   no-inversion path — e.g. property sparse_chi_squared returning None by default; chi_squared (L209)
   returns it when the dataset is array-free (self.dataset.data is None) and the hook is not None, before the
   _require raise. Keep _require for the maps. log_likelihood (fit_dataset.py:102) and
   noise_normalization (fit_interferometer.py:219-241, already returns the operator scalar when
   noise_map is dataset.noise_map, i.e. None is None) then work unchanged.
4. Optional small helper in inversion_interferometer_util.py next to operated_matrix_slim_from (L1588):
   sparse_profile_terms_from(operator, image_slim, extent_index_for_masked_pixel, xp) -> (operated_image, sparse_dirty_image, data_term) so ag/al share one implementation and autoarray can unit-test it against the
   dense Σ|d − F i|²/σ² on interferometer_7. (Alternatively keep it in ag — decide in implementation; prefer
   autoarray so the identity is tested where W̃ lives.)
5. Tests test_autoarray/inversion/inversion/interferometer/test_interferometer.py: modelled on
   ..._sparse_dirty_image_override__used_by_data_vector (L607) and ..._data_none_uses_the_sparse_operator_data_term
   (L1250): data_term override used when given, operator scalar used when absent, numpy == jax (template L1349);
   helper identity vs dense at rel 1e-10. test_autoarray/fit/test_fit_interferometer.py: chi_squared via the hook on
   an array-free fit; maps still raise match="array-free" (L440 contract preserved).

B. PyAutoGalaxy (galaxy/PyAutoGalaxy/autogalaxy/interferometer/fit_interferometer.py)

1. Replace sparse_dirty_image_from (L58-111) internals with the shared helper: new module function
   sparse_profile_terms_from(dataset, galaxies, image, xp) → (sparse_dirty_image, data_term) or (None, None)
   when no sparse operator / no ordinary light; sparse_dirty_image_from wraps it (keeps the existing callers).
2. uses_precomputed_data_term_from (L114-157): the "no ordinary light" clause stays for the in-memory path only
   when no override is available; simplest rule: sparse operator with data_term present and
   data is dataset.data / noise_map is dataset.noise_map → precomputed, with ordinary light handled by the
   corrected data_term override. Document the rule in the docstring.
3. profile_visibilities (L233-269): on dataset.transformer is None return None (no raise).
   profile_subtracted_visibilities (L271-285): return None when self.data is None without evaluating
   profile_visibilities first.
4. galaxies_to_inversion (L302-338): compute (sparse_dirty_image, data_term) once; pass data=None,
   sparse_dirty_image=…, data_term=… to DatasetInterface when _uses_precomputed_data_term.
5. No-inversion: implement the aa hook (sparse_chi_squared) returning the identity from profile_image
   when the dataset is array-free and there is no inversion; None otherwise.
6. model_data (L384-401) and galaxy_model_visibilities_dict (~L487): raise
   aa.exc.DatasetException("… array-free … use model_image_natural / dirty_model_image_natural") when
   dataset.transformer is None. inversion_with_data (L352) unchanged.
7. Tests test_autogalaxy/interferometer/test_fit_interferometer.py: delete
   test__fit_figure_of_merit__array_free_dataset__light_profile__raises (L947); add, each using
   _array_free_dataset_from (L860) + _assert_sparse_fit_matches_dense (L515) at rel 1e-8:
   LP-only (figure_of_merit == log_likelihood), LP+pixelization (log_evidence), LP+linear-light; jit variants
   modelled on L910 (coefficient-parameterised galaxy); a transformer/profile_visibilities spy (monkeypatch
   TransformerNUFFT.visibilities_from to raise) proving the N_vis path is never formed; model_data typed raise;
   the natural-dirty-image tests (L1001, L1074) pass unchanged. Also assert the in-memory sparse fit with ordinary
   light is byte-unchanged vs today (..._light_profile__unchanged_vs_data_passed L794 stays green).

C. PyAutoLens (lens/PyAutoLens/autolens/interferometer/fit_interferometer.py)

Mirror B using the ag helpers it already imports (L28-32): profile_visibilities (L129-165) returns None
array-free; profile_subtracted_visibilities (L167-182) short-circuits; tracer_to_inversion (L199-236) passes
data_term; sparse_chi_squared hook from tracer.image_2d_from; model_data (L302) and the planes/galaxy
visibility helpers (~L430/452) raise typed. Tests in test_autolens/interferometer/test_fit_interferometer.py:
replace ..._array_free_dataset__lens_light_profile__raises (L816) with lens-light-only, lens-light+pixelized-source
(_pixelized_source_tracer(lens_light=True) L554), lens-light+MGE source parity at rel 1e-8, numpy + jit
(template L779).

Analysis / plotting

No changes expected: ag model/analysis.py:89, al model/visualizer.py:54 and the three fit_interferometer_plots
modules already branch on is_array_free and use model_image_natural (= profile_image + mapped_reconstructed_data,
shipped P3). Implementation verifies no array-free branch touches fit.model_data.

Trade-offs / decisions

• Override lives on DatasetInterface, not a mutated operator — the operator is a frozen shared dataclass; the
  per-fit scalar is fit state. Matches how sparse_dirty_image already flows (decision (a)/(d) lineage).
• Precision: the identity subtracts two large numbers (data_term ≫ χ² at high S/N); in float64 this is fine at
  rel 1e-8 on logL (phase-1 parity was 1e-16 on a comparable identity). JAX tests run with x64 enabled as the
  existing sparse tests do; no fp32 path is in scope.
• In-memory sparse + ordinary light now also goes through the identity (phase 1 already did this for the dirty
  image). Behaviour change is expected to be at machine precision; the L794 unchanged-test is the guard.
• Out of scope (ledger notes): per-chunk JAX recompile → profiling campaign; visualize_before_fit switch →
  PyAutoFit draft; phase 5 cubes/phase-centre.

Key Files

• autoarray/inversion/inversion/dataset_interface.py — DatasetInterface gains data_term
• autoarray/inversion/inversion/interferometer/abstract.py — fast_chi_squared term-3 override
• autoarray/fit/fit_interferometer.py — array-free chi_squared hook for no-inversion fits
• autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py — shared identity helper next to operated_matrix_slim_from
• autogalaxy/interferometer/fit_interferometer.py — sparse_profile_terms_from, profile_visibilities, galaxies_to_inversion, model_data
• autolens/interferometer/fit_interferometer.py — mirror
• tests: test_autoarray/inversion/inversion/interferometer/test_interferometer.py, test_autoarray/fit/test_fit_interferometer.py, test_autogalaxy/interferometer/test_fit_interferometer.py, test_autolens/interferometer/test_fit_interferometer.py

Verification

• pytest test_autoarray/inversion/inversion/interferometer test_autoarray/fit/test_fit_interferometer.py
• pytest test_autogalaxy/interferometer test_autogalaxy/interferometer/model and the al mirrors; full suites per repo
  (baseline counts P3: aa 1898, ag 1293, al 776+1 xfail).
• Witness (prompt header): LP+pixelization and LP-only fits on a from_stream dataset match the in-memory fit at
  rel 1e-8 in numpy and under jax.jit; the visibilities_from spy is never called. Run also the phase-2/3 analysis
  round-trip tests (ag test_analysis_interferometer.py:171/271/318, al :220) to confirm save/reload/visualize
  still pass with an ordinary light profile in the model.

Original Prompt

Click to expand starting prompt

Streaming phase 4: non-linear light profiles array-free via the data-term identity

Type: feature
Target: PyAutoArray
Repos:

• PyAutoArray
• PyAutoGalaxy
• PyAutoLens
  Themes:
• interferometer
• sparse-operator
• memory
  Autonomy: supervised
  Priority: medium
  Status: draft
  Epic: streaming-visibilities
  Phase: 4
  Difficulty: medium
  Consequence: judge
  Witness: fits with a non-linear light profile plus a pixelization, and with non-linear light profiles only, run on an array-free dataset with chi-squared computed as data_term − 2 i_pᵀd̃ + i_pᵀW̃i_p (one operated_matrix_slim_from product shared with sparse_dirty_image_from), matching the in-memory fit at rel 1e-8 in numpy and under jax.jit; profile_visibilities is never formed.
  Review-minutes: 8
  Unattended: ready
  Parent: draft/feature/autoarray/interferometer_from_stream_array_free_dataset.md
  Blocked-by: none (phase 2 merged 2026-09-30)

Source: https://github.com/orgs/PyAutoLabs/discussions/13 phase 2, sliced 2026-09-30.

What

1. aa.DatasetInterface gains an optional precomputed data term (e.g. sparse_data_term) that fast_chi_squared term 3 reads when set
   (abstract.py ~L218-225), replacing the profile_subtracted_visibilities reduction.
2. autogalaxy computes it next to sparse_dirty_image_from (fit_interferometer.py ~L58-111) from profile_image; extend
   uses_precomputed_data_term_from so non-linear light profiles also pass data=None; autolens mirror.
3. No-inversion fits: override chi_squared in ag/al FitInterferometer on the same identity (figure_of_merit → log_likelihood).
4. Tests mirror the phase-1 sparse-vs-dense set for LP+pix and LP-only, numpy + jit; profile_visibilities spy stays empty.

[Jammy2211]

Library PRs created (streaming P4)

• feat: per-interface data_term override and sparse profile-term identity (streaming P4) #599
• feat: array-free interferometer fits with ordinary light profiles (streaming P4) PyAutoGalaxy#642
• feat: array-free interferometer fits with lens light profiles (streaming P4) PyAutoLens#762

All three on feature/streaming-p4-light-profile-identity, labelled pending-release. Merge order: Array → Galaxy → Lens (separate human /prm).

Heart RED development override (human-authorized, 2026-10-01)

Live human authorization in the CLI session (Fable 5.1), quoted verbatim: "Yes, ship (push + PR-open)" for this task. Exact RED reasons at ship time (pyauto-heart readiness):

• release validation FAILED (stage integrate)
• workspace validation not passing (0 failed, 1 timeout, cloud#36404726969: autolens_test scripts/multi_dataset/rectangular.py)
• manifest drift: public front-door organ tables (generated) — 1 mismatch(es) vs PyAutoMind/repos.yaml

None relate to this branch. Scope: push + PR-open only; merge requires a separate human /prm with every required check green; not a release; no claim that Heart is healthy.

Branch gates

• test_autoarray 1903 passed (main 1898); test_autogalaxy 1305 passed (main 1293); test_autolens 793 passed + 1 xfailed (main 776 + 1).
• Red-check of the plan's trap: with the array-free raise removed but no data_term override, fast_chi_squared fell back to the operator's unsubtracted scalar — LP+pixelization log_evidence −37.0125 vs dense −37.9117 (−2.4 %), LP+linear −24.3192 vs −25.2183 (−3.6 %), silently. With the override: rel 1.8e-15.
• Codex gpt-6-astra review (3 independent runs): F1/F2 (high, introduced — noise-map/data override on an array-free fit silently mixed subtracted dirty image with the raw data term / ignored the override without an inversion) → typed DatasetException guard; F3 (medium — array-free log_likelihood with an inversion carried eps·Σs², −0.25 nats at packaged eps=1e-3 and ×100 data) → identity on the total model image, rel 1e-15 vs dense; F4 (medium, pre-existing on main — galaxy_image_dict merge drops a mixed galaxy's ordinary light) → filed draft/bug/autogalaxy/galaxy_image_dict_mixed_galaxy_overwrites_ordinary_light.md.

Workspace impact

API changes are additive (plus the removed array-free raise). No workspace script uses from_stream / from_sparse_terms / sparse_dirty_image_from / is_array_free yet → option (iii), no migration; smoke runs in CI on the same-named branches.

Deviations from the plan

• In-memory sparse fits with ordinary light keep the subtracted-visibilities path (bit-for-bit unchanged); only array-free datasets take the identity.
• sparse_chi_squared applies the identity to the total model image (after review F3), not inversion.fast_chi_squared.

[Jammy2211]

Shipped

Merged 2026-10-01 via human /prm with every CI leg green (Tests: unittest 3.12 / 3.13 / nojax; Docs where present), mergeability CLEAN:

• PyAutoArray#599 → merge 176f61e
• PyAutoGalaxy#642 → merge 574743ab
• PyAutoLens#762 → merge bcf2ab6f

All three remain pending release (labelled pending-release; last tag 2026.9.27.2). Mind record: complete/2026/10/streaming-p4-light-profile-identity.md. Next in the epic: the Discussion #13 follow-up post, then phase 5 (draft/feature/autoarray/streaming_p5_cubes_phase_centre.md). Pre-existing finding filed: draft/bug/autogalaxy/galaxy_image_dict_mixed_galaxy_overwrites_ordinary_light.md.