feat: seedable multi-start draws, alive-versus-step curve, and optional momentum reset

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

@@ -43,10 +43,12 @@ def __init__(
         start_lower_limit: float = 0.15,
         start_upper_limit: float = 0.85,
         resurrect: bool = False,
+        seed: Optional[int] = None,
         convergence: Optional[MultiStartGradientConvergence] = None,
         iterations_per_log: int = 10,
         initializer: Optional[AbstractInitializer] = None,
         clipper: Optional[AbstractClipper] = None,
+        reset_momentum_on_clip: bool = False,
         iterations_per_full_update: int = None,
         iterations_per_quick_update: int = None,
         silence: bool = False,
@@ -146,6 +148,41 @@ def __init__(
             searchable at all. (Even so, on such landscapes a nested sampler
             still wins decisively — resurrection makes gradient MAP *viable*
             there, not competitive.)
+        seed
+            Seeds the two random draws this search owns: the broad starting
+            points (``_broad_starts``) and, when ``resurrect`` is on, the redraw
+            of dead lanes. Default ``None`` reproduces the historical fixed
+            seeds exactly, so an existing fit is bit-identical and this argument
+            is purely additive.
+
+            It exists because without it the search cannot be varied *or*
+            genuinely repeated: both draws were hardcoded, so every run of the
+            same model produced the **same** starting population no matter what
+            the caller seeded. Seeding ``random`` / ``numpy`` globally does not
+            reach them — that only perturbs the initializer — which makes a
+            multi-seed study silently a single-seed one. Note the scope: this
+            seeds *this search's* draws, not the whole framework (the
+            initializer and any sampler-owned generator are separate, still-open
+            sources of non-reproducibility).
+        reset_momentum_on_clip
+            Zero the optimizer's momentum on the coordinates a ``clipper`` just
+            clipped. Default ``False``, so the clipping path is unchanged unless
+            asked for.
+
+            It exists because projection alone leaves a lane holding the exact
+            velocity that carried it out of the prior box, so the next step
+            drives it into the same wall and it is re-projected onto the same
+            bound indefinitely — counted as alive, permanently pinned, and still
+            paying a full likelihood-and-gradient evaluation every step. The
+            reset is per-coordinate: a lane clipped in one parameter keeps its
+            momentum in all the others.
+
+            Note that a pinned lane is not automatically a failure. Where the
+            likelihood genuinely prefers a value outside the prior, sitting on
+            the bound is the correct MAP answer under the declared prior, and
+            this flag would then be discarding useful state. It is a knob for
+            the case where pinning is an artefact of momentum rather than a
+            statement about the data.
         convergence
             Auto-convergence (early-stopping) settings. When
             ``check_for_convergence`` is ``True`` (the default) the search stops
@@ -204,6 +241,8 @@ def __init__(
         self.start_lower_limit = start_lower_limit
         self.start_upper_limit = start_upper_limit
         self.resurrect = resurrect
+        self.seed = seed
+        self.reset_momentum_on_clip = reset_momentum_on_clip
         self.convergence = (
             convergence if convergence is not None else MultiStartGradientConvergence()
         )
@@ -698,6 +737,11 @@ def batched_value_and_grad(params):
             best_params = np.asarray(search_internal["best_params"])
             best_fom = float(search_internal["best_fom"])
             fom_history = list(search_internal["fom_history"])
+            # ``.get``: a ``search_internal`` written before the alive history
+            # existed has no such key, and a resumed run must not KeyError on
+            # it. An older run's curve is unrecoverable, so it resumes empty
+            # rather than being back-filled with a fabricated population.
+            alive_history = list(search_internal.get("alive_history", []))
             total_steps = int(search_internal["total_steps"])
             n_resurrections = int(search_internal.get("n_resurrections", 0))
             # ``.get`` with a default: a ``search_internal`` written before the
@@ -761,6 +805,7 @@ def batched_value_and_grad(params):
             best_params = np.asarray(params[0])
             best_fom = np.inf
             fom_history = []
+            alive_history = []
             total_steps = 0
             n_resurrections = 0
             n_value_nan_lane_steps = 0
@@ -776,7 +821,7 @@ def batched_value_and_grad(params):
 
         # 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)
+        resurrect_rng = np.random.default_rng(self._seed_for(1))
 
         stop_reason = self._stop_reason_on_resume(stop_reason)
 
@@ -846,6 +891,17 @@ def batched_value_and_grad(params):
 
                 fom_history.append(best_fom)
 
+                # The size of the living population at this step. Recorded as a
+                # history because the cumulative lane counters above are
+                # survival INTEGRALS: a dead lane keeps adding to them every
+                # subsequent step, so the same death curve reads ~60% at 150
+                # steps and ~75% at 300 and two runs at different budgets cannot
+                # be compared on the scalar at all. The curve is the
+                # budget-independent quantity, and until now it existed only in
+                # the progress log at ``iterations_per_log`` cadence — visible
+                # to a human reading stdout, unavailable to any analysis.
+                alive_history.append(int(np.count_nonzero(alive)))
+
                 # 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
@@ -890,6 +946,17 @@ def batched_value_and_grad(params):
                         np.count_nonzero(np.asarray(jnp.any(clipped_mask, axis=-1)))
                     )
 
+                    # Optionally zero the optimizer's momentum on the clipped
+                    # coordinates. Without it a clipped lane keeps the velocity
+                    # that pushed it out of the box, so it is re-projected onto
+                    # the same bound every step: alive in the counters, pinned
+                    # to the wall, and still spending a full likelihood-and-
+                    # gradient evaluation per step.
+                    if self.reset_momentum_on_clip:
+                        opt_state = self._reset_clipped_momentum(
+                            opt_state=opt_state, clipped_mask=clipped_mask, jnp=jnp
+                        )
+
                 total_steps += 1
 
                 if not self.resurrect and self.convergence.check_if_converged(
@@ -937,6 +1004,7 @@ def batched_value_and_grad(params):
                 "best_params": best_params,
                 "best_fom": best_fom,
                 "fom_history": np.asarray(fom_history),
+                "alive_history": np.asarray(alive_history),
                 "total_steps": total_steps,
                 "n_resurrections": n_resurrections,
                 "n_value_nan_lane_steps": n_value_nan_lane_steps,
@@ -1084,6 +1152,70 @@ def merge(old, fresh):
 
         return params, opt_state
 
+    # The optax moment accumulators — the "momentum" a clip should forget.
+    # Named explicitly rather than matched by shape, because shape matching is
+    # actively wrong here: Prodigy's ``params0`` and ``grad_sum`` carry the same
+    # ``(n_starts, n_params)`` shape as the moments, and ``params0`` is the
+    # reference point of its learning-rate estimate. Zeroing that would not
+    # reset momentum, it would corrupt the step-size estimate for the rest of
+    # the run. Anything not named here (``params0``, ``grad_sum``, ``estim_lr``,
+    # ``numerator_weighted``, ``count``) is deliberately left intact.
+    _MOMENTUM_FIELDS = frozenset({"mu", "nu", "exp_avg", "exp_avg_sq", "trace"})
+
+    @classmethod
+    def _reset_clipped_momentum(cls, opt_state, clipped_mask, jnp):
+        """Zero the optimizer moments wherever a coordinate was clipped.
+
+        ``clipped_mask`` is ``(n_starts, n_params)``, so the reset is
+        per-coordinate: a lane clipped in one parameter keeps its momentum in
+        every other parameter and only forgets the direction that took it out of
+        the box.
+        """
+
+        def rebuild(node):
+            fields = getattr(node, "_fields", None)
+            if fields is not None:
+                return type(node)(
+                    **{
+                        name: (
+                            jnp.where(clipped_mask, jnp.zeros_like(value), value)
+                            if name in cls._MOMENTUM_FIELDS
+                            and getattr(value, "shape", None) == clipped_mask.shape
+                            else rebuild(value)
+                        )
+                        for name, value in zip(fields, node)
+                    }
+                )
+            if isinstance(node, tuple):
+                return tuple(rebuild(child) for child in node)
+            if isinstance(node, list):
+                return [rebuild(child) for child in node]
+            return node
+
+        return rebuild(opt_state)
+
+    def _seed_for(self, stream: int):
+        """Seed material for one of this search's random draws.
+
+        ``stream`` identifies the draw: ``0`` the broad starting points, ``1``
+        the resurrection redraws. They must not share a stream — resurrection
+        would otherwise replay the starting population.
+
+        ``seed=None`` returns the bare stream index, which is exactly the
+        historical hardcoded seed at each site (``default_rng(0)`` /
+        ``default_rng(1)``), so the default path is bit-identical.
+
+        With a seed set, the pair is derived through ``SeedSequence`` rather
+        than by offsetting the seed. ``seed + stream`` looks equivalent and is
+        not: it makes seed 0's resurrection stream the same sequence as seed 1's
+        start stream, so nominally independent seeds in a multi-seed study share
+        draws. ``SeedSequence`` spreads each ``(seed, stream)`` pair over the
+        full state space instead.
+        """
+        if self.seed is None:
+            return stream
+        return np.random.SeedSequence([self.seed, stream])
+
     def _broad_starts(self, model, value_and_grad_single, jnp):
         """
         Draw ``n_starts`` broad starting points in the unit cube, map them to
@@ -1097,7 +1229,7 @@ def _broad_starts(self, model, value_and_grad_single, jnp):
         cost per draw, which on a multi-band ``FactorGraphModel`` objective
         dominated the whole fit (~13 minutes for 16 draws, cache or no cache).
         """
-        rng = np.random.default_rng(0)
+        rng = np.random.default_rng(self._seed_for(0))
 
         starts = []
         max_tries = self.n_starts * 30
@@ -1206,9 +1338,16 @@ def samples_via_internal_from(
             # current process's share — the same reasoning as the ``.get``
             # defaults above.
             "clipper": type(self.clipper).__name__,
+            "reset_momentum_on_clip": self.reset_momentum_on_clip,
             "n_clipped_lane_steps": int(
                 search_internal.get("n_clipped_lane_steps", 0)
             ),
+            # The seed this search's own draws used, so a result file says which
+            # member of a multi-seed study produced it. ``None`` records the
+            # default (historical fixed seeds) rather than being omitted, so a
+            # seeded and an unseeded run are distinguishable downstream instead
+            # of both reading as "no seed key".
+            "seed": self.seed,
             # Auto-convergence outcome: whether the run stopped on the plateau
             # check ("converged") or exhausted the ``n_steps`` ceiling
             # ("max_steps"), the settings that produced it, and the global-best

test_autofit/non_linear/search/mle/test_multi_start_gradient.py

@@ -1,4 +1,5 @@
 import inspect
+from typing import NamedTuple
 
 import numpy as np
 import pytest
@@ -278,6 +279,129 @@ def test__dict_round_trip__resurrect():
     assert restored.n_starts == 6
 
 
+def test__dict_round_trip__seed():
+    # The seed must survive serialisation, or a resumed member of a multi-seed
+    # study silently continues on the default draw.
+    restored = from_dict(to_dict(af.MultiStartAdam(seed=7, n_starts=6)))
+
+    assert isinstance(restored, af.MultiStartAdam)
+    assert restored.seed == 7
+
+
+def test__seed__default_is_bit_identical_to_the_historical_fixed_seeds():
+    # ``seed`` is additive: the default path must draw exactly what the
+    # hardcoded ``default_rng(0)`` / ``default_rng(1)`` drew before it existed,
+    # so no existing fit changes.
+    search = af.MultiStartAdam(seed=None)
+
+    assert search._seed_for(0) == 0
+    assert search._seed_for(1) == 1
+
+
+def test__seed__is_reproducible_and_varies_between_seeds():
+    def draw(seed, stream):
+        search = af.MultiStartAdam(seed=seed)
+        return np.random.default_rng(search._seed_for(stream)).uniform(size=8)
+
+    # Same seed reproduces; different seeds diverge. Without this the campaign's
+    # "at least two seeds per arm" is silently a single-seed study.
+    assert np.array_equal(draw(3, 0), draw(3, 0))
+    assert not np.array_equal(draw(3, 0), draw(4, 0))
+
+
+def test__seed__start_and_resurrect_streams_never_coincide():
+    # Guards the bug a naive ``seed + stream`` offset would introduce: seed 0's
+    # resurrection stream would be seed 1's starting stream, so nominally
+    # independent seeds would share draws and resurrection would replay the
+    # starting population.
+    def draw(seed, stream):
+        search = af.MultiStartAdam(seed=seed)
+        return np.random.default_rng(search._seed_for(stream)).uniform(size=8)
+
+    assert not np.array_equal(draw(0, 0), draw(0, 1))
+    assert not np.array_equal(draw(0, 1), draw(1, 0))
+
+
+class _AdamLikeState(NamedTuple):
+    count: np.ndarray
+    mu: np.ndarray
+    nu: np.ndarray
+
+
+class _ProdigyLikeState(NamedTuple):
+    exp_avg: np.ndarray
+    exp_avg_sq: np.ndarray
+    grad_sum: np.ndarray
+    params0: np.ndarray
+    estim_lr: np.ndarray
+
+
+def test__reset_clipped_momentum__zeroes_only_the_clipped_coordinates():
+    # (2 starts, 3 params); only start 0's middle coordinate was clipped.
+    mask = np.array([[False, True, False], [False, False, False]])
+    state = _AdamLikeState(
+        count=np.array([7, 7]),
+        mu=np.ones((2, 3)),
+        nu=np.full((2, 3), 2.0),
+    )
+
+    out = af.MultiStartAdam._reset_clipped_momentum(
+        opt_state=state, clipped_mask=mask, jnp=np
+    )
+
+    assert out.mu.tolist() == [[1.0, 0.0, 1.0], [1.0, 1.0, 1.0]]
+    assert out.nu.tolist() == [[2.0, 0.0, 2.0], [2.0, 2.0, 2.0]]
+    # A lane keeps its momentum in every coordinate that was not clipped, and
+    # non-moment state is untouched.
+    assert out.count.tolist() == [7, 7]
+
+
+def test__reset_clipped_momentum__never_touches_prodigys_reference_point():
+    # ``params0`` and ``grad_sum`` carry the SAME shape as the moments, so a
+    # shape-matched reset would zero them. ``params0`` anchors Prodigy's
+    # learning-rate estimate: zeroing it corrupts the step size for the rest of
+    # the run rather than resetting momentum.
+    mask = np.array([[True, True, True]])
+    state = _ProdigyLikeState(
+        exp_avg=np.ones((1, 3)),
+        exp_avg_sq=np.ones((1, 3)),
+        grad_sum=np.full((1, 3), 5.0),
+        params0=np.full((1, 3), 9.0),
+        estim_lr=np.array([0.5]),
+    )
+
+    out = af.MultiStartProdigy._reset_clipped_momentum(
+        opt_state=state, clipped_mask=mask, jnp=np
+    )
+
+    assert out.exp_avg.tolist() == [[0.0, 0.0, 0.0]]
+    assert out.exp_avg_sq.tolist() == [[0.0, 0.0, 0.0]]
+    assert out.params0.tolist() == [[9.0, 9.0, 9.0]]
+    assert out.grad_sum.tolist() == [[5.0, 5.0, 5.0]]
+    assert out.estim_lr.tolist() == [0.5]
+
+
+def test__reset_clipped_momentum__recurses_through_optax_chain_tuples():
+    # optax.adam's state arrives wrapped in a chain tuple, so the reset has to
+    # descend through plain tuples as well as NamedTuples.
+    mask = np.array([[True, False]])
+    nested = (_AdamLikeState(count=np.array([1]), mu=np.ones((1, 2)), nu=np.ones((1, 2))),)
+
+    out = af.MultiStartAdam._reset_clipped_momentum(
+        opt_state=nested, clipped_mask=mask, jnp=np
+    )
+
+    assert out[0].mu.tolist() == [[0.0, 1.0]]
+
+
+def test__dict_round_trip__reset_momentum_on_clip():
+    restored = from_dict(to_dict(af.MultiStartAdam(reset_momentum_on_clip=True)))
+
+    assert restored.reset_momentum_on_clip is True
+    # Default stays off, so the clipping path is unchanged unless asked for.
+    assert af.MultiStartAdam().reset_momentum_on_clip is False
+
+
 def test__samples_via_internal_from():
     model = af.Model(example.Gaussian)