diff --git a/scripts/scan_q_gauge_enumeration.py b/scripts/scan_q_gauge_enumeration.py new file mode 100644 index 0000000..9af2de0 --- /dev/null +++ b/scripts/scan_q_gauge_enumeration.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +"""Max I[S+0 : X0 : S-1 | S-0, S+1] over topological ε-machines. + +Uniform-outdegree topological machines from Johnson et al. (2010). +Workers use SIGALRM + a low mixed-state cap so hard reverses fail fast. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import signal +import time +from collections import defaultdict +from multiprocessing import Pool +from pathlib import Path + +import numpy as np + +from sofic.generators.topological_epsilon_enumeration import ( + idfa_string_to_epsilon_machine, + iter_topological_epsilon_strings, +) + +# Cap reverse MSP size so pathological machines fail immediately. +_MAX_MIXED_STATES = 128 +_WORKER_TIMEOUT_S = 2.0 + + +def _q_gauge(em) -> float: + d = em.to_bidirectional().step_distribution() + pmf: dict[tuple, float] = defaultdict(float) + for outcome, p in zip(d.outcomes, d.pmf, strict=True): + pmf[tuple(outcome)] += float(p) + + def marg(idxs: tuple[int, ...]) -> dict[tuple, float]: + out: dict[tuple, float] = defaultdict(float) + for outcome, p in pmf.items(): + out[tuple(outcome[i] for i in idxs)] += p + return out + + def H(idxs: tuple[int, ...]) -> float: + return -sum(p * math.log(p, 2) for p in marg(idxs).values() if p > 0) + + def I2(a: tuple[int, ...], b: tuple[int, ...], g: tuple[int, ...]) -> float: + abg = tuple(sorted(set(a) | set(b) | set(g))) + return H(a + g) + H(b + g) - H(abg) - (H(g) if g else 0.0) + + return I2((0,), (2,), (1, 3)) - I2((0,), (2,), (1, 3, 4)) + + +class _AlarmError(Exception): + pass + + +def _alarm_handler(signum, frame): # noqa: ARG001 + raise _AlarmError("worker timeout") + + +def _init_worker(max_mixed: int, timeout_s: float) -> None: + global _MAX_MIXED_STATES, _WORKER_TIMEOUT_S + _MAX_MIXED_STATES = max_mixed + _WORKER_TIMEOUT_S = timeout_s + # Patch mixed-state construction default for this process. + import sofic.generators.epsilon_construction as ec + import sofic.generators.mixed_state_construction as msc + + _orig_build_msp = msc.build_mixed_state_presentation + + def _capped_msp(hmm, initial_mixed_state=None, max_states=_MAX_MIXED_STATES, **kwargs): + return _orig_build_msp( + hmm, + initial_mixed_state=initial_mixed_state, + max_states=max_states, + **kwargs, + ) + + msc.build_mixed_state_presentation = _capped_msp # type: ignore[assignment] + + # epsilon_construction.build_epsilon_machine -> MSP; ensure it picks up patched symbol + if hasattr(ec, "build_mixed_state_presentation"): + ec.build_mixed_state_presentation = _capped_msp # type: ignore[attr-defined] + + +def eval_one(payload: tuple[tuple[int, ...], int, int]) -> tuple[str, float | None, object]: + transitions, n, k = payload + signal.signal(signal.SIGALRM, _alarm_handler) + signal.setitimer(signal.ITIMER_REAL, _WORKER_TIMEOUT_S) + try: + em = idfa_string_to_epsilon_machine(transitions, n=n, k=k) + value = float(_q_gauge(em)) + return ("ok", value, transitions) + except _AlarmError: + return ("timeout", None, transitions) + except Exception as exc: # noqa: BLE001 + return ("fail", None, f"{type(exc).__name__}: {exc}"[:100]) + finally: + signal.setitimer(signal.ITIMER_REAL, 0.0) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--n", type=int, required=True) + parser.add_argument("--k", type=int, required=True) + parser.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 4) - 1)) + parser.add_argument("--timeout", type=float, default=2.0) + parser.add_argument("--max-mixed", type=int, default=128) + parser.add_argument("--chunksize", type=int, default=32) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + n, k = args.n, args.k + out = args.out or Path(f"/tmp/q_gauge_n{n}_k{k}.json") + t0 = time.time() + + print(f"listing topological strings n={n} k={k} ...", flush=True) + strings = list(iter_topological_epsilon_strings(k=k, n=n)) + total = len(strings) + print( + f"count={total} workers={args.workers} timeout={args.timeout}s max_mixed={args.max_mixed}", + flush=True, + ) + + best_v = -np.inf + best_str = None + ok = fail = timeout = 0 + top: list[tuple[float, tuple[int, ...]]] = [] + + def save(status: str) -> None: + out.write_text( + json.dumps( + { + "status": status, + "n": n, + "k": k, + "total": total, + "ok": ok, + "fail": fail, + "timeout": timeout, + "processed": ok + fail + timeout, + "best": best_v if np.isfinite(best_v) else None, + "best_string": list(best_str) if best_str is not None else None, + "top": [{"value": v, "string": list(s)} for v, s in top[:25]], + "elapsed_s": time.time() - t0, + "max_mixed": args.max_mixed, + "timeout_s": args.timeout, + }, + indent=2, + ) + ) + + payloads = ((s, n, k) for s in strings) + with Pool( + processes=args.workers, + initializer=_init_worker, + initargs=(args.max_mixed, args.timeout), + maxtasksperchild=200, + ) as pool: + for i, (status, value, meta) in enumerate( + pool.imap_unordered(eval_one, payloads, chunksize=args.chunksize), + start=1, + ): + if status == "ok" and value is not None: + ok += 1 + if value > best_v: + best_v = value + best_str = meta + print(f" new best {best_v:.8f} string={best_str}", flush=True) + top.append((value, meta)) # type: ignore[arg-type] + top.sort(key=lambda t: -t[0]) + del top[50:] + elif status == "timeout": + timeout += 1 + else: + fail += 1 + + if i % 500 == 0 or i == total: + rate = i / max(time.time() - t0, 1e-9) + eta_m = (total - i) / max(rate, 1e-9) / 60.0 + print( + f" {i}/{total} ok={ok} fail={fail} timeout={timeout} " + f"best={best_v if np.isfinite(best_v) else float('nan'):.6f} " + f"{rate:.1f}/s ETA={eta_m:.1f}m", + flush=True, + ) + save("running") + + save("done") + print( + f"DONE max={best_v} string={best_str} ok={ok} fail={fail} timeout={timeout} " + f"elapsed={time.time() - t0:.1f}s -> {out}", + flush=True, + ) + + +if __name__ == "__main__": + main()