1import json, os, random
  2import numpy as np
  3
  4SEED = 2139
  5np.random.seed(SEED); random.seed(SEED)
  6
  7
  8def parts(V):
  9    return (V + V.T) / 2.0, (V - V.T) / 2.0
 10
 11
 12def lift(V, m=4):
 13    d = V.shape[0]
 14    S, A = parts(V)
 15    L = np.zeros((m*d, m*d))
 16    for i in range(m):
 17        L[i*d:(i+1)*d, i*d:(i+1)*d] = A
 18        if i + 1 < m:
 19            L[i*d:(i+1)*d, (i+1)*d:(i+2)*d] = S
 20            L[(i+1)*d:(i+2)*d, i*d:(i+1)*d] = -S
 21    return L, S, A
 22
 23
 24def embed(x, m):
 25    z = np.zeros(m * len(x)); z[:len(x)] = x; z[len(x):2*len(x)] = x
 26    return z
 27
 28
 29def cayley(L, eta):
 30    I = np.eye(len(L))
 31    return np.linalg.solve(I - .5*eta*L, I + .5*eta*L)
 32
 33
 34def run_checks():
 35    # Fixed non-normal, damped, transient-amplifying generator.
 36    V0 = np.array([[-.35, 2.0], [-.15, -.10]])
 37    x = np.array([.8, -.4])
 38    m = 5
 39    L, S, A = lift(V0, m)
 40    skew_err = np.linalg.norm(L + L.T, ord=np.inf)
 41    proj_err = np.linalg.norm((L[:2,:] @ embed(x,m)) - V0 @ x)
 42    z0 = embed(x,m)
 43    cayley_rows = []
 44    euler_rows = []
 45    for gamma in [0.0, .25, .5, 1.0, 2.0]:
 46        V = gamma * V0
 47        Lg, _, _ = lift(V,m)
 48        lam = np.linalg.norm(Lg, 2)
 49        for eta in [.05, .2, .5]:
 50            z = embed(x,m)
 51            C = cayley(Lg,eta)
 52            # ten steps: predicted exact preservation for every eta and gamma
 53            for _ in range(10): z = C @ z
 54            rel_norm_error = abs(np.linalg.norm(z)/np.linalg.norm(embed(x,m))-1)
 55            cayley_rows.append([gamma, eta, rel_norm_error])
 56        # Euler one-step identity and accumulated growth at eta=.2
 57        eta=.2; z=embed(x,m); z1=(np.eye(len(Lg))+eta*Lg)@z
 58        exact_increment = np.linalg.norm(z1)**2-np.linalg.norm(z)**2
 59        predicted_increment = eta**2*np.linalg.norm(Lg@z)**2
 60        for _ in range(10): z=(np.eye(len(Lg))+eta*Lg)@z
 61        euler_rows.append([gamma, exact_increment, predicted_increment,
 62                           np.linalg.norm(z)/np.linalg.norm(embed(x,m))])
 63
 64    # Projected transient gain sweep: ||P exp(tL) E|| / ||E|| <= 1.
 65    gain_rows=[]
 66    from scipy.linalg import expm
 67    for gamma in [0.0, .25, .5, 1., 2.]:
 68        Lg,_,_=lift(gamma*V0,m)
 69        gains=[]
 70        for t in np.linspace(0,4,81):
 71            gains.append(np.linalg.norm(expm(t*Lg)[:2,:2],2) / np.sqrt(2))
 72        gain_rows.append([gamma, float(max(gains)), float(np.sqrt(2)), float(gains[-1])])
 73    return {
 74        'skew_inf_error': float(skew_err),
 75        'projection_derivative_error': float(proj_err),
 76        'cayley_rows_gamma_eta_rel_norm_error': cayley_rows,
 77        'euler_rows_gamma_increment_predicted_increment_norm_ratio_after_10': euler_rows,
 78        'projected_gain_rows_gamma_max_gain_bound_final_gain': gain_rows,
 79    }
 80
 81
 82def tiny_task():
 83    # A matched-parameter linear memory task, trained with torch if available.
 84    try:
 85        import torch
 86        torch.manual_seed(SEED)
 87        device = 'cuda' if torch.cuda.is_available() else 'cpu'
 88        if device == 'cuda':
 89            try: torch.cuda.set_per_process_memory_fraction(.15)
 90            except Exception: pass
 91        T, N, d = 40, 256, 4
 92        rng=np.random.default_rng(SEED)
 93        u=torch.tensor(rng.normal(size=(N,T,2)),dtype=torch.float32,device=device)
 94        # target is a stable but non-normal recurrence, final-state regression
 95        V=torch.tensor([[-.18,1.2,0,0],[-.05,-.12,0,0],[0,0,-.3,.7],[0,0,-.1,-.25]],dtype=torch.float32,device=device)
 96        B=torch.tensor(rng.normal(scale=.4,size=(d,2)),dtype=torch.float32,device=device)
 97        target=torch.zeros(N,d,device=device)
 98        for k in range(T): target=target + .08*(target@V.T + u[:,k]@B.T)
 99        target=target.detach()
100        # baseline vanilla RNN, small and deliberately matched by hidden dimension
101        class Base(torch.nn.Module):
102            def __init__(self):
103                super().__init__(); self.W=torch.nn.Parameter(torch.randn(d,d,device=device)*.15); self.U=torch.nn.Parameter(torch.randn(d,2,device=device)*.15); self.q=torch.nn.Parameter(torch.randn(d,d,device=device)*.15)
104            def forward(self, x):
105                h=torch.zeros(x.shape[0],d,device=device)
106                for k in range(T): h=torch.tanh(h@self.W.T+x[:,k]@self.U.T)
107                return h@self.q.T
108        # lift has 2d state and linear readout; stable Cayley, train only input/readout
109        class Lift(torch.nn.Module):
110            def __init__(self):
111                super().__init__(); self.rawS=torch.nn.Parameter(torch.randn(d,d,device=device)*.12); self.rawA=torch.nn.Parameter(torch.randn(d,d,device=device)*.12); self.U=torch.nn.Parameter(torch.randn(2*d,2,device=device)*.15); self.q=torch.nn.Parameter(torch.randn(d,d,device=device)*.15)
112            def forward(self,x):
113                S=self.rawS+self.rawS.T; A=self.rawA-self.rawA.T; z=torch.zeros(x.shape[0],2*d,device=device); eta=.08
114                L=torch.cat([torch.cat([A,S],1),torch.cat([-S,A],1)],0); I=torch.eye(2*d,device=device); C=torch.linalg.solve(I-eta*L/2,I+eta*L/2)
115                for k in range(T): z=z@C.T + x[:,k]@self.U.T
116                return z[:,:d]@self.q.T
117        def fit(model):
118            opt=torch.optim.Adam(model.parameters(),lr=.02); losses=[]
119            for _ in range(120):
120                pred=model(u); loss=((pred-target)**2).mean(); opt.zero_grad(); loss.backward(); opt.step(); losses.append(float(loss.detach().cpu()))
121            return losses[0], losses[-1]
122        b=fit(Base()); l=fit(Lift())
123        return {'device':device,'baseline_initial_final_mse':b,'lift_initial_final_mse':l}
124    except Exception as e:
125        return {'error':repr(e)}
126
127if __name__ == '__main__':
128    out={'checks':run_checks(),'tiny_task':tiny_task()}
129    with open('results.json','w') as f: json.dump(out,f,indent=2)
130    print(json.dumps(out,indent=2))