1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4from sklearn.linear_model import LogisticRegression
  5from sklearn.metrics import accuracy_score
  6
  7SEED = 1021
  8rng = np.random.default_rng(SEED)
  9random.seed(SEED)
 10
 11
 12def certificate_sweep():
 13    # A recurrence whose exact tail has the paper's kappa^(l/3) form.
 14    eps = 1e-3
 15    C = 1.0
 16    rows = []
 17    for kappa in [0.05, 0.15, 0.30, 0.50, 0.70, 0.85, 0.95]:
 18        rho = kappa ** (1.0 / 3.0)
 19        # E_l is the remaining geometric effect after depth l.
 20        exact = lambda l: rho ** (l + 1) / (1.0 - rho)
 21        observed_ratio = exact(20) / exact(19)
 22        predicted_ratio = rho
 23        observed_l = next(l for l in range(10000) if exact(l) <= eps)
 24        # Formula in the proposal omits the geometric prefactor; use it literally.
 25        predicted_l = math.ceil(3 * math.log(C / eps) / math.log(1.0 / kappa) - 1)
 26        rows.append({"kappa": kappa, "predicted_tail_ratio": predicted_ratio,
 27                     "observed_tail_ratio": observed_ratio,
 28                     "predicted_depth": predicted_l, "observed_depth": observed_l})
 29    return rows
 30
 31
 32def boundary_sweep():
 33    # x_l=(gamma*Lambda)^l: predicted transition at gamma*Lambda=1.
 34    rows = []
 35    for q in [0.70, 0.90, 0.99, 1.00, 1.01, 1.10, 1.30]:
 36        x = q ** np.arange(30, dtype=float)
 37        # Estimate asymptotic multiplier from consecutive magnitudes.
 38        estimate = float(np.median(x[1:] / x[:-1]))
 39        rows.append({"gamma_times_Lambda": q, "observed_multiplier": estimate,
 40                     "predicted_regime": "decay" if q < 1 else ("critical" if q == 1 else "growth")})
 41    return rows
 42
 43
 44def make_sbm(n=600, p_in=0.045, p_out=0.006):
 45    y = np.repeat([0, 1], n // 2)
 46    A = np.zeros((n, n), dtype=np.float32)
 47    for i in range(n):
 48        same = y == y[i]
 49        prob = np.where(same, p_in, p_out)
 50        draw = rng.random(n) < prob
 51        draw[i] = False
 52        A[i, draw] = 1.0
 53    A = np.maximum(A, A.T)
 54    deg = A.sum(1)
 55    # symmetric normalized adjacency with self loops
 56    B = A + np.eye(n, dtype=np.float32)
 57    d = B.sum(1)
 58    P = B / np.sqrt(d[:, None] * d[None, :])
 59    X = np.zeros((n, 4), dtype=np.float32)
 60    X[:, 0] = (y * 2 - 1) + rng.normal(0, 1.0, n)
 61    X[:, 1] = (1 - 2 * y) + rng.normal(0, 1.0, n)
 62    X[:, 2:] = rng.normal(0, 1.0, (n, 2))
 63    return P.astype(np.float32), X, y
 64
 65
 66def gnn_experiment():
 67    P, X, y = make_sbm()
 68    n = len(y)
 69    perm = rng.permutation(n)
 70    train, val, test = perm[:240], perm[240:420], perm[420:]
 71    # Gamma is an attenuation estimate; Delta is the normalized propagation bound.
 72    gamma, Delta, eps, tau, max_depth = 0.72, 1.0, 0.02, 0.012, 8
 73    kappa = gamma * gamma * Delta
 74    # Calibrate C from a shallow-vs-deep validation error gap, as proposed.
 75    H = [X]
 76    for _ in range(max_depth):
 77        H.append(P @ H[-1])
 78    models = []
 79    val_errors = []
 80    for l in range(max_depth + 1):
 81        clf = LogisticRegression(C=2.0, max_iter=300, random_state=SEED)
 82        clf.fit(H[l][train], y[train])
 83        models.append(clf)
 84        val_errors.append(1 - accuracy_score(y[val], clf.predict(H[l][val])))
 85    gap = abs(val_errors[2] - val_errors[4])
 86    C = max(gap, 0.05)
 87    cert_depth = max(0, math.ceil(3 * math.log(C / eps) / math.log(1 / kappa) - 1))
 88    fixed_depth = min(max_depth, max(2, cert_depth))
 89    fixed_pred = models[fixed_depth].predict(H[fixed_depth][test])
 90    # Node-wise halting: two consecutive small logit changes plus certificate.
 91    logits = [m.predict_proba(H[l])[:, 1] for l, m in enumerate(models)]
 92    stop = np.full(n, max_depth, dtype=int)
 93    active = np.ones(n, dtype=bool)
 94    small_run = np.zeros(n, dtype=int)
 95    active_counts = []
 96    for l in range(1, max_depth + 1):
 97        d = np.abs(logits[l] - logits[l - 1])
 98        small_run = np.where(d < tau, small_run + 1, 0)
 99        tail = C * kappa ** ((l + 1) / 3)
100        halt = active & (small_run >= 2) & (tail < eps)
101        stop[halt] = l
102        active[halt] = False
103        active_counts.append(int(active.sum()))
104    adaptive_prob = np.array([logits[stop[i]][i] for i in range(n)])
105    adaptive_pred = (adaptive_prob >= 0.5).astype(int)
106    # Full batched propagation does not realize node savings itself; this is the
107    # equivalent active-root message count, useful for measuring the policy.
108    fixed_work = len(test) * fixed_depth
109    adaptive_work = int(stop[test].sum())
110    return {
111        "kappa": kappa, "calibrated_C": C, "certificate_depth": cert_depth,
112        "fixed_depth": fixed_depth, "fixed_test_accuracy": float(accuracy_score(y[test], fixed_pred)),
113        "adaptive_test_accuracy": float(accuracy_score(y[test], adaptive_pred[test])),
114        "adaptive_mean_depth": float(stop[test].mean()),
115        "work_ratio_adaptive_over_fixed": adaptive_work / fixed_work,
116        "active_nodes_after_each_layer": active_counts,
117        "validation_errors_by_depth": val_errors,
118    }
119
120
121def main():
122    out = {"seed": SEED, "certificate_sweep": certificate_sweep(),
123           "boundary_sweep": boundary_sweep(), "gnn": gnn_experiment()}
124    Path("ks_results.json").write_text(json.dumps(out, indent=2))
125    print(json.dumps(out, indent=2))
126
127if __name__ == "__main__":
128    main()