1import sys, json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, make_report, sweep_baseline, evaluate
  8
  9SEEDS = tuple(range(8))
 10SWEEP_SEEDS = (0, 1, 2, 3)
 11LRS = [1e-3, 3e-3, 6e-3]
 12EPOCHS = 18
 13BATCH = 128
 14H = 64
 15STEP = 0.12
 16LAMBDA = 0.10
 17
 18
 19def seed_all(seed):
 20    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 21    try:
 22        if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 23    except Exception:
 24        pass
 25
 26
 27def tangent_project(x, q):
 28    p = torch.where(x <= -1.0, torch.clamp(q, min=0.0), q)
 29    return torch.where(x >= 1.0, torch.clamp(p, max=0.0), p)
 30
 31
 32class ResidualRNN(nn.Module):
 33    def __init__(self, out_dim=1):
 34        super().__init__()
 35        self.inp = nn.Linear(3, H)
 36        self.W = nn.Linear(H, H, bias=False)
 37        self.bias = nn.Parameter(torch.zeros(H))
 38        self.head = nn.Linear(H, out_dim)
 39        nn.init.normal_(self.W.weight, std=0.08 / math.sqrt(H))
 40
 41    def hidden(self, x, z0=None):
 42        s = x.view(x.shape[0], -1, 3)
 43        z = torch.zeros(x.shape[0], H, device=x.device, dtype=x.dtype) if z0 is None else z0
 44        for t in range(s.shape[1]):
 45            z = z + STEP * torch.tanh(self.W(z) + self.inp(s[:, t]) + self.bias)
 46        return z
 47
 48    def forward(self, x):
 49        return self.head(self.hidden(x))
 50
 51
 52class ProjectedDissipativeRNN(nn.Module):
 53    def __init__(self, out_dim=1):
 54        super().__init__()
 55        self.inp = nn.Linear(3, H)
 56        self.L = nn.Parameter(torch.randn(H, H) * (0.08 / math.sqrt(H)))
 57        self.bias = nn.Parameter(torch.zeros(H))
 58        self.head = nn.Linear(H, out_dim)
 59
 60    def matrix(self):
 61        return self.L.T @ self.L + LAMBDA * torch.eye(H, device=self.L.device, dtype=self.L.dtype)
 62
 63    def hidden(self, x, z0=None):
 64        s = x.view(x.shape[0], -1, 3)
 65        z = torch.zeros(x.shape[0], H, device=x.device, dtype=x.dtype) if z0 is None else z0
 66        M = self.matrix()
 67        for t in range(s.shape[1]):
 68            q = -(z @ M.T) + self.inp(s[:, t]) + self.bias
 69            z = torch.clamp(z + STEP * tangent_project(z, q), -1.0, 1.0)
 70        return z
 71
 72    def forward(self, x):
 73        return self.head(self.hidden(x))
 74
 75
 76def train_one(kind, seed, lr, keep=False):
 77    seed_all(seed)
 78    ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
 79    model = ResidualRNN(ds['out_dim']) if kind == 'baseline' else ProjectedDissipativeRNN(ds['out_dim'])
 80    net, metric, history = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
 81    if net is None:
 82        raise RuntimeError('bench training failed')
 83    return (float(metric), net, ds) if keep else float(metric)
 84
 85
 86def baseline_factory(cfg):
 87    return lambda seed: train_one('baseline', seed, cfg['lr'])
 88
 89
 90def idea_factory(cfg):
 91    return lambda seed: train_one('idea', seed, cfg['lr'])
 92
 93
 94def idea_sweep():
 95    tried = []
 96    for cfg in [{'lr': lr, 'epochs': EPOCHS, 'batch': BATCH} for lr in LRS]:
 97        r = evaluate(idea_factory(cfg), SWEEP_SEEDS)
 98        tried.append({'cfg': cfg, 'mean': r['mean']})
 99    best = min(tried, key=lambda a: a['mean'])
100    return {'best_cfg': best['cfg'], 'sweep': tried}
101
102
103def signature(base, idea):
104    ratios_b, ratios_i, boxes = [], [], []
105    for (mb, db), (mi, di) in zip(base, idea):
106        x = db['xte'][:32]
107        device = next(mb.parameters()).device
108        x = x.to(device)
109        zb = torch.zeros(x.shape[0], H, device=device)
110        zi = torch.zeros(x.shape[0], H, device=device)
111        with torch.no_grad():
112            delta = 1e-3 * torch.randn_like(zb)
113            rb = (mb.hidden(x, zb + delta) - mb.hidden(x, zb)).norm(dim=1) / (delta.norm(dim=1) + 1e-12)
114            ri = (mi.hidden(x, zi + delta) - mi.hidden(x, zi)).norm(dim=1) / (delta.norm(dim=1) + 1e-12)
115            ratios_b.extend(rb.cpu().numpy().tolist()); ratios_i.extend(ri.cpu().numpy().tolist())
116            boxes.append(float(mi.hidden(x).abs().max().cpu()))
117    mean_b, mean_i = float(np.mean(ratios_b)), float(np.mean(ratios_i))
118    return {'prediction': 'projected dissipative dynamics reduces initial-state perturbation amplification and keeps hidden states in [-1,1]', 'baseline_perturbation_ratio': mean_b, 'idea_perturbation_ratio': mean_i, 'baseline_to_idea_ratio': mean_i / (mean_b + 1e-12), 'idea_max_abs_hidden': float(max(boxes)), 'confirmed': bool(mean_i < mean_b and max(boxes) <= 1.00001)}
119
120
121def main():
122    grid = [{'lr': lr, 'epochs': EPOCHS, 'batch': BATCH} for lr in LRS]
123    base = sweep_baseline(baseline_factory, grid, seeds=SWEEP_SEEDS)
124    isweep = idea_sweep()
125    bcfg, icfg = base['best_cfg'], isweep['best_cfg']
126    bfull = evaluate(baseline_factory(bcfg), SEEDS)
127    if icfg['lr'] == bcfg['lr']:
128        ifull = evaluate(idea_factory(icfg), SEEDS)
129    else:
130        ifull = evaluate(idea_factory(icfg), SEEDS)
131    bmodels = [train_one('baseline', s, bcfg['lr'], keep=True)[1:] for s in SEEDS]
132    imodels = [train_one('idea', s, icfg['lr'], keep=True)[1:] for s in SEEDS]
133    bblock = {'sweep': base['sweep'], 'best_config': bcfg, 'full': bfull}
134    rep = make_report('dynamics', 'rnn_small_matched_residual', bblock, {'config': icfg, **ifull}, {'mechanism_signature': signature(bmodels, imodels), 'track_rationale': 'Dynamics is the built-in structural match for stability/control and Lyapunov-style contraction.'})
135    rep['idea_sweep'] = isweep
136    rep['parameterization'] = {'hidden': H, 'step': STEP, 'lambda': LAMBDA, 'same_input_and_head': True, 'lr_union_tested_on_both': LRS}
137    with open('bench_report.json', 'w') as f: json.dump(rep, f, indent=2)
138    print(json.dumps(rep, indent=2))
139
140if __name__ == '__main__':
141    main()