1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, train_model, sweep_baseline, evaluate, make_report
  9
 10SEEDS = tuple(range(8))
 11TRACE = []
 12
 13class PositiveGainRNN(nn.Module):
 14    def __init__(self, hidden=64, mode='fixed', gain=1.0, a=1.0, b=0.2, dt=0.2):
 15        super().__init__()
 16        self.hidden, self.mode, self.gain = hidden, mode, gain
 17        self.a, self.b, self.dt = a, b, dt
 18        self.inp = nn.Linear(3, hidden)
 19        self.rec = nn.Linear(hidden, hidden)
 20        self.bias = nn.Parameter(torch.zeros(hidden))
 21        self.head = nn.Linear(hidden, 1)
 22        self.last_trace = None
 23
 24    def forward(self, x):
 25        seq = x.view(x.shape[0], -1, 3)
 26        h = torch.zeros(x.shape[0], self.hidden, device=x.device, dtype=x.dtype)
 27        w = torch.full((x.shape[0], 1), float(self.gain), device=x.device, dtype=x.dtype)
 28        ys, ws = [], []
 29        for k in range(seq.shape[1]):
 30            # Positive state and positive output projection, matching the intended plant.
 31            h = torch.relu((1.0 - self.dt) * h + self.dt * (self.inp(seq[:, k]) + w * self.rec(h) + self.bias))
 32            y = h.mean(dim=1, keepdim=True)
 33            if self.mode == 'pito':
 34                w = torch.relu(w + self.dt * (self.b - self.a * y * w))
 35            ys.append(y.detach())
 36            ws.append(w.detach())
 37        self.last_trace = (torch.stack(ys, 1), torch.stack(ws, 1))
 38        return self.head(h)
 39
 40def set_seed(seed):
 41    np.random.seed(seed); torch.manual_seed(seed)
 42    if torch.cuda.is_available():
 43        try: torch.cuda.manual_seed_all(seed)
 44        except Exception: pass
 45
 46def run_one(seed, cfg, mode, keep_trace=False):
 47    set_seed(seed)
 48    ds = get_dataset('dynamics', seed, n_train=400, n_test=400)
 49    if mode == 'fixed':
 50        net = PositiveGainRNN(mode='fixed', gain=cfg['gain'])
 51    else:
 52        net = PositiveGainRNN(mode='pito', gain=cfg.get('gain', 1.0), a=cfg['a'], b=cfg['b'])
 53    net, metric, hist = train_model(net, ds, epochs=20, lr=cfg['lr'], batch=128, log=lambda *_: None)
 54    if net is None or metric is None: return float('nan')
 55    if keep_trace:
 56        dev = next(net.parameters()).device
 57        with torch.no_grad(): net(ds['xte'][:128].to(dev))
 58        if net.last_trace is not None: TRACE.append((net.a, net.b, net.last_trace))
 59    return float(metric)
 60
 61def baseline_factory(cfg):
 62    return lambda seed: run_one(seed, cfg, 'fixed')
 63
 64def idea_factory(cfg):
 65    return lambda seed: run_one(seed, cfg, 'pito', keep_trace=True)
 66
 67def signature():
 68    # Re-test the theorem's prediction on traces of trained idea models. Select
 69    # a high-output tail threshold, then fit log(w - b/(aV)) where positive.
 70    rates, preds, thresholds = [], [], []
 71    for a, b, (ys, ws) in TRACE:
 72        y = ys.cpu().numpy().reshape(-1); w = ws.cpu().numpy().reshape(-1)
 73        V = float(np.quantile(y, .75))
 74        if V <= 1e-5: continue
 75        floor = b/(a*V); pred = a*V
 76        mask = (y >= V) & (w > floor + 1e-5)
 77        if mask.sum() >= 8:
 78            t = np.arange(len(w), dtype=float)[mask] * .2
 79            slope = -float(np.polyfit(t, np.log(w[mask]-floor), 1)[0])
 80            if np.isfinite(slope) and slope > 0:
 81                rates.append(slope); preds.append(pred); thresholds.append(V)
 82    if not rates:
 83        return {'prediction': 'decay slope = a*V when y>=V', 'predicted_rate': None,
 84                'observed_rate': None, 'relative_error': None, 'n_traces': 0,
 85                'confirmed': False}
 86    rel = abs(float(np.mean(rates))-float(np.mean(preds)))/float(np.mean(preds))
 87    return {'prediction': 'decay slope = a*V when y>=V',
 88            'predicted_rate': float(np.mean(preds)), 'observed_rate': float(np.mean(rates)),
 89            'relative_error': float(rel), 'threshold_V': float(np.mean(thresholds)),
 90            'n_traces': len(rates), 'confirmed': bool(rel <= .20)}
 91
 92def main():
 93    # Union of all learning rates is shared by both methods. Baseline's central
 94    # method knob (fixed recurrent gain) is swept as well.
 95    lrs = [1e-3, 3e-3, 1e-2]
 96    base_grid = [{'lr': lr, 'gain': g} for lr in lrs for g in [0.5, 1.0, 1.5]]
 97    base = sweep_baseline(baseline_factory, base_grid, seeds=(0,1,2,3))
 98    idea_grid = [{'lr': base['best_cfg']['lr'], 'a': 1.0, 'b': .2, 'gain': 1.0},
 99                 {'lr': 1e-3, 'a': 1.0, 'b': .2, 'gain': 1.0},
100                 {'lr': 1e-2, 'a': 1.0, 'b': .2, 'gain': 1.0}]
101    idea_runs = []
102    for cfg in idea_grid:
103        TRACE.clear()
104        r = evaluate(idea_factory(cfg), seeds=SEEDS)
105        idea_runs.append((cfg, r, signature()))
106    best_cfg, idea, sig = min(idea_runs, key=lambda z: z[1]['mean'])
107    # Ensure signature corresponds to selected configuration by rerunning it.
108    TRACE.clear(); idea = evaluate(idea_factory(best_cfg), seeds=SEEDS); sig = signature()
109    report = make_report('dynamics', 'positive_rnn_fixed_gain_vs_pito', base, idea,
110                         {'prediction': sig, 'selected_idea_cfg': best_cfg,
111                          'matched_structure': 'dynamics/control'})
112    report['idea_sweep'] = [{'cfg': c, 'result': r, 'mechanism_signature': s}
113                            for c, r, s in idea_runs]
114    report['custom_track'] = None
115    Path('bench_report.json').write_text(json.dumps(report, indent=2))
116    print(json.dumps(report, indent=2))
117
118if __name__ == '__main__': main()