1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn.functional as F
  6
  7ROOT = Path(__file__).resolve().parent
  8sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  9from bench import make_model, train_model, sweep_baseline, evaluate, make_report
 10from pde_boundary_track import get_dataset, C
 11
 12SEEDS = tuple(range(8))
 13EPOCHS = 25
 14BATCH = 128
 15# The baseline and idea share this union of learning rates and Adam weight decay.
 16CONFIGS = [
 17    {'lr': 1e-3, 'weight_decay': 0.0},
 18    {'lr': 3e-3, 'weight_decay': 0.0},
 19    {'lr': 6e-3, 'weight_decay': 0.0},
 20]
 21
 22
 23def seed_all(seed):
 24    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 25    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 26
 27
 28def ds_for(seed):
 29    d = get_dataset(seed, 400, 400)
 30    for k in ('xtr','ytr','xte','yte'):
 31        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
 32    d['ytr'] = d['ytr'].reshape(-1,1); d['yte'] = d['yte'].reshape(-1,1)
 33    d['input_shape'] = tuple(d['xtr'].shape[1:]); d['out_dim'] = 1
 34    return d
 35
 36
 37def baseline_fn(cfg):
 38    def run(seed):
 39        seed_all(seed); d=ds_for(seed)
 40        net=make_model('mlp_tiny', d['input_shape'], 1)
 41        _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH,
 42                                   weight_decay=cfg['weight_decay'], log=lambda *_: None)
 43        return metric
 44    return run
 45
 46
 47def boundary_residuals(net, x, device):
 48    # Values are measured from the trained model, not an analytical identity.
 49    z=x.to(device).detach().clone().requires_grad_(True)
 50    u=net(z)
 51    g=torch.autograd.grad(u.sum(), z, create_graph=False)[0]
 52    xx, yy=z[:,0], z[:,1]
 53    left=xx < 1e-7
 54    horiz=(yy < 1e-7) | (yy > 1-1e-7)
 55    ob=torch.abs(g[:,0] + C*g[:,1])[left]
 56    no=torch.abs(g[:,1])[horiz]
 57    return float(ob.mean().detach().cpu()) if ob.numel() else 0., float(no.mean().detach().cpu()) if no.numel() else 0.
 58
 59
 60def idea_train(seed, cfg, return_net=False):
 61    seed_all(seed); d=ds_for(seed)
 62    device='cuda' if torch.cuda.is_available() else 'cpu'
 63    try:
 64        net=make_model('mlp_tiny', d['input_shape'], 1).to(device)
 65        opt=torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
 66        xtr,ytr=d['xtr'].to(device),d['ytr'].to(device)
 67        # Boundary penalty is the PDE intervention; all architecture/data/budget terms match.
 68        for _ in range(EPOCHS):
 69            net.train(); perm=torch.randperm(len(xtr),device=device)
 70            for i in range(0,len(xtr),BATCH):
 71                x=xtr[perm[i:i+BATCH]].detach().clone().requires_grad_(True)
 72                y=ytr[perm[i:i+BATCH]]
 73                pred=net(x); loss=F.mse_loss(pred,y)
 74                grad=torch.autograd.grad(pred.sum(),x,create_graph=True)[0]
 75                left=x[:,0] < 1e-7; horiz=(x[:,1] < 1e-7)|(x[:,1] > 1-1e-7)
 76                bc=0.0
 77                if left.any(): bc=bc+(grad[left,0]+C*grad[left,1]).square().mean()
 78                if horiz.any(): bc=bc+grad[horiz,1].square().mean()
 79                total=loss + cfg['bc_weight']*bc
 80                opt.zero_grad(); total.backward(); opt.step()
 81        net.eval()
 82        with torch.no_grad(): metric=float(((net(d['xte'].to(device))-d['yte'].to(device))**2).mean().cpu())
 83        if return_net: return net, d, device, metric
 84        return metric
 85    except RuntimeError:
 86        # Explicit CPU fallback for shared GPU failures.
 87        torch.cuda.empty_cache() if torch.cuda.is_available() else None
 88        old=torch.cuda.is_available
 89        # Re-run with CPU by temporarily using an equivalent forced implementation.
 90        if device != 'cpu':
 91            torch.set_default_device('cpu')
 92            try: return idea_train(seed, cfg, return_net)
 93            finally: torch.set_default_device('cpu')
 94        raise
 95
 96
 97def idea_fn(cfg):
 98    return lambda seed: idea_train(seed, cfg)
 99
100
101def main():
102    # Baseline sweep uses the same lr union and the same eight-seed final evaluation.
103    base=sweep_baseline(baseline_fn, CONFIGS)
104    idea_configs=[dict(CONFIGS[CONFIGS.index(base['best_cfg'])], bc_weight=w) for w in (0.1, 1.0, 5.0)]
105    idea_results=[]
106    for cfg in idea_configs:
107        r=evaluate(idea_fn(cfg), SEEDS); r['cfg']=cfg; idea_results.append(r)
108    best=min(idea_results, key=lambda r:r['mean'])
109    # Re-test signature on all boundary points from the trained systems.
110    sig=[]
111    for s in SEEDS:
112        bcfg=base['best_cfg']; seed_all(s); d=ds_for(s)
113        bnet=make_model('mlp_tiny',d['input_shape'],1)
114        bnet,_,_=train_model(bnet,d,epochs=EPOCHS,lr=bcfg['lr'],batch=BATCH,weight_decay=bcfg['weight_decay'],log=lambda *_: None)
115        # exact boundary probe, with model derivatives measured after training
116        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)]])
117        bd=torch.as_tensor(bx, dtype=torch.float32)
118        bdev='cuda' if next(bnet.parameters()).is_cuda else 'cpu'
119        bo,bn=boundary_residuals(bnet,bd,bdev)
120        inet,_,idev,_=idea_train(s,best['cfg'],True)
121        io,inn=boundary_residuals(inet,bd,idev)
122        sig.append({'seed':s,'baseline_oblique_abs_derivative':bo,'idea_oblique_abs_derivative':io,'baseline_normal_abs_derivative':bn,'idea_normal_abs_derivative':inn})
123    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]))
124    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)}
125    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]}
126    report=make_report('anisotropic_oblique_bvp','mlp_tiny',base,best,extra)
127    report['idea']['selected_cfg']=best['cfg']; report['protocol_note']='Custom PDE track required: no built-in track contains boundary-value/PDE structure.'
128    Path('bench_report.json').write_text(json.dumps(report,indent=2))
129    print(json.dumps(report,indent=2))
130
131if __name__=='__main__': main()