1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4
  5SEED=1427
  6
  7def seed_all(seed=SEED):
  8    random.seed(seed); np.random.seed(seed)
  9
 10def spectral_norm(M):
 11    return float(np.linalg.svd(M, compute_uv=False)[0])
 12
 13def normalize(M, target):
 14    n=spectral_norm(M)
 15    return M*(target/n) if n else M.copy()
 16
 17class TwoColumnMemory:
 18    """Small MPS-inspired memory. A is the single-site transfer channel;
 19    B is a second transfer channel driven by adjacent embedded features."""
 20    def __init__(self,p=4,chi=8,alpha=.9,beta=.9,seed=SEED):
 21        rng=np.random.default_rng(seed)
 22        self.p,self.chi=p,chi
 23        self.A=np.stack([normalize(rng.normal(size=(chi,chi)),alpha) for _ in range(p)])
 24        self.B=np.stack([normalize(rng.normal(size=(chi,chi)),beta) for _ in range(2*p)])
 25    def run(self,emb):
 26        h1=np.zeros(self.chi); h2=np.zeros(self.chi); prev=np.zeros(emb.shape[1]); out=[]
 27        for e in emb:
 28            At=np.tensordot(e,self.A,axes=(0,0))
 29            z=np.concatenate([prev,e]); Bt=np.tensordot(z,self.B,axes=(0,0))
 30            h1=At@h1+e[0]*np.ones(self.chi)/math.sqrt(self.chi)
 31            h2=Bt@h2+np.dot(prev,e)*np.ones(self.chi)/math.sqrt(self.chi)
 32            out.append(np.concatenate([e,h1,h2])); prev=e
 33        return np.asarray(out)
 34
 35def contraction_sweep():
 36    """Use A=alpha*orthogonal, so the formula predicts exactly alpha**r."""
 37    rng=np.random.default_rng(SEED); chi=12; r=20
 38    Q,_=np.linalg.qr(rng.normal(size=(chi,chi)))
 39    rows=[]
 40    for a in [.6,.8,.95,1.0,1.1,1.25]:
 41        A=a*Q; v=rng.normal(size=chi); v/=np.linalg.norm(v)
 42        norms=[]
 43        for k in range(r+1): norms.append(np.linalg.norm(v)); v=A@v
 44        ratio=norms[-1]/norms[0]
 45        predicted=a**r
 46        # fitted per-step amplification from the whole trajectory
 47        slope=(math.log(norms[-1]+1e-30)-math.log(norms[0]))/r
 48        rows.append({'alpha':a,'r':r,'observed_ratio':ratio,'predicted_ratio':predicted,
 49                     'observed_per_step':math.exp(slope),'predicted_per_step':a})
 50    return rows
 51
 52def pair_scaling_sweep():
 53    """For a two-column perturbation, scale B by lambda. One-step response is linear."""
 54    rng=np.random.default_rng(SEED+1); chi=10; d=6
 55    B=normalize(rng.normal(size=(chi,chi)),.9); z=rng.normal(size=d); h=rng.normal(size=chi)
 56    base=B@h; base_norm=np.linalg.norm(base)
 57    rows=[]
 58    for lam in [0,.1,.25,.5,1.,1.5]:
 59        got=np.linalg.norm(lam*B@h)
 60        rows.append({'lambda':lam,'observed_norm':got,'predicted_norm':lam*base_norm,
 61                     'relative_error':abs(got-lam*base_norm)/(1e-12+lam*base_norm) if lam else 0.0})
 62    return rows
 63
 64def boundary_sweep():
 65    """Prediction: pair perturbations decay if gamma*lambda<1 and grow if >1."""
 66    rows=[]; r=12; chi=10
 67    for gamma in [.7,.9,1.0,1.1]:
 68        for lam in [.5,.9,1.0,1.1,1.4]:
 69            gain=gamma*lam
 70            # orthogonal transfer makes the finite-horizon prediction exact
 71            rng=np.random.default_rng(SEED+int(100*gamma)+int(10*lam))
 72            Q,_=np.linalg.qr(rng.normal(size=(chi,chi)))
 73            v=rng.normal(size=chi); v/=np.linalg.norm(v)
 74            observed=np.linalg.norm(np.linalg.matrix_power(gain*Q,r)@v)
 75            predicted=gain**r
 76            rows.append({'gamma':gamma,'lambda':lam,'product':gain,'r':r,
 77                         'observed_ratio':observed,'predicted_ratio':predicted,
 78                         'predicted_regime':'decay' if gain<1 else ('neutral' if gain==1 else 'growth'),
 79                         'observed_regime':'decay' if observed<1 else ('neutral' if abs(observed-1)<1e-10 else 'growth')})
 80    return rows
 81
 82def delayed_learning(seed=SEED):
 83    # Tiny nonlinear delayed-correlation task: label is XOR of bits 5 steps apart.
 84    import torch
 85    def run(device):
 86        torch.manual_seed(seed); np.random.seed(seed)
 87        T,N=18,1200; rng=np.random.default_rng(seed)
 88        x=rng.integers(0,2,size=(N,T,1)).astype('float32')
 89        y=(x[:,2,0] != x[:,7,0]).astype('float32')
 90        xt=torch.tensor(x,device=device); yt=torch.tensor(y,device=device)
 91        class GRU(torch.nn.Module):
 92            def __init__(self):
 93                super().__init__(); self.g=torch.nn.GRU(1,16,batch_first=True); self.o=torch.nn.Linear(16,1)
 94            def forward(self,x): return self.o(self.g(x)[0][:,-1,:]).squeeze(-1)
 95        class TC(torch.nn.Module):
 96            def __init__(self):
 97                super().__init__(); self.e=torch.nn.Linear(1,4); self.a=torch.nn.Parameter(torch.randn(4,8,8)*.08); self.b=torch.nn.Parameter(torch.randn(8,8,8)*.08); self.o=torch.nn.Linear(20,1)
 98            def forward(self,x):
 99                e=torch.tanh(self.e(x)); h1=torch.zeros(x.size(0),8,device=x.device); h2=h1.clone(); prev=torch.zeros_like(e[:,0])
100                for t in range(x.size(1)):
101                    h1=torch.bmm(torch.einsum('bp,pij->bij',e[:,t],self.a),h1.unsqueeze(-1)).squeeze(-1)+e[:,t,0:1]
102                    z=torch.cat([prev,e[:,t]],-1); h2=torch.bmm(torch.einsum('bp,pij->bij',z,self.b),h2.unsqueeze(-1)).squeeze(-1)+torch.sum(prev*e[:,t],-1,keepdim=True)
103                    prev=e[:,t]
104                return self.o(torch.cat([e[:,-1],h1,h2],-1)).squeeze(-1)
105        def train(m):
106            m.to(device); opt=torch.optim.Adam(m.parameters(),lr=.01); lossfn=torch.nn.BCEWithLogitsLoss(); t0=time.time()
107            for _ in range(80):
108                opt.zero_grad(); loss=lossfn(m(xt),yt); loss.backward(); opt.step()
109            with torch.no_grad(): acc=((torch.sigmoid(m(xt))>.5)==yt.bool()).float().mean().item()
110            return {'accuracy':acc,'loss':float(loss.item()),'seconds':time.time()-t0,'device':str(device)}
111        return {'GRU':train(GRU()),'TwoColumn':train(TC())}
112    requested='cuda' if torch.cuda.is_available() else 'cpu'
113    try:
114        return run(requested)
115    except Exception as first_error:
116        if requested == 'cuda':
117            try:
118                return run('cpu') | {'fallback_reason':str(first_error)}
119            except Exception as second_error:
120                return {'error':str(second_error),'cuda_error':str(first_error)}
121        return {'error':str(first_error)}
122
123def main():
124    seed_all(); result={'contraction':contraction_sweep(),'pair_scaling':pair_scaling_sweep(),'boundary':boundary_sweep(),'learning':delayed_learning()}
125    Path('results.json').write_text(json.dumps(result,indent=2))
126    print(json.dumps(result,indent=2))
127if __name__=='__main__': main()