1import json, random, sys
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, sweep_baseline, make_report
  8
  9TAU = 0.20
 10ALPHA = (0.5, 0.5)
 11HIDDEN = 64
 12SEEDS = tuple(range(8))
 13SWEEP_SEEDS = tuple(range(4))
 14
 15class SwitchedSSM(nn.Module):
 16    def __init__(self, input_shape, out_dim, comm_lambda=0.0):
 17        super().__init__()
 18        self.d = HIDDEN
 19        self.lam = float(comm_lambda)
 20        eye = torch.eye(self.d)
 21        self.A = nn.Parameter(-0.7 * eye[None].repeat(2, 1, 1) + 0.03 * torch.randn(2, self.d, self.d))
 22        self.U = nn.Parameter(0.08 * torch.randn(2, 3, self.d))
 23        self.head = nn.Linear(self.d, out_dim)
 24
 25    def forward(self, x):
 26        z = x.reshape(x.shape[0], -1, 3)
 27        h = torch.zeros(x.shape[0], self.d, device=x.device, dtype=x.dtype)
 28        for t in range(z.shape[1]):
 29            for i, a in enumerate(ALPHA):
 30                E = torch.matrix_exp(a * TAU * self.A[i])
 31                h = h @ E.T + (a * z[:, t, :]) @ self.U[i]
 32        return self.head(h)
 33
 34    def comm_penalty(self):
 35        B0 = ALPHA[0] * TAU * self.A[0]
 36        B1 = ALPHA[1] * TAU * self.A[1]
 37        C = B0 @ B1 - B1 @ B0
 38        return (C * C).sum()
 39
 40def seed_all(s):
 41    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 42
 43def train_idea(model, ds, epochs, lr):
 44    try:
 45        device = 'cuda' if torch.cuda.is_available() else 'cpu'
 46        model = model.to(device)
 47        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 48        opt = torch.optim.Adam(model.parameters(), lr=lr)
 49        lossf = nn.MSELoss()
 50        for _ in range(epochs):
 51            model.train()
 52            for ix in torch.randperm(len(x), device=device).split(128):
 53                loss = lossf(model(x[ix]), y[ix]) + model.lam * model.comm_penalty()
 54                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
 55        model.eval()
 56        with torch.no_grad(): metric = float(lossf(model(ds['xte'].to(device)), ds['yte'].to(device)).cpu())
 57        return model, metric
 58    except Exception:
 59        if torch.cuda.is_available():
 60            torch.cuda.empty_cache()
 61        model = model.cpu(); x, y = ds['xtr'], ds['ytr']; opt = torch.optim.Adam(model.parameters(), lr=lr); lossf = nn.MSELoss()
 62        for _ in range(epochs):
 63            for ix in torch.randperm(len(x)).split(128):
 64                loss = lossf(model(x[ix]), y[ix]) + model.lam * model.comm_penalty()
 65                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
 66        with torch.no_grad(): metric = float(lossf(model(ds['xte']), ds['yte']))
 67        return model, metric
 68
 69def train_base(cfg, seed, keep=False):
 70    seed_all(seed); ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
 71    model = SwitchedSSM(ds['input_shape'], ds['out_dim'], 0.0)
 72    # Uses the standard harness path for the matched no-penalty switched baseline.
 73    model, metric, _ = train_model(model, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128, weight_decay=0.0)
 74    return (model, metric) if keep else metric
 75
 76def train_idea_cfg(cfg, seed, keep=False):
 77    seed_all(seed); ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
 78    model = SwitchedSSM(ds['input_shape'], ds['out_dim'], cfg['comm_lambda'])
 79    model, metric = train_idea(model, ds, cfg['epochs'], cfg['lr'])
 80    return (model, metric) if keep else metric
 81
 82def base_factory(cfg):
 83    return lambda seed: train_base(cfg, seed)
 84
 85def metric_and_models(fn, cfg):
 86    vals=[]; models=[]
 87    for s in SEEDS:
 88        m,v = fn(cfg, s, True); vals.append(v); models.append(m)
 89    return {'config':cfg, 'per_seed':vals, 'mean':float(np.mean(vals)), 'models':models}
 90
 91def diagnostics(models):
 92    comm=[]; gain=[]
 93    for m in models:
 94        with torch.no_grad():
 95            A=m.A.detach().cpu(); B0=ALPHA[0]*TAU*A[0]; B1=ALPHA[1]*TAU*A[1]
 96            C=B0@B1-B1@B0; Phi=torch.matrix_exp(B1)@torch.matrix_exp(B0)
 97            comm.append(float((C*C).sum())); gain.append(float(torch.linalg.svdvals(Phi).max()))
 98    return {'observed_commutator_mean':float(np.mean(comm)), 'observed_cycle_gain_mean':float(np.mean(gain)), 'n_models':len(models)}
 99
100def main():
101    # Equal union: every idea lr is included in baseline tuning.
102    grid=[{'lr':lr,'epochs':20} for lr in (0.003,0.006,0.009)]
103    tuned=sweep_baseline(base_factory, grid, seeds=SWEEP_SEEDS)
104    best_cfg=tuned['best_cfg']
105    base_full=metric_and_models(lambda c,s,k: train_base(c,s,k), best_cfg)
106    idea_runs=[]
107    for lam in (0.0, 0.01, 0.05):
108        cfg=dict(best_cfg, comm_lambda=lam)
109        idea_runs.append(metric_and_models(lambda c,s,k: train_idea_cfg(c,s,k), cfg))
110    idea=min(idea_runs, key=lambda r:r['mean'])
111    bd=diagnostics(base_full['models']); id=diagnostics(idea['models'])
112    report=make_report('dynamics','rnn_small',
113        {'best_cfg':best_cfg, 'sweep':tuned['sweep'], 'full':{'per_seed':base_full['per_seed'],'mean':base_full['mean']}},
114        {'config':idea['config'],'per_seed':idea['per_seed'],'mean':idea['mean']},
115        {'track_justification':'Dynamics directly tests switched latent stability/control and Lyapunov contraction.',
116         'mechanism_signature':{'prediction':'commutator regularization lowers trained ordered-cycle noncommutativity and cycle gain',
117             'baseline_observed':bd, 'idea_observed':id,
118             'commutator_ratio':id['observed_commutator_mean']/max(bd['observed_commutator_mean'],1e-12),
119             'cycle_gain_delta':id['observed_cycle_gain_mean']-bd['observed_cycle_gain_mean'],
120             'confirmed': bool(id['observed_commutator_mean'] < bd['observed_commutator_mean'] and id['observed_cycle_gain_mean'] <= bd['observed_cycle_gain_mean']*1.02)}})
121    report['idea_sweep']=[{k:r[k] for k in ('config','mean','per_seed')} for r in idea_runs]
122    Path('bench_report.json').write_text(json.dumps(report, indent=2, allow_nan=False))
123    print(json.dumps(report, indent=2))
124
125if __name__=='__main__': main()