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

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

Add a resurrect=False knob to af.AbstractMultiStartGradient. 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) via a jnp.where mask
over the vmapped state pytree, leaving alive starts untouched. Keeps the
population alive on likelihoods with broad non-finite regions (pixelized
sources), where apply_if_finite alone latches a start at the cliff edge.
Default off leaves the MGE-cell behaviour unchanged. Tracks an n_resurrections
diagnostic through search_internal (resume-safe) and samples_info.

Numpy-only unit tests; JAX validation follows in autofit_workspace_test.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: 2 people

autofit/non_linear/search/mle/multi_start_gradient/search.py

@@ -37,6 +37,7 @@ def __init__(
         max_consecutive_nan: int = 8,
         start_lower_limit: float = 0.15,
         start_upper_limit: float = 0.85,
+        resurrect: bool = False,
         initializer: Optional[AbstractInitializer] = None,
         iterations_per_full_update: int = None,
         iterations_per_quick_update: int = None,
@@ -97,6 +98,20 @@ def __init__(
             The unit-cube bounds broad starts are drawn uniformly from. The
             interior default ``(0.15, 0.85)`` avoids the prior edges where many
             transforms (e.g. ``arctan2`` / ``sqrt`` at exactly 0) are singular.
+        resurrect
+            Restart-on-death. When ``True``, any start whose objective goes
+            non-finite is redrawn each step (fresh params from the start band +
+            its per-start optimizer state reinitialised), leaving alive starts
+            untouched. Default ``False`` — the parametric (MGE-class) cell has
+            only the measure-zero singularity, so the ``apply_if_finite`` guard
+            suffices and behaviour/results are unchanged. Turn it on for
+            likelihoods with broad non-finite regions (pixelized sources), where
+            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 at all. (Even so, on such landscapes a nested sampler
+            still wins decisively — resurrection makes gradient MAP *viable*
+            there, not competitive.)
         """
 
         super().__init__(
@@ -120,6 +135,7 @@ def __init__(
         )
         self.start_lower_limit = start_lower_limit
         self.start_upper_limit = start_upper_limit
+        self.resurrect = resurrect
 
         self.logger.debug(f"Creating {self.optax_method} MultiStartGradient Search")
 
@@ -206,6 +222,7 @@ def batched_value_and_grad(params):
             best_fom = float(search_internal["best_fom"])
             fom_history = list(search_internal["fom_history"])
             total_steps = int(search_internal["total_steps"])
+            n_resurrections = int(search_internal.get("n_resurrections", 0))
 
             self.logger.info(
                 "Resuming MultiStartGradient search (previous samples found)."
@@ -228,12 +245,17 @@ def batched_value_and_grad(params):
             best_fom = np.inf
             fom_history = []
             total_steps = 0
+            n_resurrections = 0
 
             self.logger.info(
                 f"Starting new {self.optax_method} MultiStartGradient search "
                 f"({self.n_starts} starts, no previous samples found)."
             )
 
+        # Deterministic RNG for redrawing dead starts (only used when
+        # ``resurrect`` is on); seeded independently of the broad-start draw.
+        resurrect_rng = np.random.default_rng(1)
+
         while total_steps < self.n_steps:
 
             steps_remaining = self.n_steps - total_steps
@@ -244,16 +266,36 @@ def batched_value_and_grad(params):
             for _ in range(iterations):
                 foms, grads = batched_value_and_grad(params)
 
-                foms_np = np.where(
-                    np.isfinite(np.asarray(foms)), np.asarray(foms), np.inf
-                )
+                alive = np.isfinite(np.asarray(foms))
+                foms_np = np.where(alive, np.asarray(foms), np.inf)
                 best_index = int(np.argmin(foms_np))
                 if foms_np[best_index] < best_fom:
                     best_fom = float(foms_np[best_index])
                     best_params = np.asarray(params[best_index])
 
                 fom_history.append(best_fom)
 
+                # Restart-on-death: redraw any start whose objective went
+                # non-finite (fresh params + reinitialised per-start optimizer
+                # state), leaving alive starts untouched. best_* is captured
+                # above, from the pre-redraw alive population. The (old,
+                # non-finite) grads of a dead start are zeroed for one step by
+                # apply_if_finite on its fresh state, so a redrawn start simply
+                # waits a step before descending.
+                if self.resurrect and not alive.all():
+                    dead_idx = np.flatnonzero(~alive)
+                    n_resurrections += int(dead_idx.size)
+                    params, opt_state = self._reinit_dead_starts(
+                        params=params,
+                        opt_state=opt_state,
+                        dead_idx=dead_idx,
+                        model=model,
+                        optimizer=optimizer,
+                        jax=jax,
+                        jnp=jnp,
+                        rng=resurrect_rng,
+                    )
+
                 updates, opt_state = step_update(grads, opt_state, params, foms)
                 params = optax.apply_updates(params, updates)
 
@@ -266,6 +308,7 @@ def batched_value_and_grad(params):
                 "best_fom": best_fom,
                 "fom_history": np.asarray(fom_history),
                 "total_steps": total_steps,
+                "n_resurrections": n_resurrections,
             }
             self.paths.save_search_internal(obj=search_internal)
 
@@ -344,6 +387,46 @@ def step_update(grads, opt_state, params, values):
 
         return optimizer, step_update
 
+    def _reinit_dead_starts(
+        self, params, opt_state, dead_idx, model, optimizer, jax, jnp, rng
+    ):
+        """
+        Redraw the dead starts (``dead_idx``) and reinitialise their per-start
+        optimizer state, leaving the alive starts untouched — the restart-on-
+        death recovery step.
+
+        Each dead row's params are redrawn uniformly from the start band and
+        mapped to physical parameters; a fresh vmapped optimizer state is built
+        and merged into the live state pytree with a boolean mask (``jnp.where``
+        per leaf). ``np.asarray`` of a JAX array is read-only, so the redraw
+        happens in an ``np.array`` copy.
+        """
+        n = params.shape[0]
+
+        params_np = np.array(params)  # writable copy (jax arrays are read-only)
+        for k in dead_idx:
+            unit_vector = rng.uniform(
+                self.start_lower_limit, self.start_upper_limit, size=model.prior_count
+            )
+            params_np[k] = np.asarray(
+                model.vector_from_unit_vector(unit_vector=list(unit_vector), xp=jnp)
+            )
+        params = jnp.asarray(params_np)
+
+        fresh_state = jax.vmap(optimizer.init)(params)
+
+        mask = np.zeros(n, dtype=bool)
+        mask[dead_idx] = True
+        mask_j = jnp.asarray(mask)
+
+        def merge(old, fresh):
+            m = mask_j.reshape((n,) + (1,) * (fresh.ndim - 1))
+            return jnp.where(m, fresh, old)
+
+        opt_state = jax.tree.map(merge, opt_state, fresh_state)
+
+        return params, opt_state
+
     def _broad_starts(self, model, fitness, batched_value_and_grad, jnp):
         """
         Draw ``n_starts`` broad starting points in the unit cube, map them to
@@ -429,6 +512,8 @@ def samples_via_internal_from(
             "optax_method": self.optax_method,
             "learning_rate": self.learning_rate,
             "max_consecutive_nan": self.max_consecutive_nan,
+            "resurrect": self.resurrect,
+            "n_resurrections": int(search_internal.get("n_resurrections", 0)),
             "time": self.timer.time if self.timer else None,
         }
 

test_autofit/non_linear/search/mle/test_multi_start_gradient.py

@@ -64,6 +64,23 @@ def test__per_rule_defaults():
     assert default.batch_size is None
     # apply_if_finite per-start rejected-step budget (the in-step NaN guard).
     assert default.max_consecutive_nan == 8
+    # restart-on-death is opt-in — default off keeps the MGE-cell behaviour.
+    assert default.resurrect is False
+
+
+def test__resurrect_defaults_off_and_is_a_carried_knob():
+    """``resurrect`` (restart-on-death) is a shared base-class knob, off by
+    default and dict-serialised so a resumed search keeps it. The redraw /
+    per-start state reinit itself is JAX and is validated in
+    autofit_workspace_test (the library suite is NumPy-only)."""
+    for cls in (
+        af.MultiStartAdam,
+        af.MultiStartADABelief,
+        af.MultiStartLion,
+        af.MultiStartProdigy,
+    ):
+        assert cls().resurrect is False
+        assert cls(resurrect=True).resurrect is True
 
 
 def test__max_consecutive_nan_is_a_carried_knob():
@@ -119,7 +136,16 @@ def test__dict_round_trip__prodigy_is_learning_rate_free():
     assert restored.n_starts == 5
     assert restored.optax_method == "prodigy"
     assert restored.learning_rate is None
-    assert restored.max_consecutive_nan == 4
+
+
+def test__dict_round_trip__resurrect():
+    # The restart-on-death flag must survive serialisation so a resumed search
+    # keeps redrawing dead starts.
+    restored = from_dict(to_dict(af.MultiStartAdam(resurrect=True, n_starts=6)))
+
+    assert isinstance(restored, af.MultiStartAdam)
+    assert restored.resurrect is True
+    assert restored.n_starts == 6
 
 
 def test__samples_via_internal_from():
@@ -141,9 +167,10 @@ def test__samples_via_internal_from():
         "best_fom": best_fom,
         "fom_history": np.asarray([-4.0, -8.0, best_fom]),
         "total_steps": 42,
+        "n_resurrections": 7,
     }
 
-    search = af.MultiStartAdam(n_starts=2, n_steps=42)
+    search = af.MultiStartAdam(n_starts=2, n_steps=42, resurrect=True)
     samples = search.samples_via_internal_from(
         model=model, search_internal=search_internal
     )
@@ -169,3 +196,6 @@ def test__samples_via_internal_from():
     assert samples.samples_info["optax_method"] == "adam"
     assert samples.samples_info["n_starts"] == 2
     assert samples.samples_info["total_steps"] == 42
+    # restart-on-death diagnostics flow through to samples_info.
+    assert samples.samples_info["resurrect"] is True
+    assert samples.samples_info["n_resurrections"] == 7