diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index eb2ad4059..7cbaa3784 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/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 diff --git a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py index 4458c71ac..55b1eba79 100644 --- a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py +++ b/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)