feat(autofit): multi-start gradient auto-convergence (phase 1)

[Jammy2211]

Overview

Give the multi-start gradient searches (af.MultiStartProdigy / MultiStartAdam / MultiStartADABelief / MultiStartLion) a first-class auto-convergence (early-stopping) mode, default ON, so users no longer hand-tune n_steps. Today AbstractMultiStartGradient runs a fixed n_steps budget with no convergence criterion. n_steps becomes a hard ceiling. Phase 1 of 2 (phase 2 = results-contract + aggregator hardening, issued as this nears shipping). Successor to the multi-start gradient v2 work (Fit#1398/#1400).

Plan

• Add a MultiStartGradientConvergence settings object mirroring the AutoCorrelationsSettings precedent (Emcee/Zeus early-stop): check_for_convergence=True default, plateau on the GLOBAL best figure-of-merit over a trailing window (rtol/atol + min_steps), with an is_test_mode() shrink hook.
• Wire the check into the search loop at the iterations_per_full_update boundary; n_steps stays a hard ceiling (never runs forever). Carry a stop_reason local (surfacing into results is phase 2).
• Scope: PARAMETRIC SOURCES ONLY (MGE/Sersic). Pixelized is out of scope — when resurrect=True the auto behavior stays conservative/off and leans on the ceiling.
• Verify + guard the JAX compile-cache path: recall/resume must reuse the disk-cached compiled value_and_grad (already default-on via autonerves), not recompile. Prove the resumed HLO is stable so the persistent cache hits.
• Tests: numpy-only unit tests for the plateau logic in PyAutoFit; JAX early-stop + resume-no-recompile validation in autofit_workspace_test.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)
• autofit_workspace_test (JAX validation)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./autofit_workspace_test	main	6 regenerated dataset artifacts (inconsequential)

Suggested branch: feature/multistart-gradient-auto-convergence-phase-1

Implementation Steps

1. New file autofit/non_linear/search/mle/multi_start_gradient/convergence.py: MultiStartGradientConvergence(check_for_convergence=True, window, rtol, atol, min_steps) + check_if_converged(fom_history) (plateau: |best_fom[-1] - best_fom[-window]| <= atol + rtol*|best_fom[-window]| once len >= max(min_steps, window)); is_test_mode() shrinks window/min_steps.
2. search.py AbstractMultiStartGradient.__init__ (~line 28): add convergence: Optional[MultiStartGradientConvergence] = None (default -> auto ON); export the class via autofit/__init__.py.
3. search.py loop (~line 302-321, after fom_history extends + save_search_internal, before perform_update): break when check_if_converged passes AND NOT resurrect; carry stop_reason ('converged'|'max_steps') so during_analysis/final state reflect early stop.
4. Add convergence defaults to packaged autofit/config/.
5. JAX cache guarantee (verify, do NOT build): confirm batched_value_and_grad (search.py:198-208) is built once per fit; on resume (search.py:216-253) prove the rebuilt closure's HLO is byte-identical so the autonerves persistent cache (JAX_COMPILATION_CACHE_DIR) HITS — no recompile on recall.
6. Tests: numpy-only unit tests (plateau boundaries, min_steps, resurrect-off, ceiling honoured); autofit_workspace_test JAX script (MGE early-stop + resume asserts cache HIT / near-zero recompile).

Key Files

• autofit/non_linear/search/mle/multi_start_gradient/search.py — loop + __init__ + samples_info
• autofit/non_linear/search/mle/multi_start_gradient/convergence.py — NEW settings + check
• autofit/non_linear/search/mcmc/auto_correlations.py — precedent to mirror
• autonerves/jax_wrapper.py — existing persistent compile cache (reference, not edited)
• autofit/__init__.py, autofit/config/ — export + defaults

Testing

• python -m pytest test_autofit/non_linear/.../multi_start_gradient/ (numpy-only)
• autofit_workspace_test JAX validation (out of curated smoke subset)

Original Prompt

Click to expand starting prompt

Auto-convergence (early stopping) for the multi-start gradient searches — Phase 1: convergence loop + settings

Type: feature
Target: PyAutoFit
Repos:

• PyAutoFit
• autofit_workspace_test
  Difficulty: large
  Autonomy: supervised
  Priority: normal
  Status: formalised

Phase 1 of 2 (see multistart_gradient_auto_convergence_phase_2.md for the
results-contract + aggregator-hardening follow-up; issue phase 2 only as this one
nears shipping — no bulk-issuing). Successor to the multi-start gradient v2 work
(contrib rules + resurrection, Fit#1398/#1400).

Goal

Give af.MultiStartProdigy / MultiStartAdam / MultiStartADABelief / MultiStartLion
a first-class auto-convergence (early-stopping) mode (default on) so users no
longer hand-tune n_steps. Today AbstractMultiStartGradient runs a FIXED
n_steps budget with NO convergence criterion. n_steps stays a HARD
CEILING / max budget (never runs forever).

Scope

• In scope: PARAMETRIC SOURCES ONLY (MGE / Sersic — the smooth, well-behaved
  regime where a global-best plateau genuinely means converged).
• Out of scope: pixelized. Its best-fom climbs in long plateaus punctuated by
  breakthrough jumps (resurrection churn), so plateau detection false-stops there.
  So when resurrect=True the auto behavior stays conservative / off and leans
  on the ceiling. (The real pixelized goal is smoothing the LIKELIHOOD, not
  band-aiding the optimizer — tracked separately.)

Design (mirror the AutoCorrelationsSettings precedent)

Emcee/Zeus terminate early via AutoCorrelationsSettings
(autofit/non_linear/search/mcmc/auto_correlations.py:
check_for_convergence=True, a companion check_if_converged, and an
is_test_mode() shrink hook). Mirror that shape:

1. New file autofit/non_linear/search/mle/multi_start_gradient/convergence.py
   — MultiStartGradientConvergence(check_for_convergence=True, window, rtol, atol, min_steps) + a check_if_converged(fom_history) method. Plateau on the GLOBAL
   best figure-of-merit over the trailing window:
   |best_fom[-1] - best_fom[-window]| <= atol + rtol * |best_fom[-window]|,
   only once len(fom_history) >= max(min_steps, window). is_test_mode()
   shrinks window/min_steps so tests terminate fast.
2. search.py __init__ (AbstractMultiStartGradient, ~line 28): add a
   convergence: Optional[MultiStartGradientConvergence] = None param; default
   constructs the settings (auto ON). Store on self. Export
   MultiStartGradientConvergence via autofit/__init__.py.
3. search.py loop (~line 302–321, at the iterations_per_full_update
   boundary, right after fom_history extends + save_search_internal, before
   perform_update): if convergence enabled AND NOT resurrect, break when
   check_if_converged passes. n_steps remains the hard ceiling. Carry a
   stop_reason local ('converged' | 'max_steps') so during_analysis and
   the final state reflect the early stop. (stop_reason / converged surfacing
   into samples_info + result artifacts is phase 2 — this phase only needs the
   local so the loop terminates correctly.)
4. Config: add convergence defaults to the packaged autofit/config/ so
   downstream workspaces inherit them.

JAX compile-cache guarantee (recall must NOT recompile)

Requirement: when a likelihood function is recalled (resume / restart), the
JIT/grad-compiled value_and_grad must be reused, not recompiled.

The stack ALREADY ships the durable lever: autonerves/jax_wrapper.py sets
JAX_COMPILATION_CACHE_DIR (→ ~/.cache/pyauto_jax) +
JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=1 for every user (compile-time arc,
PyAutoConf#128). XLA serializes each compiled executable (jit AND grad) to disk
keyed on its HLO hash, so a fresh process loads the binary instead of recompiling.
Do NOT build redundant caching machinery (prefer the lean existing lever).

Phase-1 work is to verify + guard, not to build:

• In _fit (search.py:198–208) batched_value_and_grad is built once per fit
  and reused across all loop iterations — confirm no per-iteration re-tracing.
• On resume (search.py:216–253) the closure is rebuilt from a fresh
  fitness.call; the on-disk cache only HITS if the HLO is byte-identical across
  processes. Auto-convergence introduces variable-length/resume paths — prove the
  resumed value_and_grad HLO is stable so the persistent cache hits (no recompile
  on recall). If early-stop perturbs traced shapes and busts the cache, keep the
  traced computation shape-stable (the actual work, if any).
• Add a JAX-validation script in autofit_workspace_test that runs an MGE
  multi-start fit to early-stop, resumes in a fresh process, and asserts the
  resume does NOT recompile (cache HIT / near-zero compile time).

Testing

• Unit tests (numpy-only, PyAutoFit): plateau boundaries — converges on a flat
  history, does NOT on a climbing history, respects min_steps, disabled when
  resurrect=True, n_steps ceiling still honoured. NEVER use JAX in library
  unit tests.
• JAX validation (autofit_workspace_test): end-to-end MGE early-stop +
  resume-no-recompile check (above). Cluster-scale scripts stay out of the curated
  smoke subset.

Deliverable

Library PR (PyAutoFit) shipped first so the workspace can consume the API-change
summary, then the JAX-validation follow-up in autofit_workspace_test.

[Jammy2211]

Library PR Created — phase 1 (--auto)

PR: #1407 (pending-release) · commit 03f7fb5

What shipped

MultiStartGradientConvergence settings object (mirrors AutoCorrelationsSettings) + a convergence= param on all four multi-start gradient searches, default auto-convergence ON. Global-best plateau check runs per-step; n_steps stays a hard ceiling; skipped when resurrect=True (pixelized regime unchanged); stop_reason persisted for resume. New convergence.py, wired into search.py, exported from autofit/__init__.py, 6 numpy-only unit tests.

Four-leg gate

• Tests: ✅ full test_autofit/ suite — 1521 passed / 1 skipped (worktree). No public-API removal → downstream dependents n/a.
• Smoke: ✅ autofit_workspace/scripts/searches/mle.py (the one affected curated smoke script) runs end-to-end through the new code; early-stop fires (converged early at step 1 in test-mode); exit 0.
• Review: ✅ self-CLEAN (additive + one default-behaviour change, mirrors the MCMC precedent); the review faculty runs on the committed diff.
• Heart: ⛔ RED — two reasons, both pre-existing and unrelated to this change ("PyAutoFit: 2 commit(s) behind origin", "PyAutoLens: 1 uncommitted source change(s)"), human-waived contemporaneously to authorize the ship. Merge stays a separate human act.

API impact & workspace recommendation

Additive + one behaviour-default change (searches now early-stop by default; n_steps still a ceiling). No migration needed — restore the old behaviour with convergence=af.MultiStartGradientConvergence(check_for_convergence=False). Grep found only guide/tutorial constructions of MultiStartAdam/MultiStartProdigy; no script asserts on n_steps, and mle.py smoke passes.

Recommended next (option i, on the next launch): the phase-1 workspace follow-up in autofit_workspace_test — a JAX validation script that runs an MGE multi-start fit to early-stop, resumes in a fresh process, and asserts the resume does not recompile (the persistent JAX compile cache hits). Optionally a one-line note in the searches guide scripts about the new default. Routes via /start_workspace after this PR merges.

Phase 2 (results contract: converged/stop_reason in samples_info, the fom_history convergence-trace artifact, aggregator hardening for variable-length runs) stays a draft — issued as this phase nears merge.

[Jammy2211]

Library PR merged ✅

PyAutoFit#1407 merged to main (merge commit 63d4e3a). Auto-convergence early-stopping is now the default for the multi-start gradient searches (parametric regime; n_steps hard ceiling; resurrect=True unchanged).

Remaining phase-1 work: the autofit_workspace_test JAX validation — MGE multi-start fit to early-stop, resume in a fresh process, assert no recompile (persistent JAX compile cache hits). Routing to /start_workspace now. Phase 1 records complete once that validation lands; phase 2 (results contract + aggregator) stays a draft, issued as this nears completion.

[Jammy2211]

Phase 1 complete ✅

Both PRs merged:

• Library: PyAutoFit#1407 (63d4e3a) — MultiStartGradientConvergence + default-on auto-convergence for the multi-start gradient searches.
• Workspace: autofit_workspace_test#60 (45a2dbe) — JAX validation: early-stop lands at the truth basin (converged step 158/300) + deterministic no-recompile guard (byte-identical value_and_grad HLO → persistent compile cache hits on recall).

Completion record: PyAutoMind complete/2026/07/multistart-gradient-auto-convergence-phase-1.md.

Phase 2 (results contract: converged/stop_reason in samples_info, fom_history convergence-trace artifact, aggregator hardening for variable-length runs) stays a draft — issue it when picked up.