1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED = 1102
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(min(8, torch.get_num_threads()))
 10try:
 11    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 12    if device.type == 'cuda':
 13        torch.zeros(1, device=device)
 14except Exception:
 15    device = torch.device('cpu')
 16
 17
 18def toy_verification():
 19    # A two-node positive cycle K has rho(K)=sqrt(ab), and tr(K^2)=2ab.
 20    a, b = 1.8, 0.8
 21    K = np.array([[0., a], [b, 0.]])
 22    rho = math.sqrt(a*b)
 23    trace2 = np.trace(K @ K)
 24    # Prediction 1: trace exactly equals the sum of the two length-2 cycles.
 25    p1_err = abs(trace2 - 2*a*b)
 26    # Prediction 2: repeated positive feedback changes from decay to growth at gamma*rho=1.
 27    gammas = np.array([0.35, 0.55, 0.70, 0.74, 0.76, 0.90, 1.10])
 28    n = 30
 29    growth = []
 30    for g in gammas:
 31        x = np.array([1., 0.])
 32        for _ in range(n): x = g*K @ x
 33        growth.append(float(np.linalg.norm(x)))
 34    predicted_threshold = 1.0/rho
 35    # Fit log(norm) vs steps at representative stable/unstable settings.
 36    slopes = []
 37    for g in [0.7, 0.9]:
 38        x=np.array([1., 0.]); vals=[]
 39        for _ in range(1, 16):
 40            x=g*K@x; vals.append(np.log(np.linalg.norm(x)+1e-30))
 41        slopes.append(float(np.polyfit(np.arange(1,16), vals, 1)[0]))
 42    # Prediction 3: normalized r=2 closed-walk term scales exactly as tau^-2.
 43    taus=np.array([0.5, 1., 2., 4.])
 44    normalized = np.array([trace2/(t*t*2.) for t in taus])
 45    slope_tau=float(np.polyfit(np.log(taus), np.log(normalized), 1)[0])
 46    # A disconnected/feed-forward matrix has no 2-cycle and zero trace(K^2).
 47    dag=np.array([[0., 2.0],[0., 0.]])
 48    return {
 49        'device': str(device),
 50        'predictions': {
 51            'trace_identity': {'predicted': 'tr(K^2)=2ab', 'observed': trace2, 'expected': 2*a*b, 'abs_error': p1_err},
 52            'threshold': {'predicted_gamma': predicted_threshold, 'observed_rho': rho,
 53                          'gamma_sweep': gammas.tolist(), 'norm_after_30': growth,
 54                          'classification': ['decay' if x < 1 else 'growth' for x in growth],
 55                          'theory_log_slopes_at_gamma_0.7_0.9': [math.log(g*rho) for g in [0.7,0.9]],
 56                          'observed_log_slopes': slopes},
 57            'temperature_scaling': {'predicted_loglog_slope': -2.0, 'observed_loglog_slope': slope_tau,
 58                                    'taus': taus.tolist(), 'normalized_trace': normalized.tolist()},
 59            'cycle_vanishes_in_dag': {'cycle_trace': float(np.trace(dag@dag)), 'predicted': 0.0}
 60        }
 61    }
 62
 63class IterCell(nn.Module):
 64    def __init__(self):
 65        super().__init__()
 66        self.A=nn.Parameter(torch.tensor([[0.15, 0.85],[0.55, 0.10]], dtype=torch.float32))
 67        self.B=nn.Parameter(torch.randn(2,1)*0.25)
 68        self.out=nn.Linear(2,1)
 69    def forward(self,x, return_jac=False):
 70        h=torch.zeros(x.shape[0],2,device=x.device)
 71        for _ in range(5): h=torch.tanh(h@self.A.T + x@self.B.T)
 72        return self.out(h)
 73    def cycle_penalty(self, x, tau=1.0):
 74        # Exact per-example 2x2 block Jacobian for this small scalar-block MVP.
 75        h=torch.zeros(x.shape[0],2,device=x.device)
 76        for _ in range(4): h=torch.tanh(h@self.A.T + x@self.B.T)
 77        z=h@self.A.T + x@self.B.T
 78        d=(1-torch.tanh(z)**2).detach() # intermittent/stop-gradient Jacobian estimate
 79        # J_ij = |d tanh(z_i)/d h_j|; batch average gives nonnegative block matrix.
 80        J=d.mean(0)[:,None]*self.A.abs()
 81        return torch.trace(J@J)/(tau*tau*2), float(torch.trace(J@J).detach().cpu())
 82    def spectral_proxy(self):
 83        # Same auxiliary cost class: penalize total squared Jacobian entries, not cycles.
 84        return (self.A*self.A).mean()
 85
 86def train(kind, seed=SEED, steps=500):
 87    torch.manual_seed(seed); np.random.seed(seed)
 88    model=IterCell().to(device)
 89    opt=torch.optim.Adam(model.parameters(),lr=0.015)
 90    x=torch.randn(96,1,device=device)
 91    y=2*x + 0.35*torch.sin(3*x)
 92    t0=time.time(); losses=[]; penalties=[]
 93    for step in range(steps):
 94        # fixed stream makes variants directly comparable
 95        xb=x[(step*13)%96:((step*13)%96)+32] if (step*13)%96+32<=96 else x[:32]
 96        yb=y[(step*13)%96:((step*13)%96)+32] if (step*13)%96+32<=96 else y[:32]
 97        pred=model(xb); loss=((pred-yb)**2).mean()
 98        if kind=='cycle': pen,_=model.cycle_penalty(xb); loss=loss+0.04*pen
 99        elif kind=='spectral': pen=model.spectral_proxy(); loss=loss+0.04*pen
100        else: pen=torch.tensor(0.,device=device)
101        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10.); opt.step()
102        losses.append(float(loss.detach().cpu())); penalties.append(float(pen.detach().cpu()))
103    with torch.no_grad():
104        pred=model(x); mse=float(((pred-y)**2).mean().cpu())
105        cp, raw=model.cycle_penalty(x)
106        sp=float(model.spectral_proxy().cpu())
107    return {'final_mse':mse,'final_loss':losses[-1],'mean_aux':float(np.mean(penalties[-100:])),
108            'closed_walk_trace2':raw,'spectral_proxy':sp,'seconds':time.time()-t0}
109
110def main():
111    out=toy_verification()
112    train_results={k:train(k) for k in ['none','spectral','cycle']}
113    out['mini_experiment']=train_results
114    Path('results.json').write_text(json.dumps(out,indent=2))
115    print(json.dumps(out,indent=2))
116
117if __name__=='__main__': main()