1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, train_model, sweep_baseline, make_report
  9
 10SEED = 449
 11EPOCHS = 12
 12BATCH = 128
 13WIDTH = 16
 14MODES = 2
 15TAU = 0.7
 16LR_GRID = [1e-3, 3e-3, 6e-3]
 17WD_GRID = [0.0, 1e-4]
 18IDEA_GRID = [dict(lr=x, weight_decay=0.0, lambda_c=0.02)
 19             for x in LR_GRID]
 20
 21
 22def seed_all(seed):
 23    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 24    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 25
 26
 27class SwitchedCore(nn.Module):
 28    def __init__(self, width=WIDTH, modes=MODES):
 29        super().__init__()
 30        self.width, self.modes = width, modes
 31        eye = torch.eye(width).unsqueeze(0).repeat(modes, 1, 1)
 32        self.A = nn.Parameter(-0.35 * eye + 0.04 * torch.randn(modes, width, width))
 33        self.U = nn.Parameter(0.10 * torch.randn(modes, width))
 34        self.head = nn.Linear(width, 1)
 35
 36    def matrices(self):
 37        B = [TAU / self.modes * self.A[i] for i in range(self.modes)]
 38        E = [torch.matrix_exp(b) for b in B]
 39        return B, E
 40
 41    def forward(self, x):
 42        # dynamics input is [batch, 8*3], ordered as (theta, omega, u)
 43        z = x.view(x.shape[0], -1, 3)
 44        h = torch.zeros(x.shape[0], self.width, device=x.device, dtype=x.dtype)
 45        B, E = self.matrices()
 46        for t in range(z.shape[1]):
 47            inp = z[:, t, 2:3]
 48            for i in range(self.modes):
 49                h = h @ E[i].T + (TAU / self.modes) * inp * self.U[i]
 50        return self.head(h)
 51
 52    def commutator_penalty(self):
 53        B, _ = self.matrices()
 54        total = torch.zeros((), device=self.A.device)
 55        for i in range(self.modes):
 56            for j in range(i):
 57                C = B[i] @ B[j] - B[j] @ B[i]
 58                total = total + (C * C).sum()
 59        return total
 60
 61
 62def train_switched(ds, seed, lr, weight_decay=0.0, lambda_c=0.0):
 63    seed_all(seed + 10000)
 64    model = SwitchedCore()
 65    # This custom loop is necessary because the intervention is a new loss.
 66    try:
 67        device = 'cuda' if torch.cuda.is_available() else 'cpu'
 68        model = model.to(device)
 69        xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device)
 70        xte, yte = ds['xte'].to(device), ds['yte'].to(device)
 71        opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
 72        hist = []
 73        for _ in range(EPOCHS):
 74            model.train(); perm = torch.randperm(len(xtr), device=device); total = 0.0
 75            for k in range(0, len(xtr), BATCH):
 76                ix = perm[k:k+BATCH]
 77                pred = model(xtr[ix])
 78                loss = ((pred - ytr[ix]) ** 2).mean() + lambda_c * model.commutator_penalty()
 79                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
 80                total += float(loss.detach()) * len(ix)
 81            hist.append(total / len(xtr))
 82        model.eval()
 83        with torch.no_grad(): metric = float(((model(xte) - yte) ** 2).mean().cpu())
 84        return model, metric, hist
 85    except RuntimeError:
 86        if device == 'cuda':
 87            torch.cuda.empty_cache()
 88            return train_switched_cpu(ds, seed, lr, weight_decay, lambda_c)
 89        raise
 90
 91
 92def train_switched_cpu(ds, seed, lr, weight_decay=0.0, lambda_c=0.0):
 93    seed_all(seed + 10000)
 94    model = SwitchedCore().cpu()
 95    xtr, ytr, xte, yte = ds['xtr'], ds['ytr'], ds['xte'], ds['yte']
 96    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
 97    for _ in range(EPOCHS):
 98        perm = torch.randperm(len(xtr))
 99        for k in range(0, len(xtr), BATCH):
100            ix = perm[k:k+BATCH]; pred = model(xtr[ix])
101            loss = ((pred-ytr[ix])**2).mean() + lambda_c*model.commutator_penalty()
102            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
103    model.eval()
104    with torch.no_grad(): metric = float(((model(xte)-yte)**2).mean())
105    return model, metric, []
106
107
108def baseline_train(ds, seed, cfg):
109    # Standard practice: same switched recurrent architecture, no regularizer.
110    _, metric, _ = train_switched(ds, seed, cfg['lr'], cfg['weight_decay'], 0.0)
111    return metric
112
113
114def idea_train(ds, seed, cfg):
115    _, metric, _ = train_switched(ds, seed, cfg['lr'], cfg['weight_decay'], cfg['lambda_c'])
116    return metric
117
118
119def bch_check():
120    from scipy.linalg import expm, norm
121    rng = np.random.default_rng(SEED); X = rng.normal(size=(3,3)); Y = rng.normal(size=(3,3))
122    C = X@Y-Y@X; rows=[]
123    for s in [0.02,0.04,0.08,0.16]:
124        exact=expm(s*Y)@expm(s*X)
125        corr=expm(s*(X+Y)-0.5*s*s*C)
126        rows.append([s,norm(exact-expm(s*(X+Y)),'fro'),norm(exact-corr,'fro')])
127    a=np.polyfit(np.log(np.array(rows)[-3:,0]),np.log(np.array(rows)[-3:,1]),1)[0]
128    b=np.polyfit(np.log(np.array(rows)[-3:,0]),np.log(np.array(rows)[-3:,2]),1)[0]
129    return {'rows':rows,'uncorrected_slope':float(a),'corrected_slope':float(b),'confirmed':bool(b>a+0.5)}
130
131
132def signature(models, ds):
133    vals=[]
134    for label, model in models:
135        with torch.no_grad():
136            B,E=model.matrices(); phi=E[1]@E[0]; avg=torch.matrix_exp(B[0]+B[1])
137            comm=(B[1]@B[0]-B[0]@B[1]).norm().item()**2
138            mismatch=(phi-avg).norm().item()/max(avg.norm().item(),1e-12)
139            gain=torch.linalg.matrix_norm(phi,2).item()
140            vals.append({'system':label,'commutator':comm,'ordered_avg_mismatch':mismatch,'cycle_gain':gain})
141    return {'prediction':'lower commutator should accompany lower ordered-vs-averaged mismatch',
142            'observed':vals,'predicted_sign': 'positive','confirmed': bool(vals[1]['commutator'] < vals[0]['commutator'] and vals[1]['ordered_avg_mismatch'] < vals[0]['ordered_avg_mismatch']) if len(vals)==2 else False}
143
144
145def main():
146    track='dynamics'; model_name='rnn_small'
147    # Benchmark-compatible dataset; model is a matched end-to-end switched recurrent replacement.
148    datasets={s:get_dataset(track, s, n_train=400, n_test=400) for s in range(8)}
149    base_grid=[{'lr':lr,'weight_decay':wd} for lr in LR_GRID for wd in WD_GRID]
150    base_block=sweep_baseline(lambda cfg: (lambda s: baseline_train(datasets[s],s,cfg)), base_grid)
151    idea_results=[]; best_cfg=min(IDEA_GRID, key=lambda c: np.mean([idea_train(datasets[s],s,c) for s in range(4)]))
152    for s in range(8): idea_results.append(idea_train(datasets[s],s,best_cfg))
153    idea_res={'mean':float(np.mean(idea_results)),'std':float(np.std(idea_results)), 'per_seed':idea_results,'n':8}
154    models=[]
155    for s in range(2):
156        models.append(('baseline', train_switched(datasets[s],s, base_block['best_cfg']['lr'],base_block['best_cfg']['weight_decay'],0.0)[0]))
157        models.append(('idea', train_switched(datasets[s],s, best_cfg['lr'],0.0,best_cfg['lambda_c'])[0]))
158    report=make_report(track, model_name, base_block, idea_res, {'mechanism_signature':signature(models,datasets[0]),'bch_check':bch_check(),'idea_cfg':best_cfg})
159    report['structural_match']='Dynamics/control track matches stability and Lyapunov structure.'
160    Path('bench_report.json').write_text(json.dumps(report,indent=2))
161    print(json.dumps(report,indent=2))
162
163if __name__=='__main__': main()