feat(search): add multi-start gradient MAP searches (Adam/ADABelief/Lion)

[Jammy2211]

Overview

Phase 1 of promoting the benchmark-winning multi-start gradient MAP optimizer to first-class PyAutoFit searches. The wide multi-start first-order gradient method reliably recovers the truth basin where every single-start, line-search and second-order method fails (autolens_workspace_developer PR#96+#98, Phase-3 complete: nothing beats multi-start Adam on the NNLS-kinked lens likelihood). This grows the mle/ search library alongside Drawer / BFGS / LBFGS.

Plan

• Add AbstractMultiStartGradient(AbstractMLE) under autofit/non_linear/search/mle/multi_start_gradient/, owning the shared N-broad-start, vmapped value_and_grad, optax-update MAP loop against the library-agnostic af.Fitness seam.
• Ship three concrete rule-classes — af.MultiStartAdam (lr 1e-2, certified best), af.MultiStartADABelief (lr 1e-2, tied), af.MultiStartLion (lr 1e-3) — each a thin subclass fixing its optax factory + default lr; export in autofit/__init__.py.
• Return best-basin MAP + per-start basin diagnostics through the standard Samples/Result contract, with search-internal save/resume.
• Add optax as a jax-extra dependency with a lazy, clear-error import.
• Library unit tests numpy-only (plumbing); JAX end-to-end truth-basin recovery in autofit_workspace_test behind the library-first gate.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)
• autofit_workspace_test (JAX end-to-end validation, follows library merge)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./autofit_workspace_test	main	clean

Suggested branch: feature/multi-start-gradient-search

Note: PyAutoFit is softly claimed by database-latent-wheel-load (PyAutoFit#1368, merged; only corrective release-validation remains). Human-approved to proceed in a parallel worktree — no real source collision.

Implementation Steps

1. autofit/non_linear/search/mle/multi_start_gradient/{__init__.py,search.py}:
   • AbstractMultiStartGradient(AbstractMLE) with __init__(n_starts=48, n_steps=300, learning_rate=None→_default_lr, start_low=0.15, start_high=0.85, initializer, iterations_per_*_update, silence, session, **kwargs) and class attrs _optax_factory, _default_lr.
   • _fit(model, analysis): Fitness(model, analysis, fom_is_log_likelihood=False) with use_jax_vmap when analysis._use_jax; lazy import optax; broad unit-space starts → model.vector_from_unit_vector(u, xp=jnp), keep finite-gradient starts; jax.jit(jax.vmap(jax.value_and_grad(neg_log_posterior))); optax loop; track global-best + per-start finals; save_search_internal for resume — mirroring mle/bfgs/search.py.
   • samples_via_internal_from: best-basin MAP as max-lik Sample, per-start finals as diagnostic samples; samples_info carries per-start basin stats.
   • MultiStartAdam / MultiStartADABelief / MultiStartLion subclasses.
2. autofit/__init__.py: export the three classes after LBFGS.
3. pyproject.toml: add optax to the jax extra.
4. test_autofit/non_linear/search/mle/test_multi_start_gradient.py: numpy-only — start generation/filtering, Samples build, resume round-trip, config parsing (simple analytic objective; no JAX in the library suite).
5. autofit_workspace_test: end-to-end script — MultiStartAdam recovers a known truth basin on a JAX likelihood; asserts best-basin params.

Key Files

• autofit/non_linear/search/mle/bfgs/search.py — structural analogue (Abstract + concrete subclasses, samples_via_internal_from, resume).
• autofit/non_linear/fitness.py — the af.Fitness seam (_jit/_vmap/_grad, use_jax_vmap).
• autofit/mapper/prior_model/abstract.py::vector_from_unit_vector — jax-traceable unconstrained transform.
• autolens_workspace_developer/searches_minimal/{gpu_multi_start_adam,_grad_setup}.py — reference algorithm (port algorithm only; PyAutoFit must not import autoarray/autolens).

Out of scope

Line-search / second-order methods (L-BFGS/BFGS/NCG/LM/Gauss-Newton) — all failed the benchmark on the NNLS-kinked objective.

Follow-up phases

• Phase 2 — config + packaged defaults (optax rule / N / lr / step-count as first-class packaged config).
• Phase 3 — autofit_workspace examples on an imaging lens fit (Opus-authored prose).

Original Prompt

Click to expand starting prompt

Phase 1 of promoting the multi-start gradient MAP optimizer (benchmark winner, autolens_workspace_developer PR#96+#98, Phase-3 complete) to first-class PyAutoFit searches. Adds AbstractMultiStartGradient(AbstractMLE) + af.MultiStartAdam / af.MultiStartADABelief / af.MultiStartLion, running N broad multi-starts vmapped over the af.Fitness seam with a fixed self-normalised optax update per start, returning best-basin MAP + per-start diagnostics through the standard samples/result contract. optax added as a jax-extra dep. Library unit tests numpy-only; JAX end-to-end truth-basin validation in autofit_workspace_test. Line-search/second-order methods explicitly out of scope. Follow-ups: Phase 2 config/defaults, Phase 3 workspace examples.

[Jammy2211]

Implementation complete + verified — autonomous ship parked (Heart RED)

Launched --auto (effective level supervised = min of prompt header supervised, feature-large cap supervised). The Phase-1 library work is done and verified on the branch, but the four-leg autonomous-ship gate stops at leg 4 (Heart), so nothing was committed, pushed, or opened as a PR.

Branch legs (1–3) — green

• Tests (leg 1): full test_autofit/ — 1479 passed, 1 skipped (146s), incl. 4 new test_multi_start_gradient plumbing tests; mle+optimize subtree 7/7.
• End-to-end (JAX): MultiStartAdam recovers a truth basin exactly on a convex toy JAX analysis (jax 0.10.2 + optax 0.2.5) — proving broad finite-gradient start generation, batched value_and_grad over af.Fitness, the optax loop, and samples_via_internal_from all run under real JAX. (Seed for the autofit_workspace_test validation.)
• Serialization: to_dict/from_dict round-trips.
• Review (leg 3): self-review CLEAN — new search module + 3-line __init__ export + 1-line pyproject dep + new test file; scope-clean, mirrors mle/bfgs, no unrelated changes.

Leg 4 — Heart RED (blocks ship)

Heart verdict at ship time (score 0), all pre-existing release-tail reasons, unrelated to this feature:

• PyAutoGalaxy: 1 commit behind origin
• PyAutoLens: 1 commit behind origin
• release validation FAILED (stage integrate)

Per the autonomy contract, Heart RED forbids commit/push/PR-open at every autonomy level, and --auto never acknowledges RED. The corrective-PR exception does not apply — this is a new feature that repairs no named RED reason. No code was modified to pass any leg.

Resume

The change sits complete + uncommitted in the worktree ~/Code/PyAutoLabs-wt/multi-start-gradient-search. When Heart clears to GREEN/STALE (or a human acks a YELLOW reason set), run /ship_library to commit/push and open the pending-release PR — branch legs 1–3 are already green. Follow-ons: autofit_workspace_test JAX validation (library-first gate), then Phase 2 (config/defaults) and Phase 3 (workspace examples).

[Jammy2211]

Library shipped + merged ✅

Heart cleared RED→YELLOW; the YELLOW reason set was human-acknowledged and the merge human-authorized in-session. All four ship legs passed (tests 1479p/1s, smoke n/a for an additive library class, review self-CLEAN, Heart YELLOW-acked).

• PR feat(search): add multi-start gradient MAP searches (Adam/ADABelief/Lion) #1370 — squash-merged to main (63cd4e22); branch deleted.
• af.MultiStartAdam / af.MultiStartADABelief / af.MultiStartLion are now on main, additive (no migration).

Next (library-first gate now satisfied): the autofit_workspace_test JAX end-to-end truth-basin validation script (the remaining Phase-1 leg). Then Phase 2 (config/packaged defaults) and Phase 3 (autofit_workspace examples), issued as they come up. Issue kept open until the workspace_test validation lands.

[Jammy2211]

Phase 1 complete ✅ — both PRs merged

• Library — PyAutoFit#1370 squash-merged (63cd4e22): af.MultiStartAdam / af.MultiStartADABelief / af.MultiStartLion on main.
• Validation — autofit_workspace_test#43 squash-merged (0a81c457): scripts/searches/MultiStartAdam.py recovers the 1D Gaussian truth basin under real JAX (centre=50.156, normalization=25.196, sigma=9.858).

The full Phase-1 scope (core search + samples/result + JAX end-to-end validation, library-first) is landed and verified. Closing.

Follow-ups tracked separately, not yet issued (per the don't-bulk-issue rule):

• Phase 2 — config + packaged defaults (optax rule / N / lr / step-count as first-class packaged config).
• Phase 3 — autofit_workspace examples on an imaging lens fit (Opus-authored prose).