Skip to content

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

Closed

Description

@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.

Activity

  1. Jammy2211 commented on Jul 27, 2026

    @Jammy2211
    CollaboratorAuthor

    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.
  2. Jammy2211 commented on Jul 27, 2026

    @Jammy2211
    CollaboratorAuthor

    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.

  3. Jammy2211 commented on Jul 27, 2026

    @Jammy2211
    CollaboratorAuthor

    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.

  4. Jammy2211 commented on Jul 27, 2026

    @Jammy2211
    CollaboratorAuthor

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

    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.

  5. Jammy2211 commented on Jul 27, 2026

    @Jammy2211
    CollaboratorAuthor

    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).

  6. added a commit that references this issue on Jul 27, 2026
  7. Jammy2211 commented on Jul 27, 2026

    @Jammy2211
    CollaboratorAuthor

    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.

  8. Jammy2211 commented on Jul 30, 2026

    @Jammy2211
    CollaboratorAuthor

    Phase 3 shipped — 2026-07-30

    PRs (all merged)

    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).
  9. Jammy2211 commented on Jul 31, 2026

    @Jammy2211
    CollaboratorAuthor

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

    PR

    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.
  10. Jammy2211 commented on Jul 31, 2026

    @Jammy2211
    CollaboratorAuthor

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

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

    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.
  11. added a commit that references this issue on Jul 31, 2026
  12. Jammy2211 commented on Jul 31, 2026

    @Jammy2211
    CollaboratorAuthor

    Gradient-search benchmark — the gradients work in a real search (2026-07-31)

    PR: PyAutoLabs/autolens_profiling#99 (truth-anchored Prodigy-vs-Nautilus on the galaxy-scale point-source dataset; full table + six result JSONs inside).

    Headline: MultiStartProdigy on PointSolved + FitPositionsImagePairAllSolved converges to the truth through the lens-equation solver — er 1.574, centre (0.00, 0.01) vs truth (1.6, (0,0)), logL +2.37 vs truth anchor +7.74, in ~16 min with the default 64×300 budget. This is the implicit-diff custom_jvp doing real work end-to-end.

    Supporting findings:

    • Free-centre image-plane Prodigy needs 256 starts to find the same basin (64 misses it; polish plateaus) — consistent with the filed PairAll -inf-underflow plateau issue and the solved variant's dimensionality edge.
    • The scalar-µ² free-centre source-plane likelihood is spectacularly biased here: truth logL −33788 while both samplers find −110…−313 at wrong parameters. Objective, not optimizer (Prodigy actually out-optimizes Nautilus 8.7× faster on it) — direct evidence for the tensor-weighted solved recommendation.

    Follow-ups filed in Mind per human decisions: cluster examples move demonstrated default to PointSolved (galaxy-scale stays free-centre); cluster-scale gradient-search benchmark next after point-source wraps.

  13. Jammy2211 commented on Aug 6, 2026

    @Jammy2211
    CollaboratorAuthor

    Closing as shipped: design fully implemented across all 5 phases — PyAutoArray#414, PyAutoGalaxy#531, PyAutoLens#659/#677, autolens_workspace#425 all merged (2026-08-06 tracker reconciliation audit).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions