1import json, math, random, time
 2import numpy as np
 3import torch
 4from fractional_optimizer import FractionalMemory, LateTimeOrderEstimator, AdamW
 5
 6SEED = 1729
 7
 8def seed_all(seed=SEED):
 9    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
10    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
11
12
13def get_device():
14    try:
15        d = torch.device("cuda" if torch.cuda.is_available() else "cpu")
16        if d.type == "cuda": torch.zeros(1, device=d).sum().item()
17        return d
18    except Exception:
19        return torch.device("cpu")
20
21
22def verify_power_law():
23    # M(t)=C t^(-m alpha), so the estimator should recover alpha for m=1.
24    true_alpha, m, rho = 0.68, 1, 1.8
25    est = LateTimeOrderEstimator(alpha0=.5, amin=.2, amax=.95,
26                                 beta=0.0, warmup=10, lag=10,
27                                 check_every=1, monotone_checks=2)
28    vals = []
29    for k in range(1, 401):
30        # Positive, monotonically decreasing late-time signal with tiny noise.
31        t = float(k + 1)
32        loss = t ** (-m * true_alpha) * (1 + 0.0002 * math.sin(k))
33        vals.append(est.update(k, loss))
34    raw = -math.log((401.0/391.0)**(-m*true_alpha)) / (m*math.log(401.0/391.0))
35    return {"true_alpha": true_alpha, "raw_asymptotic_estimate": raw,
36            "final_estimate": vals[-1], "absolute_error": abs(vals[-1]-true_alpha),
37            "passed": abs(vals[-1]-true_alpha) < .06}
38
39
40def make_problem(n=24, d=20, device=None):
41    g = torch.Generator(device="cpu").manual_seed(SEED)
42    X = torch.randn(n, d, generator=g).to(device)
43    true_w = torch.randn(d, 1, generator=g).to(device)
44    y = X @ true_w + .05 * torch.randn(n, 1, generator=g).to(device)
45    return X, y
46
47
48def train(kind, device, steps=180):
49    seed_all(SEED)
50    X, y = make_problem(device=device)
51    w = torch.nn.Parameter(torch.zeros(20, 1, device=device))
52    if kind == "adamw":
53        opt = AdamW([w], lr=.08, weight_decay=0.0)
54    elif kind.startswith("fixed"):
55        alpha = float(kind.split("_")[1])
56        opt = FractionalMemory([w], lr=.08, alpha=alpha, history=32)
57    else:
58        # The estimator intentionally starts conservatively and adapts only
59        # after warmup; loss itself is the positive scalar observation.
60        e = LateTimeOrderEstimator(alpha0=.5, amin=.2, amax=.95, beta=.8,
61                                   warmup=25, lag=12, check_every=3,
62                                   monotone_checks=2)
63        opt = FractionalMemory([w], lr=.08, alpha=.5, history=32,
64                               adaptive=True, estimator=e)
65    losses, alphas = [], []
66    start = time.perf_counter()
67    for step in range(1, steps + 1):
68        pred = X @ w
69        loss = ((pred-y)**2).mean()
70        loss.backward()
71        if kind == "adamw":
72            opt.step()
73        else:
74            opt.step(loss_value=float(loss.detach()))
75        opt.zero_grad(set_to_none=True)
76        losses.append(float(loss.detach()))
77        if kind == "adaptive": alphas.append(opt.estimator.alpha)
78    elapsed = time.perf_counter() - start
79    tail = float(np.mean(losses[-20:]))
80    return {"final_loss": losses[-1], "tail20_loss": tail,
81            "loss_at_40": losses[39], "seconds": elapsed,
82            "alpha_final": (alphas[-1] if alphas else None),
83            "alpha_min": (min(alphas) if alphas else None),
84            "alpha_max": (max(alphas) if alphas else None)}
85
86
87def main():
88    device = get_device()
89    check = verify_power_law()
90    results = {"device": str(device), "math_check": check, "runs": {}}
91    for kind in ["adamw", "fixed_0.3", "fixed_0.5", "fixed_0.7", "fixed_0.9", "adaptive"]:
92        results["runs"][kind] = train(kind, device)
93    print(json.dumps(results, indent=2))
94
95if __name__ == "__main__": main()