feat(autofit): restart-on-death (resurrection) for multi-start gradient search

[Jammy2211]

Overview

Phase 2 of the multi-start gradient v2 promotion (autolens_workspace_developer#101). Builds on Phase 1 (PyAutoFit#1398, merged) — the per-start jax.vmaped optimizer state. Adds restart-on-death (resurrection) to af.AbstractMultiStartGradient: each step, any start whose objective went non-finite is redrawn fresh (params from the start band + its per-start optimizer state reinitialised), leaving alive starts untouched. On likelihoods with hard non-finite regions (pixelized sources) every trajectory dies within ~25–50 steps and apply_if_finite alone latches the start at the cliff edge; resurrection keeps the population alive and converts the pixelized landscape from unsearchable to searchable (adam −51201 → +1718 over 3000 steps, job 330598).

Plan

• Add a resurrect: bool = False search knob (default off → existing MGE-cell behaviour and results are byte-for-byte unchanged; resurrection is only load-bearing on inversion-heavy landscapes).
• In the fit loop, after computing per-start objectives: if resurrect and any start is non-finite, redraw those starts' params from (start_lower_limit, start_upper_limit) and reinit their vmapped optimizer state, merging into the alive population via a boolean mask (jnp.where over the state pytree).
• Track a n_resurrections diagnostic through search_internal (resume-safe) and samples_info.
• Honest caveat in docstring/PR: even with resurrection Nautilus still wins the pixelized cell decisively; this makes gradient MAP viable there, not competitive.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)
• autofit_workspace_test (JAX validation, follow-on)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main (Phase 1 merged)	clean

Suggested branch: feature/multistart-resurrection

Implementation Steps

1. autofit/non_linear/search/mle/multi_start_gradient/search.py — add resurrect: bool = False constructor arg (store + docstring).
2. Add _reinit_dead_starts(self, params, opt_state, dead_idx, model, optimizer, jax, jnp, rng): copy params into a writable np.array (jax arrays are read-only), redraw dead rows from the start band via model.vector_from_unit_vector, jax.vmap(optimizer.init) fresh state, merge dead rows via a jnp.where mask over the vmapped state pytree (jax.tree.map). Mask size = params.shape[0] (broad-start collection may yield < n_starts).
3. _fit loop: compute alive = np.isfinite(foms); select best from alive-only (unchanged); if self.resurrect and not alive.all(), reinit dead starts (grads for dead rows are the old non-finite ones — apply_if_finite on the fresh per-start state zeroes them for one step, so redrawn params wait one step, matching the reference). Initialise resurrect_rng and n_resurrections on both the fresh and resume paths.
4. Persist n_resurrections in search_internal and surface resurrect + n_resurrections in samples_info.
5. Unit tests (numpy-only): resurrect default False + dict round-trip; _reinit_dead_starts mask logic on plain numpy arrays.
6. autofit_workspace_test — extend/adds a JAX validation exercising resurrect=True on a landscape with non-finite regions (or assert the knob is inert when off).

Key Files

• autofit/non_linear/search/mle/multi_start_gradient/search.py
• test_autofit/non_linear/search/mle/test_multi_start_gradient.py
• Reference (read-only): autolens_workspace_developer/searches_minimal/pix_lr_free.py (PIX_RESURRECT, reinit_starts), lr_free_findings.md (Phase 2).

Original Prompt

Click to expand starting prompt

(draft/feature/autofit/multistart_resurrection_restart_on_death.md — the Phase-2 prompt, filed at split time.)

Generated by the PyAutoLabs agent workflow.

[Jammy2211]

Shipped — Phase 2 complete. Both PRs merged: library #1400 and workspace PyAutoLabs/autofit_workspace_test#59. The resurrect (restart-on-death) knob is on main (default off; redraws dead starts + reinit per-start vmapped optimizer state). Validated: unit 9-pass, _reinit_dead_starts isolation, resurrect on/off identical basin. Completes the two-phase multi-start gradient v2 promotion (#1397 was Phase 1). Full pixelized-landscape validation remains A100-only (ref job 330598). Closing.