scripts/forge.py

#!/usr/bin/env python3
"""Forge a REAL trained surrogate for szl-formulas.
Kernel = ground truth. Surrogate = a 21-class formula IDENTIFIER: given the
OBSERVABLE trace of a single formula call (argument-shape observables + the
formula's own output observables, produced by the kernel's REAL evaluators),
predict WHICH of the 21 canonical formulas produced it. The registry count (21)
is read from the kernel itself (registry_count()); we assert it == 21 and say so.
Labels come from the kernel: each sample is generated BY calling the real
REGISTRY[name] evaluator, so the label is definitionally the producing formula.
A sample of traces is re-audited by replaying the kernel call and asserting the
recorded output matches. Seeded, receipted, reproducible."""
import json, os, random, sys, time, hashlib, platform, math
_here = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if os.path.isdir(os.path.join(_here, "build", "torch-universal")):
    sys.path.insert(0, os.path.join(_here, "build", "torch-universal"))  # in-repo run
else:
    sys.path.insert(0, "/tmp/kernel-probe/szl-formulas/build/torch-universal")  # forge-dev run
import szl_formulas as fx
import numpy as np
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, recall_score
import joblib

SEED = 20260721
random.seed(SEED); np.random.seed(SEED)
T0 = time.time()

REG_COUNT = fx.registry_count()
assert REG_COUNT == 21, f"expected 21 canonical formulas, kernel reports {REG_COUNT}"
CLASSES = sorted(fx.REGISTRY.keys())
assert len(CLASSES) == 21

# ---- per-formula argument generators (produce VALID args for the REAL kernel) ----
def _simplex(k):
    v = [random.random() + 1e-3 for _ in range(k)]
    s = sum(v); return [x / s for x in v]

def _axes(k):
    return [random.uniform(0.05, 1.0) for _ in range(k)]

def gen_args(name):
    if name == "lambda_aggregate":
        k = random.randint(2, 6); ax = _axes(k)
        if random.random() < 0.5:
            return [ax, _simplex(k)]
        return [ax]
    if name == "lambda_homogeneous":
        k = random.randint(2, 5)
        return [random.uniform(0.0, 3.0), _axes(k)]
    if name == "lambda_bounded":
        return [_axes(random.randint(2, 6))]
    if name == "pac_bayes_mcallester":
        return [random.uniform(0.0, 0.5), random.uniform(0.0, 5.0),
                random.randint(10, 5000), random.uniform(0.01, 0.5)]
    if name == "bekenstein_cascade":
        return [random.uniform(0.0, 10.0), random.uniform(0.0, 1e-10)]
    if name == "reidemeister_invariant":
        alphabet = "abAB"
        s = "".join(random.choice(alphabet) for _ in range(random.randint(2, 8)))
        return [s, random.choice(["R1", "R2", "R3"])]
    if name == "khipu_merkle_root":
        n = random.randint(1, 5)
        return [[{"decision_id": f"d{i}", "value": random.randint(0, 100)} for i in range(n)]]
    if name == "dsse_envelope":
        blen = random.randint(1, 24)
        return [bytes(random.randint(0, 255) for _ in range(blen)), f"key-{random.randint(0,9)}"]
    if name == "gleason_quantum_lambda":
        n = random.randint(2, 4)
        return [[[random.uniform(-1, 1) for _ in range(n)] for _ in range(n)]]
    if name == "hoeffding_tail":
        return [random.uniform(0.0, 1.0), random.randint(1, 5000)]
    if name == "pinsker_kl_bound":
        k = random.randint(2, 5); return [_simplex(k), _simplex(k)]
    if name == "fisher_rao_distance":
        k = random.randint(2, 5); return [_simplex(k), _simplex(k)]
    if name == "bohr_complementarity_floor":
        return [random.uniform(0.0, 1.0), random.uniform(0.0, 1.0)]
    if name == "kochen_specker_18vector_witness":
        rows = random.randint(2, 7)
        return [[[1 if random.random() < 0.4 else 0 for _ in range(random.randint(2, 5))] for _ in range(rows)]]
    if name == "two_witness_ks18_soundness":
        return [random.random() < 0.5, random.random() < 0.5]
    if name == "shor_codeword_distance":
        rows = random.randint(1, 5); cols = random.randint(2, 6)
        return [[[random.randint(0, 1) for _ in range(cols)] for _ in range(rows)]]
    if name == "css_ingress_verify":
        payload = bytes(random.randint(0, 255) for _ in range(random.randint(0, 20)))
        env = fx.dsse_envelope(payload, "signer")
        from hashlib import sha256
        commit = sha256(bytes.fromhex(env["payload"]) if env["payload"] else b"").digest()
        # half the time provide the matching root, half a random one
        css_root = commit if random.random() < 0.5 else bytes(random.randint(0, 255) for _ in range(4))
        return [env, css_root]
    if name == "kitaev_surface_correct":
        return [[random.randint(0, 3) for _ in range(random.randint(2, 8))]]
    if name == "reed_solomon_singleton":
        n = random.randint(2, 255); k = random.randint(1, n)
        return [n, k]
    if name == "madhava_series":
        return [random.uniform(-1.0, 1.0), random.randint(1, 60)]
    if name == "schur_concave_lambda_two_axis":
        return [random.uniform(0.0, 1.0), random.uniform(0.0, 1.0)]
    raise KeyError(name)

# ---- trace observables (features) — derived ONLY from args + REAL kernel output ----
def _flatten_num(x):
    """Yield numeric leaves from an arbitrarily nested arg/output structure."""
    if isinstance(x, bool):
        yield float(x)
    elif isinstance(x, (int, float)):
        yield float(x)
    elif isinstance(x, (bytes, bytearray)):
        for b in x: yield float(b)
    elif isinstance(x, str):
        yield float(len(x))
    elif isinstance(x, dict):
        for v in x.values():
            yield from _flatten_num(v)
    elif isinstance(x, (list, tuple)):
        for v in x:
            yield from _flatten_num(v)

def _shape_stats(obj):
    """Structural stats for a single object: (n_items, depth, is_str, is_bytes, is_dict, is_list, is_bool)."""
    def depth(o):
        if isinstance(o, (list, tuple)) and o:
            return 1 + max(depth(v) for v in o)
        if isinstance(o, dict) and o:
            return 1 + max(depth(v) for v in o.values())
        return 0
    n = 0
    if isinstance(obj, (list, tuple, dict, str, bytes, bytearray)):
        n = len(obj)
    return [float(n), float(depth(obj)),
            float(isinstance(obj, str)), float(isinstance(obj, (bytes, bytearray))),
            float(isinstance(obj, dict)), float(isinstance(obj, (list, tuple))),
            float(isinstance(obj, bool))]

def observe(args, out):
    """Produce a FIXED-LENGTH numeric feature vector from the call trace."""
    arg_leaves = list(_flatten_num(args))
    out_leaves = list(_flatten_num(out))
    def agg(leaves):
        if not leaves:
            return [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        a = np.array(leaves, dtype=np.float64)
        finite = a[np.isfinite(a)]
        if finite.size == 0:
            finite = np.array([0.0])
        return [float(finite.size), float(finite.mean()), float(finite.std()),
                float(finite.min()), float(finite.max()),
                float(np.mean(np.abs(finite) < 1e-6))]
    feats = []
    feats += [float(len(args))]                       # arity
    feats += _shape_stats(args)                       # arg container shape
    feats += agg(arg_leaves)                          # arg numeric aggregate
    feats += _shape_stats(out)                        # output shape
    feats += agg(out_leaves)                          # output numeric aggregate
    # output-type one-hots (the kernel's real return type is a strong signal)
    feats += [float(isinstance(out, bool)),
              float(isinstance(out, int) and not isinstance(out, bool)),
              float(isinstance(out, float)),
              float(isinstance(out, (bytes, bytearray))),
              float(isinstance(out, dict)),
              float(isinstance(out, (list, tuple)))]
    # magnitude of a scalar output (log-abs), else 0
    if isinstance(out, (int, float)) and not isinstance(out, bool):
        feats += [math.log1p(abs(float(out)))]
    else:
        feats += [0.0]
    return feats

FEATURE_NAMES = (
    ["arity"]
    + [f"argshape_{s}" for s in ["n", "depth", "is_str", "is_bytes", "is_dict", "is_list", "is_bool"]]
    + [f"argnum_{s}" for s in ["count", "mean", "std", "min", "max", "frac_zero"]]
    + [f"outshape_{s}" for s in ["n", "depth", "is_str", "is_bytes", "is_dict", "is_list", "is_bool"]]
    + [f"outnum_{s}" for s in ["count", "mean", "std", "min", "max", "frac_zero"]]
    + ["out_is_bool", "out_is_int", "out_is_float", "out_is_bytes", "out_is_dict", "out_is_list"]
    + ["out_scalar_logabs"]
)

# ---- generate ----
PER_CLASS = 900
X, y, audited = [], [], 0
sample_bank = {}  # name -> list of (args, recorded_out) for audit
for name in CLASSES:
    for _ in range(PER_CLASS):
        args = gen_args(name)
        out = fx.REGISTRY[name](*args)   # REAL kernel evaluator == ground truth label
        X.append(observe(args, out)); y.append(name)
        if len(sample_bank.get(name, [])) < 6:
            sample_bank.setdefault(name, []).append((args, out))

# ---- kernel-replay audit: re-call the real evaluator, assert output agrees ----
def _eq(a, b):
    if isinstance(a, float) or isinstance(b, float):
        try:
            return abs(float(a) - float(b)) <= 1e-9 * max(1.0, abs(float(a)), abs(float(b)))
        except Exception:
            return a == b
    return a == b

for name, samples in sample_bank.items():
    for args, recorded in samples:
        replay = fx.REGISTRY[name](*args)
        assert _eq(replay, recorded), f"kernel replay disagreement for {name}: {replay!r} != {recorded!r}"
        audited += 1

X = np.array(X, dtype=np.float64); y = np.array(y)
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.2, random_state=SEED, stratify=y)
clf = HistGradientBoostingClassifier(random_state=SEED, max_iter=300, early_stopping=True)
clf.fit(Xtr, ytr)
pred = clf.predict(Xte)
acc = accuracy_score(yte, pred)
per_class_recall = {c: float(recall_score(yte == c, pred == c)) for c in CLASSES}
# honest blind-spot detection: classes whose recall is materially below the mean
mean_recall = float(np.mean(list(per_class_recall.values())))
weak = {c: round(r, 4) for c, r in per_class_recall.items() if r < 0.90}

out_dir = os.path.dirname(os.path.abspath(__file__))
joblib.dump(clf, f"{out_dir}/model.joblib")
model_sha = hashlib.sha256(open(f"{out_dir}/model.joblib", "rb").read()).hexdigest()
receipt = {
  "artifact": "SZLHOLDINGS/szl-formulas surrogate v1",
  "role": "21-class formula identifier from call-trace observables — kernel remains ground truth",
  "generator": {"script": "scripts/forge.py", "seed": SEED, "kernel_version": fx.__version__,
                 "kernel_labelled": True, "kernel_registry_count": REG_COUNT,
                 "kernel_replay_audited_samples": audited,
                 "label_source": "each sample is produced by calling the REAL REGISTRY[name] evaluator; label == producing formula"},
  "data": {"rows": int(len(y)), "n_classes": len(CLASSES), "classes": CLASSES,
            "per_class_samples": PER_CLASS,
            "class_counts": {c: int((y == c).sum()) for c in CLASSES},
            "split": "80/20 stratified", "features": FEATURE_NAMES,
            "feature_policy": "trace observables ONLY: argument-shape + kernel-output shape/type/aggregate stats; no formula name leaked into features"},
  "model": {"type": "sklearn.HistGradientBoostingClassifier",
             "params": {"max_iter": 300, "early_stopping": True, "random_state": SEED},
             "file": "model.joblib", "sha256": model_sha},
  "metrics_MEASURED": {"test_accuracy_21_class": round(float(acc), 4),
                        "mean_per_class_recall": round(mean_recall, 4),
                        "per_class_recall": {k: round(v, 4) for k, v in per_class_recall.items()},
                        "blind_spots": {"policy": "classes with held-out recall < 0.90 reported honestly",
                          "classes": weak,
                          "statement": "formulas that share an argument+output OBSERVABLE signature (e.g. several bool-returning axiom checks over similar inputs) are confusable from trace shape alone; the kernel's evaluators remain authoritative for identity"}},
  "environment": {"python": platform.python_version(), "sklearn": __import__("sklearn").__version__,
                   "numpy": np.__version__, "host": "replit 2-vCPU container", "wall_seconds": round(time.time()-T0, 1)},
  "honesty": "Every number above is MEASURED by this run. The surrogate identifies formulas from trace shape; it never replaces the kernel's real evaluators. Λ untouched = Conjecture 1.",
  "trained_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
}
with open(f"{out_dir}/TRAINING_RECEIPT.json", "w") as f: json.dump(receipt, f, indent=2)
print(json.dumps({"acc": receipt["metrics_MEASURED"]["test_accuracy_21_class"],
                  "mean_recall": mean_recall, "weak": weak}, indent=2))
print(f"rows={len(y)} classes={len(CLASSES)} audited={audited} wall={receipt['environment']['wall_seconds']}s")