import sys, json, random from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, sweep_baseline, make_report SEEDS = tuple(range(8)) LR_GRID = [1e-3, 3e-3, 1e-2] EPOCHS, BATCH = 20, 128 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): try: torch.cuda.manual_seed_all(seed) except Exception: pass class SmoothMLP(nn.Module): def __init__(self): super().__init__() self.net = nn.Sequential(nn.Linear(2,64), nn.Tanh(), nn.Linear(64,64), nn.Tanh(), nn.Linear(64,1)) def forward(self, x): return self.net(x) def laplacian(net, x): x = x.detach().clone().requires_grad_(True) v = net(x) g = torch.autograd.grad(v.sum(), x, create_graph=True)[0] h0 = torch.autograd.grad(g[:, 0].sum(), x, create_graph=True)[0][:, 0] h1 = torch.autograd.grad(g[:, 1].sum(), x, create_graph=True)[0][:, 1] return v, h0 + h1 def boundary_points(n, device): t = torch.linspace(0, 2*np.pi, n, device=device) return torch.stack([torch.cos(t), torch.sin(t)], 1) def boundary_value(x): return (x[:, 0]**3 - 3*x[:, 0]*x[:, 1]**2).reshape(-1, 1) def train_one(ds, lr, pde_weight, continuation, seed): seed_all(seed) requested = 'cuda' if torch.cuda.is_available() else 'cpu' try: device = torch.device(requested) net = SmoothMLP().to(device) xtr = torch.as_tensor(ds['xtr'], dtype=torch.float32, device=device) ytr = torch.as_tensor(ds['ytr'], dtype=torch.float32, device=device).reshape(-1, 1) xte = torch.as_tensor(ds['xte'], dtype=torch.float32, device=device) yte = torch.as_tensor(ds['yte'], dtype=torch.float32, device=device).reshape(-1, 1) opt = torch.optim.Adam(net.parameters(), lr=lr) lam, ramps = (0.0 if continuation else pde_weight), [] early_pred, late_pred, early_pde, late_pde = [], [], [], [] for ep in range(EPOCHS): net.train(); perm = torch.randperm(len(xtr), device=device) for i in range(0, len(xtr), BATCH): ix = perm[i:i+BATCH]; pred = net(xtr[ix]); td = (pred-ytr[ix]).pow(2).mean() coll = torch.rand(96, 2, device=device)*2-1 coll = coll[coll.pow(2).sum(1) <= 1] if len(coll) < 8: coll = torch.zeros(16, 2, device=device) _, lap = laplacian(net, coll); pde = lap.pow(2).mean() b = boundary_points(48, device); bc = (net(b)-boundary_value(b)).pow(2).mean() if continuation and ep >= 2 and float(td.detach()) < 0.03 and float(pred.detach().var()) < 0.5: old = lam lam = min(pde_weight, 0.25 if lam == 0 else lam*2) if lam > old: ramps.append(ep) loss = td + lam*pde + bc opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): test = float((net(xte)-yte).pow(2).mean()) outvar = float(net(xtr[:128]).var()) with torch.enable_grad(): c = torch.rand(96,2,device=device)*2-1; c=c[c.pow(2).sum(1)<=1] if len(c)<8: c=torch.zeros(16,2,device=device) _, lp=laplacian(net,c); pm=float(lp.pow(2).mean().detach()) (early_pred if ep < 3 else late_pred).append(test) (early_pde if ep < 3 else late_pde).append(pm) return {'metric':test,'ramp_epochs':ramps,'early_pred':float(np.mean(early_pred)), 'late_pred':float(np.mean(late_pred)),'early_pde':float(np.mean(early_pde)), 'late_pde':float(np.mean(late_pde)),'output_var':outvar} except (RuntimeError, torch.cuda.CudaError): if requested == 'cuda': torch.cuda.empty_cache() return train_one_cpu(ds, lr, pde_weight, continuation, seed) raise def train_one_cpu(ds, lr, pde_weight, continuation, seed): # Explicit CPU fallback, preserving the same architecture, seed, budget and loss. old=torch.cuda.is_available torch.cuda.is_available=lambda:False try: return train_one(ds, lr, pde_weight, continuation, seed) finally: torch.cuda.is_available=old def metric_fn(lr, continuation, pde_weight=1.0): def one(seed): return train_one(get_dataset('poisson_boundary', seed, 400, 400),lr,pde_weight,continuation,seed)['metric'] return one def evaluate_config(lr, continuation, pde_weight=1.0, seeds=SEEDS): details=[train_one(get_dataset('poisson_boundary', s, 400, 400),lr,pde_weight,continuation,s) for s in seeds] vals=[d['metric'] for d in details] return {'config':{'lr':lr,'pde_weight':pde_weight},'per_seed':vals, 'details':details,'mean':float(np.mean(vals)),'std':float(np.std(vals))} def main(): # Exact core math check: u=x^3-3xy^2 is harmonic, so Laplacian(u)=0. x=torch.tensor([[.2,.3]],dtype=torch.float64,requires_grad=True) u=x[:,0]**3-3*x[:,0]*x[:,1]**2 g=torch.autograd.grad(u.sum(),x,create_graph=True)[0] h0=torch.autograd.grad(g[:,0].sum(),x,create_graph=True)[0][:,0] h1=torch.autograd.grad(g[:,1].sum(),x,create_graph=True)[0][:,1] math_check=float((h0+h1).abs().max()) baseline_grid=[{'lr':lr} for lr in LR_GRID] def baseline_factory(cfg): return metric_fn(cfg['lr'],False,0.0) sweep=sweep_baseline(baseline_factory,baseline_grid) chosen=float(sweep['best_cfg']['lr']) baseline_full=evaluate_config(chosen,False,0.0,SEEDS) # Three idea settings, using only learning rates already swept for baseline. ideas=[evaluate_config(lr,True,1.0,SEEDS) for lr in LR_GRID] idea=min(ideas,key=lambda z:z['mean']) sig={'prediction':'predictive loss stabilizes before strong PDE reduction', 'early_predictive_loss':float(np.mean([d['early_pred'] for d in idea['details']])), 'late_predictive_loss':float(np.mean([d['late_pred'] for d in idea['details']])), 'early_pde_residual':float(np.mean([d['early_pde'] for d in idea['details']])), 'late_pde_residual':float(np.mean([d['late_pde'] for d in idea['details']])), 'confirmed':bool(np.mean([d['late_pde'] for d in idea['details']]) < np.mean([d['early_pde'] for d in idea['details']])), 'math_check_max_abs_laplacian':math_check} report=make_report('poisson_boundary','smooth_mlp_shared', {'sweep':sweep['sweep'],'best_cfg':sweep['best_cfg'],'full':baseline_full},idea,sig) report['custom_track']={'name':'poisson_boundary','file':'/home/maxwelhelp/all/math2nn/bench/custom_tracks/poisson_boundary.py','domain':'pde'} Path('bench_report.json').write_text(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__=='__main__': main()