1import sys, json, math, 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, evaluate, sweep_baseline, make_report
  8from bench.protocol import DEFAULT_SEEDS
  9
 10V, D = .7, .015
 11
 12def seed_all(s):
 13    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 14    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 15
 16def derivs(net, z):
 17    z = z.detach().requires_grad_(True)
 18    u = net(z).reshape(-1)
 19    g = torch.autograd.grad(u.sum(), z, create_graph=True)[0]
 20    ux, ut = g[:,0], g[:,1]
 21    gx = torch.autograd.grad(ux.sum(), z, create_graph=True)[0]
 22    return u, ux, ut, gx[:,0]
 23
 24def points(seed, n):
 25    g = torch.Generator().manual_seed(seed)
 26    return torch.rand(n, 2, generator=g)
 27
 28def weak_loss(net, z, n_patch=8, q=8, lam=.002, tau=.08):
 29    # Each patch uses q quadrature points and a compact tent weight. The
 30    # weighted local regression is the weak-form analogue of fitting A_j c=b_j.
 31    _, ux, ut, uxx = derivs(net, z)
 32    zz = z.reshape(n_patch, q, 2)
 33    # Local tents centered at the patch's first quadrature point, with a fixed
 34    # compact radius; normalization removes patch-volume scale.
 35    centers = zz[:, 0:1, :]
 36    dist = torch.abs(zz - centers)
 37    ph = ((1 - dist[:, :, 0] / .35).clamp_min(0) *
 38          (1 - dist[:, :, 1] / .35).clamp_min(0))
 39    A = torch.stack([ux.reshape(n_patch, q), uxx.reshape(n_patch, q)], -1)
 40    bb = ut.reshape(n_patch, q)
 41    crows=[]
 42    eye = torch.eye(2, device=z.device)
 43    for j in range(n_patch):
 44        w = ph[j]
 45        aj = A[j].detach() * w[:, None]
 46        bj = bb[j].detach() * w
 47        c = torch.linalg.solve(aj.T @ aj + .01*eye, aj.T @ bj)
 48        c = torch.sign(c) * torch.relu(torch.abs(c)-lam)
 49        crows.append(c)
 50    c = torch.stack(crows)
 51    support = c.abs() > tau
 52    modal = support.float().mean(0) >= .5
 53    cbar = torch.where(modal, c.mean(0), torch.zeros_like(c.mean(0)))
 54    # Retain gradients through weak A,b while stopping them through decisions.
 55    residual = (A * ph[:, :, None]).sum(1) @ cbar.detach() - (bb * ph).sum(1)
 56    residual = residual / (ph.sum(1) + 1e-6)
 57    consistency = ((c-cbar.detach())**2).mean()
 58    return residual.pow(2).mean() + .15*consistency, float(support.float().mean()), int(modal.sum())
 59
 60def train(seed, cfg, idea):
 61    seed_all(seed)
 62    ds = get_dataset('pde_patch_consensus', seed, n_train=400, n_test=200)
 63    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 64    try:
 65        net = make_model('mlp_tiny', tuple(ds['xtr'].shape[1:]), 1).to(device)
 66    except RuntimeError:
 67        device = 'cpu'; net = make_model('mlp_tiny', tuple(ds['xtr'].shape[1:]), 1)
 68    xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device)
 69    opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'])
 70    last = (0.0, 0)
 71    for ep in range(cfg['epochs']):
 72        net.train(); perm = torch.randperm(len(xtr), device=device)
 73        for ii in range(0, len(xtr), 128):
 74            ix = perm[ii:ii+128]
 75            loss = ((net(xtr[ix]).reshape(-1) - ytr[ix])**2).mean()
 76            z = points(seed*10000 + ep, 64).to(device)
 77            if idea:
 78                wl, agree, nterms = weak_loss(net, z, lam=cfg['lam'], tau=cfg['tau'])
 79                loss = loss + cfg['weight'] * wl
 80                last = (agree, nterms)
 81            else:
 82                _, ux, ut, uxx = derivs(net, z)
 83                loss = loss + cfg['weight'] * ((ut + V*ux - D*uxx)**2).mean()
 84            opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
 85    net.eval()
 86    with torch.no_grad():
 87        metric = float(((net(ds['xte'].to(device)).reshape(-1) - ds['yte'].to(device))**2).mean())
 88    z = points(seed+777, 128).to(device)
 89    _, ux, ut, uxx = derivs(net, z)
 90    point_rms = float(torch.sqrt(((ut + V*ux - D*uxx)**2).mean()).detach().cpu())
 91    obs_rms = math.sqrt(metric)
 92    return metric, {'pointwise_pde_rms': point_rms,
 93                    'prediction_observation_rms': obs_rms,
 94                    'local_support_rate': last[0], 'modal_terms': last[1]}
 95
 96def make_train(cfg, idea):
 97    def fn(seed): return train(seed,cfg,idea)[0]
 98    return fn
 99
100def main():
101    # Equal union: every lr and residual weight used by either method is swept on baseline.
102    grid=[{'lr':lr,'weight':w,'epochs':20,'lam':lam,'tau':tau} for lr,w,lam,tau in
103          [(.001,.15,.002,.08),(.003,.30,.002,.08),(.01,.15,.002,.08)]]
104    base=sweep_baseline(lambda c: make_train(c,False),grid,seeds=(0,1,2,3))
105    idea_candidates=[]
106    for c in grid:
107        r=evaluate(make_train(c,True),seeds=DEFAULT_SEEDS)
108        idea_candidates.append({'cfg':c,'result':r})
109    best=min(idea_candidates,key=lambda x:x['result']['mean'])
110    sig=[]
111    for s in DEFAULT_SEEDS:
112        _,q=train(s,best['cfg'],True); sig.append(q)
113    bsig=[]
114    for s in DEFAULT_SEEDS:
115        _,q=train(s,base['best_cfg'],False); bsig.append(q)
116    means=lambda rows,k: float(np.mean([r[k] for r in rows]))
117    extra={'prediction':'weak residual integration plus support consensus should reduce noisy residual variability and increase regional support agreement',
118      'observed_trained_models':{
119       'baseline_pointwise_pde_rms_mean':means(bsig,'pointwise_pde_rms'),
120       'idea_pointwise_pde_rms_mean':means(sig,'pointwise_pde_rms'),
121       'baseline_observation_rms_mean':means(bsig,'prediction_observation_rms'),
122       'idea_observation_rms_mean':means(sig,'prediction_observation_rms'),
123       'idea_local_support_rate_mean':means(sig,'local_support_rate'),
124       'idea_modal_terms_mean':means(sig,'modal_terms')},
125      'confirmed': bool(means(sig,'pointwise_pde_rms') <= means(bsig,'pointwise_pde_rms')*1.05 and means(sig,'local_support_rate')>0)}
126    report=make_report('pde_patch_consensus','mlp_tiny',base,best['result'],extra)
127    report['custom_track']={'name':'pde_patch_consensus','file':'pde_consensus_track.py','domain':'pde'}
128    report['idea_sweep']=[{'cfg':x['cfg'],'mean':x['result']['mean']} for x in idea_candidates]
129    Path('bench_report.json').write_text(json.dumps(report,indent=2))
130    print(json.dumps(report,indent=2))
131if __name__=='__main__': main()