import sys, json, random from pathlib import Path import numpy as np import torch import torch.nn.functional as F ROOT = Path(__file__).resolve().parent sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import make_model, train_model, sweep_baseline, evaluate, make_report from pde_boundary_track import get_dataset, C SEEDS = tuple(range(8)) EPOCHS = 25 BATCH = 128 # The baseline and idea share this union of learning rates and Adam weight decay. CONFIGS = [ {'lr': 1e-3, 'weight_decay': 0.0}, {'lr': 3e-3, 'weight_decay': 0.0}, {'lr': 6e-3, 'weight_decay': 0.0}, ] def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def ds_for(seed): d = get_dataset(seed, 400, 400) for k in ('xtr','ytr','xte','yte'): d[k] = torch.as_tensor(d[k], dtype=torch.float32) d['ytr'] = d['ytr'].reshape(-1,1); d['yte'] = d['yte'].reshape(-1,1) d['input_shape'] = tuple(d['xtr'].shape[1:]); d['out_dim'] = 1 return d def baseline_fn(cfg): def run(seed): seed_all(seed); d=ds_for(seed) net=make_model('mlp_tiny', d['input_shape'], 1) _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *_: None) return metric return run def boundary_residuals(net, x, device): # Values are measured from the trained model, not an analytical identity. z=x.to(device).detach().clone().requires_grad_(True) u=net(z) g=torch.autograd.grad(u.sum(), z, create_graph=False)[0] xx, yy=z[:,0], z[:,1] left=xx < 1e-7 horiz=(yy < 1e-7) | (yy > 1-1e-7) ob=torch.abs(g[:,0] + C*g[:,1])[left] no=torch.abs(g[:,1])[horiz] return float(ob.mean().detach().cpu()) if ob.numel() else 0., float(no.mean().detach().cpu()) if no.numel() else 0. def idea_train(seed, cfg, return_net=False): seed_all(seed); d=ds_for(seed) device='cuda' if torch.cuda.is_available() else 'cpu' try: net=make_model('mlp_tiny', d['input_shape'], 1).to(device) opt=torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay']) xtr,ytr=d['xtr'].to(device),d['ytr'].to(device) # Boundary penalty is the PDE intervention; all architecture/data/budget terms match. for _ in range(EPOCHS): net.train(); perm=torch.randperm(len(xtr),device=device) for i in range(0,len(xtr),BATCH): x=xtr[perm[i:i+BATCH]].detach().clone().requires_grad_(True) y=ytr[perm[i:i+BATCH]] pred=net(x); loss=F.mse_loss(pred,y) grad=torch.autograd.grad(pred.sum(),x,create_graph=True)[0] left=x[:,0] < 1e-7; horiz=(x[:,1] < 1e-7)|(x[:,1] > 1-1e-7) bc=0.0 if left.any(): bc=bc+(grad[left,0]+C*grad[left,1]).square().mean() if horiz.any(): bc=bc+grad[horiz,1].square().mean() total=loss + cfg['bc_weight']*bc opt.zero_grad(); total.backward(); opt.step() net.eval() with torch.no_grad(): metric=float(((net(d['xte'].to(device))-d['yte'].to(device))**2).mean().cpu()) if return_net: return net, d, device, metric return metric except RuntimeError: # Explicit CPU fallback for shared GPU failures. torch.cuda.empty_cache() if torch.cuda.is_available() else None old=torch.cuda.is_available # Re-run with CPU by temporarily using an equivalent forced implementation. if device != 'cpu': torch.set_default_device('cpu') try: return idea_train(seed, cfg, return_net) finally: torch.set_default_device('cpu') raise def idea_fn(cfg): return lambda seed: idea_train(seed, cfg) def main(): # Baseline sweep uses the same lr union and the same eight-seed final evaluation. base=sweep_baseline(baseline_fn, CONFIGS) idea_configs=[dict(CONFIGS[CONFIGS.index(base['best_cfg'])], bc_weight=w) for w in (0.1, 1.0, 5.0)] idea_results=[] for cfg in idea_configs: r=evaluate(idea_fn(cfg), SEEDS); r['cfg']=cfg; idea_results.append(r) best=min(idea_results, key=lambda r:r['mean']) # Re-test signature on all boundary points from the trained systems. sig=[] for s in SEEDS: bcfg=base['best_cfg']; seed_all(s); d=ds_for(s) bnet=make_model('mlp_tiny',d['input_shape'],1) bnet,_,_=train_model(bnet,d,epochs=EPOCHS,lr=bcfg['lr'],batch=BATCH,weight_decay=bcfg['weight_decay'],log=lambda *_: None) # exact boundary probe, with model derivatives measured after training t=np.linspace(0.01,0.99,32,dtype=np.float32); bx=np.concatenate([np.c_[np.zeros(32),t],np.c_[t,np.zeros(32)],np.c_[t,np.ones(32)]]) bd=torch.as_tensor(bx, dtype=torch.float32) bdev='cuda' if next(bnet.parameters()).is_cuda else 'cpu' bo,bn=boundary_residuals(bnet,bd,bdev) inet,_,idev,_=idea_train(s,best['cfg'],True) io,inn=boundary_residuals(inet,bd,idev) sig.append({'seed':s,'baseline_oblique_abs_derivative':bo,'idea_oblique_abs_derivative':io,'baseline_normal_abs_derivative':bn,'idea_normal_abs_derivative':inn}) mean_b=float(np.mean([q['baseline_oblique_abs_derivative'] for q in sig])); mean_i=float(np.mean([q['idea_oblique_abs_derivative'] for q in sig])) signature={'prediction':'boundary-aware PDE training reduces measured oblique/no-flux residuals','mean_baseline_oblique_residual':mean_b,'mean_idea_oblique_residual':mean_i,'mean_baseline_normal_residual':float(np.mean([q['baseline_normal_abs_derivative'] for q in sig])),'mean_idea_normal_residual':float(np.mean([q['idea_normal_abs_derivative'] for q in sig])),'per_seed':sig,'confirmed': bool(mean_i < mean_b)} extra={'mechanism_signature':signature,'custom_track':{'name':'anisotropic_oblique_bvp','file':'pde_boundary_track.py','domain':'pde'},'idea_sweep':[{'cfg':r['cfg'],'mean':r['mean'],'std':r['std'],'per_seed':r['per_seed']} for r in idea_results]} report=make_report('anisotropic_oblique_bvp','mlp_tiny',base,best,extra) report['idea']['selected_cfg']=best['cfg']; report['protocol_note']='Custom PDE track required: no built-in track contains boundary-value/PDE structure.' Path('bench_report.json').write_text(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__=='__main__': main()