1import json, math, time, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7
  8SEED=274
  9random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 10torch.set_num_threads(4)
 11DEVICE='cuda' if torch.cuda.is_available() else 'cpu'
 12try:
 13    if DEVICE=='cuda': torch.cuda.get_device_properties(0)
 14except Exception:
 15    DEVICE='cpu'
 16
 17# -------- Core finite-dimensional quotient verification --------
 18def quotient(A,u,v,eps=1e-12):
 19    return float(v @ ((np.eye(A.shape[0])-A)@u) / (v@u+eps))
 20
 21def verify_math():
 22    rng=np.random.default_rng(19)
 23    # A positive, non-symmetric matrix has positive right/left Perron vectors.
 24    A=rng.uniform(.02,.9,(8,8)); np.fill_diagonal(A, .35)
 25    eig, vr=np.linalg.eig(A); j=np.argmax(eig.real)
 26    mu=float(eig[j].real); u=np.abs(vr[:,j].real); u/=np.linalg.norm(u)
 27    eigT, vl=np.linalg.eig(A.T); k=np.argmin(np.abs(eigT-mu))
 28    v=np.abs(vl[:,k].real); v/=np.linalg.norm(v)
 29    q=quotient(A,u,v); rr=np.linalg.norm((np.eye(8)-A)@u-(1-mu)*u)
 30    rl=np.linalg.norm((np.eye(8)-A.T)@v-(1-mu)*v)
 31    # Independent positive probes optimized by alternating gradient descent on residuals.
 32    U=np.exp(rng.normal(size=8)); V=np.exp(rng.normal(size=8)); U/=np.linalg.norm(U); V/=np.linalg.norm(V)
 33    before=(np.linalg.norm((np.eye(8)-A)@U-quotient(A,U,V)*U)**2+
 34            np.linalg.norm((np.eye(8)-A.T)@V-quotient(A,U,V)*V)**2)
 35    for _ in range(300):
 36        lam=quotient(A,U,V)
 37        # projected gradient-like updates for the two residuals, followed by positivity.
 38        r=(np.eye(8)-A)@U-lam*U
 39        s=(np.eye(8)-A.T)@V-lam*V
 40        U=np.maximum(U-0.08*((np.eye(8)-A).T@r-lam*r),1e-8); U/=np.linalg.norm(U)
 41        V=np.maximum(V-0.08*((np.eye(8)-A.T).T@s-lam*s),1e-8); V/=np.linalg.norm(V)
 42    after=(np.linalg.norm((np.eye(8)-A)@U-quotient(A,U,V)*U)**2+
 43           np.linalg.norm((np.eye(8)-A.T)@V-quotient(A,U,V)*V)**2)
 44    return {'positive_matrix_mu':mu,'quotient_at_eigenpair':q,
 45            'quotient_abs_error':abs(q-(1-mu)), 'right_residual':rr,
 46            'left_residual':rl,'probe_residual_before':before,
 47            'probe_residual_after':after,'probe_residual_ratio':after/before}
 48
 49class RNN(nn.Module):
 50    def __init__(self,h=32):
 51        super().__init__(); self.W=nn.Parameter(torch.randn(h,h)*.18)
 52        self.U=nn.Parameter(torch.randn(h,2)*.25); self.b=nn.Parameter(torch.zeros(h))
 53        self.out=nn.Linear(h,1)
 54    def forward(self,x):
 55        h=torch.zeros(x.shape[0],self.W.shape[0],device=x.device)
 56        for t in range(x.shape[1]): h=torch.tanh(h@self.W.T+x[:,t]@self.U.T+self.b)
 57        return self.out(h).squeeze(-1)
 58
 59def cone_penalty(model, steps=5, mu_max=.98):
 60    # Local Jacobian at the zero-input fixed point, which is the recurrent
 61    # transition relevant during the long distractor portion of this task.
 62    W=model.W; h=W.shape[0]
 63    with torch.no_grad(): A=W.detach().clone()
 64    a=torch.zeros(h,device=W.device,requires_grad=True); b=torch.zeros(h,device=W.device,requires_grad=True)
 65    opt=torch.optim.SGD([a,b],lr=.35)
 66    I=torch.eye(h,device=W.device)
 67    for _ in range(steps):
 68        u=F.softplus(a)+1e-5; v=F.softplus(b)+1e-5
 69        u=u/(u.norm()+1e-12); v=v/(v.norm()+1e-12)
 70        lam=((u-A@u)*v).sum()/(u@v+1e-6)
 71        loss=((u-A@u-lam*u)**2).sum()+((v-A.T@v-lam*v)**2).sum()
 72        opt.zero_grad(); loss.backward(); opt.step()
 73    u=(F.softplus(a).detach()+1e-5); u=u/u.norm()
 74    v=(F.softplus(b).detach()+1e-5); v=v/v.norm()
 75    # Reattach only the model matrix for the outer update.
 76    lam=((u-W@u)*v).sum()/(u@v+1e-6)
 77    mu=1-lam
 78    return F.softplus(mu-mu_max)**2, float(mu.detach()), float(lam.detach())
 79
 80def spectral_penalty(model, target=.98):
 81    # Standard differentiable power estimate of the recurrent spectral norm.
 82    s=torch.linalg.matrix_norm(model.W,ord=2)
 83    return F.softplus(s-target)**2, float(s.detach())
 84
 85def make_batch(n=64,T=30,device='cpu'):
 86    x=torch.zeros(n,T,2,device=device); y=torch.empty(n,device=device).uniform_(-1,1)
 87    x[:,0,0]=1.; x[:,0,1]=y
 88    return x,y
 89
 90def train(kind,seed,steps=350):
 91    torch.manual_seed(seed); model=RNN().to(DEVICE); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 92    losses=[]; penalties=[]; t0=time.time()
 93    for step in range(steps):
 94        x,y=make_batch(device=DEVICE); pred=model(x); task=F.mse_loss(pred,y); reg=pred.new_zeros(())
 95        if kind=='cone': reg,mu,_=cone_penalty(model); penalties.append(mu)
 96        elif kind=='spectral': reg,s=spectral_penalty(model); penalties.append(s)
 97        loss=task+(.12 if kind!='baseline' else 0.)*reg
 98        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step()
 99        losses.append(float(task.detach()))
100    # long-horizon heldout loss and gradient statistic, same fixed batch.
101    torch.manual_seed(1000+seed); x,y=make_batch(256,30,DEVICE)
102    with torch.no_grad(): test=float(F.mse_loss(model(x),y))
103    return {'final_train_loss':float(np.mean(losses[-30:])),'test_loss':test,
104            'penalty_metric':float(np.mean(penalties[-30:])) if penalties else None,
105            'seconds':time.time()-t0}
106
107def main():
108    mathcheck=verify_math(); results={'device':DEVICE,'math_check':mathcheck,'runs':{}}
109    for kind in ['baseline','spectral','cone']:
110        results['runs'][kind]=[train(kind,s) for s in [7,17,27]]
111    for kind,rs in list(results['runs'].items()):
112        if kind.endswith('_summary'): continue
113        results['runs'][kind+'_summary']={k:float(np.mean([r[k] for r in rs])) for k in ['final_train_loss','test_loss','seconds']}
114        if kind!='baseline': results['runs'][kind+'_summary']['penalty_metric']=float(np.mean([r['penalty_metric'] for r in rs]))
115    Path('results.json').write_text(json.dumps(results,indent=2))
116    print(json.dumps(results,indent=2))
117if __name__=='__main__': main()