Delay-Kernel Bifurcation Scheduler / delay_scheduler_experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4from scipy.special import lambertw
  5import torch
  6from torch import nn
  7
  8SEED = 1082
  9random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 10
 11# x'(t)=a x(t)+b x(t-tau), with characteristic roots
 12# lambda = a + W_k(b*tau*exp(-a*tau))/tau.
 13def roots_discrete(a, b, tau, branches=range(-40, 41)):
 14    if tau == 0:
 15        return np.array([a+b], dtype=complex)
 16    z = b*tau*np.exp(-a*tau)
 17    return np.array([a + lambertw(z, k)/tau for k in branches], dtype=complex)
 18
 19def rightmost(a, b, tau):
 20    rr = roots_discrete(a,b,tau)
 21    return rr[np.argmax(rr.real)]
 22
 23def hopf_prediction(a,b,k=0):
 24    w = math.sqrt(b*b-a*a)
 25    # Direct transcription of the supplied atan2 formula, normalized to tau>0.
 26    phase = math.atan2(-w/b, -a/b) + 2*math.pi*k
 27    while phase <= 0: phase += 2*math.pi
 28    return phase/w, w
 29
 30def math_check():
 31    # Parameter sweep: each case has |b|>|a| and therefore a Hopf candidate.
 32    cases=[(-1.0,-2.0),(-0.5,-1.5),(-1.5,-2.5),(-0.8,-1.8)]
 33    rows=[]
 34    for a,b in cases:
 35        tau_pred,w_pred=hopf_prediction(a,b)
 36        taus=np.linspace(0.02, tau_pred*2.0, 700)
 37        rs=np.array([rightmost(a,b,t).real for t in taus])
 38        ix=np.where(np.sign(rs[:-1]) != np.sign(rs[1:]))[0][0]
 39        tc=taus[ix] - rs[ix]*(taus[ix+1]-taus[ix])/(rs[ix+1]-rs[ix])
 40        rc=rightmost(a,b,tc)
 41        eps=max(.001,tau_pred*0.005)
 42        slope=(rightmost(a,b,tc+eps).real-rightmost(a,b,tc-eps).real)/(2*eps)
 43        freq=float(abs(rc.imag))
 44        below=rightmost(a,b,tau_pred*.8).real
 45        above=rightmost(a,b,tau_pred*1.2).real
 46        rows.append({"a":a,"b":b,"predicted_tau_c":tau_pred,
 47          "observed_tau_c":float(tc),"tau_error_pct":float(abs(tc-tau_pred)/tau_pred*100),
 48          "predicted_omega":w_pred,"observed_omega":freq,
 49          "omega_error_pct":float(abs(freq-w_pred)/w_pred*100),
 50          "dr_dtau_observed":float(slope),"stable_side_r":float(below),
 51          "unstable_side_r":float(above),"confirmed":bool(abs(tc-tau_pred)/tau_pred<.02 and
 52          abs(freq-w_pred)/w_pred<.02 and slope>0 and below<0<above)})
 53    return {"cases":rows,"all_predictions_confirmed":all(x["confirmed"] for x in rows),
 54            "max_tau_error_pct":max(x["tau_error_pct"] for x in rows),
 55            "max_omega_error_pct":max(x["omega_error_pct"] for x in rows)}
 56
 57class DelayedRNN(nn.Module):
 58    def __init__(self, d, hidden=24):
 59        super().__init__(); self.d=d; self.hidden=hidden
 60        self.inp=nn.Linear(1,hidden); self.rec=nn.Linear(hidden,hidden,bias=False)
 61        self.out=nn.Linear(hidden,2)
 62    def forward(self,x):
 63        B,T,_=x.shape; hs=[torch.zeros(B,self.hidden,device=x.device) for _ in range(self.d+1)]
 64        for t in range(T):
 65            h=torch.tanh(self.inp(x[:,t])+self.rec(hs[-self.d]))
 66            hs.append(h); hs.pop(0)
 67        return self.out(hs[-1])
 68
 69def data(n, T=32):
 70    # Two frequencies, with random phase and modest noise.
 71    y=np.random.randint(0,2,n); phase=np.random.rand(n)*2*np.pi
 72    t=np.arange(T)[None,:]
 73    f=np.where(y[:,None]==0,.12,.25)
 74    x=np.sin(2*np.pi*f*t+phase[:,None])+0.08*np.random.randn(n,T)
 75    return torch.tensor(x[:,:,None],dtype=torch.float32), torch.tensor(y,dtype=torch.long)
 76
 77def train(scheduled, device):
 78    torch.manual_seed(SEED+int(scheduled)); np.random.seed(SEED+int(scheduled))
 79    tr_y=data(768); va_y=data(256)
 80    tr=(tr_y[0].to(device),tr_y[1].to(device)); va=(va_y[0].to(device),va_y[1].to(device))
 81    model=DelayedRNN(1,24).to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3); lossfn=nn.CrossEntropyLoss()
 82    a,b=-1.,-2.; target=-.10; tau=.25; history=[]
 83    # In this tiny implementation integer buffer delay is the exposed curriculum variable.
 84    for step in range(120):
 85        if scheduled and step%10==0:
 86            r=rightmost(a,b,tau).real
 87            # supplied control law, with clipping and a stable warm-up target
 88            tau=float(np.clip(tau+0.30*(target-r),.05,2.0))
 89            history.append((step,tau,float(r)))
 90        d=max(1,min(8,int(round(tau*4))))
 91        model.d=d
 92        ix=torch.randint(0,768,(64,),device=device)
 93        opt.zero_grad(); loss=lossfn(model(tr[0][ix]),tr[1][ix]); loss.backward(); opt.step()
 94    with torch.no_grad():
 95        pred=model(va[0]).argmax(1); acc=float((pred==va[1]).float().mean().cpu())
 96        vl=float(lossfn(model(va[0]),va[1]).cpu())
 97    return {"val_loss":vl,"accuracy":acc,"final_tau":tau,"final_integer_delay":max(1,min(8,int(round(tau*4)))),"schedule":history}
 98
 99def main():
100    check=math_check()
101    device='cuda' if torch.cuda.is_available() else 'cpu'
102    try:
103        baseline=train(False,device); idea=train(True,device)
104    except Exception as e:
105        device='cpu'; baseline=train(False,device); idea=train(True,device)
106    result={"math_check":check,"device":device,"baseline":baseline,"idea":idea}
107    Path('results.json').write_text(json.dumps(result,indent=2))
108    print(json.dumps(result,indent=2))
109if __name__=='__main__': main()