diff --git a/docs/generators/generative_models.rst b/docs/generators/generative_models.rst index 59a01cf..384750d 100644 --- a/docs/generators/generative_models.rst +++ b/docs/generators/generative_models.rst @@ -141,6 +141,85 @@ only, so ``gap`` can be positive even when some non-lifted generator attains :math:`\op{G}{S^+ : S^-}` (the interleaved process above). The gap vanishes for first-order Markov processes. +Coupled generators +================== + +:mod:`sofic.generators.coupled_generators` drops the lifting restriction. A +generator with :math:`k` states is described, jointly with the causal states, +by a state law :math:`\pi`, beliefs :math:`Q_w(s^+)` and :math:`R_w(s^-)`, and a +free kernel :math:`K(x, w' \mid w)`; the product +:math:`\pi_w Q_w(s^+_0) K(x_0, w_1 \mid w_0) R_{w_1}(s^-_1)` must reproduce the +bidirectional edge law :cite:`Ellison2009`, be stationary on the past side and +consistent on the future side. :func:`~sofic.generators.coupled_generators.coupled_search` +minimizes :math:`\H{\pi}` subject to these equations (an augmented Lagrangian +with analytic gradients) from the ε-machines, the bidirectional machine, extra +seeds and random starts; :func:`~sofic.generators.coupled_generators.coupled_generative_model` +returns the best candidate that reproduces the source process. For the +interleaved process above it reaches :math:`2 h(\tfrac49) = 1.98215`, below the +best lifted generator (1.99108); for the :math:`(2, \infty)` process it finds +the ε-machine (1.5). + +Two narrower tools share the same equations: + +* :func:`~sofic.generators.coupled_generators.memory_lift` builds the + *memory-1 lift* of a channel :math:`\nu(v \mid s^+, s^-)`: draw + :math:`V_t \sim \nu(\cdot \mid B_t)` independently at each time and take + :math:`W_t = (V_t, V_{t+1})`, so consecutive states share randomness. The + interleaved process's minimal generator is the memory-1 lift of a single + flag, and :func:`~sofic.generators.coupled_generators.memory_lift_search` + finds it. +* :func:`~sofic.generators.coupled_generators.generator_kernel` decides whether + a given channel's variable is the state of *any* generator: for a fixed channel + the equations are linear in :math:`K`, and a linear program returns the + smallest violation and the kernel. The exact-common-information optimum of the + :math:`(2, \infty)` process has violation 0.0625, so no generator has it as + its state. :func:`~sofic.generators.coupled_generators.bilevel_search` + (experimental) searches channels against this violation. + +All three searches are local and return upper bounds on :math:`C_g` for the +given number of states. + +Searching for a generator +========================= + +:func:`~sofic.generators.generative_search.generative_model` combines these +tools into one pipeline and returns the best generator it can verify together +with a rigorous bracket on :math:`C_g`: + +* lower bounds: the excess entropy and the support-type bound below; +* the persistent-crypticity criteria; when both vanish every separating variable + is a generator state and :math:`C_g = \op{G}{S^+ : S^-}`; +* the exact-common-information optima (dit, several starts, since the optimum + need not be unique and only some optima generate), each polished to exact + separation with :func:`~sofic.generators.coupled_generators.polish_separating_channel` + and tested with :func:`~sofic.generators.coupled_generators.generator_kernel`; +* upper bounds over nested classes, cheapest first: the ε-machines and + bidirectional machine, memory-1 lifts, deterministic and lifted channels (for + small machines), and the coupled search seeded with everything found so far. + +.. code-block:: python + + result = generative_model(bidir) + result.lower, result.upper, result.certified # bracket on C_g + result.upper_source, result.lower_source # which stage produced each bound + result.model # the generator, verified against the source + +The searches are local; ``certified`` reports when the bracket closes (for +example the golden mean, where the support bound meets :math:`C_\mu`). + +Lower bounds +============ + +:mod:`sofic.generators.generator_bounds` bounds :math:`\op{G}{S^+ : S^-}`, and +hence :math:`C_g` (Löhr & Ay :cite:`LohrAy2009`), from below using only the +support of the causal-state joint distribution. Each value of a separating +variable has a rectangular support; typing values by maximal rectangles gives a +polytope of type masses on which entropy is minimized at a vertex +(:func:`~sofic.generators.generator_bounds.support_type_bound`). The bound is +exact for the :math:`(2, \infty)` process (:math:`2 - \tfrac38 \log_2 3`) and +for the golden mean process (:math:`C_\mu`, proving :math:`C_g = C_\mu` there); +it is uninformative when the support is the full product. + API === @@ -152,6 +231,22 @@ Constructors .. autofunction:: functional_generative_model .. autofunction:: gacs_korner_generative_model +Coupled generators and bounds +----------------------------- + +.. autofunction:: sofic.generators.generative_search.generative_model +.. autoclass:: sofic.generators.generative_search.GenerativeModelResult + :members: certified, gap +.. autofunction:: sofic.generators.coupled_generators.coupled_generative_model +.. autofunction:: sofic.generators.coupled_generators.polish_separating_channel +.. autofunction:: sofic.generators.coupled_generators.coupled_search +.. autofunction:: sofic.generators.coupled_generators.memory_lift +.. autofunction:: sofic.generators.coupled_generators.memory_lift_search +.. autofunction:: sofic.generators.coupled_generators.generator_kernel +.. autofunction:: sofic.generators.coupled_generators.bilevel_search +.. autofunction:: sofic.generators.generator_bounds.support_type_bound +.. autofunction:: sofic.generators.generator_bounds.generator_support_bound + Measures -------- diff --git a/sofic/generators/coupled_generators.py b/sofic/generators/coupled_generators.py new file mode 100644 index 0000000..2e489ac --- /dev/null +++ b/sofic/generators/coupled_generators.py @@ -0,0 +1,684 @@ +r"""Minimal-state-entropy generators without the lifting restriction. + +A generator of a stationary process with :math:`k` hidden states is described, +jointly with the forward and reverse causal states, by a state distribution +:math:`\pi`, the state's beliefs :math:`Q_w(s^+) = \Pr(S^+ = s^+ \mid W = w)` +and :math:`R_w(s^-) = \Pr(S^- = s^- \mid W = w)`, and a kernel +:math:`K(x, w' \mid w)`. Its Markov property makes the joint law of +:math:`(S^+_0, W_0, X_0, W_1, S^-_1)` the product +:math:`\pi_w Q_w(s^+_0) K(x_0, w_1 \mid w_0) R_{w_1}(s^-_1)`, which must + +* reproduce the edge law :math:`\Pr(S^+_0, X_0, S^-_1)` of the bidirectional + ε-machine (Ellison, Mahoney & Crutchfield :cite:`Ellison2009`); +* be stationary on the past side: :math:`\pi_{w'} Q_{w'}(s^+_1) = + \sum \pi_w Q_w(s^+_0) K(x_0, w' \mid w)` over :math:`s^+_1 = f(s^+_0, x_0)`; +* be consistent on the future side: :math:`R_w(s^-_0) = + \sum K(x_0, w' \mid w) R_{w'}(s^-_1)` over :math:`s^-_0 = g(x_0, s^-_1)`. + +Unlike the lifted search of +:mod:`~sofic.generators.lifted_common_information`, the kernel :math:`K` is free, +so consecutive states may share randomness. :func:`coupled_generative_model` +minimizes :math:`\H{\pi}` subject to these equations (penalty continuation, +analytic gradients) from the ε-machine, reverse ε-machine, bidirectional +machine, optional extra starts (e.g. a lifted optimum) and random starts, and +checks every candidate against the source process's word probabilities. The +result bounds the generative complexity :math:`C_g` of Löhr & Ay +:cite:`LohrAy2009` from above for the given number of states. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, cast + +import numpy as np +from scipy.optimize import minimize + +from sofic.exceptions import StochasticValidationError +from sofic.generators.lifted_common_information import _LiftedProblem +from sofic.generators.mealy import MealyHMM +from sofic.generators.minimal_generative_model import _normalized_joint_distribution +from sofic.graph import ATTR_EMISSION, ATTR_PROB, TransitionGraph + +if TYPE_CHECKING: + from sofic.generators.bidirectional_epsilon_machine import BidirectionalEpsilonMachine + +_LN2 = np.log(2.0) +_ALM_MU = 1e3 +_ALM_MU_MAX = 1e7 +_ALM_ROUNDS = 60 +_ALM_TOL = 1e-9 + + +class _GeneratorProblem: + """Edge law and causal-state maps of a bidirectional ε-machine, for k-state generator equations.""" + + def __init__(self, bidir: BidirectionalEpsilonMachine, *, cutoff: float = 0.0) -> None: + joint = _normalized_joint_distribution(bidir, cutoff=cutoff) + lifted = _LiftedProblem(bidir, joint) + self.bidir = bidir + self.lifted = lifted + self.symbols = sorted(bidir.observation_alphabet, key=repr) + plus, minus, edges = lifted.plus, lifted.minus, lifted.edges + # Pe[a, x, b'] = Pr(S+_0 = a, X_0 = x, S-_1 = b') + self.Pe = np.einsum("pa,pxq,qb->axb", plus, edges, minus) + # F[a, x, a'] = 1 when s+_1 = a' follows s+_0 = a on x; G[x, b', b] = 1 when s-_0 = b precedes s-_1 = b'. + self.F = (np.einsum("pa,pxq,qc->axc", plus, edges, plus) > 0).astype(float) + self.G = (np.einsum("pb,pxq,qc->xcb", minus, edges, minus) > 0).astype(float) + self.p_pair = lifted.p + self.p_plus = plus.T @ lifted.p + self.p_minus = minus.T @ lifted.p + self.n_plus, self.n_minus, self.n_sym = plus.shape[1], minus.shape[1], len(self.symbols) + + # ---- parametrization --------------------------------------------------------------------------- + def sizes(self, k: int) -> list[tuple[int, ...]]: + return [(k,), (k, self.n_plus), (k, self.n_minus), (k, self.n_sym * k)] + + def unpack(self, theta: np.ndarray, k: int) -> tuple[np.ndarray, ...]: + parts, i = [], 0 + for shape in self.sizes(k): + n = int(np.prod(shape)) + z = theta[i : i + n].reshape(shape) + z = np.exp(z - z.max(axis=-1, keepdims=True)) + parts.append(z / z.sum(axis=-1, keepdims=True)) + i += n + pi, Q, R, K = parts + return pi, Q, R, K.reshape(k, self.n_sym, k) + + @staticmethod + def logits(pi: np.ndarray, Q: np.ndarray, R: np.ndarray, K: np.ndarray, floor: float = 1e-9) -> np.ndarray: + return np.concatenate([np.log(np.maximum(a, floor)).ravel() for a in (pi, Q, R, K.reshape(len(pi), -1))]) + + # ---- residuals ------------------------------------------------------------------------------------ + def residuals(self, pi, Q, R, K) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + A = pi[:, None] * Q + M = np.einsum("wa,wxv->axv", A, K) + rE = np.einsum("axv,vb->axb", M, R) - self.Pe + rF = np.einsum("axv,axc->vc", M, self.F) - A + N = np.einsum("wxv,vc,xcb->wb", K, R, self.G) + rB = pi[:, None] * (N - R) + return rE, rF, rB + + def residual(self, pi, Q, R, K) -> float: + return float(max(np.abs(r).max() for r in self.residuals(pi, Q, R, K))) + + def loss_and_grad( + self, theta: np.ndarray, k: int, mu: float, shifts: tuple[np.ndarray, ...] | None = None, weight: float = 1.0 + ) -> tuple[float, np.ndarray]: + """``weight * H[pi] + mu/2 * |r + shift|^2`` and its gradient (augmented Lagrangian when shifts = lambda/mu).""" + pi, Q, R, K = self.unpack(theta, k) + A = pi[:, None] * Q + M = np.einsum("wa,wxv->axv", A, K) + rE = np.einsum("axv,vb->axb", M, R) - self.Pe + rF = np.einsum("axv,axc->vc", M, self.F) - A + N = np.einsum("wxv,vc,xcb->wb", K, R, self.G) + D = N - R + rB = pi[:, None] * D + if shifts is not None: + rE, rF, rB = rE + shifts[0], rF + shifts[1], rB + shifts[2] + loss = 0.5 * float((rE**2).sum() + (rF**2).sum() + (rB**2).sum()) + + # Gradients of the penalty in (pi, Q, R, K). + gM = np.einsum("axb,vb->axv", rE, R) + np.einsum("vc,axc->axv", rF, self.F) + gR = np.einsum("axb,axv->vb", rE, M) + gA = np.einsum("wxv,axv->wa", K, gM) - rF + gK = np.einsum("wa,axv->wxv", A, gM) + gpi = (gA * Q).sum(axis=1) + (rB * D).sum(axis=1) + gQ = gA * pi[:, None] + wB = pi[:, None] * rB + gK += np.einsum("wb,vc,xcb->wxv", wB, R, self.G) + gR += np.einsum("wb,wxv,xcb->vc", wB, K, self.G) - wB + + # State entropy. + logp = np.log2(np.maximum(pi, 1e-300)) + h = float(-(pi * logp).sum()) + gpi_total = -weight * (logp + 1.0 / _LN2) + mu * gpi + + def through_softmax(p: np.ndarray, g: np.ndarray) -> np.ndarray: + return p * (g - (g * p).sum(axis=-1, keepdims=True)) + + grads = [ + through_softmax(pi, gpi_total), + through_softmax(Q, mu * gQ), + through_softmax(R, mu * gR), + through_softmax(K.reshape(k, -1), mu * gK.reshape(k, -1)), + ] + return weight * h + mu * loss, np.concatenate([g.ravel() for g in grads]) + + # ---- seeds ---------------------------------------------------------------------------------------- + def epsilon_seed(self) -> tuple[np.ndarray, ...]: + """W = S+: the forward ε-machine as a generator.""" + pa = self.p_plus + joint_ab = self.lifted.plus.T @ (self.p_pair[:, None] * self.lifted.minus) + R = joint_ab / np.maximum(pa[:, None], 1e-300) + px_a = self.Pe.sum(axis=2) / np.maximum(pa[:, None], 1e-300) + K = px_a[:, :, None] * self.F + return pa, np.eye(self.n_plus), R, K + + def reverse_seed(self) -> tuple[np.ndarray, ...]: + """W = S-: the reverse ε-machine run forward in time as a generator.""" + pb = self.p_minus + joint_ab = self.lifted.plus.T @ (self.p_pair[:, None] * self.lifted.minus) + Q = (joint_ab / np.maximum(pb[None, :], 1e-300)).T + # Pr(X_0 = x, S-_1 = b' | S-_0 = b) for g(x, b') = b + K = np.einsum("axc,xcb->bxc", self.Pe, self.G) / np.maximum(pb[:, None, None], 1e-300) + return pb, Q, np.eye(self.n_minus), K + + def pair_seed(self) -> tuple[np.ndarray, ...]: + """W = (S+, S-): the bidirectional machine.""" + return self.channel_seed(np.eye(len(self.p_pair))) + + def channel_seed(self, c: np.ndarray) -> tuple[np.ndarray, ...]: + """The lift of a channel ``c[pair, w]`` (e.g. a lifted or unconstrained optimum).""" + lifted = self.lifted + joint = lifted.p[:, None] * c + pi = joint.sum(axis=0) + safe = np.maximum(pi, 1e-300) + Q = (lifted.plus.T @ joint).T / safe[:, None] + R = (lifted.minus.T @ joint).T / safe[:, None] + K = np.einsum("pw,pxq,qv->wxv", c, lifted.edges, c) / safe[:, None, None] + return pi, Q, R, K + + +def _as_problem(bidir: BidirectionalEpsilonMachine | _GeneratorProblem) -> _GeneratorProblem: + return bidir if isinstance(bidir, _GeneratorProblem) else _GeneratorProblem(bidir) + + +def memory_lift(bidir: BidirectionalEpsilonMachine | _GeneratorProblem, nu: np.ndarray) -> tuple[np.ndarray, ...]: + r"""Generator equations' variables for the memory-1 lift of a channel ``nu[pair, v]``. + + Draw :math:`V_t \sim \nu(\cdot \mid B_t)` independently at each time given the bidirectional states + :math:`B_t`, and take :math:`W_t = (V_t, V_{t+1})`; consecutive states share :math:`V_{t+1}`. With + :math:`m` values of :math:`V` this gives :math:`m^2` states, ordered ``(v, v') -> v * m + v'``. + ``nu`` is indexed by the causal-state pairs in sorted order. Returns ``(pi, Q, R, K)``. + """ + problem = _as_problem(bidir) + E = problem.lifted.edges # E[p, x, q] = Pr(B_0 = p, X_0 = x, B_1 = q) + p_pair = problem.lifted.p + m = nu.shape[1] + # Pr(W_0 = (v, v'), B_0 = p) and Pr(B_1 = q -> B_2 = r) transition probabilities + joint_wp = np.einsum("pxq,pv,qu->vup", E, nu, nu).reshape(m * m, -1) + pi = joint_wp.sum(axis=1) + safe = np.maximum(pi, 1e-300) + Q = joint_wp @ problem.lifted.plus / safe[:, None] + R = joint_wp @ problem.lifted.minus / safe[:, None] + step = E.sum(axis=1) / np.maximum(p_pair[:, None], 1e-300) # Pr(B_2 = r | B_1 = q) + # Pr(W_0 = (v, v'), X_0 = x, W_1 = (v', v'')) + flow = np.einsum("pxq,qr,pv,qu,rw->vuxw", E, step, nu, nu, nu) + K = np.zeros((m * m, problem.n_sym, m * m)) + for v in range(m): + for u in range(m): + K[v * m + u, :, u * m : (u + 1) * m] = flow[v, u] + K /= safe[:, None, None] + return pi, Q, R, K + + +def memory_lift_search( + bidir: BidirectionalEpsilonMachine, + m: int, + *, + starts: int = 10, + maxiter: int = 400, + 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. + """ + rng = rng or np.random.default_rng() + problem = _GeneratorProblem(bidir) + n = len(problem.lifted.p) + + def channel(z: np.ndarray) -> np.ndarray: + z = z.reshape(n, m) + z = np.exp(z - z.max(axis=1, keepdims=True)) + return z / z.sum(axis=1, keepdims=True) + + def value(z: np.ndarray, mu: float) -> float: + 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)) + return float(-(p * np.log2(p)).sum()) + mu * penalty + + out = [] + for _ in range(starts): + 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 + nu = channel(z) + pi, Q, R, K = memory_lift(problem, nu) + p = pi[pi > 1e-15] + out.append((float(-(p * np.log2(p)).sum()), problem.residual(pi, Q, R, K), nu)) + return sorted(out, key=lambda t: (t[1] > 1e-6, t[0])) + + +def generator_kernel(bidir: BidirectionalEpsilonMachine | _GeneratorProblem, c: np.ndarray) -> tuple[float, np.ndarray]: + """Best kernel for a fixed channel ``c[pair, w]``: the generator equations are linear in ``K``. + + Solves the linear program minimizing the largest violation ``t`` over stochastic kernels; ``t = 0`` means + the channel's variable is the state of a generator, whatever the coupling. ``c`` is indexed by the + causal-state pairs in sorted order. Returns ``(t, K)``. + """ + from scipy.optimize import linprog + + problem = _as_problem(bidir) + + pi, Q, R, _ = problem.channel_seed(c) + k, X = len(pi), problem.n_sym + A = pi[:, None] * Q + nK = k * X * k + idx = lambda w, x, v: (w * X + x) * k + v # noqa: E731 + rows, rhs = [], [] + + def add(terms, target: float) -> None: + row = np.zeros(nK) + for j, val in terms: + row[j] += val + rows.append(row) + rhs.append(target) + + for a in range(problem.n_plus): # edge law + for x in range(X): + for b in range(problem.n_minus): + add(((idx(w, x, v), A[w, a] * R[v, b]) for w in range(k) for v in range(k)), problem.Pe[a, x, b]) + for v in range(k): # past-side stationarity + for c_ in range(problem.n_plus): + add( + ( + (idx(w, x, v), A[w, a] * problem.F[a, x, c_]) + for w in range(k) + for x in range(X) + for a in range(problem.n_plus) + ), + A[v, c_], + ) + for w in range(k): # future-side consistency + for b in range(problem.n_minus): + add( + ( + (idx(w, x, v), pi[w] * R[v, c_] * problem.G[x, c_, b]) + for x in range(X) + for v in range(k) + for c_ in range(problem.n_minus) + ), + pi[w] * R[w, b], + ) + m = len(rows) + dense = np.array(rows) + target = np.array(rhs) + # |dense K - rhs| <= t ; rows of K sum to one; K >= 0 + A_ub = np.block([[dense, -np.ones((m, 1))], [-dense, -np.ones((m, 1))]]) + b_ub = np.concatenate([target, -target]) + A_eq = np.zeros((k, nK + 1)) + for w in range(k): + A_eq[w, w * X * k : (w + 1) * X * k] = 1.0 + cost = np.zeros(nK + 1) + cost[-1] = 1.0 + res = linprog(cost, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=np.ones(k), bounds=[(0, None)] * (nK + 1), method="highs") + if not res.success: + return np.inf, np.zeros((k, X, k)) + return float(res.x[-1]), res.x[:-1].reshape(k, X, k) + + +def _decompose( + P: np.ndarray, lam0: np.ndarray, Q0: np.ndarray, R0: np.ndarray, maxiter: int = 500 +) -> tuple[np.ndarray, np.ndarray, np.ndarray, float]: + """Minimize ``H[lam]`` over exact decompositions ``P = sum_w lam_w Q_w R_w^T`` with SLSQP from ``(lam0, Q0, R0)``. + + Returns ``(lam, Q, R, residual)``, the residual being the largest entry of ``|sum_w lam_w Q_w R_w^T - P|``. + """ + k, na, nb = len(lam0), P.shape[0], P.shape[1] + n = k + k * na + k * nb + + def split(x): + return x[:k], x[k : k + k * na].reshape(k, na), x[k + k * na :].reshape(k, nb) + + def H(x): + lam = np.maximum(x[:k], 1e-300) + return float(-(lam * np.log2(lam)).sum()) + + def dH(x): + g = np.zeros(n) + g[:k] = -(np.log2(np.maximum(x[:k], 1e-15)) + 1 / _LN2) + return g + + def dec(x): + lam, A, B = split(x) + return (np.einsum("w,wa,wb->ab", lam, A, B) - P).ravel()[:-1] + + def ddec(x): + lam, A, B = split(x) + J = np.zeros((na, nb, n)) + J[:, :, :k] = np.einsum("wa,wb->abw", A, B) + for w in range(k): + for a in range(na): + J[a, :, k + w * na + a] = lam[w] * B[w] + for b in range(nb): + J[:, b, k + k * na + w * nb + b] = lam[w] * A[w] + return J.reshape(na * nb, n)[:-1] + + S = np.zeros((1 + 2 * k, n)) + S[0, :k] = 1 + for w in range(k): + S[1 + w, k + w * na : k + (w + 1) * na] = 1 + S[1 + k + w, k + k * na + w * nb : k + k * na + (w + 1) * nb] = 1 + ones = np.ones(1 + 2 * k) + res = minimize( + H, + np.concatenate([lam0, Q0.ravel(), R0.ravel()]), + jac=dH, + method="SLSQP", + bounds=[(0, 1)] * n, + constraints=[ + {"type": "eq", "fun": dec, "jac": ddec}, + {"type": "eq", "fun": lambda x: S @ x - ones, "jac": lambda x: S}, + ], + options={"maxiter": maxiter, "ftol": 1e-14}, + ) + lam, A, B = split(np.clip(res.x, 0, None)) + residual = float(np.abs(np.einsum("w,wa,wb->ab", lam, A, B) - P).max()) + return lam, A, B, residual + + +def _pair_joint(problem: _GeneratorProblem) -> np.ndarray: + lifted = problem.lifted + return lifted.plus.T @ (lifted.p[:, None] * lifted.minus) + + +def _channel_from_decomposition(problem: _GeneratorProblem, lam, A, B) -> np.ndarray: + lifted = problem.lifted + joint = np.einsum("w,pa,pb,wa,wb->pw", lam, lifted.plus, lifted.minus, A, B) + return joint / np.maximum(joint.sum(axis=1, keepdims=True), 1e-300) + + +def polish_separating_channel( + bidir: BidirectionalEpsilonMachine | _GeneratorProblem, c: np.ndarray, *, maxiter: int = 500 +) -> np.ndarray: + """Re-solve a separating channel's decomposition ``p(s+, s-) = sum_w lam_w Q_w(s+) R_w(s-)`` exactly. + + Starts from the channel's own ``(lam, Q, R)`` and minimizes ``H[lam]`` with SLSQP under the exact + decomposition constraints (unconstrained optimizers such as dit's satisfy them only to a tolerance, which + is enough to make an otherwise valid generator state fail :func:`generator_kernel`). Returns the polished + channel ``c'[pair, w] = lam_w Q_w(s+) R_w(s-) / p(s+, s-)``. + """ + problem = _as_problem(bidir) + pi, Q, R, _ = problem.channel_seed(c) + keep = pi > 1e-12 + lam, A, B, _ = _decompose(_pair_joint(problem), pi[keep], Q[keep], R[keep], maxiter) + return _channel_from_decomposition(problem, lam, A, B) + + +def separating_optima( + bidir: BidirectionalEpsilonMachine | _GeneratorProblem, + *, + starts: int = 60, + sizes: Sequence[int] | None = None, + atol: float = 1e-7, + rng: np.random.Generator | None = None, +) -> tuple[float, list[np.ndarray]]: + """Distinct optima of the exact common information ``G[S+ : S-]`` found by multistart exact decomposition. + + Runs :func:`_decompose` (SLSQP with exact constraints) from ``starts`` random points for each number of terms in + ``sizes`` (default ``2 .. min(|S+| |S-|, 8)``, within the Kumar-Li-El Gamal bound :cite:`Kumar2014`), keeps the + solutions within ``atol`` of the best value, removes duplicates (equal sets of terms up to relabeling) and + returns ``(best value, channels)``. The optimum need not be unique, and only some optima may be generator + states; :func:`generator_kernel` decides which. + """ + problem = _as_problem(bidir) + rng = rng or np.random.default_rng() + P = _pair_joint(problem) + na, nb = P.shape + sizes = sizes or range(2, min(na * nb, 8) + 1) + found = [] + for k in sizes: + for _ in range(starts): + lam, A, B, residual = _decompose( + P, rng.dirichlet(np.ones(k)), rng.dirichlet(np.ones(na), size=k), rng.dirichlet(np.ones(nb), size=k) + ) + if residual <= 1e-10: + live = lam > 1e-10 + p = lam[live] + found.append((float(-(p * np.log2(p)).sum()), lam[live], A[live], B[live])) + if not found: + return np.inf, [] + best = min(f[0] for f in found) + optima, keys = [], [] + for value, lam, A, B in sorted(found, key=lambda f: f[0]): + if value > best + atol: + break + key = sorted(tuple(np.round(np.concatenate([[l], a, b]), 6)) for l, a, b in zip(lam, A, B, strict=True)) + if key not in keys: + keys.append(key) + optima.append(_channel_from_decomposition(problem, lam, A, B)) + return best, optima + + +def bilevel_search( + bidir: BidirectionalEpsilonMachine, + starts: Sequence[np.ndarray], + *, + weight: float = 1e3, + maxiter: int = 300, +) -> list[tuple[float, float, np.ndarray]]: + """Outer search over channels ``c[pair, w]`` scoring ``H[W] + weight * t(c)``, where ``t`` is the + linear-programming violation of :func:`generator_kernel`, from the given starting channels (derivative-free; + small problems). Returns ``(entropy, violation, channel)`` per start, best feasible first. + """ + problem = _GeneratorProblem(bidir) + n = len(problem.lifted.p) + out = [] + for c0 in starts: + k = c0.shape[1] + + def channel(z: np.ndarray, k: int = k) -> np.ndarray: + z = z.reshape(n, k) + z = np.exp(z - z.max(axis=1, keepdims=True)) + return z / z.sum(axis=1, keepdims=True) + + def score(z: np.ndarray, channel=channel) -> float: + c = channel(z) + pi = problem.lifted.p @ c + p = pi[pi > 1e-15] + return float(-(p * np.log2(p)).sum()) + weight * generator_kernel(problem, c)[0] + + z0 = np.log(np.maximum(c0, 1e-6)).ravel() + z = minimize(score, z0, method="Powell", options={"maxiter": maxiter, "xtol": 1e-6, "ftol": 1e-10}).x + c = channel(z) + pi = problem.lifted.p @ c + p = pi[pi > 1e-15] + out.append((float(-(p * np.log2(p)).sum()), generator_kernel(problem, c)[0], c)) + return sorted(out, key=lambda t: (t[1] > 1e-7, t[0])) + + +Seed = tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] + + +def _pad(seed: Seed, k: int, rng: np.random.Generator, noise: float) -> Seed: + """Fit a seed with fewer or more states into ``k`` states, adding small mass to empty states.""" + pi, Q, R, K = seed + m = len(pi) + if m > k: + keep = np.argsort(-pi)[:k] + pi, Q, R, K = pi[keep], Q[keep], R[keep], K[np.ix_(keep, np.arange(K.shape[1]), keep)] + m = k + pi2 = np.full(k, noise) + pi2[:m] = pi + Q2 = rng.dirichlet(np.ones(Q.shape[1]), size=k) + Q2[:m] = Q + R2 = rng.dirichlet(np.ones(R.shape[1]), size=k) + R2[:m] = R + K2 = rng.dirichlet(np.ones(K.shape[1] * k), size=k).reshape(k, K.shape[1], k) * noise + K2[:m, :, :m] += K + K2 /= K2.reshape(k, -1).sum(axis=1)[:, None, None] + return pi2 / pi2.sum(), Q2, R2, K2 + + +@dataclass +class CoupledSearchResult: + """One optimized k-state generator: state law, beliefs, kernel, entropy and residuals.""" + + pi: np.ndarray + Q: np.ndarray + R: np.ndarray + K: np.ndarray + entropy: float + residual: float + start: str + + +def _optimize(problem: _GeneratorProblem, theta: np.ndarray, k: int, maxiter: int) -> np.ndarray: + """Feasibility first, then an augmented Lagrangian for min H[pi] subject to the generator equations.""" + theta = minimize( + problem.loss_and_grad, + theta, + args=(k, 1.0, None, 0.0), + jac=True, + method="L-BFGS-B", + options={"maxiter": maxiter}, + ).x + mu = _ALM_MU + lam = [np.zeros_like(r) for r in problem.residuals(*problem.unpack(theta, k))] + for _ in range(_ALM_ROUNDS): + shifts = tuple(l / mu for l in lam) + theta = minimize( + problem.loss_and_grad, + theta, + args=(k, mu, shifts), + jac=True, + method="L-BFGS-B", + options={"maxiter": maxiter}, + ).x + res = problem.residuals(*problem.unpack(theta, k)) + lam = [l + mu * r for l, r in zip(lam, res, strict=False)] + if max(np.abs(r).max() for r in res) < _ALM_TOL: + break + mu = min(mu * 2.0, _ALM_MU_MAX) + return theta + + +def coupled_search( + bidir: BidirectionalEpsilonMachine, + k: int, + *, + starts: int = 10, + extra: Sequence[tuple[str, tuple[np.ndarray, ...]]] = (), + maxiter: int = 3000, + residual_atol: float = 1e-6, + 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. + """ + rng = rng or np.random.default_rng() + problem = _GeneratorProblem(bidir) + named = [("eps", problem.epsilon_seed()), ("reverse", problem.reverse_seed()), ("bidir", problem.pair_seed())] + named += list(extra) + thetas = [(name, problem.logits(*_pad(cast(Seed, seed), k, rng, 1e-3))) for name, seed in named] + thetas += [ + (f"random{i}", rng.normal(scale=1.5, size=sum(int(np.prod(s)) for s in problem.sizes(k)))) + for i in range(starts) + ] + out = [] + for name, theta in thetas: + theta = _optimize(problem, theta, k, maxiter) + pi, Q, R, K = problem.unpack(theta, k) + residual = problem.residual(pi, Q, R, K) + if residual <= residual_atol: + p = pi[pi > 1e-12] + out.append(CoupledSearchResult(pi, Q, R, K, float(-(p * np.log2(p)).sum()), residual, name)) + return sorted(out, key=lambda r: r.entropy) + + +class CoupledGenerativeModel(MealyHMM): + """A generator found by :func:`coupled_generative_model`; ``generative_complexity`` is its state entropy.""" + + markov_residual: float + + def __init__(self, *, markov_residual: float = 0.0, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.markov_residual = float(markov_residual) + + def generative_complexity(self) -> float: + return self.state_entropy() + + +def _model_from_kernel( + problem: _GeneratorProblem, result: CoupledSearchResult, cutoff: float +) -> CoupledGenerativeModel: + active = [w for w in range(len(result.pi)) if result.pi[w] > cutoff] + names = {w: f"G{i}" for i, w in enumerate(active)} + graph = TransitionGraph() + for w in active: + graph.add_state(names[w]) + for w in active: + row = result.K[w][:, active] + row = np.where(row > cutoff, row, 0.0) + row /= row.sum() + for xi, x in enumerate(problem.symbols): + for j, v in enumerate(active): + if row[xi, j] > 0: + graph.add_transition(names[w], names[v], **{ATTR_PROB: float(row[xi, j]), ATTR_EMISSION: x}) + mass = result.pi[active] / result.pi[active].sum() + model = CoupledGenerativeModel( + graph=graph, + initial_distribution={names[w]: float(m) for w, m in zip(active, mass, strict=False)}, + observation_alphabet=problem.bidir.observation_alphabet, + markov_residual=result.residual, + ) + model.validate() + return model + + +def coupled_generative_model( + bidir: BidirectionalEpsilonMachine, + *, + states: int | Sequence[int] | None = None, + starts: int = 10, + extra: Sequence[tuple[str, tuple[np.ndarray, ...]]] = (), + maxiter: int = 3000, + residual_atol: float = 1e-6, + reproduction_atol: float = 1e-5, + cutoff: float = 1e-9, + rng: np.random.Generator | None = None, +) -> CoupledGenerativeModel: + """Least-state-entropy generator found over the given numbers of states (default 2 .. #causal-state pairs). + + Every candidate is checked against the source's word probabilities; raises if none reproduces the process. + """ + from sofic.generators.minimal_generative_model import _reproduction_error, _reproduction_length + + rng = rng or np.random.default_rng() + problem = _GeneratorProblem(bidir) + if states is None: + states = range(2, len(problem.p_pair) + 1) + elif isinstance(states, int): + states = [states] + candidates = [] + for k in states: + candidates += coupled_search( + bidir, k, starts=starts, extra=extra, maxiter=maxiter, residual_atol=residual_atol, rng=rng + ) + reference = bidir.forward_machine + length = _reproduction_length(reference) + for result in sorted(candidates, key=lambda r: r.entropy): + model = _model_from_kernel(problem, result, cutoff) + if _reproduction_error(model, reference, max_length=length) <= reproduction_atol: + return model + raise StochasticValidationError("coupled generative model: no candidate reproduced the source process") + + +__all__ = [ + "CoupledGenerativeModel", + "CoupledSearchResult", + "bilevel_search", + "coupled_generative_model", + "coupled_search", + "generator_kernel", + "memory_lift", + "memory_lift_search", + "polish_separating_channel", + "separating_optima", +] diff --git a/sofic/generators/generative_search.py b/sofic/generators/generative_search.py new file mode 100644 index 0000000..132d3f7 --- /dev/null +++ b/sofic/generators/generative_search.py @@ -0,0 +1,280 @@ +r"""Search for minimal-state-entropy generators, bracketing the generative complexity. + +:func:`generative_model` runs a pipeline over nested model classes and returns the best generator it can +verify together with a rigorous bracket :math:`[\text{lower}, \text{upper}]` on the generative complexity +:math:`C_g` of Löhr & Ay :cite:`LohrAy2009`: + +1. *Lower bounds.* The excess entropy :math:`\mathbf{E} = \op{I}{S^+ : S^-} \le \op{G}{S^+ : S^-} \le C_g` + and, when the causal-state joint has few maximal rectangles, the support-type and flow bounds of + :mod:`~sofic.generators.generator_bounds`. +2. *Theorem check.* If :math:`\H{S^+_{t+1} \mid S^-_t, X_t} = 0` and :math:`\H{S^-_t \mid S^+_{t+1}, X_t} = 0` + (both persistent crypticities vanish), every variable separating :math:`S^+` from :math:`S^-` is a generator + state, so :math:`C_g = \op{G}{S^+ : S^-}`; this is reported but needs no special handling, since the + unconstrained optimum then passes the next test. +3. *Exact test of the unconstrained optima.* The optima of the exact common information (Kumar, Li & El Gamal + :cite:`Kumar2014`) are enumerated by multistart exact decomposition + (:func:`~sofic.generators.coupled_generators.separating_optima`; there can be several, of which only some + generate), dit's polished optimum is added, and each is checked with the linear program of + :func:`~sofic.generators.coupled_generators.generator_kernel`; if a kernel exists, :math:`C_g` equals the + exact-common-information value found. +4. *Upper bounds over nested classes:* the ε-machines and bidirectional machine; deterministic channels + (:func:`~sofic.generators.minimal_generative_model.functional_generative_model`); lifted channels + (:func:`~sofic.generators.minimal_generative_model.minimal_generative_model`); memory-:math:`m` lifts + (:func:`~sofic.generators.coupled_generators.memory_lift_search`, :math:`m = 1`); and general + :math:`k`-state generators (:func:`~sofic.generators.coupled_generators.coupled_search`), seeded with every + earlier solution. + +Every candidate is checked against the source's word probabilities before it counts. The searches are local, +so ``upper`` is an upper bound; ``certified`` reports whether it meets ``lower``. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import numpy as np + +from sofic.exceptions import StochasticValidationError +from sofic.generators.coupled_generators import ( + CoupledGenerativeModel, + CoupledSearchResult, + _GeneratorProblem, + _model_from_kernel, + coupled_search, + generator_kernel, + memory_lift, + memory_lift_search, + polish_separating_channel, + separating_optima, +) +from sofic.generators.generator_bounds import flow_bound, maximal_rectangles, support_type_bound + +if TYPE_CHECKING: + from sofic.generators.bidirectional_epsilon_machine import BidirectionalEpsilonMachine + +_MAX_RECTANGLES = 7 +_FLOW_THRESHOLDS = ((0.6,), (0.55, 0.65)) + + +def _entropy(p: np.ndarray) -> float: + p = np.asarray(p, float) + p = p[p > 1e-15] + return float(-(p * np.log2(p)).sum()) + + +@dataclass +class GenerativeModelResult: + """Best verified generator and a bracket on the generative complexity.""" + + model: CoupledGenerativeModel + upper: float + lower: float + upper_source: str + lower_source: str + exact_common_information_estimate: float | None + separating_variables_generate: bool + stages: list[dict[str, Any]] = field(default_factory=list) + + @property + def certified(self) -> bool: + """Whether the upper and lower bounds meet (to 1e-6 bits).""" + return self.upper - self.lower <= 1e-6 + + @property + def gap(self) -> float: + return self.upper - self.lower + + +def _persistent_criteria(problem: _GeneratorProblem) -> tuple[float, float]: + """``H[S+_1 | S-_0, X_0]`` and ``H[S-_0 | S+_1, X_0]`` in bits (both zero iff both persistent crypticities + vanish).""" + lifted = problem.lifted + J = np.einsum("pb,pxq,qa->bxa", lifted.minus, lifted.edges, lifted.plus) + h_joint = _entropy(J.ravel()) + return h_joint - _entropy(J.sum(axis=2).ravel()), h_joint - _entropy(J.sum(axis=0).ravel()) + + +def _channel_matrix(problem: _GeneratorProblem, channel: dict) -> np.ndarray: + states = sorted({w for row in channel.values() for w in row}, key=repr) + return np.array([[channel[pair].get(w, 0.0) for w in states] for pair in problem.lifted.pairs]) + + +def generative_model( + bidir: BidirectionalEpsilonMachine, + *, + max_states: int | None = None, + starts: int = 20, + memory_starts: int = 10, + residual_atol: float = 1e-7, + reproduction_atol: float = 1e-5, + stages: Sequence[str] = ("machines", "exact", "memory", "functional", "lifted", "coupled"), + large: int = 9, + exact_starts: int = 30, + dit_fallback: bool = True, + rng: np.random.Generator | None = None, +) -> GenerativeModelResult: + """Run the generator-search pipeline of the module docstring. + + ``max_states`` caps the coupled search (default: the number of causal-state pairs, at most 8); + ``starts`` and ``memory_starts`` are random starts per state count for the coupled and memory-lift + searches; ``stages`` selects which upper-bound stages run. The functional and lifted stages are skipped + when there are more than ``large`` causal-state pairs (they dominate the run time there and, on the examples + tried, add nothing the memory-lift and coupled stages miss). Near-feasible candidates (e.g. an unconstrained + optimum whose kernel program has a small violation) seed the coupled search for repair. + """ + from sofic.generators.lifted_common_information import _unconstrained_channel + from sofic.generators.minimal_generative_model import ( + _normalized_joint_distribution, + _reproduction_error, + _reproduction_length, + functional_generative_model, + minimal_generative_model, + ) + + rng = rng or np.random.default_rng() + problem = _GeneratorProblem(bidir) + lifted = problem.lifted + log: list[dict[str, Any]] = [] + reference = bidir.forward_machine + length = _reproduction_length(reference) + + # ---- lower bounds ------------------------------------------------------------------------------- + excess = _entropy(problem.p_plus) + _entropy(problem.p_minus) - _entropy(problem.p_pair) + lower, lower_source = excess, "excess entropy" + joint = {pair: float(m) for pair, m in zip(lifted.pairs, lifted.p, strict=True)} + if len(maximal_rectangles(set(joint))) <= _MAX_RECTANGLES: + support, _, _ = support_type_bound(joint) + log.append({"stage": "support bound", "value": support}) + if support > lower: + lower, lower_source = support, "support-type bound" + for thresholds in _FLOW_THRESHOLDS: + flow = flow_bound(bidir, thresholds=thresholds) + log.append({"stage": f"flow bound {thresholds}", "value": flow}) + if flow > lower + 1e-9: + lower, lower_source = flow, f"flow bound {thresholds}" + log.append({"stage": "excess entropy", "value": excess}) + + f_crit, r_crit = _persistent_criteria(problem) + theorem = f_crit < 1e-10 and r_crit < 1e-10 + log.append({"stage": "persistent criteria", "value": (f_crit, r_crit), "theorem applies": theorem}) + + # ---- upper bounds --------------------------------------------------------------------------------- + candidates: list[tuple[float, str, tuple[np.ndarray, ...]]] = [] + repair: list[tuple[str, tuple[np.ndarray, ...]]] = [] + small = len(lifted.pairs) <= large + + def offer(name: str, seed: tuple[np.ndarray, ...]) -> None: + residual = problem.residual(*seed) + value = _entropy(seed[0]) + log.append({"stage": name, "value": value, "residual": residual}) + if residual <= residual_atol: + candidates.append((value, name, seed)) + + if "machines" in stages: + offer("epsilon machine", problem.epsilon_seed()) + offer("reverse epsilon machine", problem.reverse_seed()) + offer("bidirectional machine", problem.pair_seed()) + + g_estimate = None + if "exact" in stages: + # The unconstrained problem can have several optima, only some of them generator states (12 and 2 for the + # interleaved chains), so enumerate them by multistart exact decomposition and test each with the kernel + # program; dit's optimum, polished, is added as one more candidate. + best, optima = separating_optima(problem, starts=exact_starts, rng=rng) + if np.isfinite(best): + g_estimate = best + if dit_fallback: + try: + joint_n = _normalized_joint_distribution(bidir, cutoff=1e-10) + c = _unconstrained_channel( + lifted, + joint_n, + "exact", + niter=10, + maxiter=1000, + polish=1e-8, + backend="numpy", + cutoff=1e-10, + rng=rng, + ) + if c is not None: + c = polish_separating_channel(problem, c) + value = _entropy(lifted.p @ c) + if g_estimate is None or value < g_estimate - 1e-9: + g_estimate, optima = value, [c] + elif value <= g_estimate + 1e-7: + optima.append(c) + except (StochasticValidationError, ImportError): + pass + accepted = 0 + for i, c in enumerate(optima): + violation, K = generator_kernel(problem, c) + pi, Q, R, _ = problem.channel_seed(c) + log.append( + { + "stage": f"exact common information, optimum {i}", + "value": _entropy(pi), + "kernel violation": violation, + } + ) + if violation <= 1e-9: + offer("exact-common-information optimum", (pi, Q, R, K)) + accepted += 1 + else: + repair.append((f"exact-common-information optimum {i} (repair)", (pi, Q, R, K))) + log.append( + { + "stage": "exact common information", + "value": g_estimate, + "optima": len(optima), + "generating optima": accepted, + } + ) + + if "memory" in stages: + for _value, residual, nu in memory_lift_search(bidir, 2, starts=memory_starts, rng=rng)[:3]: + if residual <= residual_atol: + offer("memory-1 lift", memory_lift(problem, nu)) + + for name, builder in (("functional", functional_generative_model), ("lifted", minimal_generative_model)): + if name not in stages or not small: + if name in stages: + log.append({"stage": f"{name} channel", "skipped": f"more than {large} causal-state pairs"}) + continue + try: + kwargs = {"rng": rng} if name == "lifted" else {} + model = builder(bidir, **kwargs) + offer(f"{name} channel", problem.channel_seed(_channel_matrix(problem, model.pair_state_channel))) + except (StochasticValidationError, ImportError) as exc: + log.append({"stage": f"{name} channel", "error": str(exc)}) + + if "coupled" in stages: + k_max = max_states or min(len(lifted.pairs), 8) + extra = [(name, seed) for _v, name, seed in sorted(candidates, key=lambda t: t[0])[:6]] + repair + for k in range(2, k_max + 1): + for res in coupled_search(bidir, k, starts=starts, extra=extra, residual_atol=residual_atol, rng=rng)[:2]: + offer(f"coupled k={k} ({res.start})", (res.pi, res.Q, res.R, res.K)) + + # ---- verify and choose ------------------------------------------------------------------------------- + for value, name, seed in sorted(candidates, key=lambda t: t[0]): + pi, Q, R, K = seed + result = CoupledSearchResult(pi, Q, R, K, value, problem.residual(pi, Q, R, K), name) + model = _model_from_kernel(problem, result, 1e-9) + if _reproduction_error(model, reference, max_length=length) <= reproduction_atol: + upper = model.state_entropy() + return GenerativeModelResult( + model=model, + upper=float(upper), + lower=float(min(lower, upper)), + upper_source=name, + lower_source=lower_source, + exact_common_information_estimate=g_estimate, + separating_variables_generate=theorem, + stages=log, + ) + raise StochasticValidationError("generative_model: no candidate reproduced the source process") + + +__all__ = ["GenerativeModelResult", "generative_model"] diff --git a/sofic/generators/generator_bounds.py b/sofic/generators/generator_bounds.py new file mode 100644 index 0000000..c1f2434 --- /dev/null +++ b/sofic/generators/generator_bounds.py @@ -0,0 +1,270 @@ +r"""Lower bounds on exact common information and generative complexity from the support of p(s+, s-). + +Every value :math:`w` of a variable separating the forward and reverse causal states has +:math:`p(s^+, s^- \mid w) = q_w(s^+)\, r_w(s^-)`, so its support is a rectangle inside the support of the +causal-state joint. Assign each value to a maximal rectangle containing its support (a function of +:math:`W`); the type masses :math:`m_T` then satisfy, for every family :math:`\mathcal T'` of maximal rectangles, + +.. math:: + + p\Bigl(\text{cells covered only by } \mathcal T'\Bigr) \;\le\; \sum_{T \in \mathcal T'} m_T + \;\le\; p\Bigl(\textstyle\bigcup_{T \in \mathcal T'} T\Bigr), + +and :math:`\H{W} \ge \H{m}`. Entropy is concave, so its minimum over this polytope is attained at a vertex. +The bound applies to the exact common information :math:`\op{G}{S^+ : S^-}` (Kumar, Li & El Gamal +:cite:`Kumar2014`) and therefore to the generative complexity :math:`C_g` (Löhr & Ay :cite:`LohrAy2009`), +whose generator states separate the causal states. It is informative only when the support is sparse; for +the :math:`(2, \infty)` run-length-limited process it equals the exact common information. +""" + +from __future__ import annotations + +import itertools +from collections.abc import Hashable +from collections.abc import Set as AbstractSet +from fractions import Fraction +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from sofic.generators.bidirectional_epsilon_machine import BidirectionalEpsilonMachine + + +def maximal_rectangles(support: AbstractSet[tuple[Hashable, Hashable]]) -> list[tuple[frozenset, frozenset]]: + """Maximal rectangles ``U x V`` contained in a set of cells (maximal bicliques), by closure of row sets.""" + rows = sorted({a for a, _ in support}, key=repr) + cols_of = {a: frozenset(b for a2, b in support if a2 == a) for a in rows} + found = set() + for size in range(1, len(rows) + 1): + for subset in itertools.combinations(rows, size): + V = frozenset.intersection(*(cols_of[a] for a in subset)) + if not V: + continue + U = frozenset(a for a in rows if cols_of[a] >= V) + found.add((U, V)) + return sorted(found, key=lambda t: (sorted(map(repr, t[0])), sorted(map(repr, t[1])))) + + +def _solve(M: list[list[Fraction]], r: list[Fraction]) -> list[Fraction] | None: + n = len(M) + A = [row[:] + [r[i]] for i, row in enumerate(M)] + for col in range(n): + piv = next((i for i in range(col, n) if A[i][col] != 0), None) + if piv is None: + return None + A[col], A[piv] = A[piv], A[col] + p = A[col][col] + A[col] = [v / p for v in A[col]] + for i in range(n): + if i != col and A[i][col] != 0: + f = A[i][col] + A[i] = [a - f * b for a, b in zip(A[i], A[col], strict=False)] + return [A[i][n] for i in range(n)] + + +def support_type_bound( + joint: dict[tuple[Hashable, Hashable], float], *, max_denominator: int = 10**6 +) -> tuple[float, list[tuple[frozenset, frozenset]], tuple[Fraction, ...]]: + """Lower bound on ``G[S+ : S-]`` (hence on ``C_g``) from the rectangle types of the support. + + ``joint`` maps causal-state pairs to probabilities (rationalized with ``max_denominator``). Returns the bound + (bits), the maximal rectangles, and a minimizing vertex of type masses. + """ + p = {cell: Fraction(v).limit_denominator(max_denominator) for cell, v in joint.items() if v > 0} + total = sum(p.values()) + p = {cell: v / total for cell, v in p.items()} + types = maximal_rectangles(set(p)) + n = len(types) + contains = [{cell for cell in p if cell[0] in U and cell[1] in V} for U, V in types] + ineq = [([-Fraction(int(j == i)) for j in range(n)], Fraction(0)) for i in range(n)] + for size in range(1, n + 1): + for fam in itertools.combinations(range(n), size): + union = set().union(*(contains[i] for i in fam)) + only = {cell for cell in union if all(i in fam for i in range(n) if cell in contains[i])} + row = [Fraction(int(j in fam)) for j in range(n)] + ineq.append((row, sum((p[c] for c in union), Fraction(0)))) + ineq.append(([-x for x in row], -sum((p[c] for c in only), Fraction(0)))) + best = (np.inf, None) + for S in itertools.combinations(range(len(ineq)), n - 1): + M = [[Fraction(1)] * n] + [ineq[i][0] for i in S] + r = [Fraction(1)] + [ineq[i][1] for i in S] + x = _solve(M, r) + if x is None or not all(sum(a * xi for a, xi in zip(row, x, strict=False)) <= v for row, v in ineq): + continue + h = -sum(float(m) * np.log2(float(m)) for m in x if m > 0) + if h < best[0]: + best = (h, tuple(x)) + if best[1] is None: + raise ValueError("support_type_bound: the type polytope has no vertex") + return float(best[0]), types, best[1] + + +def _concave_entropy_min( + A_ub: np.ndarray, b_ub: np.ndarray, A_eq: np.ndarray, b_eq: np.ndarray, *, tol: float = 1e-7, max_nodes: int = 20000 +) -> float: + """Minimum of the (concave) entropy of x over a polytope, by branch-and-bound with secant LP bounds. + + On a box ``lo <= x <= hi`` each term ``-x_i log2 x_i`` is at least its secant, so the LP with the secant + objective bounds the minimum from below; the entropy at the LP solution bounds it from above. + """ + from scipy.optimize import linprog + + n = A_eq.shape[1] + + def f(t: np.ndarray) -> np.ndarray: + t = np.clip(t, 0.0, 1.0) + return np.where(t > 0, -t * np.log2(np.where(t > 0, t, 1.0)), 0.0) + + def solve(lo: np.ndarray, hi: np.ndarray): + width = np.maximum(hi - lo, 1e-15) + slope = (f(hi) - f(lo)) / width + res = linprog( + slope, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=b_eq, bounds=list(zip(lo, hi, strict=True)), method="highs" + ) + if not res.success: + return None + x = res.x + return float((f(lo) + slope * (x - lo)).sum()), x + + lo, hi = np.zeros(n), np.ones(n) + root = solve(lo, hi) + if root is None: + return np.inf + best_upper = float(f(root[1]).sum()) + heap = [(root[0], 0, lo, hi, root[1])] + import heapq + + counter = 1 + nodes = 0 + while heap: + bound, _, lo, hi, x = heapq.heappop(heap) + if bound >= best_upper - tol or nodes > max_nodes: + return min(bound, best_upper) if nodes > max_nodes else best_upper + nodes += 1 + width = np.maximum(hi - lo, 1e-15) + gap = f(x) - (f(lo) + (f(hi) - f(lo)) / width * (x - lo)) + i = int(np.argmax(gap)) + split = x[i] if lo[i] + 1e-12 < x[i] < hi[i] - 1e-12 else (lo[i] + hi[i]) / 2 + for new_lo, new_hi in ( + (lo, np.where(np.arange(n) == i, split, hi)), + (np.where(np.arange(n) == i, split, lo), hi), + ): + child = solve(new_lo, new_hi) + if child is None: + continue + best_upper = min(best_upper, float(f(child[1]).sum())) + if child[0] < best_upper - tol: + heapq.heappush(heap, (child[0], counter, new_lo, new_hi, child[1])) + counter += 1 + return best_upper + + +def flow_bound(bidir: BidirectionalEpsilonMachine, *, thresholds: tuple[float, ...] = (0.55, 0.65)) -> float: + r"""Lower bound on the generative complexity :math:`C_g` from support types plus generator flows. + + Generators satisfy, beyond separation, the flow identities of their Markov property: if the forward causal + state :math:`a'` follows only :math:`a` (on symbol :math:`x`), then + :math:`p(w') q_{w'}(a') = \sum_w p(w)\, q_w(a)\, K(x, w' \mid w)`, and symmetrically for reverse causal states + with a unique successor. States of a type that must hold a *forced* cell (one no other type covers) in row + :math:`a'` are therefore entered from states that give :math:`a` high probability. Binning each type by such a + belief at ``thresholds`` turns these facts into linear constraints on the bin masses (the inflow bound in its + LP-dual form); since the bin label is a function of the state, :math:`\H{W}` is at least the minimum entropy of + the bin masses over the resulting polytope, computed by branch-and-bound. This automates the certificate + :math:`C_g \ge 1.49702` for the :math:`(2, \infty)` process. + """ + from sofic.generators.coupled_generators import _GeneratorProblem + + problem = _GeneratorProblem(bidir) + lifted = problem.lifted + P = lifted.plus.T @ (lifted.p[:, None] * lifted.minus) + na, nb = P.shape + support = {(a, b) for a in range(na) for b in range(nb) if P[a, b] > 1e-14} + types = maximal_rectangles(support) + contains = [{c for c in support if c[0] in U and c[1] in V} for U, V in types] + forced = [{c for c in contains[t] if sum(c in s for s in contains) == 1} for t in range(len(types))] + + # unique-predecessor forward sources (a, a') and unique-successor reverse sources (b, b') + fwd = [(int(np.argwhere(problem.F[:, :, a2])[0][0]), a2) for a2 in range(na) if problem.F[:, :, a2].sum() == 1] + rev = [(b, int(np.argwhere(problem.G[:, :, b])[0][1])) for b in range(nb) if problem.G[:, :, b].sum() == 1] + # candidate binning variables: ("+", a) for forward sources whose target row has forced cells, ("-", b') mirror + cands = [("+", a) for a, a2 in fwd if any(c[0] == a2 for t in forced for c in t)] + cands += [("-", b2) for b, b2 in rev if any(c[1] == b for t in forced for c in t)] + + def holds(t: int, var: tuple[str, int]) -> bool: + U, V = types[t] + return var[1] in (U if var[0] == "+" else V) + + binned: dict[int, tuple[str, int]] = {} + for var in sorted(set(cands), key=lambda v: sum(holds(t, v) for t in range(len(types)))): + for t in range(len(types)): + if t not in binned and holds(t, var): + binned[t] = var + edges = [0.0, *thresholds, 1.0] + columns = [] # (type, lower bound, upper bound) of the belief binned on, or (type, 0, 1) when unbinned + for t in range(len(types)): + if t in binned: + columns += [(t, edges[i], edges[i + 1]) for i in range(len(edges) - 1)] + else: + columns.append((t, 0.0, 1.0)) + n = len(columns) + A_ub, b_ub = [], [] + + def type_row(fam): + return np.array([1.0 if col[0] in fam else 0.0 for col in columns]) + + for size in range(1, len(types) + 1): + for fam in itertools.combinations(range(len(types)), size): + union = set().union(*(contains[t] for t in fam)) + only = {c for c in union if all(t in fam for t in range(len(types)) if c in contains[t])} + A_ub.append(type_row(fam)) + b_ub.append(sum(P[c] for c in union)) + A_ub.append(-type_row(fam)) + b_ub.append(-sum(P[c] for c in only)) + + def coeffs(var: tuple[str, int], which: int) -> np.ndarray: + """Per column: bin bound (which = 1 upper, 0 lower) of the belief var, 1/0 if the type holds var unbinned.""" + out = np.zeros(n) + for j, (t, lb, ub) in enumerate(columns): + if not holds(t, var): + continue + if binned.get(t) == var: + out[j] = ub if which else lb + else: + out[j] = 1.0 if which else 0.0 + return out + + for var in set(binned.values()): + marginal = float(P[var[1], :].sum()) if var[0] == "+" else float(P[:, var[1]].sum()) + A_ub.append(coeffs(var, 0)) + b_ub.append(marginal) + A_ub.append(-coeffs(var, 1)) + b_ub.append(-marginal) + + def flow(var: tuple[str, int], t0: int, need: float) -> None: + u = coeffs(var, 1) + own = type_row({t0}) + for lam in sorted(set(u) | {1.0}): + A_ub.append(-(lam * own + np.maximum(u - lam, 0.0))) + b_ub.append(-need) + + for a, a2 in fwd: + for t0 in range(len(types)): + need = sum(P[c] for c in forced[t0] if c[0] == a2) + if need > 0: + flow(("+", a), t0, need) + for b, b2 in rev: + for t0 in range(len(types)): + need = sum(P[c] for c in forced[t0] if c[1] == b) + if need > 0: + flow(("-", b2), t0, need) + + return _concave_entropy_min(np.array(A_ub), np.array(b_ub), np.ones((1, n)), np.array([1.0])) + + +def generator_support_bound(bidir: BidirectionalEpsilonMachine) -> float: + """:func:`support_type_bound` of a bidirectional ε-machine's causal-state joint distribution.""" + return support_type_bound({k: float(v) for k, v in bidir.joint_distribution().items()})[0] + + +__all__ = ["flow_bound", "generator_support_bound", "maximal_rectangles", "support_type_bound"] diff --git a/tests/test_coupled_generators.py b/tests/test_coupled_generators.py new file mode 100644 index 0000000..ce64106 --- /dev/null +++ b/tests/test_coupled_generators.py @@ -0,0 +1,214 @@ +"""Tests for generators without the lifting restriction, and support-type lower bounds.""" + +from __future__ import annotations + +from itertools import product + +import numpy as np +import pytest + +from sofic.examples.epsilon_machines import from_symbol_matrices, golden_mean_bidirectional +from sofic.generators import coupled_generators as cg +from sofic.generators.bidirectional_epsilon_machine import BidirectionalEpsilonMachine +from sofic.generators.generator_bounds import generator_support_bound, maximal_rectangles, support_type_bound +from sofic.generators.minimal_generative_model import _reproduction_error, _reproduction_length + + +def _h(p: float) -> float: + return float(-(p * np.log2(p) + (1 - p) * np.log2(1 - p))) + + +def _rll() -> BidirectionalEpsilonMachine: + machine = from_symbol_matrices( + ("A", "B", "C"), + (0, 1), + {0: np.array([[0.5, 0, 0], [0, 0, 1.0], [1.0, 0, 0]]), 1: np.array([[0, 0.5, 0], [0, 0, 0], [0, 0, 0]])}, + ) + return BidirectionalEpsilonMachine.from_forward(machine) + + +def _interleaved(eps: float) -> BidirectionalEpsilonMachine: + T = np.array([[1 - eps, eps], [eps, 1 - eps]]) + names = ["00", "01", "10", "11"] + mats = {x: np.zeros((4, 4)) for x in (0, 1)} + for i, (a, b) in enumerate(names): + for c in "01": + mats[int(c)][i, names.index(b + c)] = T[int(a), int(c)] + return BidirectionalEpsilonMachine.from_forward(from_symbol_matrices(names, (0, 1), mats)) + + +def _rll_exact_channel(problem: cg._GeneratorProblem) -> np.ndarray: + """The deterministic exact-common-information optimum W* of the (2, inf) process as a channel.""" + pairs = problem.lifted.pairs + partner: dict = {} + for a, b in pairs: + partner.setdefault(b, set()).add(a) + r00 = next(b for b, s in partner.items() if "B" in s) + labels = [0 if b == r00 else 2 if a == "C" else 1 for a, b in pairs] + return np.eye(3)[labels] + + +def _flag_channel(problem: cg._GeneratorProblem, eps: float) -> np.ndarray: + """V_t = flag of the pair (z_{t-2}, z_t) of the interleaved process, as a channel of the bidirectional state.""" + T = np.array([[1 - eps, eps], [eps, 1 - eps]]) + joint = dict(zip(problem.lifted.pairs, problem.lifted.p, strict=False)) + first_future = {} + for _, m in problem.lifted.pairs: + for c, d in product((0, 1), repeat=2): + if all( + abs(joint.get((s, m), 0.0) - 0.25 * T[int(s[0]), c] * T[int(s[1]), d]) < 1e-12 + for s in ("00", "01", "10", "11") + ): + first_future[m] = c + alpha = (1 - 2 * eps) / (2 * (1 - eps)) + p_flag = alpha / ((1 - eps) / 2) + nu = np.zeros((len(problem.lifted.pairs), 2)) + for i, (s, m) in enumerate(problem.lifted.pairs): + f = p_flag if (s[0], first_future[m]) == ("0", 0) else 0.0 + nu[i] = (f, 1 - f) + return nu + + +def test_loss_gradient_matches_finite_differences() -> None: + problem = cg._GeneratorProblem(_interleaved(0.1)) + k = 3 + rng = np.random.default_rng(0) + theta = rng.normal(size=sum(int(np.prod(s)) for s in problem.sizes(k))) + shifts = tuple(0.01 * rng.normal(size=r.shape) for r in problem.residuals(*problem.unpack(theta, k))) + _, grad = problem.loss_and_grad(theta, k, 10.0, shifts, 0.7) + step = 1e-6 + numeric = np.array( + [ + ( + problem.loss_and_grad(theta + step * e, k, 10.0, shifts, 0.7)[0] + - problem.loss_and_grad(theta - step * e, k, 10.0, shifts, 0.7)[0] + ) + / (2 * step) + for e in np.eye(len(theta)) + ] + ) + assert np.max(np.abs(numeric - grad)) <= 1e-6 * np.max(np.abs(grad)) + + +@pytest.mark.parametrize("bidir", [_rll(), _interleaved(0.1), golden_mean_bidirectional(0.4)]) +def test_machine_seeds_solve_the_generator_equations(bidir) -> None: + problem = cg._GeneratorProblem(bidir) + for seed in (problem.epsilon_seed(), problem.reverse_seed(), problem.pair_seed()): + assert problem.residual(*seed) < 1e-12 + + +def test_rll_coupled_search_reaches_the_epsilon_machine() -> None: + results = cg.coupled_search(_rll(), 3, starts=10, rng=np.random.default_rng(2)) + assert results[0].entropy == pytest.approx(1.5, abs=1e-5) + + +def test_memory_lift_of_the_flag_is_the_minimal_interleaved_generator() -> None: + """The coupled generator of the interleaved process is the memory-1 lift of a single flag.""" + problem = cg._GeneratorProblem(_interleaved(0.1)) + pi, Q, R, K = cg.memory_lift(problem, _flag_channel(problem, 0.1)) + assert problem.residual(pi, Q, R, K) < 1e-12 + assert float(-(pi * np.log2(pi)).sum()) == pytest.approx(2 * _h(4 / 9), abs=1e-12) + + +def test_generator_kernel_decides_generator_states() -> None: + rll = cg._GeneratorProblem(_rll()) + assert cg.generator_kernel(rll, np.eye(len(rll.lifted.pairs)))[0] == pytest.approx(0.0, abs=1e-10) + assert cg.generator_kernel(rll, _rll_exact_channel(rll))[0] > 0.05 # no generator has W* as its state + z = cg._GeneratorProblem(_interleaved(0.1)) + nu = _flag_channel(z, 0.1) + joint = np.einsum("pxq,pv,qu->vup", z.lifted.edges, nu, nu).reshape(4, -1) + product_flag = (joint / z.lifted.p[None, :]).T + violation, kernel = cg.generator_kernel(z, product_flag) + assert violation == pytest.approx(0.0, abs=1e-9) + assert np.allclose(kernel.reshape(4, -1).sum(axis=1), 1.0) + + +def test_support_bound_is_the_exact_common_information_of_rll() -> None: + value, types, masses = support_type_bound({k: float(v) for k, v in _rll().joint_distribution().items()}) + assert value == pytest.approx(2 - 3 / 8 * np.log2(3), abs=1e-12) + assert len(types) == 3 + assert sorted(masses) == sorted([0.5, 0.375, 0.125]) + + +def test_support_bound_proves_golden_mean_generative_complexity() -> None: + bidir = golden_mean_bidirectional(0.5) + assert generator_support_bound(bidir) == pytest.approx(_h(1 / 3), abs=1e-12) + + +def test_maximal_rectangles_of_a_full_support_is_one_rectangle() -> None: + cells = {(a, b) for a in "xy" for b in "uv"} + assert maximal_rectangles(cells) == [(frozenset("xy"), frozenset("uv"))] + + +def test_coupled_generative_model_reproduces_the_process() -> None: + bidir = golden_mean_bidirectional(0.5) + model = cg.coupled_generative_model(bidir, states=2, starts=4, rng=np.random.default_rng(0)) + reference = bidir.forward_machine + assert _reproduction_error(model, reference, max_length=_reproduction_length(reference)) < 1e-5 + assert model.generative_complexity() == pytest.approx(_h(1 / 3), abs=1e-5) + + +def test_polishing_makes_the_flag_channel_exact() -> None: + """A slightly perturbed product-flag channel is polished back to an exact generator state.""" + z = cg._GeneratorProblem(_interleaved(0.1)) + nu = _flag_channel(z, 0.1) + joint = np.einsum("pxq,pv,qu->vup", z.lifted.edges, nu, nu).reshape(4, -1) + c = (joint / z.lifted.p[None, :]).T + rng = np.random.default_rng(0) + noisy = np.clip(c + 1e-4 * rng.random(c.shape), 0, None) + noisy /= noisy.sum(axis=1, keepdims=True) + assert cg.generator_kernel(z, noisy)[0] > 1e-6 + polished = cg.polish_separating_channel(z, noisy) + assert cg.generator_kernel(z, polished)[0] == pytest.approx(0.0, abs=1e-9) + pi = z.lifted.p @ polished + assert float(-(pi[pi > 0] * np.log2(pi[pi > 0])).sum()) == pytest.approx(2 * _h(4 / 9), abs=1e-7) + + +def test_generative_model_certifies_the_golden_mean() -> None: + from sofic.generators.generative_search import generative_model + + result = generative_model(golden_mean_bidirectional(0.5), starts=3, max_states=3, rng=np.random.default_rng(0)) + assert result.certified + assert result.upper == pytest.approx(_h(1 / 3), abs=1e-6) + assert result.separating_variables_generate + + +def test_generative_model_brackets_the_rll_process() -> None: + from sofic.generators.generative_search import generative_model + + result = generative_model( + _rll(), stages=("machines", "coupled"), starts=6, max_states=3, rng=np.random.default_rng(1) + ) + h11 = -(0.5 * np.log2(0.5) + (3 / 11) * np.log2(3 / 11) + (5 / 22) * np.log2(5 / 22)) + assert result.lower == pytest.approx(h11, abs=1e-6) + assert result.lower_source.startswith("flow bound") + assert result.upper == pytest.approx(1.5, abs=1e-5) + assert not result.certified and not result.separating_variables_generate + + +def test_separating_optima_finds_the_generating_flags() -> None: + """The interleaved process has several optima of G[S+ : S-]; exactly two of them are generator states.""" + problem = cg._GeneratorProblem(_interleaved(0.1)) + best, optima = cg.separating_optima(problem, starts=30, sizes=(4,), rng=np.random.default_rng(0)) + assert best == pytest.approx(2 * _h(4 / 9), abs=1e-7) + generating = [c for c in optima if cg.generator_kernel(problem, c)[0] <= 1e-9] + assert len(optima) > len(generating) == 2 + + +def test_separating_optima_of_rll_are_not_generator_states() -> None: + problem = cg._GeneratorProblem(_rll()) + best, optima = cg.separating_optima(problem, starts=20, sizes=(3, 4), rng=np.random.default_rng(0)) + assert best == pytest.approx(2 - 3 / 8 * np.log2(3), abs=1e-7) + assert len(optima) == 2 + assert all(cg.generator_kernel(problem, c)[0] > 1e-3 for c in optima) + + +def test_flow_bound_reproduces_the_rll_certificates() -> None: + from sofic.generators.generator_bounds import flow_bound + + rll = _rll() + h5 = -(2 * (5 / 12) * np.log2(5 / 12) + (1 / 6) * np.log2(1 / 6)) + h11 = -(0.5 * np.log2(0.5) + (3 / 11) * np.log2(3 / 11) + (5 / 22) * np.log2(5 / 22)) + assert flow_bound(rll, thresholds=(0.6,)) == pytest.approx(h5, abs=1e-6) + assert flow_bound(rll, thresholds=(0.55, 0.65)) == pytest.approx(h11, abs=1e-6) + assert flow_bound(golden_mean_bidirectional(0.5)) == pytest.approx(_h(1 / 3), abs=1e-6)