1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4from scipy.linalg import expm, norm, eigvals
  5import torch
  6import torch.nn as nn
  7
  8SEED = 449
  9np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 10torch.set_num_threads(4)
 11DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 12
 13
 14def bch_check():
 15    # For two generators, the residual after retaining the commutator term
 16    # should scale cubically with the step size.
 17    rng = np.random.default_rng(SEED)
 18    X = rng.normal(size=(3,3)); Y = rng.normal(size=(3,3))
 19    X -= 1.7*np.eye(3); Y -= 1.7*np.eye(3)
 20    C = X @ Y - Y @ X
 21    rows = []
 22    for s in [0.02, 0.04, 0.08, 0.16, 0.32]:
 23        exact = expm(s*Y) @ expm(s*X)
 24        avg = expm(s*(X+Y) - 0.5*s*s*C)
 25        first = expm(s*(X+Y))
 26        rows.append((s, norm(exact-first, 'fro'), norm(exact-avg, 'fro')))
 27    x = np.log([r[0] for r in rows[-3:]])
 28    y1 = np.log([r[1] for r in rows[-3:]])
 29    y2 = np.log([r[2] for r in rows[-3:]])
 30    return {"rows": [[float(v) for v in r] for r in rows],
 31            "first_order_loglog_slope": float(np.polyfit(x,y1,1)[0]),
 32            "commutator_corrected_loglog_slope": float(np.polyfit(x,y2,1)[0])}
 33
 34
 35def matrix_metrics(A, tau=0.7, alpha=(0.5,0.5)):
 36    B = [alpha[i]*tau*A[i] for i in range(2)]
 37    Phi = expm(B[1]) @ expm(B[0])
 38    avg = expm(B[0]+B[1])
 39    comm = norm(B[0]@B[1]-B[1]@B[0], 'fro')**2
 40    mismatch = norm(Phi-avg, 'fro') / max(norm(avg,'fro'), 1e-12)
 41    return float(comm), float(mismatch), Phi
 42
 43
 44def regularizer_optimization():
 45    # Directly optimize two matrices toward a fixed noncommuting task matrix,
 46    # with and without the proposed commutator penalty.
 47    torch.manual_seed(SEED+1)
 48    target = torch.tensor(np.array([[.72,.35],[-.15,.61]], dtype=np.float32))
 49    out = {}
 50    for name, lam in [("unregularized",0.0),("commutator_regularized",2.0)]:
 51        A = nn.Parameter(torch.randn(2,2)*.7)
 52        D = nn.Parameter(torch.randn(2,2)*.7)
 53        opt = torch.optim.Adam([A,D], lr=.04)
 54        for _ in range(500):
 55            B0=.35*A; B1=.35*D
 56            phi=torch.matrix_exp(B1)@torch.matrix_exp(B0)
 57            loss=(phi-target).pow(2).mean()+lam*(B0@B1-B1@B0).pow(2).sum()
 58            opt.zero_grad(); loss.backward(); opt.step()
 59        a=A.detach().numpy(); d=D.detach().numpy()
 60        comm, mismatch, _=matrix_metrics([a,d],tau=.7)
 61        fit=float(((expm(.35*d)@expm(.35*a))-target.numpy())**2 .mean()) if False else float(np.mean((expm(.35*d)@expm(.35*a)-target.numpy())**2))
 62        out[name]={"task_fit_mse":fit,"commutator":comm,"ordered_averaged_relative_error":mismatch}
 63    return out
 64
 65
 66class SwitchedSSM(nn.Module):
 67    def __init__(self, d=8, modes=2, regularized=False, single=False):
 68        super().__init__(); self.d=d; self.m=modes; self.regularized=regularized; self.single=single
 69        n=1 if single else modes
 70        self.A=nn.Parameter(torch.randn(n,d,d)*.08 - .35*torch.eye(d).unsqueeze(0).repeat(n,1,1))
 71        self.U=nn.Parameter(torch.randn(modes,d)*.12)
 72        self.v=nn.Parameter(torch.randn(d)*.12)
 73        self.readout=nn.Linear(d,1)
 74    def forward(self,x):
 75        h=torch.zeros(x.shape[0],self.d,device=x.device)
 76        alpha=1.0/self.m; tau=.7
 77        E=[]
 78        for i in range(self.m):
 79            A=self.A[0] if self.single else self.A[i]
 80            E.append(torch.matrix_exp(alpha*tau*A))
 81        for t in range(x.shape[1]):
 82            for i in range(self.m):
 83                h=(h.unsqueeze(1) @ E[i].T).squeeze(1)
 84                h=h+alpha*x[:,t:t+1]*self.U[i]
 85        return self.readout(h).squeeze(-1)
 86    def penalties(self):
 87        if self.single: return torch.tensor(0.,device=self.A.device)
 88        B=[.5*.7*a for a in self.A]
 89        return sum(((B[i]@B[j]-B[j]@B[i])**2).sum() for i in range(self.m) for j in range(i))
 90    def diagnostics(self):
 91        with torch.no_grad():
 92            As=self.A.detach().cpu().numpy(); As=As if not self.single else np.repeat(As,2,axis=0)
 93            comm, mismatch, phi=matrix_metrics(As,tau=.7)
 94            gain=max(abs(eigvals(phi)))
 95        return {"commutator":comm,"ordered_averaged_relative_error":mismatch,"cycle_spectral_gain":float(gain)}
 96
 97
 98def sequence_experiment():
 99    torch.manual_seed(SEED+2)
100    dev=DEVICE
101    ntrain, ntest, T=192,96,18
102    rng=np.random.default_rng(SEED+2)
103    X=rng.normal(size=(ntest+ntrain,T)).astype('float32')
104    # Stable long-memory target: exponentially discounted sequence sum.
105    weights=(.88**np.arange(T-1,-1,-1)).astype('float32')
106    y=(X*weights).sum(1).astype('float32')
107    xt=torch.tensor(X[:ntrain],device=dev); yt=torch.tensor(y[:ntrain],device=dev)
108    xe=torch.tensor(X[ntrain:],device=dev); ye=torch.tensor(y[ntrain:],device=dev)
109    results={}
110    for name, reg, single in [("single_generator",0.,True),("switched_unregularized",0.,False),("switched_commutator_regularized",.08,False)]:
111        torch.manual_seed(SEED+10+len(results)); model=SwitchedSSM(8,2,reg>0,single).to(dev)
112        opt=torch.optim.Adam(model.parameters(),lr=.025)
113        for step in range(100):
114            pred=model(xt); loss=((pred-yt)**2).mean()+reg*model.penalties()
115            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
116        with torch.no_grad(): val=float(((model(xe)-ye)**2).mean().cpu())
117        results[name]={"validation_mse":val,**model.diagnostics()}
118    return results
119
120
121def main():
122    global DEVICE
123    try:
124        b=bch_check(); ro=regularizer_optimization(); seq=sequence_experiment()
125    except Exception as e:
126        # CUDA failures are explicitly allowed to fall back to CPU.
127        if DEVICE=="cuda":
128            DEVICE="cpu"; torch.cuda.empty_cache(); b=bch_check(); ro=regularizer_optimization(); seq=sequence_experiment()
129        else: raise
130    result={"device":DEVICE,"bch":b,"regularizer_optimization":ro,"sequence_experiment":seq}
131    Path("results.json").write_text(json.dumps(result,indent=2))
132    print(json.dumps(result,indent=2))
133if __name__=="__main__": main()