1import math, random, json
 2import numpy as np
 3import torch
 4
 5SEED = 2947
 6np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 7
 8def poles_from_corr(c, rank=2, ridge=1e-6):
 9    c = np.asarray(c, dtype=np.float64)
10    m = min(rank, (len(c)-1)//2)
11    H0 = np.array([[c[i+j] for j in range(m)] for i in range(m)])
12    H1 = np.array([[c[i+j+1] for j in range(m)] for i in range(m)])
13    R = np.linalg.solve(H0.T @ H0 + ridge*np.eye(m), H0.T @ H1)
14    return np.linalg.eigvals(R), H0
15
16def math_check():
17    true = np.array([.93, .62]); weights = np.array([.7, .3]); t=np.arange(40)
18    c=sum(w*r**t for w,r in zip(weights,true))
19    est,H=poles_from_corr(c,2); est=est[np.argsort(-np.abs(est))]
20    mods=np.abs(est); dominant=float(mods[0])
21    half_pred=math.log(.5)/math.log(dominant)
22    half_emp=int(np.where(np.abs(c/c[0])<=.5)[0][0])
23    return {"true_poles":true.tolist(), "estimated_moduli":mods.tolist(),
24      "dominant_tau":float(-1/math.log(dominant)), "predicted_half_life":float(half_pred),
25      "empirical_half_life_steps":half_emp, "hankel_condition":float(np.linalg.cond(H)),
26      "pole_recovery_max_error":float(np.max(np.abs(np.sort(mods)[::-1]-true)))}
27
28class LinearMemory(torch.nn.Module):
29    def __init__(self,n=8):
30        super().__init__()
31        self.W=torch.nn.Parameter(.98*torch.eye(n)+.04*torch.randn(n,n))
32        self.inp=torch.nn.Parameter(.2*torch.randn(n)); self.out=torch.nn.Parameter(.2*torch.randn(n))
33    def forward(self,x):
34        h=torch.zeros(x.shape[0],self.W.shape[0],device=x.device); ys=[]
35        for k in range(x.shape[1]):
36            h=torch.tanh(h+x[:,k,0:1]*self.inp); ys.append((h*self.out).sum(1))
37        return torch.stack(ys,1)
38
39def one_task(use_res, device, steps=350):
40    torch.manual_seed(SEED + int(use_res)); model=LinearMemory().to(device)
41    opt=torch.optim.Adam(model.parameters(),lr=.012); rng=np.random.default_rng(SEED+int(use_res))
42    pole_log=[]
43    for step in range(steps):
44        x=rng.normal(size=(64,24,1)).astype('float32'); target=np.zeros((64,24),dtype='float32'); target[:,-1]=x[:,11,0]
45        xb=torch.tensor(x,device=device); yb=torch.tensor(target,device=device)
46        pred=model(xb); loss=((pred-yb)**2).mean(); total=loss
47        if use_res and step%5==0:
48            with torch.no_grad():
49                h=torch.zeros(64,8,device=device); states=[]
50                for _ in range(18):
51                    h=torch.tanh(h+.15*torch.randn(64,1,device=device)*model.inp); states.append(h[:,0])
52                cc=torch.stack([(states[0]*states[t]).mean() for t in range(18)])
53                poles,_=poles_from_corr(cc.cpu().numpy(),rank=3)
54                dominant=float(np.max(np.abs(poles))); pole_log.append(dominant)
55            # Differentiable proxy: penalize excess operator norm only when fitted pole is too slow.
56            excess=max(0.,dominant-.92)
57            total=loss + .08*excess**2*torch.linalg.matrix_norm(model.W,ord=2)**2
58        opt.zero_grad(); total.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step()
59    with torch.no_grad():
60        x=rng.normal(size=(256,24,1)).astype('float32'); target=torch.tensor(x[:,11,0],device=device)
61        pred=model(torch.tensor(x,device=device)); mse=float(((pred[:,-1]-target)**2).mean().cpu())
62        rho=float(np.max(np.abs(np.linalg.eigvals(model.W.detach().cpu().numpy()))))
63    return {"mse":mse,"spectral_radius":rho,"mean_fitted_pole":float(np.mean(pole_log)) if pole_log else None,
64            "max_fitted_pole":float(np.max(pole_log)) if pole_log else None,"device":device}
65
66def run_task(use_res):
67    preferred='cuda' if torch.cuda.is_available() else 'cpu'
68    if preferred=='cuda':
69        try: return one_task(use_res,'cuda')
70        except Exception as e:
71            torch.cuda.empty_cache(); fallback=one_task(use_res,'cpu'); fallback['cuda_fallback']=str(e); return fallback
72    return one_task(use_res,'cpu')
73
74def main():
75    result={"math_check":math_check(),"task_baseline":run_task(False),"task_resonance":run_task(True)}
76    with open('results.json','w') as f: json.dump(result,f,indent=2)
77    print(json.dumps(result,indent=2))
78if __name__=='__main__': main()