1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6SEED = 1234
  7K = 15
  8STEPS = 1800
  9LR = 0.03
 10N_COL = 96
 11DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 12
 13def seed_all(seed):
 14    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 15    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 16
 17# Formal-series recurrence for u' = 1 + u^2, u(0)=0.
 18# (n+1)c[n+1] = 1_{n=0} + sum_{i=0}^n c[i]c[n-i].
 19def formal_support(K):
 20    c = np.zeros(K + 1, dtype=np.float64)
 21    for n in range(K):
 22        rhs = (1.0 if n == 0 else 0.0) + sum(c[i] * c[n-i] for i in range(n + 1))
 23        c[n + 1] = rhs / (n + 1)
 24    return [i for i, x in enumerate(c) if abs(x) > 1e-14], c
 25
 26def polynomial_value(coef, t, exponents):
 27    # Horner is unnecessary at this small K and this makes the exponent mask explicit.
 28    powers = torch.stack([t ** int(i) for i in exponents], dim=1)
 29    return (powers * coef.unsqueeze(0)).sum(dim=1)
 30
 31def train(exponents, seed):
 32    seed_all(seed)
 33    t = torch.linspace(0.0, 1.0, N_COL, device=DEVICE)
 34    # Coefficients are the minimal coefficient-output layer described in the idea.
 35    coef = torch.nn.Parameter(0.01 * torch.randn(len(exponents), device=DEVICE))
 36    opt = torch.optim.Adam([coef], lr=LR)
 37    history = []
 38    ex = [int(i) for i in exponents]
 39    for step in range(STEPS):
 40        opt.zero_grad(set_to_none=True)
 41        u = polynomial_value(coef, t, ex)
 42        du = torch.zeros_like(t)
 43        for j, i in enumerate(ex):
 44            if i > 0:
 45                du = du + coef[j] * i * t ** (i - 1)
 46        residual = du - 1.0 - u * u
 47        loss = torch.mean(residual ** 2) + 10.0 * coef[ex.index(0)] ** 2 if 0 in ex else torch.mean(residual ** 2)
 48        loss.backward(); opt.step()
 49        if step in (0, 99, 299, 599, 1199, STEPS - 1): history.append(float(loss.detach().cpu()))
 50    with torch.no_grad():
 51        te = torch.linspace(0.0, 1.0, 401, device=DEVICE)
 52        ue = polynomial_value(coef, te, ex)
 53        exact = torch.tan(te)
 54        rel = torch.linalg.vector_norm(ue - exact) / torch.linalg.vector_norm(exact)
 55        u = polynomial_value(coef, t, ex)
 56        du = torch.zeros_like(t)
 57        for j, i in enumerate(ex):
 58            if i > 0: du += coef[j] * i * t ** (i - 1)
 59        residual = du - 1.0 - u * u
 60        return {"final_loss": float((torch.mean(residual**2)).cpu()),
 61                "relative_error": float(rel.cpu()),
 62                "coeff_l2": float(torch.linalg.vector_norm(coef).cpu()),
 63                "history": history, "parameters": len(ex)}
 64
 65def main():
 66    global DEVICE
 67    try:
 68        seed_all(SEED)
 69        _ = torch.zeros(1, device=DEVICE)
 70    except Exception:
 71        DEVICE = "cpu"
 72        seed_all(SEED)
 73    support, coeffs = formal_support(K)
 74    # Independent direct check: all even coefficients are forced to zero by recurrence.
 75    even_abs = max(abs(coeffs[i]) for i in range(0, K + 1, 2))
 76    odd_nonzero = [i for i in range(1, K + 1, 2) if abs(coeffs[i]) > 1e-14]
 77    results = {"device": DEVICE, "K": K, "collocation": N_COL, "steps": STEPS,
 78               "formal_support": support, "expected_odd_support": odd_nonzero,
 79               "max_forced_even_coefficient": even_abs, "math_check_pass": bool(even_abs < 1e-14)}
 80    runs = {"dense": [], "restricted": []}
 81    for seed in (11, 22, 33):
 82        for name, ex in (("dense", list(range(K + 1))), ("restricted", support)):
 83            try:
 84                runs[name].append(train(ex, seed))
 85            except Exception as e:
 86                if DEVICE != "cpu":
 87                    DEVICE = "cpu"
 88                    runs[name].append(train(ex, seed))
 89                else:
 90                    raise
 91    results["runs"] = runs
 92    for name in runs:
 93        for metric in ("final_loss", "relative_error", "parameters"):
 94            vals = [x[metric] for x in runs[name]]
 95            results.setdefault(name, {})[metric + "_mean"] = float(np.mean(vals))
 96            results[name][metric + "_std"] = float(np.std(vals))
 97    Path("results.json").write_text(json.dumps(results, indent=2))
 98    print(json.dumps(results, indent=2))
 99
100if __name__ == "__main__": main()