1import json, os, random
  2from collections import deque
  3import numpy as np
  4import torch
  5import sys
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7import bench
  8
  9SEEDS = tuple(range(8))
 10# Shared union: every idea learning rate is also evaluated by baseline.
 11GRID = [
 12    {'lr': 0.001}, {'lr': 0.003}, {'lr': 0.006},
 13]
 14EPOCHS = 8
 15NTRAIN, NTEST = 400, 400
 16
 17class HankelResidualObserver:
 18    def __init__(self, t_ini=8, horizon=1, max_cols=40, ridge=1e-2, tau=1.0):
 19        self.t_ini, self.horizon = t_ini, horizon
 20        self.max_cols, self.ridge, self.tau = max_cols, ridge, tau
 21        self.r = deque(maxlen=t_ini + horizon + max_cols + 8)
 22        self.last = None
 23        self.last_q = None
 24    def update(self, residual):
 25        self.r.append(float(residual))
 26    def predict(self):
 27        if len(self.r) < self.t_ini + self.horizon:
 28            return 0.0, 0.0, None
 29        a = np.asarray(self.r, dtype=np.float64)
 30        n = len(a) - self.t_ini - self.horizon + 1
 31        starts = np.arange(max(0, n-self.max_cols), n)
 32        dp = np.stack([a[j:j+self.t_ini] for j in starts], axis=1)
 33        df = np.stack([a[j+self.t_ini:j+self.t_ini+self.horizon] for j in starts], axis=1)
 34        d = a[-self.t_ini:]
 35        gram = dp.T @ dp + self.ridge*np.eye(dp.shape[1])
 36        g = np.linalg.solve(gram, dp.T @ d)
 37        recon = dp @ g
 38        q = np.linalg.norm(recon-d)/(np.linalg.norm(d)+1e-8)
 39        gamma = float(np.clip(1.0-q/self.tau, 0.0, 1.0))
 40        pred = float((df @ g)[0]) * gamma
 41        return pred, gamma, float(q)
 42
 43def seed_all(seed):
 44    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 45    if torch.cuda.is_available():
 46        torch.cuda.manual_seed_all(seed)
 47
 48def fit(seed, lr):
 49    seed_all(seed)
 50    ds = bench.get_dataset('sequence', seed, n_train=NTRAIN, n_test=NTEST)
 51    model = bench.make_model('transformer_tiny', ds['input_shape'], ds['out_dim'])
 52    net, metric, hist = bench.train_model(model, ds, epochs=EPOCHS, lr=lr,
 53                                          batch=128, weight_decay=0.0, log=lambda _: None)
 54    if net is None:
 55        raise RuntimeError('bench training failed')
 56    return net, ds
 57
 58def score(seed, lr, observer_on=False, collect_signature=False):
 59    net, ds = fit(seed, lr)
 60    net.eval()
 61    model_device = next(net.parameters()).device
 62    x, y = ds['xte'], ds['yte']
 63    with torch.no_grad():
 64        # Keep batches small and use the harness-produced trained model.
 65        pred = []
 66        for i in range(0, len(x), 128):
 67            pred.append(net(x[i:i+128].to(model_device)).detach().cpu().numpy().reshape(-1))
 68    base = np.concatenate(pred).astype(np.float64)
 69    actual = y.detach().cpu().numpy().reshape(-1).astype(np.float64)
 70    if not observer_on:
 71        return float(np.mean((base-actual)**2)), None
 72    obs = HankelResidualObserver()
 73    corrected, gammas, qs, residual_forecasts, next_residuals = [], [], [], [], []
 74    # Strictly causal: residual at sample i is appended only after predicting i.
 75    for i in range(len(actual)):
 76        dh, gamma, q = obs.predict()
 77        corrected.append(base[i] + dh)
 78        gammas.append(gamma); qs.append(q if q is not None else np.nan)
 79        if q is not None:
 80            residual_forecasts.append(dh/gamma if gamma > 1e-12 else 0.0)
 81            next_residuals.append(actual[i]-base[i])
 82        obs.update(actual[i]-base[i])
 83    corr = np.asarray(corrected)
 84    result = float(np.mean((corr-actual)**2))
 85    sig = {
 86        'residual_forecast_mse': float(np.mean((np.asarray(residual_forecasts)-np.asarray(next_residuals))**2)) if next_residuals else None,
 87        'residual_naive_mse': float(np.mean(np.asarray(next_residuals)**2)) if next_residuals else None,
 88        'correction_mse': result,
 89        'base_mse': float(np.mean((base-actual)**2)),
 90        'mean_gamma': float(np.mean(gammas)),
 91        'active_fraction': float(np.mean(np.asarray(gammas)>0)),
 92        'mean_mismatch_q': float(np.nanmean(qs)),
 93    }
 94    # Quantitative confirmation requires both causal residual predictability and task improvement.
 95    sig['confirmed'] = bool(sig['residual_forecast_mse'] < sig['residual_naive_mse'] and result < sig['base_mse'])
 96    return result, sig
 97
 98def eval_cfg(cfg, idea=False, signature=False):
 99    vals, sigs = [], []
100    for s in SEEDS:
101        v, sig = score(s, cfg['lr'], observer_on=idea, collect_signature=signature)
102        vals.append(v)
103        if sig is not None: sigs.append(sig)
104    out = {'mean': float(np.mean(vals)), 'std': float(np.std(vals)),
105           'per_seed': [float(v) for v in vals], 'n': len(vals)}
106    if sigs:
107        keys = ['residual_forecast_mse','residual_naive_mse','correction_mse','base_mse','mean_gamma','active_fraction','mean_mismatch_q']
108        out['signature_summary'] = {k: float(np.mean([z[k] for z in sigs])) for k in keys}
109    return out
110
111def baseline_factory(cfg):
112    return lambda seed: score(seed, cfg['lr'], observer_on=False)[0]
113
114def main():
115    # sweep_baseline itself is the required baseline selection mechanism; use all 8 seeds.
116    base = bench.sweep_baseline(baseline_factory, GRID, seeds=SEEDS)
117    idea_trials = []
118    for cfg in GRID:
119        r = eval_cfg(cfg, idea=True)
120        idea_trials.append({'cfg': cfg, **r})
121    best = min(idea_trials, key=lambda z: z['mean'])
122    idea = {k: best[k] for k in ('mean','std','per_seed','n')}
123    sig = eval_cfg(best['cfg'], idea=True, signature=True)['signature_summary']
124    # make_report computes paired deltas and the mandated sign-permutation p-value.
125    report = bench.make_report('sequence', 'transformer_tiny', base, idea,
126        {'mechanism_signature': {**sig, 'confirmed': bool(sig['residual_forecast_mse'] < sig['residual_naive_mse'] and sig['correction_mse'] < sig['base_mse'])},
127         'idea_sweep': [{'cfg': z['cfg'], 'mean': z['mean'], 'std': z['std']} for z in idea_trials],
128         'protocol': {'seeds': list(SEEDS), 'epochs': EPOCHS, 'n_train': NTRAIN, 'n_test': NTEST,
129                      'shared_lr_grid': GRID, 'observer': {'t_ini':8,'horizon':1,'max_cols':40,'ridge':1e-2,'tau':1.0}}})
130    with open('bench_report.json','w') as f: json.dump(report, f, indent=2)
131    print(json.dumps(report, indent=2))
132
133if __name__ == '__main__':
134    main()