feat(autofit): multi-start gradient convergence results contract (phase 2)

[Jammy2211]

Overview

Phase 2 of 2 of the multi-start gradient auto-convergence work (phase 1 shipped: PyAutoFit#1407 + autofit_workspace_test#60). The results contract the author flagged as "must ship together": make the auto-convergence outcome inspectable and keep the aggregator robust to the variable-length runs early-stopping now produces. Single cohesive PR (the generic phaser over-split it; the contract is one deliverable).

Plan

• Add converged: bool + stop_reason ('converged'|'max_steps') + the convergence settings to samples_info (the loop already persists stop_reason in search_internal).
• Surface the fom_history convergence trace in the standard result artifacts so users can verify the plateau — carry fom_history into samples_info and add an MLE plotter that draws the global-best figure-of-merit vs step.
• Harden aggregator / samples_summary for variable-length runs (different total_steps across datasets); guard the known zero-weight / NaN diagnostic-row IndexError edge (Fit#275 lineage; the multi-start searches write zero-weight rows for the non-best starts).
• Tests: numpy-only unit tests (samples_info keys, aggregator variable-length + zero-weight edge); extend the JAX validation to round-trip converged/stop_reason/trace through the results DB.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)
• autofit_workspace_test (JAX validation extension)

Branch Survey

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

Suggested branch: feature/multistart-gradient-convergence-results-phase-2

Implementation Steps

1. search.py samples_via_internal_from (~L507, samples_info): read stop_reason = search_internal.get("stop_reason"); add converged (stop_reason == "converged"), stop_reason, the convergence settings dict (check_for_convergence/window/rtol/atol/min_steps), and the fom_history trace (list) to samples_info.
2. autofit/non_linear/plot/mle_plotters.py: add figure_of_merit_vs_iteration(samples, ...) that reads samples.samples_info.get("fom_history") and plots best figure-of-merit (or -0.5*fom = log-posterior) vs step, marking the plateau; no-op when absent (LBFGS/Drawer have no fom_history).
3. autofit/non_linear/search/mle/abstract_mle.py plot_results (~L19): wire the new plotter, gated by should_plot("figure_of_merit_vs_iteration"); guard on fom_history presence so generic MLE searches are unaffected.
4. autofit/aggregator/search_output.py (~L320, except IndexError): confirm/extend the zero-weight/NaN diagnostic-row guard for variable-length multi-start results; add a regression test rather than widening a silent guard.
5. Config: mirror any new should_plot key into packaged autofit/config/visualize/plots.yaml (and note for downstream workspaces).

Key Files

• autofit/non_linear/search/mle/multi_start_gradient/search.py — samples_info
• autofit/non_linear/plot/mle_plotters.py — new convergence-trace plotter
• autofit/non_linear/search/mle/abstract_mle.py — plot_results wiring
• autofit/aggregator/search_output.py — zero-weight/NaN IndexError guard
• test_autofit/non_linear/search/mle/test_multi_start_gradient.py + test_autofit/aggregator/ — unit tests
• autofit_workspace_test/scripts/jax_assertions/multi_start_gradient_auto_convergence.py — extend round-trip check

Testing

• python -m pytest test_autofit/ (numpy-only unit tests for samples_info + aggregator)
• autofit_workspace_test JAX validation extension (converged/stop_reason/trace round-trip)

Original Prompt

Click to expand starting prompt

Auto-convergence for the multi-start gradient searches — Phase 2: results-contract + aggregator hardening

Type: feature
Target: PyAutoFit
Repos:

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

Phase 2 of 2 — the results contract the prompt author flagged as "part 2, must
ship together" with the convergence loop
(multistart_gradient_auto_convergence_phase_1.md). Do not issue this until
phase 1 nears shipping (no bulk-issuing a prompt series). Depends on the
phase-1 stop_reason / converged locals and the MultiStartGradientConvergence
settings object landing first.

Goal

Make the auto-convergence outcome inspectable and keep the aggregator robust to
the variable-length runs early stopping now produces.

Scope

1. samples_info: add converged: bool + stop_reason ('converged' |
   'max_steps') + the convergence settings to samples_info
   (search.py:507, samples_via_internal_from).
2. Result artifacts: surface the fom_history convergence trace in the
   standard result artifacts so users can visually verify the plateau (the loop
   already tracks fom_history).
3. Aggregator / samples_summary coverage: extend test coverage for
   variable-length runs (different total_steps across datasets). Guard the known
   zero-weight / NaN diagnostic-row aggregator IndexError edge — see
   PyAutoFit#275 (_quick_fit consolidation, samples-weight-threshold) — since
   the multi-start searches write zero-weight diagnostic rows for the non-best
   starts.

Testing

• Unit tests (numpy-only): samples_info carries the new keys; aggregator
  handles variable-length runs + the zero-weight diagnostic-row edge without
  IndexError.
• JAX validation (autofit_workspace_test): aggregate a small set of MGE
  early-stopped runs and confirm the convergence trace + converged/stop_reason
  round-trip through the results database.

Deliverable

Library PR (PyAutoFit) + aggregator test coverage. Workspace-facing artifact
example follows once the API lands.

[Jammy2211]

Library PR Created — phase 2 (--auto)

PR: #1410 (pending-release) · commit 704a3b9

What shipped

Results contract: converged/stop_reason/convergence/fom_history added to samples_info; new figure_of_merit_vs_iteration MLE plotter (+ default-on config toggle, KeyError-tolerant gate) drawing the global-best FoM trace; aggregator regression tests for variable-length + zero-weight/NaN rows. Plus a real fix: Samples.max_log_likelihood_index/max_log_posterior_index used np.argmax which selects the multi-start's NaN diagnostic rows — now np.nanargmax (consistent with max_log_likelihood_sample; only changes results where NaN was present).

Four-leg gate

• Tests: ✅ 1525 passed / 1 skipped (+ new samples_info, variable-length/NaN, legacy-degradation tests)
• Smoke: ✅ searches/mle.py end-to-end (exit 0)
• Review: ✅ self-CLEAN (the shared samples.py nanargmax fix is in-scope for the zero-weight/NaN edge; full suite green)
• Heart: ⛔ RED — 2 pre-existing/unrelated reasons, human-waived to authorize the ship.

Next (after merge)

Extend autofit_workspace_test/scripts/jax_assertions/multi_start_gradient_auto_convergence.py to round-trip converged/stop_reason/fom_history through the results DB (via /start_workspace). Records phase-2 complete once that lands. This completes the 2-phase auto-convergence arc.

[Jammy2211]

Phase 2 complete — auto-convergence arc DONE ✅

Both PRs merged:

• Library: PyAutoFit#1410 (5e434d1) — results contract (samples_info converged/stop_reason/convergence/fom_history), figure_of_merit_vs_iteration plotter, and the max_log_likelihood_index/max_log_posterior_index argmax→nanargmax fix.
• Workspace: autofit_workspace_test#61 (0fe6869) — JAX results-DB round-trip (converged/stop_reason/fom_history trace, length 158, survive serialisation).

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

The 2-phase multi-start gradient auto-convergence arc is complete: default-on early-stopping (phase 1, #1406) + inspectable results contract (phase 2, #1409). No further phases planned.