1import json, math, sys
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import train_model, evaluate, sweep_baseline, make_report
  9
 10META = {
 11    'name': 'burgers_periodic_split',
 12    'domain': 'pde',
 13    'description': 'Periodic underresolved Burgers one-step regression on a 32-point grid.'
 14}
 15
 16
 17def _D(v):
 18    n = v.shape[-1]
 19    # x=2*pi*j/n, so Fourier integer wavenumbers are fftfreq(n)*n.
 20    k = torch.fft.rfftfreq(n, d=1.0/n).to(v.device)
 21    return torch.fft.irfft(1j * k * torch.fft.rfft(v, dim=-1), n=n, dim=-1)
 22
 23
 24def split_quadratic(u, alpha):
 25    beta = 1.0 - float(alpha)
 26    return float(alpha) * _D(0.5 * u * u) + beta * u * _D(u)
 27
 28
 29class SplitQuadraticNet(nn.Module):
 30    def __init__(self, alpha):
 31        super().__init__()
 32        self.alpha = float(alpha)
 33        self.inp = nn.Conv1d(1, 16, 1)
 34        self.mix = nn.Conv1d(16, 16, 1)
 35        self.out = nn.Conv1d(16, 1, 1)
 36        self.gamma = nn.Parameter(torch.tensor(0.1))
 37
 38    def forward(self, x):
 39        z = torch.tanh(self.inp(x))
 40        h = torch.tanh(self.mix(z))
 41        # The only intervention is the split coefficient in this quadratic block.
 42        return self.out(h + self.gamma * split_quadratic(h, self.alpha))
 43
 44
 45def get_dataset(seed, n_train=400, n_test=160):
 46    rng = np.random.default_rng(int(seed)); n = 32
 47    x = np.arange(n) * 2 * np.pi / n
 48
 49    def make(m):
 50        xs, ys = [], []
 51        for _ in range(int(m)):
 52            u = np.zeros(n)
 53            # Deliberately weakly underresolved high modes expose product aliasing.
 54            for mode in rng.choice(np.arange(5, 16), size=4, replace=False):
 55                u += rng.normal(0, 0.16 / np.sqrt(mode)) * np.sin(mode*x + rng.uniform(0, 2*np.pi))
 56                u += rng.normal(0, 0.16 / np.sqrt(mode)) * np.cos(mode*x + rng.uniform(0, 2*np.pi))
 57            dt = 0.025
 58            uh = np.fft.rfft(u); k = np.fft.rfftfreq(n, d=1.0/n)
 59            du = np.fft.irfft(1j*k*uh, n=n)
 60            # Target is a stable, filtered Burgers Euler step generated by padded product.
 61            mgrid = 48; up = np.zeros(mgrid, complex); q=n//2
 62            up[:q] = uh[:q] * (mgrid/n); up[-q:] = uh[-q:] * (mgrid/n)
 63            vp = np.fft.irfft(up, n=mgrid)
 64            kp=np.fft.rfftfreq(mgrid,d=1.0/mgrid)
 65            flux=np.fft.irfft(1j*kp*np.fft.rfft(0.5*vp*vp),n=mgrid)
 66            fluxh=np.fft.rfft(flux); trunc=np.zeros(n,complex)
 67            trunc[:q]=fluxh[:q]*(n/mgrid); trunc[-q:]=fluxh[-q:]*(n/mgrid)
 68            target=u-dt*(np.fft.irfft(trunc,n=n))
 69            xs.append(u.astype('float32')); ys.append(target.astype('float32'))
 70        return np.asarray(xs)[:,None,:], np.asarray(ys)[:,None,:]
 71
 72    xtr,ytr=make(n_train); xte,yte=make(n_test)
 73    return {'xtr':xtr,'ytr':ytr,'xte':xte,'yte':yte,'task':'regression','metric':'mse',
 74            'input_shape':(1,n),'out_dim':n}
 75
 76
 77def run():
 78    torch.set_num_threads(4)
 79    # All idea learning rates are also evaluated for the baseline.
 80    lrs = [1e-3, 3e-3, 1e-2]
 81    grid = [{'lr': lr, 'alpha': 1.0} for lr in lrs]
 82    idea_grid = [{'lr': lr, 'alpha': 1/3} for lr in lrs]
 83
 84    def factory(cfg):
 85        def train(seed):
 86            torch.manual_seed(seed); np.random.seed(seed)
 87            d=get_dataset(seed,400,160)
 88            d={k:(torch.from_numpy(v) if isinstance(v,np.ndarray) else v) for k,v in d.items()}
 89            net,metric,_=train_model(SplitQuadraticNet(cfg['alpha']),d,epochs=18,lr=cfg['lr'],batch=128,log=lambda *_:None)
 90            return float(metric) if metric is not None else float('inf')
 91        return train
 92
 93    base=sweep_baseline(factory,grid)
 94    # Equal-sized idea sweep; its selected setting is evaluated on all eight paired seeds.
 95    idea_sweep=sweep_baseline(factory,idea_grid)
 96    idea=idea_sweep['full']
 97    rep=make_report('burgers_periodic_split','custom_split_quadratic_net',base,idea,{
 98        'predicted': {'aliasing_root_alpha_for_tau4_i1_j1': 1/3,
 99                      'prediction': 'alpha=1/3 reduces the dominant unresolved quadratic-product aliasing proxy'},
100        'observed': {'baseline_alpha': 1.0, 'idea_alpha': 1/3,
101                     'baseline_test_mse': base['full']['mean'],
102                     'idea_test_mse': idea['mean'],
103                     'idea_sweep': idea_sweep['sweep']},
104        'confirmed': False,
105        'measurement_note': 'The trained-model aliasing proxy was not retained by the generic harness metric path; confirmation is false rather than inferred from the formula.'
106    })
107    rep['custom_track']={'name':META['name'],'file':'split_bench.py','domain':'pde'}
108    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
109    print(json.dumps(rep,indent=2))
110
111if __name__ == '__main__': run()