1import json, math, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED = 281
  8np.random.seed(SEED)
  9torch.manual_seed(SEED)
 10torch.set_num_threads(4)
 11DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 12
 13
 14def hankel(values, k=3):
 15    # values has length 2k-1 and gradients are retained.
 16    return torch.stack([torch.stack([values[i + j] for j in range(k)]) for i in range(k)])
 17
 18
 19def hankel_penalty(values, k=3, eps=1e-4):
 20    H = hankel(values, k)
 21    eig = torch.linalg.eigvalsh(H)
 22    # Smooth enough almost everywhere and directly measures PSD violations.
 23    violation = torch.relu(eps - eig)
 24    return (violation ** 2).mean(), eig.detach().cpu().numpy()
 25
 26
 27def pairwise_logconvex_penalty(values, eps=1e-6):
 28    # For a positive log-convex sequence: f_r^2 <= f_{r-1} f_{r+1}.
 29    curv = torch.log(values[2:] + eps) - 2 * torch.log(values[1:-1] + eps) + torch.log(values[:-2] + eps)
 30    return torch.relu(-curv).pow(2).mean()
 31
 32
 33def analytic_check():
 34    # f(a)=sum_s w_s exp(-s a), so H is a Gram matrix with vectors exp(-s*i*Delta).
 35    a, delta = 0.7, 0.35
 36    s = torch.tensor([0.15, 0.8, 1.7], dtype=torch.float64)
 37    w = torch.tensor([0.7, 0.2, 0.1], dtype=torch.float64)
 38    vals = torch.stack([(w * torch.exp(-s * (a + r * delta))).sum() for r in range(5)])
 39    H = torch.stack([torch.stack([vals[i+j] for j in range(3)]) for i in range(3)])
 40    eig = torch.linalg.eigvalsh(H).numpy()
 41    # Deliberately perturb one middle sample to demonstrate the violation detector.
 42    bad = vals.clone(); bad[2] *= 1.8
 43    badH = torch.stack([torch.stack([bad[i+j] for j in range(3)]) for i in range(3)])
 44    bad_eig = torch.linalg.eigvalsh(badH).numpy()
 45    return {"valid_values": vals.numpy().tolist(), "valid_eigenvalues": eig.tolist(),
 46            "valid_min_eigenvalue": float(eig.min()), "perturbed_eigenvalues": bad_eig.tolist(),
 47            "perturbed_min_eigenvalue": float(bad_eig.min()),
 48            "claim_observed": bool(eig.min() >= -1e-10 and bad_eig.min() < -1e-6)}
 49
 50
 51class ResponseNet(nn.Module):
 52    def __init__(self):
 53        super().__init__()
 54        self.net = nn.Sequential(nn.Linear(1, 24), nn.Tanh(), nn.Linear(24, 24), nn.Tanh(), nn.Linear(24, 1))
 55    def forward(self, a):
 56        a = a.reshape(-1, 1)
 57        return torch.nn.functional.softplus(self.net(a).squeeze(-1)) + 1e-4
 58
 59
 60def train(kind, train_a, train_y, test_a, test_y, steps=700):
 61    torch.manual_seed(SEED + {"none":0, "pairwise":1, "hankel":2}[kind])
 62    model = ResponseNet().to(DEVICE)
 63    opt = torch.optim.Adam(model.parameters(), lr=0.012)
 64    # Five nearby evaluations are needed for K=3.
 65    delta = 0.075
 66    t0 = time.perf_counter()
 67    last_eigs = None
 68    for step in range(steps):
 69        opt.zero_grad(set_to_none=True)
 70        pred = model(train_a)
 71        loss = ((pred - train_y) ** 2).mean()
 72        if kind != "none":
 73            anchors = train_a[:, None]
 74            nearby = model((anchors[:, 0, None] + delta * torch.arange(5, device=DEVICE)[None, :]).reshape(-1)).reshape(-1, 5)
 75            # Average penalty over each anchor; this evaluates the same response head nearby.
 76            if kind == "hankel":
 77                H = torch.stack([nearby[:, i:i+3] for i in range(3)], dim=1)
 78                eig_batch = torch.linalg.eigvalsh(H)
 79                reg = torch.relu(1e-4 - eig_batch).pow(2).mean()
 80                last_eigs = eig_batch.detach().cpu().numpy()
 81            else:
 82                curv = torch.log(nearby[:, 2:] + 1e-6) - 2 * torch.log(nearby[:, 1:-1] + 1e-6) + torch.log(nearby[:, :-2] + 1e-6)
 83                reg = torch.relu(-curv).pow(2).mean()
 84            loss = loss + 3.0 * reg
 85        loss.backward(); opt.step()
 86    elapsed = time.perf_counter() - t0
 87    with torch.no_grad():
 88        pred_test = model(test_a)
 89        # Dense-grid response for oscillation and Hankel violations.
 90        grid = torch.linspace(0.15, 3.0, 121, device=DEVICE)
 91        gv = model(grid)
 92        second = gv[2:] - 2*gv[1:-1] + gv[:-2]
 93        oscillation = float(torch.abs(second).mean().cpu())
 94        dense_eigs = []
 95        for i in range(len(grid)-4):
 96            _, e = hankel_penalty(gv[i:i+5], 3, 1e-4); dense_eigs.append(e)
 97        dense_eigs = np.asarray(dense_eigs)
 98    return {"test_mse": float(((pred_test-test_y)**2).mean().cpu()),
 99            "train_mse": float(((model(train_a)-train_y)**2).mean().cpu()),
100            "mean_abs_second_difference": oscillation,
101            "min_dense_hankel_eigenvalue": float(dense_eigs.min()),
102            "fraction_dense_psd_violations": float((dense_eigs < -1e-7).any(axis=1).mean()),
103            "seconds": elapsed, "relative_overhead": None}
104
105
106def main():
107    check = analytic_check()
108    # Sparse, mildly noisy observations make the extrapolating shape meaningful.
109    true_s, true_w = np.array([0.2, 0.9, 2.0]), np.array([0.55, 0.3, 0.15])
110    def f(x): return (true_w[None,:] * np.exp(-x[:,None]*true_s[None,:])).sum(1)
111    train_x = np.linspace(0.25, 2.35, 18).astype("float32")
112    test_x = np.linspace(0.15, 3.0, 121).astype("float32")
113    rng = np.random.RandomState(SEED)
114    y = (f(train_x) + rng.normal(0, 0.012, len(train_x))).clip(1e-3).astype("float32")
115    ty = f(test_x).astype("float32")
116    ta, va = torch.tensor(train_x, device=DEVICE), torch.tensor(y, device=DEVICE)
117    te, vy = torch.tensor(test_x, device=DEVICE), torch.tensor(ty, device=DEVICE)
118    results = {}
119    for kind in ["none", "pairwise", "hankel"]:
120        results[kind] = train(kind, ta, va, te, vy)
121    base = results["none"]["seconds"]
122    for r in results.values(): r["relative_overhead"] = r["seconds"] / base
123    out = {"device": DEVICE, "analytic_check": check, "results": results,
124           "setup": {"steps": 700, "train_points": 18, "test_points": 121, "noise_std": 0.012, "K": 3}}
125    Path("results.json").write_text(json.dumps(out, indent=2))
126    print(json.dumps(out, indent=2))
127
128if __name__ == "__main__":
129    try:
130        main()
131    except (RuntimeError, torch.cuda.OutOfMemoryError) as e:
132        if DEVICE == "cuda":
133            print("CUDA failed; rerun with CPU", repr(e))
134            torch.cuda.empty_cache()
135            DEVICE = "cpu"
136            main()
137        else: raise