import sys, json, random from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, train_model, sweep_baseline, make_report SEED = 449 EPOCHS = 12 BATCH = 128 WIDTH = 16 MODES = 2 TAU = 0.7 LR_GRID = [1e-3, 3e-3, 6e-3] WD_GRID = [0.0, 1e-4] IDEA_GRID = [dict(lr=x, weight_decay=0.0, lambda_c=0.02) for x in LR_GRID] def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) class SwitchedCore(nn.Module): def __init__(self, width=WIDTH, modes=MODES): super().__init__() self.width, self.modes = width, modes eye = torch.eye(width).unsqueeze(0).repeat(modes, 1, 1) self.A = nn.Parameter(-0.35 * eye + 0.04 * torch.randn(modes, width, width)) self.U = nn.Parameter(0.10 * torch.randn(modes, width)) self.head = nn.Linear(width, 1) def matrices(self): B = [TAU / self.modes * self.A[i] for i in range(self.modes)] E = [torch.matrix_exp(b) for b in B] return B, E def forward(self, x): # dynamics input is [batch, 8*3], ordered as (theta, omega, u) z = x.view(x.shape[0], -1, 3) h = torch.zeros(x.shape[0], self.width, device=x.device, dtype=x.dtype) B, E = self.matrices() for t in range(z.shape[1]): inp = z[:, t, 2:3] for i in range(self.modes): h = h @ E[i].T + (TAU / self.modes) * inp * self.U[i] return self.head(h) def commutator_penalty(self): B, _ = self.matrices() total = torch.zeros((), device=self.A.device) for i in range(self.modes): for j in range(i): C = B[i] @ B[j] - B[j] @ B[i] total = total + (C * C).sum() return total def train_switched(ds, seed, lr, weight_decay=0.0, lambda_c=0.0): seed_all(seed + 10000) model = SwitchedCore() # This custom loop is necessary because the intervention is a new loss. try: device = 'cuda' if torch.cuda.is_available() else 'cpu' model = model.to(device) xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device) xte, yte = ds['xte'].to(device), ds['yte'].to(device) opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) hist = [] for _ in range(EPOCHS): model.train(); perm = torch.randperm(len(xtr), device=device); total = 0.0 for k in range(0, len(xtr), BATCH): ix = perm[k:k+BATCH] pred = model(xtr[ix]) loss = ((pred - ytr[ix]) ** 2).mean() + lambda_c * model.commutator_penalty() opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step() total += float(loss.detach()) * len(ix) hist.append(total / len(xtr)) model.eval() with torch.no_grad(): metric = float(((model(xte) - yte) ** 2).mean().cpu()) return model, metric, hist except RuntimeError: if device == 'cuda': torch.cuda.empty_cache() return train_switched_cpu(ds, seed, lr, weight_decay, lambda_c) raise def train_switched_cpu(ds, seed, lr, weight_decay=0.0, lambda_c=0.0): seed_all(seed + 10000) model = SwitchedCore().cpu() xtr, ytr, xte, yte = ds['xtr'], ds['ytr'], ds['xte'], ds['yte'] opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) for _ in range(EPOCHS): perm = torch.randperm(len(xtr)) for k in range(0, len(xtr), BATCH): ix = perm[k:k+BATCH]; pred = model(xtr[ix]) loss = ((pred-ytr[ix])**2).mean() + lambda_c*model.commutator_penalty() opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step() model.eval() with torch.no_grad(): metric = float(((model(xte)-yte)**2).mean()) return model, metric, [] def baseline_train(ds, seed, cfg): # Standard practice: same switched recurrent architecture, no regularizer. _, metric, _ = train_switched(ds, seed, cfg['lr'], cfg['weight_decay'], 0.0) return metric def idea_train(ds, seed, cfg): _, metric, _ = train_switched(ds, seed, cfg['lr'], cfg['weight_decay'], cfg['lambda_c']) return metric def bch_check(): from scipy.linalg import expm, norm rng = np.random.default_rng(SEED); X = rng.normal(size=(3,3)); Y = rng.normal(size=(3,3)) C = X@Y-Y@X; rows=[] for s in [0.02,0.04,0.08,0.16]: exact=expm(s*Y)@expm(s*X) corr=expm(s*(X+Y)-0.5*s*s*C) rows.append([s,norm(exact-expm(s*(X+Y)),'fro'),norm(exact-corr,'fro')]) a=np.polyfit(np.log(np.array(rows)[-3:,0]),np.log(np.array(rows)[-3:,1]),1)[0] b=np.polyfit(np.log(np.array(rows)[-3:,0]),np.log(np.array(rows)[-3:,2]),1)[0] return {'rows':rows,'uncorrected_slope':float(a),'corrected_slope':float(b),'confirmed':bool(b>a+0.5)} def signature(models, ds): vals=[] for label, model in models: with torch.no_grad(): B,E=model.matrices(); phi=E[1]@E[0]; avg=torch.matrix_exp(B[0]+B[1]) comm=(B[1]@B[0]-B[0]@B[1]).norm().item()**2 mismatch=(phi-avg).norm().item()/max(avg.norm().item(),1e-12) gain=torch.linalg.matrix_norm(phi,2).item() vals.append({'system':label,'commutator':comm,'ordered_avg_mismatch':mismatch,'cycle_gain':gain}) return {'prediction':'lower commutator should accompany lower ordered-vs-averaged mismatch', '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} def main(): track='dynamics'; model_name='rnn_small' # Benchmark-compatible dataset; model is a matched end-to-end switched recurrent replacement. datasets={s:get_dataset(track, s, n_train=400, n_test=400) for s in range(8)} base_grid=[{'lr':lr,'weight_decay':wd} for lr in LR_GRID for wd in WD_GRID] base_block=sweep_baseline(lambda cfg: (lambda s: baseline_train(datasets[s],s,cfg)), base_grid) idea_results=[]; best_cfg=min(IDEA_GRID, key=lambda c: np.mean([idea_train(datasets[s],s,c) for s in range(4)])) for s in range(8): idea_results.append(idea_train(datasets[s],s,best_cfg)) idea_res={'mean':float(np.mean(idea_results)),'std':float(np.std(idea_results)), 'per_seed':idea_results,'n':8} models=[] for s in range(2): models.append(('baseline', train_switched(datasets[s],s, base_block['best_cfg']['lr'],base_block['best_cfg']['weight_decay'],0.0)[0])) models.append(('idea', train_switched(datasets[s],s, best_cfg['lr'],0.0,best_cfg['lambda_c'])[0])) report=make_report(track, model_name, base_block, idea_res, {'mechanism_signature':signature(models,datasets[0]),'bch_check':bch_check(),'idea_cfg':best_cfg}) report['structural_match']='Dynamics/control track matches stability and Lyapunov structure.' Path('bench_report.json').write_text(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__=='__main__': main()