Universal Clock Regularization for Recurrent Dynamics / clock_experiment.py
Mechanism confirmed, baseline not beaten
1import math, random, json
2from pathlib import Path
3import numpy as np
4
5SEED=7
6np.random.seed(SEED); random.seed(SEED)
7
8def metrics(phi,dt):
9 w=np.diff(np.unwrap(phi))/dt
10 return float(w.mean()),float(w.var()),float(np.sqrt(w.var())/(abs(w.mean())+1e-12))
11
12def math_sweeps():
13 omega,dt=1.7,.002
14 t=np.arange(0,30*2*np.pi/omega,dt); th=omega*t
15 scaling=[]
16 for m in [1,2,3,5]:
17 for a in [0,.01,.03,.06]:
18 _,v,e=metrics(th+a*np.sin(m*th),dt)
19 scaling.append({'m':m,'a':a,'observed_var':v,'predicted_var':(omega*a*m)**2/2,'observed_E':e,'predicted_E':abs(a*m)/math.sqrt(2)})
20 noise=[]
21 for sig in [0,.001,.003,.01,.03]:
22 vals=[]
23 for r in range(30):
24 rng=np.random.default_rng(SEED+r)
25 vals.append(metrics(th+rng.normal(0,sig,len(th)),dt)[1])
26 pred=2*sig**2/dt**2
27 noise.append({'sigma':sig,'observed_var':float(np.mean(vals)),'predicted_var':pred})
28 boundary=[]
29 for inc in [0.5,1,2,2.8,3.0,3.13,3.15,3.3,4.0]:
30 ph=np.arange(1000)*inc
31 mw,v,e=metrics(ph,1.0)
32 boundary.append({'increment':inc,'estimated_mean':mw,'error':abs(mw-inc),'E':e})
33 return scaling,noise,boundary
34
35def gru_experiment():
36 try:
37 import torch
38 import torch.nn as nn
39 torch.manual_seed(SEED); device='cuda' if torch.cuda.is_available() else 'cpu'
40 dt=.1; omega=1.0; T=45; ntrain=96; ntest=24
41 def make(n,seed):
42 rng=np.random.default_rng(seed); x=[]
43 for i in range(n):
44 phase=rng.uniform(0,2*np.pi); amp=1+rng.normal(0,.04)
45 tt=np.arange(T+1)*dt
46 z=np.stack([amp*np.sin(omega*tt+phase),amp*np.cos(omega*tt+phase)],1)
47 z+=rng.normal(0,.025,z.shape); x.append(z)
48 return torch.tensor(np.array(x),dtype=torch.float32,device=device)
49 train,test=make(ntrain,10),make(ntest,20)
50 class Model(nn.Module):
51 def __init__(self,clock):
52 super().__init__(); self.clock=clock; self.gru=nn.GRU(2,16,batch_first=True); self.out=nn.Linear(16,2)
53 self.phase=nn.Linear(16,2) if clock else None
54 def forward(self,x):
55 h,_=self.gru(x); y=self.out(h); return y,h
56 def run(clock, dev):
57 torch.manual_seed(SEED); m=Model(clock).to(dev); opt=torch.optim.Adam(m.parameters(),lr=3e-3)
58 tr=train.to(dev); te=test.to(dev)
59 # teacher-forced one-step prediction; phase head is intentionally auxiliary.
60 for ep in range(220):
61 inp,tar=tr[:,:-1],tr[:,1:]; y,h=m(inp); loss=((y-tar)**2).mean()
62 if clock:
63 q=m.phase(h); ph=torch.atan2(q[...,0],q[...,1]); d=torch.atan2(torch.sin(torch.diff(ph,dim=1)),torch.cos(torch.diff(ph,dim=1)))/dt
64 # target natural rate is known in this synthetic benchmark; cycle term prevents zero-clock collapse.
65 loss=loss+.08*d.var(dim=1).mean()+.08*(d.mean(dim=1)-omega).pow(2).mean()
66 opt.zero_grad(); loss.backward(); opt.step()
67 with torch.no_grad():
68 y,h=m(te[:,:-1]); one=float(((y-te[:,1:])**2).mean().cpu())
69 # assess phase consistency of learned phase head; baseline uses geometric atan2.
70 if clock: ph=torch.atan2(m.phase(h)[...,0],m.phase(h)[...,1])
71 else: ph=torch.atan2(h[...,0],h[...,1])
72 ds=torch.atan2(torch.sin(torch.diff(ph,dim=1)),torch.cos(torch.diff(ph,dim=1)))/dt
73 return {'one_step_mse':one,'clock_E':float((ds.std(1)/(ds.mean(1).abs()+1e-6)).mean().cpu()),'device':str(dev)}
74 try:
75 return {'baseline':run(False,device),'clock':run(True,device)}
76 except Exception as cuda_error:
77 # Required robust fallback after any CUDA/cuDNN allocation or execution error.
78 if device != 'cpu':
79 return {'baseline':run(False,'cpu'),'clock':run(True,'cpu'),'cuda_error':str(cuda_error)}
80 raise
81 except Exception as e:
82 return {'error':str(e),'fallback':'math-only'}
83
84def main():
85 s,n,b=math_sweeps(); out={'scaling':s,'noise':n,'unwrap_boundary':b,'gru':gru_experiment()}
86 Path('results.json').write_text(json.dumps(out,indent=2))
87 print(json.dumps(out,indent=2))
88if __name__=='__main__': main()