1import sys, json, 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 get_dataset, make_model, sweep_baseline, make_report
  8
  9SEEDS = tuple(range(8))
 10LR_GRID = [1e-3, 3e-3, 1e-2]
 11EPOCHS, BATCH = 20, 128
 12
 13def seed_all(seed):
 14    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 15    if torch.cuda.is_available():
 16        try: torch.cuda.manual_seed_all(seed)
 17        except Exception: pass
 18
 19class SmoothMLP(nn.Module):
 20    def __init__(self):
 21        super().__init__()
 22        self.net = nn.Sequential(nn.Linear(2,64), nn.Tanh(), nn.Linear(64,64), nn.Tanh(), nn.Linear(64,1))
 23    def forward(self, x): return self.net(x)
 24
 25def laplacian(net, x):
 26    x = x.detach().clone().requires_grad_(True)
 27    v = net(x)
 28    g = torch.autograd.grad(v.sum(), x, create_graph=True)[0]
 29    h0 = torch.autograd.grad(g[:, 0].sum(), x, create_graph=True)[0][:, 0]
 30    h1 = torch.autograd.grad(g[:, 1].sum(), x, create_graph=True)[0][:, 1]
 31    return v, h0 + h1
 32
 33def boundary_points(n, device):
 34    t = torch.linspace(0, 2*np.pi, n, device=device)
 35    return torch.stack([torch.cos(t), torch.sin(t)], 1)
 36
 37def boundary_value(x):
 38    return (x[:, 0]**3 - 3*x[:, 0]*x[:, 1]**2).reshape(-1, 1)
 39
 40def train_one(ds, lr, pde_weight, continuation, seed):
 41    seed_all(seed)
 42    requested = 'cuda' if torch.cuda.is_available() else 'cpu'
 43    try:
 44        device = torch.device(requested)
 45        net = SmoothMLP().to(device)
 46        xtr = torch.as_tensor(ds['xtr'], dtype=torch.float32, device=device)
 47        ytr = torch.as_tensor(ds['ytr'], dtype=torch.float32, device=device).reshape(-1, 1)
 48        xte = torch.as_tensor(ds['xte'], dtype=torch.float32, device=device)
 49        yte = torch.as_tensor(ds['yte'], dtype=torch.float32, device=device).reshape(-1, 1)
 50        opt = torch.optim.Adam(net.parameters(), lr=lr)
 51        lam, ramps = (0.0 if continuation else pde_weight), []
 52        early_pred, late_pred, early_pde, late_pde = [], [], [], []
 53        for ep in range(EPOCHS):
 54            net.train(); perm = torch.randperm(len(xtr), device=device)
 55            for i in range(0, len(xtr), BATCH):
 56                ix = perm[i:i+BATCH]; pred = net(xtr[ix]); td = (pred-ytr[ix]).pow(2).mean()
 57                coll = torch.rand(96, 2, device=device)*2-1
 58                coll = coll[coll.pow(2).sum(1) <= 1]
 59                if len(coll) < 8: coll = torch.zeros(16, 2, device=device)
 60                _, lap = laplacian(net, coll); pde = lap.pow(2).mean()
 61                b = boundary_points(48, device); bc = (net(b)-boundary_value(b)).pow(2).mean()
 62                if continuation and ep >= 2 and float(td.detach()) < 0.03 and float(pred.detach().var()) < 0.5:
 63                    old = lam
 64                    lam = min(pde_weight, 0.25 if lam == 0 else lam*2)
 65                    if lam > old: ramps.append(ep)
 66                loss = td + lam*pde + bc
 67                opt.zero_grad(); loss.backward(); opt.step()
 68            net.eval()
 69            with torch.no_grad():
 70                test = float((net(xte)-yte).pow(2).mean())
 71                outvar = float(net(xtr[:128]).var())
 72            with torch.enable_grad():
 73                c = torch.rand(96,2,device=device)*2-1; c=c[c.pow(2).sum(1)<=1]
 74                if len(c)<8: c=torch.zeros(16,2,device=device)
 75                _, lp=laplacian(net,c); pm=float(lp.pow(2).mean().detach())
 76            (early_pred if ep < 3 else late_pred).append(test)
 77            (early_pde if ep < 3 else late_pde).append(pm)
 78        return {'metric':test,'ramp_epochs':ramps,'early_pred':float(np.mean(early_pred)),
 79                'late_pred':float(np.mean(late_pred)),'early_pde':float(np.mean(early_pde)),
 80                'late_pde':float(np.mean(late_pde)),'output_var':outvar}
 81    except (RuntimeError, torch.cuda.CudaError):
 82        if requested == 'cuda':
 83            torch.cuda.empty_cache()
 84            return train_one_cpu(ds, lr, pde_weight, continuation, seed)
 85        raise
 86
 87def train_one_cpu(ds, lr, pde_weight, continuation, seed):
 88    # Explicit CPU fallback, preserving the same architecture, seed, budget and loss.
 89    old=torch.cuda.is_available
 90    torch.cuda.is_available=lambda:False
 91    try:
 92        return train_one(ds, lr, pde_weight, continuation, seed)
 93    finally:
 94        torch.cuda.is_available=old
 95
 96def metric_fn(lr, continuation, pde_weight=1.0):
 97    def one(seed):
 98        return train_one(get_dataset('poisson_boundary', seed, 400, 400),lr,pde_weight,continuation,seed)['metric']
 99    return one
100
101def evaluate_config(lr, continuation, pde_weight=1.0, seeds=SEEDS):
102    details=[train_one(get_dataset('poisson_boundary', s, 400, 400),lr,pde_weight,continuation,s) for s in seeds]
103    vals=[d['metric'] for d in details]
104    return {'config':{'lr':lr,'pde_weight':pde_weight},'per_seed':vals,
105            'details':details,'mean':float(np.mean(vals)),'std':float(np.std(vals))}
106
107def main():
108    # Exact core math check: u=x^3-3xy^2 is harmonic, so Laplacian(u)=0.
109    x=torch.tensor([[.2,.3]],dtype=torch.float64,requires_grad=True)
110    u=x[:,0]**3-3*x[:,0]*x[:,1]**2
111    g=torch.autograd.grad(u.sum(),x,create_graph=True)[0]
112    h0=torch.autograd.grad(g[:,0].sum(),x,create_graph=True)[0][:,0]
113    h1=torch.autograd.grad(g[:,1].sum(),x,create_graph=True)[0][:,1]
114    math_check=float((h0+h1).abs().max())
115    baseline_grid=[{'lr':lr} for lr in LR_GRID]
116    def baseline_factory(cfg): return metric_fn(cfg['lr'],False,0.0)
117    sweep=sweep_baseline(baseline_factory,baseline_grid)
118    chosen=float(sweep['best_cfg']['lr'])
119    baseline_full=evaluate_config(chosen,False,0.0,SEEDS)
120    # Three idea settings, using only learning rates already swept for baseline.
121    ideas=[evaluate_config(lr,True,1.0,SEEDS) for lr in LR_GRID]
122    idea=min(ideas,key=lambda z:z['mean'])
123    sig={'prediction':'predictive loss stabilizes before strong PDE reduction',
124         'early_predictive_loss':float(np.mean([d['early_pred'] for d in idea['details']])),
125         'late_predictive_loss':float(np.mean([d['late_pred'] for d in idea['details']])),
126         'early_pde_residual':float(np.mean([d['early_pde'] for d in idea['details']])),
127         'late_pde_residual':float(np.mean([d['late_pde'] for d in idea['details']])),
128         'confirmed':bool(np.mean([d['late_pde'] for d in idea['details']]) < np.mean([d['early_pde'] for d in idea['details']])),
129         'math_check_max_abs_laplacian':math_check}
130    report=make_report('poisson_boundary','smooth_mlp_shared',
131        {'sweep':sweep['sweep'],'best_cfg':sweep['best_cfg'],'full':baseline_full},idea,sig)
132    report['custom_track']={'name':'poisson_boundary','file':'/home/maxwelhelp/all/math2nn/bench/custom_tracks/poisson_boundary.py','domain':'pde'}
133    Path('bench_report.json').write_text(json.dumps(report,indent=2))
134    print(json.dumps(report,indent=2))
135if __name__=='__main__': main()