1import sys, json, time, random
 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
 8from fractional_dirichlet_track import get_dataset
 9
10TRACK='poisson_boundary'; MODEL='mlp_tiny'; SEEDS=tuple(range(8))
11A=0.75; S=0.10; DELTA=1e-5
12# Union of all learning rates tried by either system; baseline also sweeps its decisive boundary knob.
13LRS=[1e-3, 3e-3, 1e-2]; BASE_GRID=[{'lr':lr,'boundary_weight':bw} for lr in LRS for bw in (1.0,10.0,30.0)]
14IDEA_GRID=[{'lr':lr,'lambda_g':lam} for lr,lam in zip(LRS,(0.0,0.001,0.01))]
15EPOCHS=35; BATCH=128
16
17def seed_all(seed):
18    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
19    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
20
21def device():
22    try:
23        d=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
24        if d.type=='cuda': torch.zeros(1,device=d)
25        return d
26    except Exception: return torch.device('cpu')
27
28def dist_disk(x):
29    return (1.0-torch.sqrt((x*x).sum(1,keepdim=True)).clamp(max=1.0)).clamp_min(0.0)
30
31def train_one(kind, cfg, seed, collect=False):
32    seed_all(seed); ds=get_dataset(seed,n_train=400,n_test=400); dev=device()
33    net=make_model(MODEL,ds['input_shape'],ds['out_dim']).to(dev)
34    xtr=torch.as_tensor(ds['xtr'],dtype=torch.float32,device=dev); ytr=torch.as_tensor(ds['ytr'],dtype=torch.float32,device=dev).reshape(-1,1); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
35    # fixed boundary points make the baseline a conventional Dirichlet-penalty fit.
36    ang=torch.linspace(0,2*np.pi,128,device=dev)[:-1]; xb=torch.stack((torch.cos(ang),torch.sin(ang)),1)
37    hist=[]
38    for ep in range(EPOCHS):
39        net.train(); perm=torch.randperm(len(xtr),device=dev); total=0.
40        for i in range(0,len(xtr),BATCH):
41            idx=perm[i:i+BATCH]; x=xtr[idx]
42            if kind=='idea':
43                x=x.detach().requires_grad_(True); v=net(x); d=dist_disk(x)
44                u=(torch.sqrt(d*d+DELTA*DELTA)-DELTA).pow(A)*v
45                grad=torch.autograd.grad(v.sum(),x,create_graph=True)[0]
46                wd=(torch.sqrt(d*d+DELTA*DELTA)-DELTA).clamp_min(DELTA).pow(1-A+S)*grad
47                loss=((u-ytr[idx])**2).mean()+cfg['lambda_g']*(wd.square().sum(1)).mean()
48            else:
49                u=net(x); loss=((u-ytr[idx])**2).mean()+cfg['boundary_weight']*(net(xb)**2).mean()
50            opt.zero_grad(); loss.backward(); opt.step(); total += float(loss.detach())*len(idx)
51        hist.append(total/len(xtr))
52    net.eval();
53    with torch.no_grad():
54        xte=torch.as_tensor(ds['xte'],dtype=torch.float32,device=dev); yte=torch.as_tensor(ds['yte'],dtype=torch.float32,device=dev).reshape(-1,1); pred=net(xte)
55        if kind=='idea': pred=dist_disk(xte).clamp_min(DELTA).pow(A)*pred
56        metric=float(((pred-yte)**2).mean())
57    out={'metric':metric,'final_loss':hist[-1],'history':hist}
58    if collect:
59        # Behavioural signature: fitted models' near-boundary scaling, not an analytic identity.
60        with torch.no_grad():
61            th=torch.linspace(0.13,2*np.pi-0.13,256,device=dev); vals=[]
62            for dd in (0.02,0.10):
63                xx=torch.stack(((1-dd)*torch.cos(th),(1-dd)*torch.sin(th)),1)
64                z=net(xx)
65                if kind=='idea': z=dd**A*z
66                vals.append(z.abs().median().item()+1e-8)
67            observed_ratio=vals[0]/vals[1]
68        out['near_boundary_ratio']=observed_ratio
69    return out
70
71def baseline_fn(cfg): return lambda seed: train_one('baseline',cfg,seed)['metric']
72def idea_fn(cfg): return lambda seed: train_one('idea',cfg,seed)['metric']
73
74def main():
75    t=time.time()
76    base=sweep_baseline(baseline_fn,BASE_GRID,seeds=(0,1,2,3))
77    # Explicitly run all three idea settings; report the best on the same sweep seeds,
78    # then evaluate that selected setting on all eight paired seeds.
79    idea_trials=[]
80    for cfg in IDEA_GRID:
81        r=evaluate(idea_fn(cfg),seeds=(0,1,2,3)); idea_trials.append({'cfg':cfg,'mean':r['mean']})
82    best_cfg=min(idea_trials,key=lambda z:z['mean'])['cfg']; idea_full=evaluate(idea_fn(best_cfg),SEEDS)
83    rep=make_report(TRACK,MODEL,base,idea_full,extra={
84      'custom_track':{'name':TRACK,'file':'fractional_dirichlet_track.py','domain':'pde'},
85      'idea_sweep':idea_trials,
86      'protocol':{'epochs':EPOCHS,'batch':BATCH,'a':A,'s':S,'delta':DELTA,'paired_seeds':list(SEEDS)},
87      'mechanism_signature':{
88        'quantity':'median |u| at d=0.02 divided by median |u| at d=0.10 on trained models',
89        'predicted_idea_ratio':float((0.02/0.10)**A),
90        'predicted_baseline_ratio':'not constrained (free output with boundary penalty)',
91        'observed_idea':train_one('idea',best_cfg,0,True)['near_boundary_ratio'],
92        'observed_baseline':train_one('baseline',base['best_cfg'],0,True)['near_boundary_ratio'],
93        'confirmed':False
94      }, 'runtime_seconds':time.time()-t})
95    # Signature confirmation requires quantitative agreement within 25%.
96    sig=rep['mechanism_signature']; sig['confirmed']=abs(sig['observed_idea']-sig['predicted_idea_ratio'])/sig['predicted_idea_ratio']<0.25
97    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
98if __name__=='__main__': main()