1import json, random, importlib.util, sys
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 7from bench import make_model, sweep_baseline, evaluate, make_report, validate_and_promote, reload_custom_tracks
 8
 9TRACK='planar_ising_autoregressive'; SEEDS=tuple(range(8))
10GRID=[{'lr':lr,'epochs':5,'batch':128} for lr in (.001,.003,.01)]
11HERE=Path(__file__).resolve().parent
12spec=importlib.util.spec_from_file_location('local_ising_track', HERE/'ising_teacher_track.py')
13MOD=importlib.util.module_from_spec(spec); spec.loader.exec_module(MOD)
14
15def seed_all(s):
16    random.seed(s); np.random.seed(s); torch.manual_seed(s)
17    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
18
19def train_one(seed,cfg,soft,capture=False):
20    seed_all(seed); raw=MOD.get_dataset(seed,400,200)
21    xtr=torch.tensor(raw['xtr']); ytr=torch.tensor(raw['ytr']); xte=torch.tensor(raw['xte'])
22    target=torch.tensor(raw['qtr'] if soft else raw['ytr'])
23    net=make_model('mlp_tiny',tuple(xtr.shape[1:]),1)
24    try:
25        dev='cuda' if torch.cuda.is_available() else 'cpu'; net.to(dev)
26        xtr,target=xtr.to(dev),target.to(dev); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
27        for _ in range(cfg['epochs']):
28            p=torch.randperm(len(xtr),device=dev)
29            for st in range(0,len(p),cfg['batch']):
30                ix=p[st:st+cfg['batch']]; loss=nn.functional.mse_loss(net(xtr[ix]),target[ix])
31                opt.zero_grad(); loss.backward(); opt.step()
32        with torch.no_grad(): pred=net(xte.to(dev)).cpu().numpy().ravel()
33    except RuntimeError:
34        seed_all(seed); net=make_model('mlp_tiny',tuple(xtr.shape[1:]),1); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
35        for _ in range(cfg['epochs']):
36            p=torch.randperm(len(xtr))
37            for st in range(0,len(p),cfg['batch']):
38                ix=p[st:st+cfg['batch']]; loss=nn.functional.mse_loss(net(xtr[ix]),target[ix]); opt.zero_grad(); loss.backward(); opt.step()
39        with torch.no_grad(): pred=net(xte).numpy().ravel()
40    metric=float(np.mean((pred-raw['yte'].ravel())**2)); qerr=float(np.mean((pred-raw['qte'].ravel())**2))
41    return {'metric':metric,'q_mse':qerr,'pred_mean':float(pred.mean())} if capture else metric
42
43def main():
44    # Promotion makes the track registered in the shared bench; never edit bench directly.
45    if TRACK not in __import__('bench').all_track_names():
46        if not validate_and_promote(HERE/'ising_teacher_track.py',TRACK,'exp2721 Kac-Ward Exact Teacher Stage-2'):
47            raise RuntimeError('custom track promotion failed')
48    reload_custom_tracks()
49    base=sweep_baseline(lambda c:(lambda s:train_one(s,c,False)),GRID)
50    idea_sweep=[]
51    for c in GRID:
52        idea_sweep.append({'cfg':c,'result':evaluate(lambda s,c=c:train_one(s,c,True),seeds=SEEDS)})
53    chosen=min(idea_sweep,key=lambda z:z['result']['mean']); idea=chosen['result']; icfg=chosen['cfg']
54    b0=train_one(0,base['best_cfg'],False,True); i0=train_one(0,icfg,True,True)
55    sig={'prediction':'exact soft conditional labels reduce trained-NN conditional probability MSE versus sampled labels','baseline_seed0':b0,'idea_seed0':i0,'q_mse_reduction':b0['q_mse']-i0['q_mse'],'confirmed':bool(i0['q_mse']<b0['q_mse'])}
56    report=make_report(TRACK,'mlp_tiny',base,idea,extra=sig)
57    report['idea_sweep']=idea_sweep
58    report['custom_track']={'name':TRACK,'file':'ising_teacher_track.py','domain':'sequence-level sampling'}
59    (HERE/'bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
60if __name__=='__main__': main()