feat: point-source likelihood variants design (arXiv:2406.15280)

[Jammy2211]

Overview

Design phase (1 of 5) for implementing the point-source likelihood variants of Lombardi 2024 (arXiv:2406.15280, Gravity.jl) in PyAutoLens: source-plane and image-plane chi-squareds with no free source-centre (y,x) parameters (analytic, Jacobian-weighted solve), tensor noise weighting for better errors, analytic flux/reference-time alternatives (magnification-first), the principled missing-image penalty, and (final phase) JAX gradients incl. solver implicit differentiation. This issue carries the design decisions; no source edits happen under it.

Plan

• Read Lombardi 2024 §§2, 5, 6.1 in full; pin exact equations for the linearized source-plane likelihood, pairing-marginalization normalization, flux/time-delay marginalization, and gradients.
• Decide the variant set: tensor-weighted solved-centre source-plane; centre-free PairRepeat + PairAll; verdict on the numpy-only Hungarian Pair; flux F* / reference-time T* alternatives; missing-image penalty disposition (PairAll already carries -log(n_permutations)).
• Decide the API: parameter-free al.ps model class + fit-side centre resolution (priors derive from __init__, so fit_positions_cls alone cannot remove the (y,x) priors); add fit_flux_cls/fit_time_delays_cls hooks (currently hard-wired); source_plane_coordinate semantics.
• Decide the JAX plan: xp threading, the two known FitPositionsSource jit blockers, and the phase-5 custom_jvp boundary for PointSolver.
• Design the phase-2 test matrix (numpy-only): analytic optima vs brute force, solved ≈ profiled likelihood parity, tensor-vs-scalar weighting case.
• Post the design here for approval; update the phase 2–5 Mind prompts to match.

Detailed implementation plan

Work Classification

Library (design phase; the 5-phase series is Both — library first, workspaces in phases 3–5).

Affected Repositories

• PyAutoLens (primary)
• PyAutoGalaxy + PyAutoArray (expected from phase 2 — al.ps model class; Grid2DIrregular xp fix)
• autolens_workspace, autolens_workspace_test, autolens_profiling, autolens_workspace_developer (phases 3–5)

Branch Survey (2026-07-27)

Repository	Current Branch	Dirty?
./PyAutoLens	main	clean
./PyAutoGalaxy	main	clean
./PyAutoArray	main	clean

Warnings: PyAutoGalaxy is claimed by in-flight task validation-searches-env-optax (hard conflict if still active at phase 2 start); PyAutoArray has unregistered feature/delaunay-nan-callback checked out in a worktree (warning only).

Suggested branch: feature/point-source-chi-squared-variants (phases 2+; phase 1 makes no source edits)
Worktree root: ~/Code/PyAutoLabs-wt/point-source-chi-squared-variants/ (created by /start_library at phase 2)

Implementation Steps

The 9 design decisions are specified in the Mind prompt (reproduced verbatim below): variant set; parameter-free al.ps class + fit-side API; name-pairing preservation; JAX plan; error-fidelity test design; analytic flux (magnification-first flux space vs paper-exact magnitude space); analytic reference time vs min-subtraction; missing-image penalty vs existing -log(n_permutations) and unmatched_model_policy; phase-5 gradient groundwork.

Key Files

• PyAutoGalaxy/autogalaxy/profiles/point_sources.py — Point/PointFlux (only classes; centre always free)
• PyAutoLens/autolens/point/fit/abstract.py:144-153 — source_plane_coordinate single funnel
• PyAutoLens/autolens/point/fit/positions/ — FitPositionsImagePair{,All,Repeat}, FitPositionsSource
• PyAutoLens/autolens/point/fit/dataset.py:34,114,130 — fit_positions_cls hook; hard-wired FitFluxes/FitTimeDelays
• PyAutoLens/autolens/point/model/analysis.py:42,181-218 — AnalysisPoint default + pytree registration
• PyAutoArray/autoarray/structures/grids/irregular_2d.py:170 — jit blocker (no xp)
• Five false "barycenter" docstrings + autolens_workspace/scripts/cluster/likelihood_function.py:312-318

Original Prompt

Click to expand starting prompt (phase 1 of 5; umbrella prompt in PyAutoMind)

Point-source chi-squared variants (arXiv:2406.15280) — Phase 1: design

Parent: point_source_chi_squared_paper_variants.md (verbatim request there).
Phase 1 of 5. Judgment-heavy — stays with the lead (Opus) session per
PyAutoBrain/skills/WORKFLOW.md; no source edits in this phase.

Goal

Produce the approved design for centre-free point-source likelihoods:
which variants from Lombardi 2024 (arXiv:2406.15280, Gravity.jl) we
implement, and the exact API. Output = design writeup on the GitHub issue

• any corrections to the phase 2–5 prompts.

Grounding (verified 2026-07-27)

• Model classes: PyAutoGalaxy/autogalaxy/profiles/point_sources.py has only
  Point(centre) and PointFlux(centre, flux) — centre is always a free
  (y,x). No centre-free class exists anywhere.
• Centre plumbing is a single funnel:
  PyAutoLens/autolens/point/fit/abstract.py:144-153
source_plane_coordinate → self.profile.centre. Every fit uses it.
• Fit classes (autolens/point/fit/positions/): FitPositionsImagePair
  (Hungarian, numpy-only, not jittable), FitPositionsImagePairAll
  (all-to-all LogSumExp mixture, from this paper, JAX-blessed),
  FitPositionsImagePairRepeat (nearest-with-repeats + unmatched-model
  policies, JAX-shaped, AnalysisPoint default), FitPositionsSource
  (source-plane, scalar µ² noise weighting).
• fit_positions_cls is the only pluggable hook (fit/dataset.py:34 default
  FitPositionsImagePair; model/analysis.py:42 default
  FitPositionsImagePairRepeat — note the defaults DIFFER).
• FALSE DOCSTRINGS: a "barycenter of ray-traced positions" option is claimed
  but not implemented in fit/positions/abstract.py:48,
  image/abstract.py:49, image/pair.py:30, image/pair_all.py:26,
  image/pair_repeat.py:18, and the workspace guide
  autolens_workspace/scripts/cluster/likelihood_function.py:312-318
  (claims FitPositionsSource(profile=None) uses a centroid — it actually
  raises PointExtractionException). Phase 2 must make the docstrings true
  or delete them; phase 4 fixes the guide.
• JAX blockers on the source-plane path (relevant to any new source-plane
  variant): Grid2DIrregular.grid_2d_via_deflection_grid_from
  (PyAutoArray/autoarray/structures/grids/irregular_2d.py:170) takes no
  xp; FitPositionsSource missing from pytree registration
  (model/analysis.py:181-218). distances_to_coordinate_from (:204) is
  already _xp-threaded.

Paper taxonomy to design against

• Image-plane (Eqs. 29–38): direct association / best-match / marginalize
  over all pairings (LogSumExp — our PairAll). Baseline forms keep the
  source position β as a free non-linear parameter.
• Source-plane linearized (Eqs. ~41–48): linearize the lens equation around
  each observed image; β enters linearly and is solved analytically as a
  precision-weighted (full local Jacobian / inverse-magnification-tensor
  weighted) mean of back-traced positions. This is BOTH the centre-free
  source-plane chi-squared AND the better-errors variant (tensor weighting
  vs our scalar µ²).
• Centre-free image-plane: profile the analytic β* (from weighted
  back-projection) into the image-plane chi-squareds — forward-solve model
  images from β* instead of a sampled profile.centre, for the pair-repeat
  and all-to-all schemes.

Decisions to make (the actual design work)

1. Variant set: (a) source-plane analytic-centre with tensor weighting;
   (b) source-plane analytic-centre with current scalar µ² weighting (for
   continuity/comparison)? (c) centre-free PairRepeat; (d) centre-free
   PairAll; (e) is a centre-free Hungarian Pair worth it given it is
   numpy-only? Recommend which becomes the documented default for cluster
   fits (Lenstool parity uses FitPositionsSource).
2. API: removing the (y,x) priors REQUIRES a parameter-free al.ps model
   class (priors derive from __init__) — fit_positions_cls alone cannot
   do it. Decide: one new class (e.g. al.ps.PointSolved, name TBD) +
   orthogonal centre-resolution on the fit side, vs new fit subclasses per
   scheme (the guide's existing class-attribute pattern, cf.
   unmatched_model_policy), vs both. Decide the semantics of
   source_plane_coordinate when the centre is solved (keep the single
   funnel; no silent None-guards — absence of a centre must be explicit).
3. Name-pairing: dataset name → profile pairing must keep working for the
   new class (fit/abstract.py:84-90).
4. JAX plan: the analytic β* (weighted mean) is fixed-shape and xp-friendly;
   decide whether phase 2 also fixes the two FitPositionsSource jit
   blockers above (recommended — new source-plane variants sit on the same
   primitives). Any PyAutoArray change makes this a 3-repo task
   (autoarray + autogalaxy + autolens); the FeatureDecision's
   "repos: autolens" is understated — record the override.
5. Error-fidelity check design: how phase 2 unit tests demonstrate (numpy
   only, no JAX in unit tests) that solved-centre likelihoods match the
   free-centre likelihood profiled over centre, and that tensor weighting
   reproduces image-plane errors better than scalar µ² on a known asymmetric
   configuration.
6. Fluxes — IN SCOPE, alternative to the current implementation. Paper
   (Eq. 39 + §6.1): magnitude-space Gaussian with lensing modulus
   LMᵢ = 2.5 log10|Aᵢ⁻¹|; intrinsic magnitude M enters linearly →
   analytically marginalized/profiled (conjugate prior), removing M as a
   free parameter. PyAutoLens is magnification-first flux-space:
   FitFluxes.model_data = |µᵢ| × profile.flux (fit/fluxes.py:110-124)
   with flux a free param on PointFlux. Decide the port: flux-space
   analytic profiling (linear least squares,
   F* = Σ µᵢ f̂ᵢ/σᵢ² / Σ µᵢ²/σᵢ² — natural fit to our flux-space noise
   maps and outputs) vs paper-exact magnitude space; and what the model
   class looks like (flux-free PointFlux sibling composing with the
   solved-centre class). NOTE: there is no fit_flux_cls hook —
   FitFluxes/FitTimeDelays are hard-wired in fit/dataset.py:114,130;
   the design must add pluggable hooks (mirroring fit_positions_cls) or
   an equivalent mechanism.
7. Time delays — IN SCOPE, alternative to current. Paper (Eqs. 24-25):
   tᵢ = T + Tᵢ with reference time T entering linearly → analytic
   marginalization, same structure as M. Current FitTimeDelays takes
   residuals relative to the MINIMUM delay (fit/times_delays.py:109-110)
   — an ad-hoc T elimination. Design the precision-weighted analytic T
   alternative and compare error behaviour vs min-subtraction.
8. Missing-image penalty — scope against what exists. The paper's
   mechanism is purely combinatorial: marginalizing over pairings with
   normalization 1/|Σ| (with repetition |Σ| = P^I; without repetition
   |Σ| = P!/(P−I)!) penalizes many-predicted/few-observed configurations;
   no explicit detection-probability term. FitPositionsImagePairAll
   ALREADY implements -log(n_permutations) (pair_all.py:151-154, cited
   to this paper). Verify it matches P^I exactly; decide whether to add
   the without-repetition variant; decide whether the paper-principled
   penalty should be offered as an unmatched_model_policy alternative to
   PairRepeat's ad-hoc no_image_residual = 1e4 floor and
   noise-normalized "penalize" policy (pair_repeat.py:93-95, 215-219).
9. JAX-gradient groundwork for phase 5 (design only, no implementation):
   confirm which variants get analytic/custom-JVP treatment vs plain
   autodiff — see the phase 5 prompt for the paper's gradient formulas.

Consults

• autolens_assistant/wiki/core/concepts/point_source.md (read; background
  current as of 2026-07-09) and skills/al_point_source.md.
• PyAutoMemory methods wiki (likelihood/jax) per FeatureDecision.
• Paper full text: https://arxiv.org/html/2406.15280 §§2, 5 (and 6.1 for the
  deferred marginalization option).

Exit criteria

Design posted to the issue; user approves variant set + API; phase 2–4
prompts updated to match. No code changes.

[Jammy2211]

Phase 1 design — point-source likelihood variants (Lombardi 2024, arXiv:2406.15280)

All code facts below verified against main on 2026-07-27.

D0. Paper results this design builds on

• Linearized source-plane likelihood (§5.1): Taylor-expand the lens equation around each observed image θ̂ᵢ, giving β ≈ β̂ᵢ + Aᵢ(θ̂ᵢ − θᵢ) with Aᵢ the lens Jacobian at the observed position. β then enters the chi-squared linearly and is solved in closed form:
  β* = (Σᵢ AᵢᵀΘᵢAᵢ)⁻¹ Σᵢ AᵢᵀΘᵢ β̂ᵢ (Θᵢ = image-plane position precision).
  Marginalizing β (flat prior) adds a −½ log|Σᵢ AᵢᵀΘᵢAᵢ|-type determinant term. The per-image precision transform AᵢᵀΘᵢAᵢ is the tensor weighting — the principled version of our scalar µ² weighting, and the "better errors" variant.
• Missing-image penalty (Eq. 32): the all-pairings marginalization normalizes by 1/|Σ|, |Σ| = P^I with repetition — a purely combinatorial Occam penalty. Verified: FitPositionsImagePairAll.chi_squared implements exactly this (n_permutations = n_finite_model^len(data), pair_all.py:145-154). Nothing new to implement; the without-repetition form (P!/(P−I)!, Eq. 37) requires factorial permutation enumeration — JAX-hostile, and rejected below.
• Analytic nuisance closed forms (§6.1): M* = Σλᵢ(m̂ᵢ+LMᵢ)/Σλᵢ (magnitude space); T* = Στᵢ(t̂ᵢ−Tᵢ)/Στᵢ.
• Honesty note: the paper never substitutes β* into its image-plane likelihoods (β stays free there). The centre-free image-plane variants below are our extension — a profile-likelihood hybrid with direct precedent in glafic's source-position-optimized image-plane χ². Docs must attribute accordingly.

D1. Variant set (decision)

#	Variant	Basis	Verdict
S1	Source-plane, solved centre, tensor weighting + det term	paper §5.1	IMPLEMENT — becomes the recommended source-plane fit (cluster default candidate)
S1b	S1 with scalar µ weighting	—	class-attribute option on S1 (weighting = "jacobian" | "magnification") for Lenstool-style comparisons; cheap, no new class
S2	Existing FitPositionsSource (free centre, scalar µ²)	—	unchanged, kept for continuity
I1	All-to-all image-plane, solved centre	our extension (glafic precedent)	IMPLEMENT
I2	Pair-repeat image-plane, solved centre	our extension	IMPLEMENT
I3	Hungarian FitPositionsImagePair, solved centre	—	SKIP — numpy-only (scipy Hungarian, not jittable), superseded by I1/I2; documented why
P	Missing-image penalty	paper Eq. 32	verification + docs only — already correct in PairAll; PairRepeat's unmatched_model_policy heuristics stay as-is (best-match is not a normalized mixture; no principled combinatorial term applies)
F	Analytic flux F*	paper §6.1, ported	IMPLEMENT, flux space: F* = Σᵢ µᵢ f̂ᵢ/σᵢ² / Σᵢ µᵢ²/σᵢ² + −½log(Σµᵢ²/σᵢ²) marginalization term. Magnitude space (paper-exact) REJECTED: our noise maps are flux-space Gaussians; converting changes the error model, not just the parametrization. Magnification-first, matching `FitFluxes.model_data =
T	Analytic reference time T*	paper §6.1	IMPLEMENT: precision-weighted T* + norm term, as alternative to the current double min-subtraction (times_delays.py:109-110)

D2. API (decision)

Model class — one new, parameter-free class in autogalaxy/profiles/point_sources.py:

class PointSolved:   # al.ps.PointSolved — zero non-linear parameters
    ...

• Removing the (y,x) priors requires this: priors derive from __init__, so no fit_positions_cls setting can do it.
• No flux sibling needed: the analytic-flux fit never reads profile.flux (F* comes from data + magnifications), so PointSolved covers positions+fluxes+delays datasets alone. FitFluxes (free flux) keeps requiring PointFlux.
• Name pairing (fit/abstract.py:84-90) unchanged — pairing is by dataset/profile name, not class.

Fit side — a SolvedCentre mixin overriding the single funnel:

• All fits already resolve the model centre through one property: AbstractFitPoint.source_plane_coordinate → profile.centre (abstract.py:144-153). The mixin overrides that property to return β* (tensor-weighted back-projection); everything downstream (incl. the PointSolver forward solve for image-plane fits, image/abstract.py:121-127) is untouched.
• Concrete classes: FitPositionsSourceSolved, FitPositionsImagePairAllSolved, FitPositionsImagePairRepeatSolved, FitFluxesSolved, FitTimeDelaysSolved (names open to bikeshed).
• Loud mismatches, both directions (no silent laundering): a centre-requiring fit given PointSolved raises with a message naming the Solved variant; a Solved fit given Point/PointFlux raises (otherwise the centre priors would be sampled but ignored — silent 2-parameter waste).
• New hooks: fit_flux_cls / fit_time_delays_cls on FitPointDataset and AnalysisPoint, mirroring fit_positions_cls; defaults = current FitFluxes/FitTimeDelays (behaviour unchanged). Currently hard-wired at fit/dataset.py:114,130.
• The five false "barycenter" docstrings (positions/abstract.py:48, image/abstract.py:49, pair.py:30, pair_all.py:26, pair_repeat.py:18) become true by pointing at the mixin; the workspace guide's false profile=None claim (cluster/likelihood_function.py:312-318) is fixed in phase 4.

Jacobian source: A = I − ∂α/∂θ assembles from ag.LensCalc.hessian_from (lens_calc.py:382, xp-threaded, JAX path exists) — same machinery as today's Hessian magnifications; no new differential-geometry code.

D3. JAX plan (decision)

• β*, F*, T* are fixed-shape linear algebra (2×2 solve / scalar ratios) — xp-threaded from day one; new fit classes registered in AnalysisPoint._register_fit_point_pytrees (model/analysis.py:181-218).
• In scope for phase 2 (approved rationale: the new source-plane variant sits on the same primitives): fix the two known FitPositionsSource jit blockers — Grid2DIrregular.grid_2d_via_deflection_grid_from gets xp (PyAutoArray/autoarray/structures/grids/irregular_2d.py:170), and FitPositionsSource joins the pytree registration. This makes the series a 3-repo task: PyAutoArray + PyAutoGalaxy + PyAutoLens (recorded override of the FeatureDecision's single-repo estimate).
• Solver gradients (custom_jvp via the fixed-point rule ∂θᵢ/∂p = Aᵢ⁻¹ ∂α/∂p) remain phase 5; phase 3's gradient scripts cover only the solver-free solved variants.

D4. Phase-2 test matrix (numpy-only; JAX via workspace_test scripts only)

1. β* equals brute-force minimizer of the tensor source-plane χ² over centre (fixed mock tracer).
2. Solved-centre log-likelihood ≈ free-centre likelihood profiled over centre, per scheme (S1, I1, I2).
3. Tensor vs scalar: an anisotropic-magnification mock where tensor weighting reproduces the image-plane χ² ordering scalar µ² gets wrong.
4. F*/T* equal brute-force scans; solved ≈ profiled parity for both.
5. Penalty regression: n_permutations == P_finite^I on a counted mock; predicted-image surplus lowers likelihood monotonically.
6. Loud-failure matrix: {Point,PointFlux,PointSolved} × {centre-requiring, solved} fits — every invalid cell raises with an actionable message.
7. Structure/__Env__ conventions mirror test_autolens/point/.

D5. Consequences for phases 3–5

• Phase 2 scope grows by: fit_flux_cls/fit_time_delays_cls hooks, FitFluxesSolved, FitTimeDelaysSolved, the PyAutoArray xp fix. Phase prompts updated accordingly.
• Recommended-default question (does S1 become the documented cluster default over FitPositionsSource?) is answered by phase 3's cluster likelihood_breakdown numbers + phase 2's error-fidelity tests; phase 4 writes the recommendation.

[Jammy2211]

Convention correction to D0/D2 above (before phase 2 reads it): with the standard lensing convention A ≡ ∂β/∂θ (so A⁻¹ is the magnification tensor, µ = 1/det A), the per-image precision on a back-traced position is Wᵢ = Aᵢ⁻ᵀ Θᵢ Aᵢ⁻¹ — inverse Jacobians, i.e. magnification-tensor weighting — and β* = (Σᵢ Wᵢ)⁻¹ Σᵢ Wᵢ β̂ᵢ. Sanity check: the isotropic limit gives Wᵢ → µᵢ²/σᵢ², matching the existing scalar weighting in FitPositionsSource.chi_squared_map. The AᵀΘA form written above is correct only if A is read as ∂θ/∂β; taken in the standard convention it would deweight high-magnification images — backwards. Phase 2 implements Wᵢ = Aᵢ⁻ᵀΘᵢAᵢ⁻¹ with A assembled from LensCalc.hessian_from, and the parity/brute-force tests in D4 guard the sign of this choice.

[Jammy2211]

Resolution of the isotropic-limit dispute (raised by the phase-2 implementer — they were right): with A = aI, µ = 1/a², so the tensor weighting's isotropic limit is W → (µ/σ²)I, not µ²/σ² as my correction comment claimed. Full picture: the eigenvalues of W = σ⁻²A⁻ᵀA⁻¹ are λ²/σ² with λ the linear stretch per eigendirection — isotropically each stretch is õ (→ µ/σ²), while near a critical curve the tangential stretch is ≈ µ (→ µ²/σ² along the arc). So the existing scalar µ²/σ² convention (FitPositionsSource, Lenstool-style) is the near-critical tangential limit, not the isotropic one, and the tensor is exact in both regimes. Implementation keeps both: weighting=\"jacobian\" (exact, default) and weighting=\"magnification\" (traditional scalar, for comparisons). Also fixed in review: the 2-D marginalization constant is now the exact (2π)² Gaussian-integral normalization. Bonus catch by the implementer: ag.LensCalc.jacobian_from returns (x,y)-ordered components against the codebase's (y,x) convention — invisible to its determinant-only existing callers, re-indexed at the new consumption site and flagged in its docstring.

[Jammy2211]

Phase 2 library PRs created (all suites green: autoarray 906, autogalaxy 1009, autolens 473; all labeled pending-release):

• feat: thread xp through Grid2DIrregular.grid_2d_via_deflection_grid_from PyAutoArray#414 — xp threading for grid_2d_via_deflection_grid_from (merge first)
• feat: add zero-parameter ps.PointSolved point-source profile PyAutoGalaxy#531 — ag.ps.PointSolved (+ jacobian_from ordering docstring warning)
• feat: analytically-solved point-source likelihood variants (#657) #659 — the five *Solved fit classes, SolvedCentre mixin, fit_flux_cls/fit_time_delays_cls hooks, pytree registrations, mismatch errors, docstring truth sweep

API changes are additions-only — no existing script is affected. Heart was YELLOW at ship (pre-existing, unrelated reasons; human-acknowledged: workspace validation 13-failed of 2026-07-21 [predates today's searches-shard fix], 33 stale parked scripts, tenant-firewall manifest drift ×5).

Next: phases 3 (workspace_test jax_likelihood + profiling likelihood_runtime/breakdown examples) and 4 (guides) via /start_workspace after these merge; phase 5 (JAX solver gradients) last.

[Jammy2211]

Independent Codex review (codex-cli 0.145.0) of the three PRs: FINDINGS — 9 items. Disposition after adversarial verification:

#	Finding	Verdict	Action
1	noise_normalization used det Wᵢ (model-dependent) instead of observed-plane det Θᵢ — spuriously favours high magnification	CONFIRMED	Fixed: normalization is now Σ log((2π)²σᵢ⁴); FitPositionsSource's magnified-noise convention intentionally unchanged, difference documented
2	jacobian_from crashes under jit for all solved fits	REFUTED empirically — solved source-plane runs under vmap+jit with exact numpy parity (−94.70750993 both backends)	none
3	xp.asarray(deflection_grid) hits __array__ on traced wrappers	CONFIRMED (reproduced: TracerArrayConversionError in the workspace_test harness)	Fixed: unwrap .array before subtraction; harness now passes where it crashed
4	Mismatch exceptions swallowed by FitPointDataset's component-skip except blocks	CONFIRMED (worse than stated — all three components)	Fixed: new exc.PointProfileMismatchException, deliberately not a PointExtractionException subclass; regression test through FitPointDataset
5/6	New kwargs inserted mid-signature break positional callers	CONFIRMED	Fixed: hooks moved to signature ends
7	xp=np default replaced the pre-existing self._xp propagation	CONFIRMED	Fixed: xp=None → self._xp
8	MockPointSolver ignores source_plane_coordinate → forwarding untested	CONFIRMED (β*-vs-truth was tested; forwarding was not)	Fixed: RecordingMockSolver test
9	Flux/time marginalization constants unpinned	CONFIRMED	Fixed: closed-form pins

Bonus from the fix cycle: the pre-existing source-plane jit blocker is now genuinely cleared — scripts/point_source/jax_likelihood/source_plane.py's vmap+jit path passes (was TracerArrayConversionError); only the separately-tracked fit_from-returns-PointSolver pytree gap remains (pre-existing, out of scope).

All suites green after fixes: autolens 477, autoarray 906, autogalaxy 1009. Fix commits pushed to PyAutoLabs/PyAutoArray#414 and #659 (PyAutoGalaxy#531 unchanged by the review).

[Jammy2211]

Phases 1+2 complete — all three library PRs MERGED (2026-07-27, post codex-review fixes; CI green on 3.12+3.13 everywhere): PyAutoLabs/PyAutoArray#414 → PyAutoLabs/PyAutoGalaxy#531 → #659. Feature branches and the task worktree are cleaned up; active.md released the library-repo claims.

Resuming here (phases 3–5): run /start_workspace on PyAutoMind/draft/feature/autolens/point_source_chi_squared_paper_variants_phase_3_workspace_examples.md (workspace_test jax_likelihood + profiling likelihood_runtime/breakdown examples — note the solved source-plane path already passes vmap+jit with exact numpy parity, so phase 3 mostly extends harnesses + records baselines), then phase 4 (guides incl. the likelihood_function.py false-claim fix), then phase 5 (PointSolver custom_jvp gradients). The phase prompts carry the binding design deltas + this thread has the weighting-convention resolution.

[Jammy2211]

Phase 3 shipped — 2026-07-30

PRs (all merged)

• Solved point-source fit coverage: jax_likelihood variants, gradients, fluxes+time delays (#657 phase 3) autolens_workspace_test#237 — solved-fit jax_likelihood coverage, gradients, fluxes+time-delays dataset
• Solved point-source runtime profiling + cluster breakdown solved steps (#657 phase 3) autolens_profiling#96 — solved runtime siblings + cluster breakdown solved steps
• Correct stale source-plane xp-blocker claims in jax_profiling (#657) autolens_workspace_developer#121 — stale xp-blocker prose corrected

Summary

Every phase-2 solved fit class now has JAX regression coverage: vmap literals with exact numpy parity for FitPositionsSourceSolved (−94.70750993), FitPositionsImagePairAllSolved (−82.33883111), FitPositionsImagePairRepeatSolved (−89.71129442), and the fluxes+time-delays combo (−137.67455417); jit(fit_from) passes exactly for all image-plane solved variants; finite-difference gradient checks pass for the solved source-plane fit. Profiling: the analytic β* step costs ~timing-noise-level overhead per likelihood (cluster image-plane net +3% across systems, source-plane +9–30% on sub-0.1 s eager totals) while removing 2 free parameters per point source from the non-linear space.

Session Notes

• The old source-plane BLOCKER (xp bug) is confirmed gone; the sole remaining jit(fit_from) gap is the fit-return pytree issue (PointSolver leaf, output component [1][1]) — tracked in PyAutoPrompt/autolens/fit_point_pytree.md, affects old and solved fits alike, vmap/log_likelihood_function unaffected.
• Eager-vs-vmap parity for time-delay fits requires identical cosmology on both sides (H0-sensitivity) — the new script documents and asserts this with an H0-matched comparison; positions-only scripts are H0-independent so their existing pattern stands.
• al.output_to_json has nondeterministic JSON key order run-to-run (value-identical output; PYTHONHASHSEED-driven) — filed in Mind ideas for a serializer fix.
• imaging/subhalo_recovery.py exceeds the 300 s smoke cap on the WSL dev laptop (solo, quiet machine) while green on main CI — margin worth watching after the parked-sweep re-measure.
• Remaining phases: 4 (guides, incl. the false profile=None-centroid claim in cluster/likelihood_function.py), 5 (PointSolver custom_jvp gradients).

[Jammy2211]

Phase 4 (guides) PR open — 2026-07-31

PR

• docs: point-source solved-variant guides (phase 4 of #657) autolens_workspace#425 (pending-release, awaiting human merge)

Upstream

Library API already on main from phase 2 (PyAutoArray#414 / PyAutoGalaxy#531 / PyAutoLens#659) — no library-first merge gate applies.

Summary

The workspace guides now document the full likelihood-option matrix (pairing scheme × free-vs-solved centre): the pairing guide carries the deep prose (solved-variant table, tensor weighting = "jacobian"|"magnification" conventions, solved fluxes/time delays, missing-image-penalty status, runnable solved-centre demo), point_source/fit.py gains a runnable __Solved Source Centre__ section plus the honest defaults story (FitPointDataset → FitPositionsImagePair vs AnalysisPoint → FitPositionsImagePairRepeat), and the false FitPositionsSource(profile=None) back-traced-barycenter claim in cluster/likelihood_function.py is replaced with the real centre-free option. The deferred cluster recommendation (phase-1 design item) is now written from the phase-3 numbers: search with FitPositionsSourceSolved, validate with the image-plane chi-squared; demonstrated defaults in the modeling scripts are unchanged.

Session Notes

• Ship ran under an explicitly human-acknowledged Heart RED: the nightly Release Integrate failure is interferometer/start_here.py OOM (~86 GB JAX allocation) in the release-wheels leg — unrelated to this docs-only change; filed as PyAutoMind/draft/bug/autolens/interferometer_release_leg_oom.md.
• Validation: local smoke suite 26/26 passed; the four code-changed scripts executed end-to-end (solved centre recovers the true (0.07, 0.07); fluxes solved model prints N=5 as documented).
• Attribution: image-plane solved variants cite the glafic precedent (Oguri 2010) and are not attributed to Lombardi 2024; the scalar µ² weighting is labelled as the Lenstool-style convention.
• Assistant-wiki refresh (point_source concept + skill pages) filed in Mind ideas for maintainer-mode once released.
• Remaining: phase 5 (PointSolver custom_jvp gradients) — note PyAutoLens is currently claimed by potential-correction-validation (fix: potential-correction end-to-end validation (parity + evidence sampling) #672); hand-check file overlap before starting.

[Jammy2211]

Phase 5 (solver gradients) PRs open — 2026-07-31

PRs (library-first; all pending-release, awaiting human merge)

• feat: PointSolver implicit-diff gradients (phase 5 of #657) #677 — the implicit-diff custom_jvp (merge FIRST)
• test: solver-gradient certification + staircase FD methodology (phase 5 of #657) autolens_workspace_test#240 — certification blocks + staircase FD methodology
• docs: point-source gradient probe status flip (phase 5 of #657) autolens_workspace_developer#123 — gradient-probe status flip
• feat: multi_start_prodigy point_source image_plane_solved cell (phase 5 of #657) autolens_profiling#98 — multi_start_prodigy image_plane_solved cell

Summary

PointSolver positions now carry gradients via an implicit fixed-point custom_jvp (autolens/point/solver/implicit_diff.py): A dθ = dα|_θ + dβ at the solved positions — the same mechanism gravity.jl itself uses (the paper never differentiates through its solver; its autodiff only verifies the analytic rule). FD-certified across the variant matrix: solved and sampled centres through the solver (≤2% vs the stair-noise floor, spike best-steps <0.4%), PairRepeat subgradients live, and nested autodiff supplies the flux/magnification gradients the paper declares analytically out of reach. af.MultiStartProdigy can now run on image-plane point-source likelihoods (first search-integration cell registered in the profiling sweep). All values are unchanged: numpy path byte-identical, test_autolens/ 492 passed, every jax_likelihood regression literal re-verified.

Session Notes

• The historical "not differentiable" verdict was about differentiating THROUGH the refinement — still true. The implicit rule differentiates AT the solution; same verdict-goes-stale shape as the Delaunay frozen-tables arc.
• Solver staircase: the forward solve quantizes at pixel_scale_precision, so naive FD reads exactly zero against a correct nonzero AD gradient (the exact-solve envelope derivative — what a gradient search wants). Certification uses a fine-precision solver + per-parameter FD step sweep; methodology encoded in the cert script prose.
• Compatibility gate: the rule engages only when the tracer flattens to JAX leaves (the af.Model/Fitness registered path); hand-built tracers keep the old zero-gradient forward — nothing existing changes behaviour.
• Known limitation (follow-up filed in Mind): free cosmology parameters cannot cross the custom_jvp boundary (Tracer keeps cosmology as no_flatten aux → stale tracer). Harmless for 2-plane positions; multi-plane free-cosmology gradient fits need cosmology flattened into the pytree.
• Also filed: FitPositionsImagePairAll mixture underflows to -inf at ≳38σ mismatch (not a stable logsumexp) — shrinks the gradient-search basin; stabilization follow-up in Mind ideas.
• Ship ran under the same human-acked Heart RED as phase 4 (nightly interferometer release-leg OOM, unrelated; re-acked for this ship). Local wst smoke 20/21 — sole failure is the known imaging/subhalo_recovery.py WSL 300s-cap margin, green on CI.
• This completes the implementation phases of the series; phases 1–5 all shipped or awaiting merge.