Anytime search: deadlines for the coupled, memory-lift, targeted and universal searches

CHANGELOG.md

@@ -31,6 +31,12 @@ brute-force reference or a closed-form value and has a regression test.
 
 ### Anytime generator search
 
+- The remaining search loops take a ``deadline`` too: `coupled_search` (the ε-machine seed always
+  runs), `memory_lift_search`, `targeted_line_split_seeds` (no new start or hop, and only the best
+  decomposition is polished, once it passes) and `universal_generative_optimizer` (no larger state
+  count once a generator is verified). The pipeline gives them 50%, 20%, 80% and 60% of the remaining
+  budget, so with ``time_budget`` set every stage either finishes or stops between iterations.
+
 - `generative_model`, `generative_model_from_bidirectional` and `minimal_wyner_model` take
   ``time_budget`` (seconds). Once it is spent, the remaining stages are skipped (each logged as
   ``skipped: time budget exhausted``) and the best verified bracket so far is returned, with

sofic/generators/coupled_generators.py

@@ -413,12 +413,14 @@ def memory_lift_search(
     *,
     starts: int = 10,
     maxiter: int = 400,
+    deadline: float | None = None,
     rng: np.random.Generator | None = None,
 ) -> list[tuple[float, float, np.ndarray]]:
     """Minimize the state entropy of memory-1 lifts with ``m`` values of V subject to the generator equations.
 
     A penalty continuation over the channel's logits (numerical gradients; meant for small problems).
-    Returns ``(entropy, residual, nu)`` for every start, best first.
+    Returns ``(entropy, residual, nu)`` for every start, best first. No new start begins after
+    ``deadline`` (a :func:`time.monotonic` value).
     """
     rng = rng or np.random.default_rng()
     problem = _GeneratorProblem(bidir)
@@ -437,6 +439,8 @@ def value(z: np.ndarray, mu: float) -> float:
 
     out = []
     for _ in range(starts):
+        if deadline is not None and out and time.monotonic() > deadline:
+            break
         z = rng.normal(scale=2.0, size=n * m)
         for mu in (1e2, 1e4, 1e6, 1e8):
             z = minimize(value, z, args=(mu,), method="L-BFGS-B", options={"maxiter": maxiter}).x
@@ -841,6 +845,7 @@ def targeted_line_split_seeds(
     hops: int = 30,
     hop_scales: Sequence[float] = (0.03, 0.05, 0.1),
     residual_atol: float = 1e-7,
+    deadline: float | None = None,
     rng: np.random.Generator | None = None,
 ) -> list[tuple[str, Seed]]:
     """Generators found by searching the line-split family directly for kernel-LP feasibility.
@@ -855,8 +860,13 @@ def targeted_line_split_seeds(
     ``hop_scales``, re-run the local search, keep improvements). The ``polish_best`` lowest-entropy
     decompositions (all if ``None``) are polished by the coupled ALM from several logit floors; seeds
     that reach the generator equations to ``residual_atol`` are returned. ``configs`` lists
-    ``(broad, axis)`` pairs.
+    ``(broad, axis)`` pairs. After ``deadline`` (a :func:`time.monotonic` value) no new start or hop
+    begins and only the best decomposition found so far is polished.
     """
+
+    def late() -> bool:
+        return deadline is not None and time.monotonic() > deadline
+
     rng = rng or np.random.default_rng()
     P = problem.lifted.plus.T @ (problem.p_pair[:, None] * problem.lifted.minus)
     found: list[tuple[float, str, np.ndarray, int, int]] = []
@@ -870,21 +880,26 @@ def targeted_line_split_seeds(
             starts = [base] + [base + rng.normal(scale=0.3, size=base.shape) for _ in range(perturbations)]
             lines = "columns" if axis == 0 else "rows"
             for s, z in enumerate(starts):
+                if late() and found:
+                    break
                 local = _line_split_local(problem, P, z, broad, axis, nm_factor=300, cobyla_iter=3000)
                 if local is not None:
                     name = f"targeted line split {i}.{s} ({broad} broad + {lines})"
                     found.append((local[0], name, local[1], broad, axis))
     found.sort(key=lambda item: item[0])
     for j, (h, name, z, broad, axis) in enumerate(found[:hop_chains]):
         for step in range(hops):
+            if late():
+                break
             scale = hop_scales[step % len(hop_scales)]
             local = _line_split_local(problem, P, z + rng.normal(scale=scale, size=z.shape), broad, axis)
             if local is not None and local[0] < h - 1e-6:
                 h, z = local
         found[j] = (h, name if h >= found[j][0] - 1e-6 else f"{name}, hops", z, broad, axis)
     found.sort(key=lambda item: item[0])
     out: list[tuple[str, Seed]] = []
-    for _h, name, z, broad, axis in found if polish_best is None else found[:polish_best]:
+    keep = found if polish_best is None else found[:polish_best]
+    for _h, name, z, broad, axis in keep[:1] if late() else keep:
         beliefs = _line_split_beliefs(P, z, broad, axis)
         _violation, K = _kernel_lp_beliefs(problem, *beliefs)
         k = len(beliefs[0])
@@ -1103,11 +1118,13 @@ def coupled_search(
     residual_atol: float = 1e-6,
     alternating: bool = False,
     dual_strategy: bool = True,
+    deadline: float | None = None,
     rng: np.random.Generator | None = None,
 ) -> list[CoupledSearchResult]:
     """Optimize k-state generators from the standard seeds, ``extra`` named seeds and random starts.
 
-    Returns every candidate whose generator equations hold to ``residual_atol``, best first.
+    Returns every candidate whose generator equations hold to ``residual_atol``, best first. No new start
+    begins after ``deadline`` (a :func:`time.monotonic` value); the ε-machine seed always runs.
     """
     rng = rng or np.random.default_rng()
     problem = _GeneratorProblem(bidir)
@@ -1119,7 +1136,9 @@ def coupled_search(
         for i in range(starts)
     ]
     out = []
-    for name, theta in thetas:
+    for index, (name, theta) in enumerate(thetas):
+        if deadline is not None and index > 0 and time.monotonic() > deadline:
+            break
         pi, Q, R, K = problem.unpack(theta, k)
         init_res = problem.residual(pi, Q, R, K)
         if init_res > residual_atol:

sofic/generators/generative_search.py

@@ -662,7 +662,9 @@ def maybe_return_early() -> GenerativeModelResult | None:
                 offer(f"{name}, floor {floor:g}", problem.unpack(theta, k))
 
     if "memory" in stages and not out_of_time("memory lifts"):
-        for _value, residual, nu in memory_lift_search(bidir, 2, starts=memory_starts, rng=rng)[:3]:
+        for _value, residual, nu in memory_lift_search(bidir, 2, starts=memory_starts, deadline=share(0.2), rng=rng)[
+            :3
+        ]:
             if residual <= residual_atol:
                 offer("memory-1 lift", memory_lift(problem, nu))
 
@@ -710,6 +712,7 @@ def maybe_return_early() -> GenerativeModelResult | None:
                     reproduction_atol=reproduction_atol,
                     alternating=coupled_alternating,
                     dual_strategy=coupled_dual_strategy,
+                    deadline=share(0.6),
                     rng=rng,
                 )
             except StochasticValidationError as exc:
@@ -741,13 +744,14 @@ def maybe_return_early() -> GenerativeModelResult | None:
                     residual_atol=residual_atol,
                     alternating=coupled_alternating,
                     dual_strategy=coupled_dual_strategy,
+                    deadline=share(0.5),
                     rng=rng,
                 )[: None if wyner else 2]:
                     offer(f"coupled k={k} ({res.start})", (res.pi, res.Q, res.R, res.K))
 
     best_value = min((c[0] for c in candidates), default=np.inf)
     if targeted and optima and best_value - lower > _TARGETED_GAP and not out_of_time("targeted line-split search"):
-        seeds = targeted_line_split_seeds(problem, optima, residual_atol=residual_atol, rng=rng)
+        seeds = targeted_line_split_seeds(problem, optima, residual_atol=residual_atol, deadline=share(0.8), rng=rng)
         log.append({"stage": "targeted line-split search", "gap before": best_value - lower, "generators": len(seeds)})
         for name, seed in seeds:
             offer(name, seed)

sofic/generators/universal_optimizer.py

@@ -9,6 +9,7 @@
 
 from __future__ import annotations
 
+import time
 from collections.abc import Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
@@ -111,9 +112,14 @@ def universal_generative_optimizer(
     alternating: bool = False,
     dual_strategy: bool = True,
     cutoff: float = 1e-9,
+    deadline: float | None = None,
     rng: np.random.Generator | None = None,
 ) -> UniversalGenerativeResult:
-    """Search over hidden state count and seeds; return the best verified coupled generator."""
+    """Search over hidden state count and seeds; return the best verified coupled generator.
+
+    Past ``deadline`` (a :func:`time.monotonic` value) no larger state count is tried once some
+    generator is verified, and each :func:`coupled_search` stops starting new seeds.
+    """
     from sofic.generators.minimal_generative_model import _reproduction_error, _reproduction_length
 
     rng = rng or np.random.default_rng()
@@ -132,6 +138,8 @@ def universal_generative_optimizer(
     n_verified = 0
 
     for k in range(2, k_max + 1):
+        if deadline is not None and best_verified is not None and time.monotonic() > deadline:
+            break
         candidates = coupled_search(
             bidir,
             k,
@@ -141,6 +149,7 @@ def universal_generative_optimizer(
             residual_atol=residual_atol,
             alternating=alternating,
             dual_strategy=dual_strategy,
+            deadline=deadline,
             rng=rng,
         )
         n_feasible += len(candidates)

tests/test_coupled_generators.py

@@ -387,3 +387,19 @@ def test_separating_optima_and_bilevel_respect_an_expired_deadline() -> None:
     assert optima and np.isfinite(best)
     out = cg.bilevel_search(_rll(), optima, deadline=time.monotonic())
     assert len(out) == 1
+
+
+def test_coupled_memory_and_targeted_searches_stop_at_an_expired_deadline() -> None:
+    import time
+
+    bidir = _rll()
+    start = time.monotonic()
+    results = cg.coupled_search(bidir, 3, starts=50, deadline=time.monotonic(), rng=np.random.default_rng(0))
+    assert results and all(r.residual <= 1e-6 for r in results)
+    lifts = cg.memory_lift_search(bidir, 2, starts=50, deadline=time.monotonic(), rng=np.random.default_rng(0))
+    assert len(lifts) == 1
+    problem = cg._GeneratorProblem(bidir)
+    _best, optima = cg.separating_optima(problem, starts=10, rng=np.random.default_rng(0))
+    seeds = cg.targeted_line_split_seeds(problem, optima, deadline=time.monotonic(), rng=np.random.default_rng(0))
+    assert len(seeds) <= 1
+    assert time.monotonic() - start < 120