1import json, math, time
  2import torch
  3from quadratic_tree import make_random_tree, sequential_postorder, rake_levels, reconstruct, EdgeFactor
  4
  5torch.set_num_threads(min(12, torch.get_num_threads()))
  6DT = torch.float64
  7
  8def clone_case(case, requires=False):
  9    parent, U, h, k, edges = case
 10    def cp(x):
 11        y = x.detach().clone()
 12        y.requires_grad_(requires)
 13        return y
 14    ee = {c: EdgeFactor(cp(e.Hpp), cp(e.Hpc), cp(e.Hcc), cp(e.gp), cp(e.gc), cp(e.k))
 15          for c, e in edges.items()}
 16    return list(parent), cp(U), cp(h), cp(k), ee
 17
 18def sync(device):
 19    if device.type == "cuda":
 20        torch.cuda.synchronize(device)
 21
 22def math_checks():
 23    # Prediction 1: the eliminated parent quadratic correction is exactly quadratic
 24    # in coupling gamma: delta H / gamma^2 is constant.
 25    u, v = 1.7, 0.8
 26    gammas = [0.05, 0.1, 0.2, 0.4, 0.8]
 27    observed = [-(g*g)/(u+v) for g in gammas]
 28    predicted = [-(g*g)/(u+v) for g in gammas]
 29    rel_quad = max(abs(a-b) for a,b in zip(observed,predicted)) / max(abs(x) for x in predicted)
 30    ratios = [observed[i]/(gammas[i]*gammas[i]) for i in range(len(gammas))]
 31    ratio_spread = max(ratios)-min(ratios)
 32
 33    # Prediction 2: damping changes the magnitude by (u+v)/(u+v+epsilon).
 34    gamma = 0.5
 35    epsilons = [0.0, 0.1, 0.3, 0.8]
 36    damp_obs = [abs(-(gamma*gamma)/(u+v+e)) for e in epsilons]
 37    damp_pred = [gamma*gamma/(u+v+e) for e in epsilons]
 38    damp_err = max(abs(a-b) for a,b in zip(damp_obs,damp_pred))
 39    monotone = all(damp_obs[i+1] < damp_obs[i] for i in range(len(damp_obs)-1))
 40
 41    # Prediction 3: a balanced binary tree has logarithmic synchronized rake depth.
 42    depth_rows = []
 43    for n in [3, 7, 15, 31, 63]:
 44        parent = [-1] + [(i-1)//2 for i in range(1,n)]
 45        d = 2
 46        eye = torch.eye(d, dtype=DT)
 47        U = eye.unsqueeze(0).repeat(n,1,1)
 48        h = torch.zeros(n,d,dtype=DT); k = torch.zeros(n,dtype=DT)
 49        edges = {c: EdgeFactor(torch.zeros(d,d,dtype=DT), torch.zeros(d,d,dtype=DT),
 50                                eye, torch.zeros(d,dtype=DT), torch.zeros(d,dtype=DT),
 51                                torch.zeros((),dtype=DT)) for c in range(1,n)}
 52        _,_,_,_,levels,_ = rake_levels(parent,U,h,k,edges)
 53        depth_rows.append({"n": n, "observed_levels": len(levels),
 54                           "predicted_levels": int(math.floor(math.log2(n)))})
 55    depth_ok = all(x["observed_levels"] == x["predicted_levels"] for x in depth_rows)
 56    return {
 57        "quadratic_scaling": {"gammas": gammas, "observed_delta_H": observed,
 58                               "predicted_delta_H": predicted, "max_relative_error": rel_quad,
 59                               "delta_H_over_gamma2_spread": ratio_spread},
 60        "damping_scaling": {"epsilons": epsilons, "observed_magnitude": damp_obs,
 61                             "predicted_magnitude": damp_pred, "max_absolute_error": damp_err,
 62                             "strictly_decreasing": monotone},
 63        "balanced_depth": {"rows": depth_rows, "prediction_confirmed": depth_ok}
 64    }
 65
 66def agreement(device):
 67    case = make_random_tree(31, 4, seed=77, device=device, dtype=DT)
 68    a = clone_case(case, requires=True)
 69    b = clone_case(case, requires=True)
 70    sr = sequential_postorder(*a, eps=1e-6)
 71    rr = rake_levels(*b, eps=1e-6)
 72    loss_s = sr[0].square().sum() + sr[1].square().sum() + sr[2].square()
 73    loss_r = rr[0].square().sum() + rr[1].square().sum() + rr[2].square()
 74    loss_s.backward(); loss_r.backward()
 75    # Inputs are cloned in the same order, so compare all differentiable tensors.
 76    max_out = max((sr[i]-rr[i]).abs().max().item() for i in range(3))
 77    max_grad = max((a[i].grad-b[i].grad).abs().max().item() for i in range(1,4))
 78    z = reconstruct(b[0], rr[5], b[1], b[2], b[3], b[4], rr[3], eps=1e-6)
 79    finite = bool(torch.isfinite(z).all())
 80    return {"max_root_coefficient_difference": max_out, "max_input_gradient_difference": max_grad,
 81            "reconstruction_finite": finite, "levels": len(rr[4])}
 82
 83def benchmark(device):
 84    rows=[]
 85    for n in [127, 511, 1023]:
 86        case = make_random_tree(n, 8, seed=100+n, device=device, dtype=DT)
 87        times=[]
 88        for mode in ["sequential", "rake"]:
 89            # Warmup and three measured runs; no autograd in timing.
 90            fn = sequential_postorder if mode == "sequential" else rake_levels
 91            fn(*case, eps=1e-6)
 92            sync(device); t0=time.perf_counter()
 93            result=None
 94            for _ in range(3): result=fn(*case, eps=1e-6)
 95            sync(device); elapsed=(time.perf_counter()-t0)/3
 96            times.append(elapsed*1000)
 97        rows.append({"n":n, "sequential_ms":times[0], "rake_ms":times[1],
 98                     "rake_over_sequential":times[1]/times[0]})
 99    return rows
100
101def main():
102    try:
103        device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
104        # CUDA errors, including shared-memory exhaustion, must not abort the run.
105        math_result=math_checks()
106        agree=agreement(device)
107        bench=benchmark(device)
108    except Exception as exc:
109        device=torch.device("cpu")
110        math_result=math_checks()
111        agree=agreement(device)
112        bench=benchmark(device)
113        fallback=str(exc)
114    else:
115        fallback=None
116    result={"device":str(device), "math_checks":math_result, "agreement":agree,
117            "benchmark":bench, "cuda_fallback_error":fallback}
118    with open("results.json","w") as f: json.dump(result,f,indent=2)
119    print(json.dumps(result,indent=2))
120
121if __name__ == "__main__": main()