1import json, math, sys
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import make_model, train_model, sweep_baseline, make_report, evaluate
  8from custom_strain_track import get_dataset
  9
 10SEEDS = tuple(range(8))
 11GRID = [dict(lr=1e-3, weight_decay=0.0), dict(lr=3e-3, weight_decay=0.0), dict(lr=1e-2, weight_decay=1e-4)]
 12EPOCHS = 18
 13BATCH = 128
 14
 15def device():
 16    return 'cuda' if torch.cuda.is_available() else 'cpu'
 17
 18def ico_axes():
 19    p=(1+math.sqrt(5))/2
 20    a=np.array([(0,1,p),(1,p,0),(p,0,1),(0,1,-p),(1,-p,0),(p,0,-1)], np.float32)
 21    return torch.tensor(a/math.sqrt(1+p*p), dtype=torch.float32)
 22AXES = ico_axes()
 23
 24def idea_features(x):
 25    # x contains row-major F (9) and six observed directional measurements.
 26    F=x[:, :9].reshape(-1,3,3); y=x[:,9:]
 27    v=AXES.to(x.device)
 28    w=torch.einsum('bij,nj->bni', F, v)
 29    A=torch.stack([w[:,:,0]*w[:,:,0], w[:,:,1]*w[:,:,1], w[:,:,2]*w[:,:,2],
 30                   math.sqrt(2)*w[:,:,0]*w[:,:,1], math.sqrt(2)*w[:,:,0]*w[:,:,2],
 31                   math.sqrt(2)*w[:,:,1]*w[:,:,2]], dim=-1)
 32    eye=torch.eye(6, device=x.device, dtype=x.dtype).expand(x.shape[0],6,6)
 33    c=torch.linalg.solve(A.transpose(1,2)@A + 1e-5*eye, A.transpose(1,2)@y.unsqueeze(-1)).squeeze(-1)
 34    # trace-free postprocessing in Frobenius-preserving coordinates.
 35    tr=(c[:,0]+c[:,1]+c[:,2])/3
 36    c=c.clone(); c[:,:3]=c[:,:3]-tr[:,None]
 37    return torch.cat([x[:,:9], c], dim=1)
 38
 39def prepare(seed, transformed):
 40    d=get_dataset(seed,400,400)
 41    out={'track':'strain_congruence','task':'regression','metric':'mse','input_shape':(15,), 'out_dim':1}
 42    for k in ('xtr','xte','ytr','yte'):
 43        out[k]=torch.as_tensor(d[k], dtype=torch.float32)
 44    if transformed:
 45        out['xtr']=idea_features(out['xtr']); out['xte']=idea_features(out['xte'])
 46    return out
 47
 48cache={}
 49def run(side,cfg,seed,keep=False):
 50    key=(side,tuple(sorted(cfg.items())),int(seed))
 51    if key in cache: return cache[key][0]
 52    torch.manual_seed(10000+int(seed)); np.random.seed(10000+int(seed))
 53    ds=prepare(seed, side=='idea')
 54    net=make_model('mlp_tiny', ds['input_shape'], 1)
 55    try:
 56        net,metric,_=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['weight_decay'])
 57    except Exception:
 58        # train_model normally performs this fallback; preserve an explicit CPU retry.
 59        ds={k:(v.cpu() if torch.is_tensor(v) else v) for k,v in ds.items()}
 60        torch.manual_seed(10000+int(seed)); net=make_model('mlp_tiny',(15,),1)
 61        net,metric,_=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=cfg['weight_decay'])
 62    cache[key]=(float(metric),net.cpu(),ds)
 63    return float(metric)
 64
 65def make_fn(side):
 66    return lambda cfg: (lambda seed: run(side,cfg,int(seed)))
 67
 68def main():
 69    # Baseline sweep uses the same union of settings that is tested for the idea.
 70    base=sweep_baseline(make_fn('baseline'),GRID,seeds=(0,1,2,3))
 71    idea_trials=[]
 72    for cfg in GRID:
 73        r=evaluate(make_fn('idea')(cfg),SEEDS)
 74        idea_trials.append({'cfg':cfg,'result':r})
 75    best=min(idea_trials,key=lambda z:z['result']['mean'])
 76    rep=make_report('strain_congruence','mlp_tiny',base,best['result'],extra={})
 77    rep['idea_sweep']=idea_trials
 78    rep['custom_track']={'name':'strain_congruence','file':'custom_strain_track.py','domain':'dynamics'}
 79    # Mechanism signature: trained-system errors stratified by observed deformation condition.
 80    cond_rows=[]
 81    bc=base['best_cfg']; ic=best['cfg']
 82    for s in SEEDS:
 83        bm=cache[('baseline',tuple(sorted(bc.items())),s)][1]
 84        im=cache[('idea',tuple(sorted(ic.items())),s)][1]
 85        ds=cache[('idea',tuple(sorted(ic.items())),s)][2]
 86        raw=prepare(s,False); F=raw['xte'][:,:9].reshape(-1,3,3)
 87        cond=torch.linalg.cond(F).numpy(); bins=[cond<1.5,(cond>=1.5)&(cond<3),(cond>=3)]
 88        with torch.no_grad():
 89            pb=bm(raw['xte']).squeeze().numpy(); pi=im(ds['xte']).squeeze().numpy(); yy=raw['yte'].squeeze().numpy()
 90        for name,mask in zip(('low','mid','high'),bins):
 91            if mask.any(): cond_rows.append((name,float(np.mean((pb[mask]-yy[mask])**2)),float(np.mean((pi[mask]-yy[mask])**2))))
 92    sig={}
 93    for name in ('low','mid','high'):
 94        z=[r for r in cond_rows if r[0]==name]
 95        if z: sig[name]={'baseline_mse':float(np.mean([r[1] for r in z])),'idea_mse':float(np.mean([r[2] for r in z]))}
 96    high=sig.get('high',{'baseline_mse':float('nan'),'idea_mse':float('nan')})
 97    rep['mechanism_signature']={'prediction':'icosahedral reconstruction should reduce learned prediction error under ill-conditioned SL(3) deformations','condition_bins':sig,'confirmed':bool(np.isfinite(high['idea_mse']) and high['idea_mse']<high['baseline_mse'])}
 98    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
 99    print(json.dumps(rep,indent=2))
100
101if __name__=='__main__': main()