1import sys, json, random
  2import numpy as np
  3import torch
  4from torch import nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
  7
  8SEEDS = tuple(range(8))
  9SWEEP_SEEDS = (0, 1, 2, 3)
 10LRS = [1e-3, 3e-3, 1e-2]
 11EPOCHS = 24
 12BATCH = 128
 13H = 4.0
 14ALPHA = 0.5
 15
 16
 17def math_check():
 18    hs = np.array([0.5, 1., 2., 4., 8., 32.])
 19    q = hs / (1.0 + ALPHA * hs)
 20    # Contractive scalar test x'=-x: amplification is |1-step|.
 21    euler = np.abs(1.0 - hs)
 22    damped = np.abs(1.0 - q)
 23    return {'h': hs.tolist(), 'q': q.tolist(),
 24            'q_bound': 1.0 / ALPHA,
 25            'euler_amplification': euler.tolist(),
 26            'damped_amplification': damped.tolist(),
 27            'predicted_bounded_effective_step': bool(np.all(q < 1.0 / ALPHA)),
 28            'damped_nonexpansive_test': bool(np.all(damped <= 1.0 + 1e-12)),
 29            'euler_nonexpansive_test': bool(np.all(euler <= 1.0 + 1e-12))}
 30
 31
 32class ResidualDynamics(nn.Module):
 33    """Eight-step residual state model; only the integration rule differs."""
 34    def __init__(self, mode, hidden=64, nominal_h=H, alpha=ALPHA):
 35        super().__init__()
 36        self.mode, self.h, self.alpha = mode, nominal_h, alpha
 37        self.inp = nn.Linear(3, hidden)
 38        self.field = nn.Sequential(nn.Linear(hidden, hidden), nn.Tanh(),
 39                                   nn.Linear(hidden, hidden))
 40        self.head = nn.Linear(hidden, 1)
 41        self.last = {}
 42
 43    def vector_field(self, z, u):
 44        # Shared feature map; input injection is identical in both systems.
 45        return self.field(z + self.inp(u))
 46
 47    def forward(self, x):
 48        seq = x.view(x.shape[0], -1, 3)
 49        z = torch.zeros(x.shape[0], self.inp.out_features, device=x.device, dtype=x.dtype)
 50        q = self.h / (1.0 + self.alpha * self.h)
 51        max_norm, max_update, max_resid = 0., 0., 0.
 52        for k in range(seq.shape[1]):
 53            u = seq[:, k]
 54            if self.mode == 'baseline':
 55                f = self.vector_field(z, u)
 56                zn = z + self.h * f
 57                resid = torch.zeros((), device=z.device)
 58            else:
 59                f1 = self.vector_field(z, u)
 60                mid = z + 0.5 * q * f1
 61                f2 = self.vector_field(mid, u)
 62                zn = z + q * f2
 63                resid = (zn - (z + q * f1)).norm(dim=1).mean()
 64            upd = (zn - z).norm(dim=1).mean()
 65            max_norm = max(max_norm, float(zn.detach().norm(dim=1).max()))
 66            max_update = max(max_update, float(upd.detach()))
 67            max_resid = max(max_resid, float(resid.detach()))
 68            z = zn
 69        self.last = {'max_activation': max_norm, 'max_update': max_update,
 70                     'stage_residual': max_resid,
 71                     'finite': bool(torch.isfinite(z).all().item())}
 72        return self.head(z)
 73
 74
 75def factory(mode, lr, seed):
 76    # Explicit deterministic pairing: same initialization for both systems.
 77    torch.manual_seed(1000 + int(seed))
 78    np.random.seed(1000 + int(seed)); random.seed(1000 + int(seed))
 79    ds = get_dataset('dynamics', seed=int(seed), n_train=400, n_test=200)
 80    model = ResidualDynamics(mode)
 81    net, metric, history = train_model(model, ds, epochs=EPOCHS, lr=lr,
 82                                       batch=BATCH, weight_decay=0.0,
 83                                       log=lambda *_: None)
 84    if net is None or metric is None:
 85        return float('inf'), {'failed': True}
 86    stats = dict(net.last)
 87    stats['test_metric'] = float(metric)
 88    stats['nan'] = not bool(stats.pop('finite', False))
 89    return float(metric), stats
 90
 91
 92def make_fn(mode, cfg):
 93    def run(seed):
 94        val, stats = factory(mode, cfg['lr'], seed)
 95        RUN_STATS.setdefault(mode, {}).setdefault(str(cfg['lr']), {})[str(seed)] = stats
 96        return val
 97    return run
 98
 99
100def signature(base_res, idea_res):
101    b = [v for v in base_res.values() if v]
102    i = [v for v in idea_res.values() if v]
103    br = float(np.mean([x['max_update'] for x in b])) if b else float('nan')
104    ir = float(np.mean([x['max_update'] for x in i])) if i else float('nan')
105    # Prediction is damping of the actual per-step update at the trained NN scale.
106    ratio = ir / br if br > 0 else float('nan')
107    return {'prediction': 'denominator block has smaller trained-state update proxy at h=4',
108            'nominal_h': H, 'alpha': ALPHA, 'observed_baseline_update': br,
109            'observed_idea_update': ir, 'observed_update_ratio': ratio,
110            'predicted_upper_ratio_from_q_over_h': (H/(1+ALPHA*H))/H,
111            'confirmed': bool(np.isfinite(ratio) and ratio < 0.85)}
112
113
114if __name__ == '__main__':
115    RUN_STATS = {}
116    check = math_check()
117    grid = [{'lr': x} for x in LRS]
118    base = sweep_baseline(lambda cfg: make_fn('baseline', cfg), grid, seeds=SWEEP_SEEDS)
119    # Evaluate all shared learning rates on the idea side; the best is selected only
120    # after the same-sized, parity-complete search.
121    idea_trials = []
122    for cfg in grid:
123        r = evaluate(make_fn('idea', cfg), SEEDS)
124        idea_trials.append({'cfg': cfg, 'result': r})
125    best = min(idea_trials, key=lambda a: a['result']['mean'])
126    rep = make_report('dynamics', 'residual_rnn_small', base, best['result'], {
127        'math_check': check,
128        'idea_config': best['cfg'],
129        'idea_sweep': idea_trials,
130        'mechanism_signature': signature(
131            RUN_STATS.get('baseline', {}).get(str(base['best_cfg']['lr']), {}),
132            RUN_STATS.get('idea', {}).get(str(best['cfg']['lr']), {}))})
133    with open('bench_report.json', 'w') as f:
134        json.dump(rep, f, indent=2)
135    print(json.dumps(rep, indent=2))