1import json, random, sys
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  8
  9SEEDS = tuple(range(8))
 10SWEEP_SEEDS = tuple(range(4))
 11EPOCHS = 24
 12BATCH = 128
 13TRACK = 'tabular'
 14MODEL = 'mlp_tiny'
 15LRS = [1e-3, 3e-3, 1e-2]
 16
 17
 18def seed_all(seed):
 19    random.seed(seed)
 20    np.random.seed(seed)
 21    torch.manual_seed(seed)
 22    if torch.cuda.is_available():
 23        try:
 24            torch.cuda.manual_seed_all(seed)
 25        except Exception:
 26            pass
 27
 28
 29def baseline_fn(cfg):
 30    def run(seed):
 31        seed_all(seed)
 32        ds = get_dataset(TRACK, seed, n_train=400, n_test=400)
 33        model = make_model(MODEL, ds['input_shape'], ds['out_dim'])
 34        _, metric, _ = train_model(model, ds, epochs=EPOCHS,
 35                                   lr=cfg['lr'], batch=BATCH,
 36                                   weight_decay=cfg.get('weight_decay', 0.0),
 37                                   log=lambda *_: None)
 38        return float(metric)
 39    return run
 40
 41
 42def idea_train(model, ds, epochs, lr, tail_weight, batch=BATCH):
 43    """Tail residual correction: first fit the cheap/global predictor, then
 44    upweight positive residuals above the empirical upper-tail boundary.
 45    The same MLP and optimizer family as the baseline are retained."""
 46    errs = []
 47    devices = []
 48    if torch.cuda.is_available():
 49        devices.append('cuda')
 50    devices.append('cpu')
 51    for device in devices:
 52        try:
 53            x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 54            net = model.to(device)
 55            opt = torch.optim.Adam(net.parameters(), lr=lr)
 56            n = x.shape[0]
 57            # Cheap surrogate phase: ordinary global MSE.
 58            warm = max(1, epochs // 2)
 59            gen = torch.Generator(device='cpu')
 60            gen.manual_seed(17000 + int(n) + int(round(lr * 1e6)))
 61            for ep in range(epochs):
 62                perm = torch.randperm(n, generator=gen)
 63                for start in range(0, n, batch):
 64                    ix = perm[start:start + batch].to(device)
 65                    pred = net(x[ix])
 66                    residual = pred - y[ix]
 67                    if ep < warm:
 68                        loss = (residual * residual).mean()
 69                    else:
 70                        # Estimated upper-tail boundary is the training beta quantile.
 71                        q = torch.quantile(y, 0.90)
 72                        # Smoothly emphasize samples in/above the tail, retaining
 73                        # ordinary MSE in the central region.
 74                        weights = 1.0 + (tail_weight - 1.0) * torch.sigmoid((y[ix] - q) / 0.15)
 75                        loss = (weights * residual * residual).mean()
 76                    opt.zero_grad(set_to_none=True)
 77                    loss.backward()
 78                    opt.step()
 79            with torch.no_grad():
 80                pred = net(ds['xte'].to(device))
 81                metric = ((pred - ds['yte'].to(device)) ** 2).mean().item()
 82            return net, float(metric), device
 83        except Exception as exc:
 84            errs.append(f'{device}:{str(exc)[:100]}')
 85            try:
 86                torch.cuda.empty_cache()
 87            except Exception:
 88                pass
 89    raise RuntimeError('idea training failed: ' + ' | '.join(errs))
 90
 91
 92def idea_fn(cfg):
 93    def run(seed):
 94        seed_all(seed)
 95        ds = get_dataset(TRACK, seed, n_train=400, n_test=400)
 96        model = make_model(MODEL, ds['input_shape'], ds['out_dim'])
 97        _, metric, _ = idea_train(model, ds, EPOCHS, cfg['lr'], cfg['tail_weight'])
 98        return float(metric)
 99    return run
100
101
102def behavior(cfg, idea=False):
103    global_vals, tail_vals, cvar_vals = [], [], []
104    for seed in SEEDS:
105        seed_all(seed)
106        ds = get_dataset(TRACK, seed, n_train=400, n_test=400)
107        model = make_model(MODEL, ds['input_shape'], ds['out_dim'])
108        if idea:
109            net, _, dev = idea_train(model, ds, EPOCHS, cfg['lr'], cfg['tail_weight'])
110        else:
111            net, _, hist = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'],
112                                       batch=BATCH, weight_decay=cfg.get('weight_decay', 0),
113                                       log=lambda *_: None)
114            dev = next(net.parameters()).device
115        with torch.no_grad():
116            p = net(ds['xte'].to(dev)).detach().cpu().numpy().ravel()
117        y = ds['yte'].numpy().ravel()
118        threshold = np.quantile(y, .90)
119        tail = y >= threshold
120        global_vals.append(float(np.mean((p-y)**2)))
121        tail_vals.append(float(np.mean((p[tail]-y[tail])**2)))
122        # Rockafellar empirical CVaR prediction versus observed test CVaR.
123        eta = np.quantile(p, .90)
124        pc = eta + np.maximum(p-eta, 0).mean()/.10
125        et = np.quantile(y, .90)
126        tc = et + np.maximum(y-et, 0).mean()/.10
127        cvar_vals.append(float(abs(pc-tc)))
128    return {'global_mse': float(np.mean(global_vals)),
129            'tail_mse': float(np.mean(tail_vals)),
130            'tail_to_global': float(np.mean(tail_vals)/max(np.mean(global_vals), 1e-12)),
131            'cvar_abs_error': float(np.mean(cvar_vals))}
132
133
134def main():
135    # Baseline is swept on the same learning-rate union used by the idea.
136    grid = [{'lr': lr, 'weight_decay': 0.0} for lr in LRS]
137    base = sweep_baseline(baseline_fn, grid, seeds=SWEEP_SEEDS)
138    idea_grid = [{'lr': lr, 'tail_weight': 4.0} for lr in LRS]
139    idea_sweep = []
140    for cfg in idea_grid:
141        r = evaluate(idea_fn(cfg), seeds=SWEEP_SEEDS)
142        idea_sweep.append({'cfg': cfg, 'result': r})
143    best = min(idea_sweep, key=lambda z: z['result']['mean'])
144    idea_full = evaluate(idea_fn(best['cfg']), seeds=SEEDS)
145    bcfg = base['best_cfg']
146    sig_b = behavior(bcfg, idea=False)
147    sig_i = behavior(best['cfg'], idea=True)
148    signature = {
149        'prediction': 'tail residual correction should preferentially reduce upper-tail prediction error relative to global error',
150        'baseline_behavior': sig_b,
151        'idea_behavior': sig_i,
152        'predicted_tail_focus': True,
153        'observed_tail_to_global_ratio_change': sig_i['tail_to_global'] - sig_b['tail_to_global'],
154        'confirmed': bool(sig_i['tail_to_global'] < sig_b['tail_to_global'])
155    }
156    report = make_report(TRACK, MODEL, base, idea_full,
157                         {'idea_config': best['cfg'], 'idea_sweep': idea_sweep,
158                          'mechanism_signature': signature})
159    report['protocol_notes'] = {'epochs': EPOCHS, 'batch': BATCH,
160        'baseline_grid': grid, 'idea_grid': idea_grid,
161        'intervention': 'same mlp_tiny trained with global-MSE warmup followed by upper-tail weighted residual correction'}
162    with open('bench_report.json', 'w') as f:
163        json.dump(report, f, indent=2)
164    print(json.dumps(report, indent=2))
165
166if __name__ == '__main__':
167    main()