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

[Jammy2211]

Summary

Phase 2 of the multi-start gradient v2 promotion (autolens_workspace_developer#101), building on the per-start jax.vmaped optimizer state from Phase 1 (#1398, merged).

Adds restart-on-death (resurrection) to af.AbstractMultiStartGradient via a resurrect: bool = False knob. When on, any start whose objective goes non-finite is redrawn each step — fresh params from the start band + its per-start optimizer state reinitialised — merged into the alive population with a jnp.where mask over the vmapped state pytree, leaving alive starts untouched. On likelihoods with broad non-finite regions (pixelized sources) every trajectory otherwise walks into a wall within ~25–50 steps and apply_if_finite alone latches the start at the cliff edge; resurrection keeps the population alive and makes the landscape searchable (adam −51201 → +1718 over 3000 steps, ref job 330598).

Default off — the parametric (MGE-class) cell is unchanged. An n_resurrections diagnostic flows through search_internal (resume-safe) and samples_info. Honest caveat (in docstring): even with resurrection a nested sampler still wins such landscapes decisively; this makes gradient MAP viable there, not competitive.

API Changes

• Added a resurrect: bool = False constructor argument on af.AbstractMultiStartGradient (inherited by all MultiStart* searches). Default False → no behaviour change.
• Added an n_resurrections diagnostic to search_internal and samples_info.

No migration required — the new argument is optional and off by default; existing MultiStart* behaviour is byte-for-byte unchanged.

Test Plan

• Library unit suite (numpy-only): pytest test_autofit/non_linear/search/mle/ — 9 passed (resurrect default-off + carried knob, dict round-trip, n_resurrections/resurrect in samples_info).
• JAX (validated locally): _reinit_dead_starts redraws only the dead rows (params + per-start optimizer state), alive rows byte-identical; resurrect on/off give the identical best basin on the 1D Gaussian (49.82/25.27/9.996) with resurrection firing incidentally (n_resurrections=1) without perturbing the winner.
• Follow-up (/ship_workspace): autofit_workspace_test JAX validation. Full pixelized-landscape validation is A100-only (the searches_minimal reference proved it; this port mirrors that exact reinit_starts logic).

Full API Changes (for automation & release notes)

Added

• af.AbstractMultiStartGradient(..., resurrect: bool = False, ...) — restart-on-death; redraws dead starts (params + per-start optimizer state) each step. Serialised in to_dict/from_dict.
• search_internal["n_resurrections"] and samples_info["n_resurrections"] / samples_info["resurrect"] — cumulative redraw diagnostic (resume-safe).
• AbstractMultiStartGradient._reinit_dead_starts(...) — the redraw + vmapped-state reinit + masked merge helper.

Migration

• None. resurrect defaults off; existing MultiStart* code and results are unchanged.

Generated by the PyAutoLabs agent workflow.

[Jammy2211]

Linked workspace PR (resurrection JAX validation): PyAutoLabs/autofit_workspace_test#59 (pending-release). Library-first: this PR merges first, then #59.