sim-numerics/sweep_controls.py

"""PREREG_controls_2026-09-22: the controls, as registered (commit 77312be).

K2 (point 6): A' with only the flat detector swapped - Lagouanelle with memory (K2L),
Chandrupatla's band (K2P), no detector (K2F) - on the six corpora of section 9, harness of
sweep_fourth.py (same problems, seeding, charging, uniform criterion).
K1 (point 5): library Chandrupatla and Brent on g = sign(f)|f|^(1/m*), m* the true order read at
the known root, on the four corpora whose roots are known.

    python -u sweep_controls.py            run -> sweep_controls.json, then the report
    python -u sweep_controls.py --smoke    K2 on corpus 4 and K1 on all four corpora, no file
    python -u sweep_controls.py --report
"""
from __future__ import annotations

import json
import math
import random
import sys
from collections import defaultdict
from multiprocessing import Pool

import sweep_fourth
from corpus_roots_jl import solved
from sweep_full import MAXFEV, XTOL
from sweep_oos import sign_p

TAIL = "+half+close+sub"
K2 = {"K2L": "contra+flatlagmempt" + TAIL, "K2P": "contra+flatbandmempt" + TAIL,
      "K2F": "contra+flatflagmempt" + TAIL}
for _k, _pre in K2.items():
    sweep_fourth.CONFIGS[_k] = sweep_fourth.spec(_pre)
SETS = ("APS", "CHA", "NEW", "RJL", "APP", "D")
K1_SETS = ("APS", "CHA", "NEW", "D")


def m_star(f, r):
    """The order at the known root: slope of ln|f| against ln|x - r|, both sides averaged,
    h = 1e-5 max(1, |r|) (registered).  1 when undefined."""
    # amendment 2 of the PREREG (after the smokes, before any full result): each side's slope
    # is rounded to 0.01 and m* is defined only when BOTH sides show a zero (slope > 0) -
    # a jump (C0.03), a flat side (aps.15.21-30) or a pole-like side (C0.04) is undefined
    h = 1e-5 * max(1.0, abs(r))
    vals = []
    for s in (1.0, -1.0):
        try:
            f1, f2 = abs(f(r + s * h)), abs(f(r + s * 2 * h))
            m = (math.log(f2) - math.log(f1)) / math.log(2.0)
        except (ValueError, OverflowError, ZeroDivisionError):
            return 1.0
        if not math.isfinite(m) or round(m, 2) <= 0.0:
            return 1.0
        vals.append(round(m, 2))
    return (vals[0] + vals[1]) / 2.0


def nonsimple(m):
    """The group 'multiple/fractional': the branches of section 8.1 where a tangent step loses to
    a halving, m >= 2 or m <= 2/3 (amendment 2)."""
    return m >= 2.0 or m <= 2.0 / 3.0


def k1_problems(which):
    from corpus_fourth import build as b4
    from corpus_new import build as b1
    from sweep_full import load_problems
    out = []
    if which in ("APS", "CHA"):
        out = [(pid, f, a, b, r) for s, pid, f, a, b, r in load_problems() if s == which]
    elif which == "NEW":
        out = [(pid, f, a, b, r) for _c, pid, f, a, b, r in b1()[0]]
    else:
        out = [(pid, f, a, b, r) for _c, pid, f, a, b, r in b4()[0]]
    return out


def k1_job(which):
    import scipy.optimize._tstutils as T

    from simroot.baselines import brent_scipy, chandrupatla_lib
    res = {"chandrupatla_g": {}, "brent_g": {}, "mstar": {}}
    for pid, f, a, b, r in k1_problems(which):
        T._f6_cache.clear()
        random.seed(20260921)
        m = m_star(f, r)
        res["mstar"][pid] = m
        inv = 1.0 / m

        def g(x, _f=f, _i=inv):
            v = _f(x)
            return 0.0 if v == 0 else math.copysign(abs(v) ** _i, v)
        for name, meth in (("chandrupatla_g", chandrupatla_lib), ("brent_g", brent_scipy)):
            T._f6_cache.clear()
            random.seed(20260921)
            try:
                x, n, conv = meth(g, a, b, XTOL, MAXFEV)
                x = float(x)
            except Exception:
                x, n, conv = float("nan"), MAXFEV, False
            ok = bool(conv and solved(f, x))
            over = ok and int(n) > 2 * math.ceil(math.log2(abs(b - a) / XTOL)) + 4
            res[name][pid] = [int(n), ok, bool(over)]
    return which, res


def k2_job(j):
    name, which = j
    return name, which, sweep_fourth.evaluate((name, which))[1]


def cost(v):
    return v[0] if v[1] else MAXFEV


def h2h(a, b, keep=None):
    ps = [p for p in a if keep is None or keep(p)]
    w = sum(cost(a[p]) < cost(b[p]) for p in ps)
    l = sum(cost(a[p]) > cost(b[p]) for p in ps)
    return w, len(ps) - w - l, l


def tot(r, keep=None):
    return sum(cost(v) for p, v in r.items() if keep is None or keep(p))


def reference():
    seen = json.load(open("sweep_fourth_seen.json", encoding="utf8"))
    four = json.load(open("sweep_fourth_fourth.json", encoding="utf8"))
    ref = {}
    for k in ("Aprime", "chandrupatla"):
        ref[k] = dict(seen[k])
        ref[k].update(four[k])
    return ref


def report(R):
    ref = reference()
    ap = ref["Aprime"]
    print("C2 - six corpora, A' against each detector (total; A' w/t/l, p); violations in ()")
    for k in K2:
        line = [f"  {k}:"]
        for s in SETS:
            w, t, l = h2h(ap[s], R["K2"][k][s])
            line.append(f"{s} {tot(R['K2'][k][s])}({sum(v[2] for v in R['K2'][k][s].values())})"
                        f" {w}/{t}/{l}")
        allk = sum(tot(R["K2"][k][s]) for s in SETS)
        print(" ".join(line) + f"  | all {allk} vs A' {sum(tot(ap[s]) for s in SETS)}")
    print("\nC1 - corpus 4 (decisive): A' total, control total, A' w/t/l, p -> verdict")
    passed = {}
    for k in K2:
        a, c = ap["D"], R["K2"][k]["D"]
        w, t, l = h2h(a, c)
        passed[k] = tot(a) < tot(c) and w > l and sign_p(w, l) < 0.05
        print(f"  {k}: A' {tot(a)}  {k} {tot(c)}  {w}/{t}/{l} p={sign_p(w, l):.2g}  -> "
              + ("PASSED" if passed[k] else "NOT passed"))
    print("  C1 overall: " + ("PASSED for every control" if all(passed.values())
                              else "NOT passed for " + ", ".join(k for k in K2 if not passed[k])))
    print("\nC4 - violations of the 2x bound, all K2 variants, six corpora:",
          sum(v[2] for k in K2 for s in SETS for v in R["K2"][k][s].values()))
    best = min(K2, key=lambda k: tot(R["K2"][k]["D"]))
    print(f"\nC5 - corpus 4 by class, A' against the best control ({best}): totals, w/t/l")
    for cls in ("D1", "D2", "D3", "D4", "D5", "D6"):
        keep = lambda p, c=cls: p.startswith(c)  # noqa: E731
        w, t, l = h2h(ap["D"], R["K2"][best]["D"], keep)
        print(f"  {cls}: A' {tot(ap['D'], keep)}  {best} {tot(R['K2'][best]['D'], keep)}  {w}/{t}/{l}")
    print("\nC3 - K1: transform with the true order (m* >= 2 or m* <= 2/3 = 'multiple/fractional')")
    for s in K1_SETS:
        k1 = R["K1"][s]
        mult = lambda p, _m=k1["mstar"]: nonsimple(_m[p])  # noqa: E731
        rest = lambda p, _m=k1["mstar"]: not nonsimple(_m[p])  # noqa: E731
        for title, keep in (("multiple/fractional", mult), ("rest", rest)):
            n = sum(1 for p in k1["mstar"] if keep(p))
            if n == 0:
                continue
            w, t, l = h2h({p: v for p, v in ap[s].items() if keep(p)}, k1["chandrupatla_g"])
            print(f"  {s:<4} {title:<20} n={n:3d}  A' {tot(ap[s], keep):5d}  "
                  f"Chandrupatla(g) {tot(k1['chandrupatla_g'], keep):5d}  "
                  f"Brent(g) {tot(k1['brent_g'], keep):5d}  "
                  f"Chandrupatla(f) {tot(ref['chandrupatla'][s], keep):5d}   "
                  f"A' vs Ch(g) {w}/{t}/{l} p={sign_p(w, l):.2g}")
    fam = defaultdict(list)
    for s in K1_SETS:
        for p, m in R["K1"][s]["mstar"].items():
            if nonsimple(m):
                fam[s].append(round(m, 2))
    print("  m* on the multiple/fractional problems (first 8 per corpus):",
          {s: v[:8] for s, v in fam.items()})


def main():
    smoke = "--smoke" in sys.argv
    if "--report" not in sys.argv:
        k2_jobs = [(k, "fourth") for k in K2] + ([] if smoke else [(k, "seen") for k in K2])
        k1_sets = K1_SETS            # the oracle must meet every corpus, also in the smoke
        with Pool(min(6, len(k2_jobs) + len(k1_sets))) as pool:
            k2 = pool.map_async(k2_job, k2_jobs)
            k1 = pool.map_async(k1_job, k1_sets)
            k2, k1 = k2.get(), k1.get()
        R = {"K2": {}, "K1": {}}
        for name, _which, res in k2:
            R["K2"].setdefault(name, {}).update(res)
        for which, res in k1:
            R["K1"][which] = res
        if smoke:
            ap = reference()["Aprime"]
            for k in K2:
                a, c = ap["D"], R["K2"][k]["D"]
                print(f"smoke {k}: corpus 4 total {tot(c)} (A' {tot(a)}), A' w/t/l {h2h(a, c)}, "
                      f"violations {sum(v[2] for v in c.values())}")
            for s in K1_SETS:
                k1c = R["K1"][s]
                print(f"smoke K1 {s}: Chandrupatla(g) {tot(k1c['chandrupatla_g'])} Brent(g) "
                      f"{tot(k1c['brent_g'])}  m* range {min(k1c['mstar'].values()):.3g}.."
                      f"{max(k1c['mstar'].values()):.3g}")
            print("controls smoke done")
            return
        with open("sweep_controls.json", "w", encoding="utf8") as fh:
            json.dump(R, fh)
    with open("sweep_controls.json", encoding="utf8") as fh:
        R = json.load(fh)
    report(R)
    print("controls done")


if __name__ == "__main__":
    main()