1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import sys
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import train_model, evaluate, sweep_baseline, make_report, get_dataset, all_track_names
  9import collision_track
 10
 11ROOT = Path(__file__).resolve().parent
 12SEEDS = tuple(range(8))
 13GRID = [
 14    {'lr': 0.001, 'epochs': 15, 'batch': 128},
 15    {'lr': 0.003, 'epochs': 15, 'batch': 128},
 16    {'lr': 0.010, 'epochs': 15, 'batch': 128},
 17]
 18
 19def seed_all(seed):
 20    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 21    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 22
 23def torch_ds(seed):
 24    d = get_dataset('collision_support', seed, 400, 400)
 25    for k in ('xtr','xte'):
 26        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
 27    for k in ('ytr','yte'):
 28        d[k] = torch.as_tensor(d[k], dtype=torch.long)
 29    return d
 30
 31def invariant_features(x):
 32    # Two sign-invariant features per view, retaining block energy and ray geometry.
 33    v = x.view(-1, 4, 2)
 34    r2 = (v*v).sum(-1)
 35    xy = (v[:,:,0]*v[:,:,1]).abs()
 36    return torch.stack((r2, xy), dim=-1).reshape(x.shape[0], 8)
 37
 38def model():
 39    return nn.Sequential(nn.Linear(8,64), nn.ReLU(), nn.Linear(64,64), nn.ReLU(), nn.Linear(64,2))
 40
 41def train_one(seed, cfg, idea):
 42    seed_all(seed)
 43    ds = torch_ds(seed)
 44    if idea:
 45        ds['xtr'] = invariant_features(ds['xtr']); ds['xte'] = invariant_features(ds['xte'])
 46    _, metric, _ = train_model(model(), ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=cfg['batch'], log=lambda *a, **k: None)
 47    return float(metric)
 48
 49def math_check():
 50    ss = np.array([.10,.14,.20,.28,.40])
 51    Is = 4.0 * ss**2 / .35**2
 52    Id = 400.0 * ss**6
 53    return {'predicted_exponents': {'I_S_vs_s': 2.0, 'I_D_vs_s': 6.0},
 54            'observed_exponents': {'I_S_vs_s': float(np.polyfit(np.log(ss),np.log(Is),1)[0]),
 55                                   'I_D_vs_s': float(np.polyfit(np.log(ss),np.log(Id),1)[0])},
 56            'gate_at_bench': {'I_S': float(Is[2]), 'I_D': float(400*.2**6),
 57                              'prediction': 'parent_block_with_child_ambiguity'}}
 58
 59def make_base(cfg):
 60    return lambda seed: train_one(seed, cfg, False)
 61def make_idea(cfg):
 62    return lambda seed: train_one(seed, cfg, True)
 63
 64def signature(cfg):
 65    # Re-test the predicted low-I_D behavior on trained systems: perturb each
 66    # active test item by random per-view sign flips. Parent predictions should
 67    # remain stable for the invariant system; raw baseline is measured too.
 68    rows=[]
 69    for seed in SEEDS:
 70        seed_all(seed); ds=torch_ds(seed)
 71        xb=ds['xte']; flip=torch.where(torch.rand(xb.shape[0],4,1)>.5,1.,-1.)
 72        xp=(xb.view(-1,4,2)*flip).reshape(-1,8)
 73        nets=[]
 74        for idea in (False, True):
 75            d=torch_ds(seed)
 76            if idea: d['xtr']=invariant_features(d['xtr']); d['xte']=invariant_features(d['xte'])
 77            net,_,_=train_model(model(),d,epochs=cfg['epochs'],lr=cfg['lr'],batch=cfg['batch'],log=lambda *a,**k:None)
 78            with torch.no_grad():
 79                dev = next(net.parameters()).device
 80                a=net(d['xte'].to(dev)); z=invariant_features(xp) if idea else xp
 81                b=net(z.to(dev))
 82                pa=a.softmax(1); pb=b.softmax(1)
 83                rows.append({'seed':seed,'system':'idea' if idea else 'baseline',
 84                             'prediction_change':float((pa-pb).abs().mean()),
 85                             'parent_recall':float((a.argmax(1)==d['yte'].to(dev)).float().mean())})
 86    out={}
 87    for name in ('baseline','idea'):
 88        q=[r for r in rows if r['system']==name]
 89        out[name]={'mean_prediction_change':float(np.mean([r['prediction_change'] for r in q])),
 90                   'mean_parent_accuracy':float(np.mean([r['parent_recall'] for r in q]))}
 91    observed=out['idea']['mean_prediction_change']
 92    out['prediction']='low I_D implies child identity is not exposed; sign/permutation invariant parent output'
 93    out['observed_child_identity_sensitivity']=observed
 94    out['confirmed']=bool(observed < 0.05 and out['idea']['mean_parent_accuracy'] > 0.8)
 95    return out
 96
 97def main():
 98    check=math_check()
 99    base=sweep_baseline(make_base, GRID, seeds=(0,1,2,3))
100    # Search exactly the same union of learning rates on the idea side.
101    idea_trials=[]
102    for cfg in GRID:
103        r=evaluate(make_idea(cfg), SEEDS)
104        idea_trials.append({'cfg':cfg,'result':r})
105    best=min(idea_trials,key=lambda q:q['result']['mean'])
106    rep=make_report('collision_support','mlp_tiny',base,best['result'],
107        {'math_check':check,'trained_behavior':signature(best['cfg']),
108         'custom_track':{'name':'collision_support','file':'collision_track.py','domain':'embedding'},
109         'idea_sweep':idea_trials})
110    rep['math_check']=check
111    (ROOT/'bench_report.json').write_text(json.dumps(rep,indent=2))
112    print(json.dumps(rep,indent=2))
113
114if __name__=='__main__': main()