1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, make_model, train_model, sweep_baseline, make_report
  7
  8SEEDS = tuple(range(8))
  9GRID = [
 10    {'lr': 0.0015, 'weight_decay': 0.0, 'epochs': 12},
 11    {'lr': 0.0030, 'weight_decay': 0.0, 'epochs': 12},
 12    {'lr': 0.0060, 'weight_decay': 0.0, 'epochs': 12},
 13    {'lr': 0.0015, 'weight_decay': 1e-4, 'epochs': 12},
 14    {'lr': 0.0030, 'weight_decay': 1e-4, 'epochs': 12},
 15    {'lr': 0.0060, 'weight_decay': 1e-4, 'epochs': 12},
 16]
 17
 18def entropy(y, z):
 19    y = np.asarray(y, int); z = np.asarray(z, int)
 20    if z.ndim == 1: z = z[:, None]
 21    keys = np.zeros(len(y), dtype=np.int64)
 22    mult = 1
 23    for c in range(z.shape[1]):
 24        keys += z[:, c] * mult; mult *= int(z[:, c].max() + 1)
 25    _, inv = np.unique(keys, return_inverse=True)
 26    ny = int(y.max()) + 1
 27    tab = np.zeros((int(inv.max()) + 1, ny), dtype=np.int64)
 28    np.add.at(tab, (inv, y), 1)
 29    p = tab / np.maximum(tab.sum(1, keepdims=True), 1)
 30    h = -(np.where(p > 0, p * np.log(np.maximum(p, 1e-300)), 0)).sum(1)
 31    return float((tab.sum(1) / len(y) * h).sum())
 32
 33def quantile_codes(a, q):
 34    a = np.asarray(a)
 35    out = np.empty_like(a, dtype=np.int64)
 36    for j in range(a.shape[1]):
 37        edges = np.unique(np.quantile(a[:, j], np.linspace(0, 1, q + 1)[1:-1]))
 38        out[:, j] = np.searchsorted(edges, a[:, j], side='right')
 39    return out
 40
 41def selector(xtr, ytr, seed):
 42    # Fixed, train-only quantile discretization for structure discovery.
 43    qx, qy = 4, 8
 44    xd = quantile_codes(xtr, qx)
 45    ye = quantile_codes(ytr.reshape(-1, 1), qy).ravel()
 46    h0 = entropy(ye, np.zeros((len(ye), 1), dtype=np.int64))
 47    sh = np.array([h0 - entropy(ye, xd[:, [j]]) for j in range(xtr.shape[1])])
 48    rng = np.random.default_rng(1000 + seed)
 49    boot = []
 50    for _ in range(24):
 51        ix = rng.integers(0, len(ye), len(ye))
 52        hb = entropy(ye[ix], np.zeros((len(ix), 1), dtype=np.int64))
 53        boot.append([hb - entropy(ye[ix], xd[ix, [j]]) for j in range(xtr.shape[1])])
 54    boot = np.asarray(boot)
 55    jstar = int(np.argmax(sh))
 56    gamma = float(sh[jstar] / (boot[:, jstar].std(ddof=1) + 1e-8))
 57    # Held-out predictive risk: polynomial ridge captures Friedman nonlinearities
 58    # while avoiding a different neural architecture during final training.
 59    n = len(xtr); cut = max(1, int(.75*n)); a, b = xtr[:cut], xtr[cut:]
 60    ya, yb = ytr[:cut], ytr[cut:]
 61    def feat(x, j):
 62        v = x[:, j:j+1]; return np.concatenate([v, v*v, np.sin(v)], axis=1)
 63    sr = []
 64    for j in range(xtr.shape[1]):
 65        z, zv = feat(a, j), feat(b, j)
 66        reg = 1e-2 * np.eye(z.shape[1]); w = np.linalg.solve(z.T@z + reg, z.T@ya)
 67        sr.append(float(np.mean((yb - zv@w)**2)))
 68    sr = -np.asarray(sr)  # larger is better, as S_R relative to common null risk
 69    mh = np.argsort(sh)[-5:][::-1]
 70    mr = np.argsort(sr)[-5:][::-1]
 71    kappa = float(h0 / math.log(qy))
 72    # qY=8 makes kappa operationally sensitive to whether bins resolve residual uncertainty.
 73    use_h = bool(kappa < 0.90 and gamma > 2.0)
 74    chosen = mh if use_h else mr
 75    return {'mask_h': mh.tolist(), 'mask_r': mr.tolist(), 'mask': chosen.tolist(),
 76            'sh': sh.tolist(), 'sr': sr.tolist(), 'kappa': kappa,
 77            'gamma': gamma, 'used_entropy': use_h, 'h0': h0}
 78
 79def run_one(track, seed, cfg, idea):
 80    torch.manual_seed(seed); np.random.seed(seed)
 81    d = get_dataset(track, seed, n_train=1200, n_test=500)
 82    meta = selector(d['xtr'].numpy(), d['ytr'].numpy(), seed) if idea else None
 83    if idea: cols = meta['mask']
 84    else: cols = list(range(d['xtr'].shape[1]))
 85    ds = dict(d)
 86    ds['xtr'], ds['xte'] = d['xtr'][:, cols], d['xte'][:, cols]
 87    ds['input_shape'] = tuple(ds['xtr'].shape[1:])
 88    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
 89    _, metric, _ = train_model(net, ds, epochs=cfg['epochs'], lr=cfg['lr'],
 90                               weight_decay=cfg['weight_decay'], log=lambda *_: None)
 91    return float(metric), meta
 92
 93def fn(track, cfg, idea):
 94    return lambda seed: run_one(track, seed, cfg, idea)[0]
 95
 96def main():
 97    track, model = 'tabular', 'mlp_tiny'
 98    baseline = sweep_baseline(lambda c: fn(track, c, False), GRID, seeds=SEEDS)
 99    # Equal-sized idea sweep over exactly the baseline union; final result is best config.
100    tried=[]
101    for cfg in GRID:
102        r = __import__('bench').evaluate(fn(track, cfg, True), SEEDS)
103        tried.append({'cfg':cfg, 'mean':r['mean']})
104    best_cfg = min(GRID, key=lambda c: next(x['mean'] for x in tried if x['cfg']==c))
105    idea = __import__('bench').evaluate(fn(track, best_cfg, True), SEEDS)
106    # Test-model signature: selector's predicted choice and the observed trained-model metric.
107    sig=[]
108    for s in SEEDS:
109        _, m = run_one(track, s, best_cfg, True)
110        full = run_one(track, s, best_cfg, False)[0]
111        chosen = run_one(track, s, best_cfg, True)[0]
112        sig.append({'seed':s, 'kappa':m['kappa'], 'gamma':m['gamma'],
113                    'used_entropy':bool(m['used_entropy']), 'predicted_fallback':bool(not m['used_entropy']),
114                    'observed_subset_minus_full_mse':chosen-full})
115    # Signature prediction is that high kappa/high stochasticity triggers fallback; verify on trained outcomes.
116    high = [x for x in sig if x['kappa'] >= .90]
117    confirmed = bool(high) and all(x['predicted_fallback'] for x in high) and np.mean([x['observed_subset_minus_full_mse'] for x in high]) <= 0.02
118    idea_block={'best_cfg':best_cfg, 'sweep':tried, 'full':idea}
119    report=make_report(track, model, baseline, idea, {
120      'mechanism_signature': {'prediction':'kappa>=0.90 routes away from entropy mask and avoids materially worse trained-model MSE',
121        'observed':sig, 'high_kappa_n':len(high), 'high_kappa_observed_mean_delta':float(np.mean([x['observed_subset_minus_full_mse'] for x in high])) if high else None,
122        'confirmed':confirmed},
123      'selection_details':'Entropy and validation-risk masks selected from train split; final systems use identical mlp_tiny/training.'})
124    Path('bench_report.json').write_text(json.dumps(report, indent=2, default=lambda o: o.item() if isinstance(o, np.generic) else str(o)))
125    print(json.dumps(report, indent=2, default=lambda o: o.item() if isinstance(o, np.generic) else str(o)))
126if __name__=='__main__': main()