1import json, time, math, random
  2import numpy as np
  3import torch
  4
  5SEED = 429
  6random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  7torch.set_num_threads(min(8, torch.get_num_threads()))
  8try:
  9    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 10except Exception:
 11    device = torch.device("cpu")
 12
 13# Finite Caratheodory-kernel quantities. The block K is assembled for scalar output.
 14def ck_quantities(A, B, C, lambdas, eps=1e-5):
 15    n = A.shape[0]
 16    dtype = torch.complex64 if A.dtype == torch.float32 else torch.complex128
 17    Ac = A.to(dtype); Bc = B.to(dtype); Cc = C.to(dtype)
 18    eye = torch.eye(n, dtype=dtype, device=A.device)
 19    phis = []
 20    for z in lambdas:
 21        z = torch.as_tensor(z, dtype=dtype, device=A.device)
 22        X = torch.linalg.solve(eye - z * Ac, Bc[:, None])
 23        phis.append((1.0 + (Cc[None, :] @ X).squeeze()).reshape(()))
 24    phis = torch.stack(phis)
 25    H = phis.real
 26    # For scalar output K is M by M; this is the exact block formula.
 27    z = torch.as_tensor(lambdas, dtype=dtype, device=A.device)
 28    K = (phis[:, None] + phis.conj()[None, :]) / (2 * (1 - z[:, None] * z.conj()[None, :]))
 29    K = (K + K.conj().T) / 2
 30    hmin = H.min()
 31    kmin = torch.linalg.eigvalsh(K).real.min()
 32    penalty = torch.relu(eps - H).square().mean() + torch.relu(eps - kmin).square()
 33    return penalty, hmin, kmin
 34
 35def make_lambdas(M=12, rho=.93):
 36    # fixed points, with a modest radial spread as suggested in the idea
 37    return np.array([(rho + .035*np.sin(2*np.pi*j/M))*np.exp(2j*np.pi*j/M) for j in range(M)])
 38
 39# Cheap direct verification: a passive constant Phi has PSD H and PSD kernel;
 40# a deliberately non-passive constant Phi violates both, and gradient descent repairs it.
 41def math_check():
 42    lam = make_lambdas(12)
 43    A = torch.zeros(1); B = torch.ones(1); Cgood = torch.ones(1)
 44    _, hg, kg = ck_quantities(A, B, Cgood, lam)
 45    Cbad = torch.tensor([-1.4], requires_grad=True)
 46    before = ck_quantities(A, B, Cbad, lam)
 47    opt = torch.optim.SGD([Cbad], lr=.25)
 48    for _ in range(30):
 49        opt.zero_grad(); loss, _, _ = ck_quantities(A, B, Cbad, lam); loss.backward(); opt.step()
 50    after = ck_quantities(A, B, Cbad, lam)
 51    return dict(good_hmin=float(hg), good_kmin=float(kg), bad_hmin=float(before[1]),
 52                bad_kmin=float(before[2]), repaired_hmin=float(after[1]),
 53                repaired_kmin=float(after[2]), repaired_C=float(Cbad.detach()))
 54
 55class LinearSSM(torch.nn.Module):
 56    def __init__(self, n, init_scale=.9):
 57        super().__init__()
 58        self.A = torch.nn.Parameter(init_scale * torch.randn(n, n) / math.sqrt(n))
 59        self.B = torch.nn.Parameter(torch.randn(n) / math.sqrt(n))
 60        self.C = torch.nn.Parameter(torch.randn(n) / math.sqrt(n))
 61    def forward(self, x):
 62        # x: batch,time; zero initial state
 63        h = torch.zeros(x.shape[0], self.A.shape[0], device=x.device)
 64        ys=[]
 65        for t in range(x.shape[1]):
 66            h = h @ self.A.T + x[:, t:t+1] * self.B[None, :]
 67            ys.append((h * self.C[None, :]).sum(1))
 68        return torch.stack(ys, 1)
 69
 70def target_data(A, B, C, batch, length, device):
 71    x = torch.randn(batch, length, device=device)
 72    h = torch.zeros(batch, A.shape[0], device=device); ys=[]
 73    for t in range(length):
 74        h = h @ A.T + x[:, t:t+1] * B[None, :]
 75        ys.append((h*C[None,:]).sum(1))
 76    return x, torch.stack(ys, 1)
 77
 78def run_one(kind, seed, device):
 79    torch.manual_seed(seed); np.random.seed(seed)
 80    n=8; T=24; Tlong=96
 81    # Stable teacher; train data and long evaluation use independent fixed draws.
 82    diag = torch.linspace(.55,.88,n,device=device)
 83    At = torch.diag(diag); Bt = torch.linspace(.2,.8,n,device=device)/math.sqrt(n); Ct = torch.sin(torch.arange(n,device=device)+1)/math.sqrt(n)
 84    xtr,ytr=target_data(At,Bt,Ct,96,T,device); xev,yev=target_data(At,Bt,Ct,96,Tlong,device)
 85    model=LinearSSM(n).to(device)
 86    opt=torch.optim.Adam(model.parameters(),lr=.025)
 87    lam=make_lambdas()
 88    t0=time.perf_counter()
 89    for step in range(350):
 90        opt.zero_grad(); pred=model(xtr); task=(pred-ytr).square().mean()
 91        if kind == 'ck':
 92            ck,_,_=ck_quantities(model.A,model.B,model.C,lam,eps=1e-4)
 93            loss=task + .15*ck
 94        else: loss=task
 95        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
 96    elapsed=time.perf_counter()-t0
 97    with torch.no_grad():
 98        train=((model(xtr)-ytr)**2).mean().item(); long=((model(xev)-yev)**2).mean().item()
 99        ck,hmin,kmin=ck_quantities(model.A,model.B,model.C,lam,eps=1e-4)
100        rho=torch.linalg.eigvals(model.A).abs().max().item()
101    return dict(train_mse=train,long_mse=long,hmin=float(hmin),kmin=float(kmin),ck_penalty=float(ck),spectral_radius=rho,seconds=elapsed)
102
103def main():
104    check=math_check(); results={}
105    for kind in ['baseline','ck']:
106        runs=[run_one(kind,s,device) for s in [SEED,SEED+1,SEED+2]]
107        results[kind]={k:float(np.mean([r[k] for r in runs])) for k in runs[0]}
108        results[kind]['runs']=runs
109    results['device']=str(device); results['math_check']=check
110    results['overhead_ratio']=results['ck']['seconds']/results['baseline']['seconds']
111    with open('results.json','w') as f: json.dump(results,f,indent=2)
112    print(json.dumps(results,indent=2))
113
114if __name__ == '__main__': main()