Skip to content

feat: analytically-solved point-source likelihood variants (#657) - #659

Merged
Jammy2211 merged 2 commits into
mainfrom
feature/point-source-chi-squared-variants
Jul 27, 2026
Merged

Jammy2211 merged 2 commits into
mainfrom
feature/point-source-chi-squared-variants

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Implements the analytically-solved point-source likelihood variants designed in #657 (phase 2 of 5; Lombardi 2024, arXiv:2406.15280). Point-source models no longer need to sample a source-plane centre (y,x) — or a flux, or a reference time — as non-linear free parameters: paired with the new zero-parameter al.ps.PointSolved profile, the *Solved fit classes solve them in closed form each likelihood call. The source-plane variant uses the paper's linearized likelihood with full Jacobian-tensor noise weighting (exact per-eigendirection stretch weighting; the traditional scalar µ²/σ² remains available as weighting="magnification") plus the analytic marginalization determinant term. The solved image-plane variants (all-to-all and pair-repeat) feed the solved centre into the existing PointSolver forward solve — an extension in the spirit of glafic's source-position optimization, not a paper result. Also: pluggable fit_flux_cls/fit_time_delays_cls hooks (previously hard-wired), pytree registration for the new classes and FitPositionsSource, informative errors replacing bare AttributeError on profile/fit mismatches, and a truth sweep of five docstrings that claimed an unimplemented "barycenter" option.

API Changes

All additions — existing classes, defaults, and behaviour are unchanged. Five new fit classes (FitPositionsSourceSolved, FitPositionsImagePairAllSolved, FitPositionsImagePairRepeatSolved, FitFluxesSolved, FitTimeDelaysSolved), the SolvedCentre mixin, and two new optional keywords (fit_flux_cls, fit_time_delays_cls) on FitPointDataset/AnalysisPoint. Invalid profile×fit pairings now raise informative PointExtractionExceptions instead of bare AttributeErrors.
See full details below.

Test Plan

  • pytest test_autolens/point/ (85 passed locally) and full pytest test_autolens/ (473 passed locally)
  • Correctness tests included: solved centre/flux/time equal brute-force optima; solved log-likelihood matches free-parameter likelihood profiled over the parameter; tensor-vs-scalar anisotropic ordering matches the image-plane chi-squared; PairAll permutation-count penalty regression; loud-failure matrix over {Point, PointFlux, PointSolved} × fit classes
Full API Changes (for automation & release notes)

Added

  • al.FitPositionsSourceSolved — source-plane chi-squared with analytically-solved centre (Lombardi 2024 §5.1); class attribute weighting = "jacobian" (default, exact tensor) or "magnification" (traditional scalar µ²/σ²); includes the analytic marginalization term -0.5*log(det(ΣWᵢ)/(2π)²).
  • al.FitPositionsImagePairAllSolved, al.FitPositionsImagePairRepeatSolved — existing image-plane chi-squareds driven by the solved centre (glafic-style extension).
  • al.FitFluxesSolved — flux solved in closed form (F* = Σµᵢf̂ᵢ/σᵢ² / Σµᵢ²/σᵢ², magnification-first flux space); works with Point or PointSolved; raises on PointFlux (its flux prior would be silently ignored).
  • al.FitTimeDelaysSolved — precision-weighted analytic reference time T* (alternative to the min-subtraction in FitTimeDelays, which is unchanged).
  • autolens.point.fit.solved.SolvedCentre mixin + precision_tensor_components_from — the shared β* linear algebra.
  • FitPointDataset(..., fit_flux_cls=FitFluxes, fit_time_delays_cls=FitTimeDelays) and the same on AnalysisPoint — previously hard-wired; defaults preserve current behaviour exactly.
  • Pytree registration for the five new classes and (new) FitPositionsSource.

Changed Behaviour

  • Profile/fit mismatches raise informative PointExtractionExceptions naming the fix (e.g. centre-bearing profile × *Solved fit; profile without centre × centre-requiring fit; PointFlux × FitFluxesSolved), instead of bare AttributeErrors.
  • Docstrings only: five files no longer claim an unimplemented "barycenter of ray-traced positions" option; FitPositionsSource's "common centroid" wording corrected (the reference point is the profile centre, or β* in the Solved class).

Migration

  • No changes required to existing scripts. To adopt: point_0 = af.Model(al.ps.PointSolved) (0 free parameters) + al.AnalysisPoint(..., fit_positions_cls=al.FitPositionsSourceSolved) (or an ImagePair*Solved class; add fit_flux_cls=al.FitFluxesSolved / fit_time_delays_cls=al.FitTimeDelaysSolved for datasets with fluxes/delays).

Depends on the companion PyAutoArray PR (xp threading through Grid2DIrregular.grid_2d_via_deflection_grid_from) and PyAutoGalaxy PR (ag.ps.PointSolved) on the same branch name — merge those first.

Closes nothing on merge — #657 stays open for phases 3–5 (workspace/profiling examples, guides, JAX gradients).

Generated by the PyAutoLabs agent workflow.

FitPositionsSourceSolved (tensor-weighted, marginalized centre; Lombardi
2024 arXiv:2406.15280 §5.1), FitPositionsImagePairAllSolved /
FitPositionsImagePairRepeatSolved (solved centre driving the existing
PointSolver forward solve), FitFluxesSolved (analytic flux, flux-space,
magnification-first), FitTimeDelaysSolved (analytic reference time);
SolvedCentre mixin; fit_flux_cls / fit_time_delays_cls hooks; pytree
registration incl. FitPositionsSource; informative mismatch errors;
docstring truth sweep.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Jul 27, 2026
…tches, kwarg order, test gaps

- FitPositionsSourceSolved.noise_normalization now uses the observed-plane
  det(Theta_i) (model-independent; Lombardi Eq. 46), not det(W_i) which
  spuriously favoured high-magnification models.
- New exc.PointProfileMismatchException (NOT a PointExtractionException
  subclass) for profile/fit mismatches, so FitPointDataset's component-skip
  handlers can no longer swallow them; regression test through
  FitPointDataset added.
- fit_flux_cls / fit_time_delays_cls moved to the end of the
  FitPointDataset / AnalysisPoint signatures — positional callers unbroken.
- Tests: RecordingMockSolver asserts beta* is actually forwarded to the
  solver; flux/time marginalization + normalization constants pinned to
  closed forms.
- Solved source-plane path verified numpy==jax under vmap+jit
  (-94.70750993 both backends on the workspace_test harness).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

@Jammy2211 Jammy2211 removed the pending-release PR queued for the next release build label Sep 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant