1import os, sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7import bench
  8
  9SEEDS = tuple(range(8))
 10# Union is used on both sides: baseline and idea see every tested setting.
 11GRID = [
 12    {'lr': 0.001, 'weight_decay': 0.0},
 13    {'lr': 0.003, 'weight_decay': 0.0},
 14    {'lr': 0.006, 'weight_decay': 0.0},
 15]
 16EPOCHS = 18
 17BATCH = 128
 18
 19class FullGRU(nn.Module):
 20    def __init__(self, hidden=45):
 21        super().__init__()
 22        self.rnn = nn.GRU(3, hidden, batch_first=True)
 23        self.head = nn.Linear(hidden, 1)
 24        self._no_cudnn = False
 25    def forward(self, x):
 26        q = x.view(x.shape[0], -1, 3)
 27        try:
 28            _, h = self.rnn(q)
 29        except RuntimeError:
 30            self._no_cudnn = True
 31        if self._no_cudnn:
 32            old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False
 33            try: _, h = self.rnn(q)
 34            finally: torch.backends.cudnn.enabled = old
 35        return self.head(h[-1])
 36    @torch.no_grad()
 37    def phase_predictions(self, x):
 38        q = x.view(x.shape[0], 8, 3)
 39        try:
 40            _, hm = self.rnn(q[:, :4]); _, hf = self.rnn(q)
 41        except RuntimeError:
 42            old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False
 43            try:
 44                _, hm = self.rnn(q[:, :4]); _, hf = self.rnn(q)
 45            finally: torch.backends.cudnn.enabled = old
 46        return self.head(hm[-1]), self.head(hf[-1])
 47
 48class ChainedGRU(nn.Module):
 49    """Two event phases; phase-2 starts from phase-1's differentiable terminal h."""
 50    def __init__(self, hidden=32):
 51        super().__init__()
 52        self.phase1 = nn.GRU(3, hidden, batch_first=True)
 53        self.phase2 = nn.GRU(3, hidden, batch_first=True)
 54        self.head = nn.Linear(hidden, 1)
 55        self._no_cudnn = False
 56    def _run(self, module, x, h=None):
 57        try: return module(x, h)
 58        except RuntimeError:
 59            self._no_cudnn = True
 60            old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False
 61            try: return module(x, h)
 62            finally: torch.backends.cudnn.enabled = old
 63    def forward(self, x):
 64        q = x.view(x.shape[0], 8, 3)
 65        _, h1 = self._run(self.phase1, q[:, :4])
 66        _, h2 = self._run(self.phase2, q[:, 4:], h1)
 67        return self.head(h2[-1])
 68    @torch.no_grad()
 69    def phase_predictions(self, x):
 70        q = x.view(x.shape[0], 8, 3)
 71        _, h1 = self._run(self.phase1, q[:, :4])
 72        # This is an observable prediction at the event, before phase 2.
 73        mid = self.head(h1[-1])
 74        _, h2 = self._run(self.phase2, q[:, 4:], h1)
 75        return mid, self.head(h2[-1])
 76
 77class InspectFull(FullGRU):
 78    @torch.no_grad()
 79    def phase_predictions(self, x):
 80        q = x.view(x.shape[0], 8, 3)
 81        _, hm = self.rnn(q[:, :4]); _, hf = self.rnn(q)
 82        return self.head(hm[-1]), self.head(hf[-1])
 83
 84def seed_all(seed):
 85    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 86
 87def train_one(kind, cfg, seed, inspect=False):
 88    seed_all(seed)
 89    ds = bench.get_dataset('dynamics', seed, n_train=400, n_test=400)
 90    model = (ChainedGRU() if kind == 'idea' else FullGRU())
 91    # train_model is the benchmark's required robust CUDA/CPU path.
 92    net, metric, hist = bench.train_model(model, ds, epochs=EPOCHS,
 93        lr=cfg['lr'], batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *_: None)
 94    if net is None: return float('nan'), None
 95    if not inspect: return float(metric), None
 96    dev = next(net.parameters()).device
 97    x = torch.as_tensor(ds['xte'], dtype=torch.float32, device=dev)
 98    y = torch.as_tensor(ds['yte'], dtype=torch.float32, device=dev).reshape(-1,1)
 99    with torch.no_grad():
100        mid, final = net.phase_predictions(x)
101    # midpoint observed theta is input's fourth theta; final observed y is target.
102    obs_mid = x.view(-1,8,3)[:,3,0:1]
103    return float(metric), {'mid_pred_mae': float((mid-obs_mid).abs().mean().cpu()),
104                           'mid_obs_mae_baseline_reference': float(obs_mid.abs().mean().cpu()),
105                           'final_pred_mae': float((final-y).abs().mean().cpu()),
106                           'final_obs_abs': float(y.abs().mean().cpu())}
107
108def make_fn(kind, cfg):
109    return lambda s: train_one(kind, cfg, s)[0]
110
111def main():
112    # Baseline sweep uses the same union of hyperparameters as the idea sweep.
113    base = bench.sweep_baseline(lambda c: make_fn('base', c), GRID, seeds=(0,1,2,3))
114    base['grid_union'] = GRID
115    idea_trials = []
116    for cfg in GRID:
117        r = bench.evaluate(make_fn('idea', cfg), SEEDS)
118        idea_trials.append({'cfg': cfg, **r})
119    best = min(idea_trials, key=lambda z: z['mean'])
120    idea_res = {'best_cfg': best['cfg'], 'sweep': idea_trials,
121                'per_seed': best['per_seed'], 'mean': best['mean'],
122                'std': best['std'], 'n': best['n']}
123    # Re-train paired best systems to measure behaviour, not an identity.
124    sig = {'baseline': [], 'idea': [], 'paired_seed': []}
125    for s in SEEDS:
126        bm, bs = train_one('base', best['cfg'], s, True)
127        im, ins = train_one('idea', best['cfg'], s, True)
128        sig['baseline'].append(bs); sig['idea'].append(ins); sig['paired_seed'].append(s)
129    sig['mean_final_mae_baseline'] = float(np.mean([v['final_pred_mae'] for v in sig['baseline']]))
130    sig['mean_final_mae_idea'] = float(np.mean([v['final_pred_mae'] for v in sig['idea']]))
131    sig['mean_mid_mae_baseline'] = float(np.mean([v['mid_pred_mae'] for v in sig['baseline']]))
132    sig['mean_mid_mae_idea'] = float(np.mean([v['mid_pred_mae'] for v in sig['idea']]))
133    # Prediction being tested: chaining should help the post-boundary dynamics.
134    sig['prediction'] = 'phase chaining improves final/post-event trajectory prediction'
135    sig['confirmed'] = bool(sig['mean_final_mae_idea'] < sig['mean_final_mae_baseline'])
136    report = bench.make_report('dynamics', 'rnn_small', base, idea_res,
137        {'mechanism_signature': sig,
138         'protocol_notes': '400 train/400 test, 8 paired seeds, 18 epochs; midpoint is the known phase boundary.'})
139    Path('bench_report.json').write_text(json.dumps(report, indent=2))
140    print(json.dumps(report, indent=2))
141
142if __name__ == '__main__':
143    try: main()
144    except Exception:
145        if torch.cuda.is_available():
146            torch.cuda.empty_cache(); os.environ['CUDA_VISIBLE_DEVICES']=''
147            main()
148        else: raise