Anytime pipeline on large joints: stage deadlines, sparse stationary solve, ladder fixes

CHANGELOG.md

@@ -7,6 +7,18 @@ brute-force reference or a closed-form value and has a regression test.
 
 ### Bidirectional builds
 
+- `stationary_distribution_hmm` solves chains with more than 2,000 states sparsely (sparse transition
+  matrix from the edges, one balance equation replaced by the normalization, `spsolve`). The dense
+  path built the full matrix and ran `matrix_rank` and `eig` on it, $O(n^3)$, which never finished on
+  a 50k-state reverse machine. Symbolic weights, several closed classes, or a failed residual check
+  fall back to the dense path and its warnings.
+- The retry ladder and `try_build_bidirectional` reuse the reverse and bidirectional machines built by
+  the preflight instead of inferring and building them a second time (about half the build time).
+- `build_bidirectional_with_ladder` and `try_build_bidirectional` take `max_reverse_states`: a ladder
+  build whose truncated reverse machine is larger is classified as an explosion, so surveys fall back
+  to window bounds. Off by default. A 4-state survey process built at the 50k cap with tens of
+  thousands of reverse states, and the pipeline then spent its whole budget setting up the problem.
+
 - `build_bidirectional_epsilon_machine` prunes joint states to a fixed point: dropping states
   outside the stationary support could leave others with incomplete outgoing rows, so builds from
   truncated reverse machines failed with "joint masses from ... sum to 0.86". Restriction to the
@@ -32,13 +44,32 @@ brute-force reference or a closed-form value and has a regression test.
 
 ### Generator caps
 
+- With a deadline, each branch-and-bound child LP of `generator_weight_cap` gets the time left as a
+  HiGHS `time_limit`; a cut-off LP keeps its parent's bound, so the cap stays valid. On a 4-state
+  survey process (4 × 36 causal states) one degenerate node LP had run more than 8 minutes past a
+  90 s stage deadline; the stage now ends on time (92.7 s).
+
 - `generator_weight_cap` adds the forward analogue of the reverse joint-moment rows: per-state forward
   stationarity times the state's own reverse belief gives
   `T_h(a', b) = sum_{x, g, a : f(a, x) = a'} Gt[x, g, h, a, b]`, which every generator satisfies.
   Same cost; two-state caps drop (e.g. 0.918 -> 0.907 at 600 nodes) and the coverage bound on the
   survey processes `((1,1),(2,0),(1,0))` and `((1,·),(2,0),(1,0))` rises 1.261 -> 1.277 and
   1.175 -> 1.184.
 
+### Anytime pipeline on large joints
+
+- `generative_model` skips the dit exact- and Wyner-common-information fallbacks when the joint has
+  more than 30 causal-state pairs (logged as a skipped `dit fallbacks` stage). dit's optimizers take no
+  deadline: on finished 4-state survey cases with 30-70 pairs they used 190-760 s of a 900 s budget,
+  and on a 95-pair case one call ran past the survey's hard timeout. Every 3-state survey joint has at
+  most 21 pairs, so the benchmark is unaffected.
+- `memory_lift_search` stops a start in progress at its `deadline` (keeping the last finished penalty
+  stage) instead of only refusing to begin new starts. Its first start used to run all four
+  finite-difference L-BFGS-B stages whatever the deadline, more than 14 minutes on a 95-pair joint.
+- The line-split seeds' ALM polish (`_optimize`, `_optimize_coupled`) stops at the pipeline deadline
+  and returns its last finished step; one polish had run more than 10 minutes past the budget on a
+  144-pair joint.
+
 ### Support-type bound speed
 
 - Exact vertex enumeration of the support-type polytope (`support_type_bound`, the certified flow

sofic/generators/bidirectional_build.py

@@ -8,7 +8,7 @@
 from __future__ import annotations
 
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Literal
 
 import numpy as np
@@ -63,40 +63,47 @@ def _pe_max_drift(a: np.ndarray, b: np.ndarray) -> float:
 
 def preflight_bidir(forward: EpsilonMachine, *, reverse_max_states: int = 10_000) -> BidirPreflightResult:
     """Classify reverse presentation and whether a bidirectional machine builds at ``reverse_max_states``."""
+    return _preflight(forward, reverse_max_states)[0]
+
+
+def _preflight(
+    forward: EpsilonMachine, reverse_max_states: int
+) -> tuple[BidirPreflightResult, BidirectionalEpsilonMachine | None]:
+    """:func:`preflight_bidir` plus the bidirectional machine it built (``None`` unless it built)."""
     twins = forward.reverse_is_finite()
     recurrence = forward.reverse_epsilon_machine_is_finite()
     try:
         reverse = infer_reverse_epsilon_machine(forward, max_states=reverse_max_states)
-        build_bidirectional_epsilon_machine(forward, reverse)
+        bidir = build_bidirectional_epsilon_machine(forward, reverse)
     except MixedStateExplosionError as exc:
         return BidirPreflightResult(
             "explosion",
             reverse_max_states=reverse_max_states,
             reverse_recurrence_finite=recurrence,
             reverse_twins_finite=twins,
             error=str(exc)[:300],
-        )
+        ), None
     except (StochasticValidationError, ValueError) as exc:
         return BidirPreflightResult(
             "failed_build",
             reverse_max_states=reverse_max_states,
             reverse_recurrence_finite=recurrence,
             reverse_twins_finite=twins,
             error=str(exc)[:300],
-        )
+        ), None
     if recurrence:
         return BidirPreflightResult(
             "finite",
             reverse_max_states=reverse_max_states,
             reverse_recurrence_finite=True,
             reverse_twins_finite=twins,
-        )
+        ), bidir
     return BidirPreflightResult(
         "truncated",
         reverse_max_states=reverse_max_states,
         reverse_recurrence_finite=False,
         reverse_twins_finite=twins,
-    )
+    ), bidir
 
 
 def build_bidirectional_with_ladder(
@@ -105,13 +112,16 @@ def build_bidirectional_with_ladder(
     caps: tuple[int, ...] = _DEFAULT_CAPS,
     edge_tol: float = 1e-8,
     time_budget: float | None = None,
+    max_reverse_states: int | None = None,
 ) -> BidirBuildResult:
     """Build a bidirectional machine, increasing reverse ``max_states`` until ``Pe`` stabilizes.
 
     With ``time_budget`` (seconds) the ladder stops before a cap whose forecast cost would overrun
     the budget, and ``budget_exhausted`` is set on the result. A single mixed-state construction
     cannot be interrupted, so the forecast scales the previous cap's time by the square of the cap
     ratio (construction time grows faster than linearly in the cap); the first cap always runs.
+    A build whose reverse machine has more than ``max_reverse_states`` states is classified as an
+    explosion: such truncated machines are far too large for any generator search.
     """
     start = time.monotonic()
     last_pe: np.ndarray | None = None
@@ -128,13 +138,16 @@ def build_bidirectional_with_ladder(
                 exhausted = True
                 break
         attempt = time.monotonic()
-        pre = preflight_bidir(forward, reverse_max_states=cap)
+        pre, bidir = _preflight(forward, cap)
+        if bidir is not None and max_reverse_states is not None:
+            size = len(list(bidir.reverse_machine.states()))
+            if size > max_reverse_states:
+                error = f"reverse machine has {size} states, more than max_reverse_states={max_reverse_states}"
+                pre, bidir = replace(pre, bidir_class="explosion", error=error), None
         last_preflight = pre
-        if pre.bidir_class in ("explosion", "failed_build"):
+        if bidir is None:
             previous = (cap, time.monotonic() - attempt)
             continue
-        reverse = infer_reverse_epsilon_machine(forward, max_states=cap)
-        bidir = build_bidirectional_epsilon_machine(forward, reverse)
         pe = _edge_tensor(bidir)
         previous = (cap, time.monotonic() - attempt)
         if last_pe is not None:
@@ -172,21 +185,22 @@ def try_build_bidirectional(
     caps: tuple[int, ...] = _DEFAULT_CAPS,
     edge_tol: float = 1e-8,
     time_budget: float | None = None,
+    max_reverse_states: int | None = None,
 ) -> BidirBuildResult:
     """Single-cap build, or ladder when ``use_ladder`` and the default cap explodes.
 
-    ``time_budget`` (seconds, including the first preflight) is passed on to
-    :func:`build_bidirectional_with_ladder`.
+    ``time_budget`` (seconds, including the first preflight) and ``max_reverse_states`` are passed on
+    to :func:`build_bidirectional_with_ladder`.
     """
     start = time.monotonic()
-    pre = preflight_bidir(forward, reverse_max_states=reverse_max_states)
+    pre, bidir = _preflight(forward, reverse_max_states)
     if pre.bidir_class == "explosion" and use_ladder:
         remaining = None if time_budget is None else max(0.0, time_budget - (time.monotonic() - start))
-        return build_bidirectional_with_ladder(forward, caps=caps, edge_tol=edge_tol, time_budget=remaining)
-    if pre.bidir_class in ("explosion", "failed_build"):
+        return build_bidirectional_with_ladder(
+            forward, caps=caps, edge_tol=edge_tol, time_budget=remaining, max_reverse_states=max_reverse_states
+        )
+    if bidir is None:
         return BidirBuildResult(None, pre)
-    reverse = infer_reverse_epsilon_machine(forward, max_states=reverse_max_states)
-    bidir = build_bidirectional_epsilon_machine(forward, reverse)
     return BidirBuildResult(bidir, pre, edge_stabilized=True, pe_drift=0.0)
 
 

sofic/generators/coupled_generators.py

@@ -407,6 +407,10 @@ def memory_lift(bidir: BidirectionalEpsilonMachine | _GeneratorProblem, nu: np.n
     return pi, Q, R, K
 
 
+class _DeadlineReachedError(Exception):
+    """Raised inside an objective to stop an optimizer at a deadline."""
+
+
 def memory_lift_search(
     bidir: BidirectionalEpsilonMachine,
     m: int,
@@ -420,7 +424,8 @@ def memory_lift_search(
 
     A penalty continuation over the channel's logits (numerical gradients; meant for small problems).
     Returns ``(entropy, residual, nu)`` for every start, best first. No new start begins after
-    ``deadline`` (a :func:`time.monotonic` value).
+    ``deadline`` (a :func:`time.monotonic` value), and a start in progress stops at it, keeping the logits
+    of its last finished penalty stage.
     """
     rng = rng or np.random.default_rng()
     problem = _GeneratorProblem(bidir)
@@ -432,6 +437,8 @@ def channel(z: np.ndarray) -> np.ndarray:
         return z / z.sum(axis=1, keepdims=True)
 
     def value(z: np.ndarray, mu: float) -> float:
+        if deadline is not None and time.monotonic() > deadline:
+            raise _DeadlineReachedError
         pi, Q, R, K = memory_lift(problem, channel(z))
         p = pi[pi > 1e-15]
         penalty = sum(float((r**2).sum()) for r in problem.residuals(pi, Q, R, K))
@@ -443,7 +450,10 @@ def value(z: np.ndarray, mu: float) -> float:
             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
+            try:
+                z = minimize(value, z, args=(mu,), method="L-BFGS-B", options={"maxiter": maxiter}).x
+            except _DeadlineReachedError:
+                break
         nu = channel(z)
         pi, Q, R, K = memory_lift(problem, nu)
         p = pi[pi > 1e-15]
@@ -997,8 +1007,12 @@ def _optimize(
     maxiter: int,
     *,
     alternating: bool = False,
+    deadline: float | None = None,
 ) -> np.ndarray:
-    """Feasibility pass, then block ALM; joint beliefs+kernel by default."""
+    """Feasibility pass, then block ALM; joint beliefs+kernel by default.
+
+    At ``deadline`` (a :func:`time.monotonic` value) the run stops and returns its last finished step.
+    """
     block_mus = (_ALM_MU, _ALM_MU, _ALM_MU)
     pi, Q, R, K = problem.unpack(theta, k)
     if problem.residual(pi, Q, R, K) <= _ALM_TOL:
@@ -1007,31 +1021,45 @@ def _optimize(
     def _block_max(res: tuple[np.ndarray, ...]) -> tuple[float, float, float]:
         return float(np.abs(res[0]).max()), float(np.abs(res[1]).max()), float(np.abs(res[2]).max())
 
+    def lbfgs(fun, x: np.ndarray, args: tuple) -> np.ndarray:
+        def timed(*a):
+            if deadline is not None and time.monotonic() > deadline:
+                raise _DeadlineReachedError
+            return fun(*a)
+
+        return minimize(timed, x, args=args, jac=True, method="L-BFGS-B", options={"maxiter": maxiter}).x
+
     if alternating:
         theta_b = problem.logits_beliefs(pi, Q, R)
         lam = [np.zeros_like(r) for r in problem.residuals(pi, Q, R, K)]
-        theta_b = minimize(
-            problem.loss_and_grad_beliefs,
-            theta_b,
-            args=(k, K, 1.0, None, 0.0, block_mus),
-            jac=True,
-            method="L-BFGS-B",
-            options={"maxiter": maxiter},
-        ).x
-        pi, Q, R = problem.unpack_beliefs(theta_b, k)
+        try:
+            theta_b = lbfgs(problem.loss_and_grad_beliefs, theta_b, (k, K, 1.0, None, 0.0, block_mus))
+            pi, Q, R = problem.unpack_beliefs(theta_b, k)
+            for _ in range(_ALM_ROUNDS):
+                _viol, K = _kernel_lp_beliefs(problem, pi, Q, R)
+                shifts = tuple(l / m for l, m in zip(lam, block_mus, strict=False))
+                theta_b = lbfgs(problem.loss_and_grad_beliefs, theta_b, (k, K, 1.0, shifts, 0.0, block_mus))
+                pi, Q, R = problem.unpack_beliefs(theta_b, k)
+                res = problem.residuals(pi, Q, R, K)
+                lam = [l + m * r for l, r, m in zip(lam, res, block_mus, strict=False)]
+                mE, mF, mB = _block_max(res)
+                if max(mE, mF, mB) < _ALM_TOL:
+                    break
+                block_mus = tuple(
+                    min(m * (2.0 if v > _ALM_TOL else 1.0), _ALM_MU_MAX)
+                    for m, v in zip(block_mus, (mE, mF, mB), strict=False)
+                )
+        except _DeadlineReachedError:
+            pass
+        return problem.logits(pi, Q, R, K)
+
+    try:
+        theta = lbfgs(problem.loss_and_grad, theta, (k, 1.0, None, 0.0, block_mus))
+        lam = [np.zeros_like(r) for r in problem.residuals(*problem.unpack(theta, k))]
         for _ in range(_ALM_ROUNDS):
-            _viol, K = _kernel_lp_beliefs(problem, pi, Q, R)
             shifts = tuple(l / m for l, m in zip(lam, block_mus, strict=False))
-            theta_b = minimize(
-                problem.loss_and_grad_beliefs,
-                theta_b,
-                args=(k, K, 1.0, shifts, 0.0, block_mus),
-                jac=True,
-                method="L-BFGS-B",
-                options={"maxiter": maxiter},
-            ).x
-            pi, Q, R = problem.unpack_beliefs(theta_b, k)
-            res = problem.residuals(pi, Q, R, K)
+            theta = lbfgs(problem.loss_and_grad, theta, (k, 1.0, shifts, 0.0, block_mus))
+            res = problem.residuals(*problem.unpack(theta, k))
             lam = [l + m * r for l, r, m in zip(lam, res, block_mus, strict=False)]
             mE, mF, mB = _block_max(res)
             if max(mE, mF, mB) < _ALM_TOL:
@@ -1040,35 +1068,8 @@ def _block_max(res: tuple[np.ndarray, ...]) -> tuple[float, float, float]:
                 min(m * (2.0 if v > _ALM_TOL else 1.0), _ALM_MU_MAX)
                 for m, v in zip(block_mus, (mE, mF, mB), strict=False)
             )
-        return problem.logits(pi, Q, R, K)
-
-    theta = minimize(
-        problem.loss_and_grad,
-        theta,
-        args=(k, 1.0, None, 0.0, block_mus),
-        jac=True,
-        method="L-BFGS-B",
-        options={"maxiter": maxiter},
-    ).x
-    lam = [np.zeros_like(r) for r in problem.residuals(*problem.unpack(theta, k))]
-    for _ in range(_ALM_ROUNDS):
-        shifts = tuple(l / m for l, m in zip(lam, block_mus, strict=False))
-        theta = minimize(
-            problem.loss_and_grad,
-            theta,
-            args=(k, 1.0, shifts, 0.0, block_mus),
-            jac=True,
-            method="L-BFGS-B",
-            options={"maxiter": maxiter},
-        ).x
-        res = problem.residuals(*problem.unpack(theta, k))
-        lam = [l + m * r for l, r, m in zip(lam, res, block_mus, strict=False)]
-        mE, mF, mB = _block_max(res)
-        if max(mE, mF, mB) < _ALM_TOL:
-            break
-        block_mus = tuple(
-            min(m * (2.0 if v > _ALM_TOL else 1.0), _ALM_MU_MAX) for m, v in zip(block_mus, (mE, mF, mB), strict=False)
-        )
+    except _DeadlineReachedError:
+        pass
     return theta
 
 
@@ -1081,18 +1082,21 @@ def _optimize_coupled(
     alternating: bool = False,
     dual_strategy: bool = True,
     residual_atol: float = 1e-6,
+    deadline: float | None = None,
 ) -> np.ndarray:
     """Run block ALM; optionally retry with alternating ``K`` LP if joint ALM stays infeasible."""
     pi, Q, R, K = problem.unpack(theta, k)
     if problem.residual(pi, Q, R, K) <= residual_atol:
         return theta
     if alternating:
-        return _optimize(problem, theta, k, maxiter, alternating=True)
-    theta_joint = _optimize(problem, theta, k, maxiter, alternating=False)
+        return _optimize(problem, theta, k, maxiter, alternating=True, deadline=deadline)
+    theta_joint = _optimize(problem, theta, k, maxiter, alternating=False, deadline=deadline)
     r_joint = problem.residual(*problem.unpack(theta_joint, k))
     if r_joint <= residual_atol or not dual_strategy:
         return theta_joint
-    theta_alt = _optimize(problem, theta, k, maxiter, alternating=True)
+    if deadline is not None and time.monotonic() > deadline:
+        return theta_joint
+    theta_alt = _optimize(problem, theta, k, maxiter, alternating=True, deadline=deadline)
     r_alt = problem.residual(*problem.unpack(theta_alt, k))
     if r_alt + 1e-15 < r_joint:
         return theta_alt