1import sys, json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
  7
  8TRACK = 'oscillatory_poisson_dirichlet'
  9WIDTH, DEPTH = 64, 2
 10
 11# Golden-ratio continued-fraction denominators: 1,1,2,3,5,8,13,21.
 12def denominators(j=8):
 13    out=[]; qm1, q = 0, 1
 14    for _ in range(j):
 15        out.append(q); qm1, q = q, q+qm1
 16    return np.asarray(out, dtype=np.float32)
 17Q = denominators(8)
 18
 19def encode_np(x, kind):
 20    x = np.asarray(x, dtype=np.float32).reshape(-1, 1)
 21    if kind == 'baseline':
 22        return x
 23    z = 2*np.pi*x*Q.reshape(1,-1)
 24    # prescribed square-root energy normalization
 25    return np.concatenate([x, np.cos(z)/np.sqrt(Q), np.sin(z)/np.sqrt(Q)], axis=1).astype(np.float32)
 26
 27class Net(nn.Module):
 28    def __init__(self, input_dim):
 29        super().__init__()
 30        layers=[]; d=input_dim
 31        for _ in range(DEPTH):
 32            layers += [nn.Linear(d, WIDTH), nn.ReLU()]; d=WIDTH
 33        layers.append(nn.Linear(d,1)); self.net=nn.Sequential(*layers)
 34    def forward(self,x): return self.net(x)
 35
 36def make_train(kind, seed, cfg, keep=False):
 37    torch.manual_seed(10000+int(seed)); np.random.seed(10000+int(seed)); random.seed(10000+int(seed))
 38    d0 = get_dataset(TRACK, seed, 400, 400)
 39    d = dict(d0)
 40    for k in ('xtr','ytr','xte','yte'):
 41        d[k] = encode_np(d0[k], kind) if k.startswith('x') else np.asarray(d0[k], dtype=np.float32).reshape(-1,1)
 42    d['input_shape'] = d['xtr'].shape[1:]
 43    for k in ('xtr','ytr','xte','yte'):
 44        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
 45    model = Net(d['xtr'].shape[1])
 46    net, metric, hist = train_model(model, d, epochs=int(cfg['epochs']), lr=float(cfg['lr']), batch=128, log=lambda *_: None)
 47    if net is None: raise RuntimeError('training failed')
 48    if keep: return float(metric), net, d
 49    return float(metric)
 50
 51def factory(kind, cfg):
 52    return lambda seed: make_train(kind, seed, cfg)
 53
 54def math_check():
 55    # Verify recursion, absolute summability, and finite encoded feature energy.
 56    qs=denominators(16)
 57    recursion=bool(np.all(qs[2:] == qs[1:-1] + qs[:-2]))
 58    inv=float(np.sum(1/qs))
 59    x=np.linspace(0,1,1001)
 60    z=2*np.pi*x[:,None]*qs[None,:]
 61    energy=np.mean((np.cos(z)/np.sqrt(qs))**2+(np.sin(z)/np.sqrt(qs))**2,axis=0)
 62    return {'q_first_8':qs[:8].astype(int).tolist(),'recursion_ok':recursion,
 63            'sum_inverse_q_16':inv,'pair_energy_mean':float(energy.mean()),
 64            'pair_energy_range':[float(energy.min()),float(energy.max())],
 65            'confirmed': recursion and inv < 4.0 and bool(np.allclose(energy,1/qs,rtol=.03,atol=.01))}
 66
 67def signature(base_cfg, idea_cfg):
 68    rows=[]
 69    for s in range(8):
 70        bm,bn,bd=make_train('baseline',s,base_cfg,True)
 71        im,inn,idata=make_train('idea',s,idea_cfg,True)
 72        dev=next(bn.parameters()).device
 73        xb_base=torch.as_tensor(bd['xte'],dtype=torch.float32, device=dev)
 74        xb_idea=torch.as_tensor(idata['xte'],dtype=torch.float32, device=next(inn.parameters()).device)
 75        y=bd['yte'].numpy().ravel()
 76        with torch.no_grad():
 77            bp=bn(xb_base).cpu().numpy().ravel()
 78            ip=inn(xb_idea).cpu().numpy().ravel()
 79        # measured high-frequency projection on the q=21 channel
 80        phase=np.sin(2*np.pi*np.asarray(bd['xte']).reshape(-1)*21); phase-=phase.mean()
 81        def proj(v):
 82            v=v-v.mean(); return float(np.dot(v,phase)/np.dot(phase,phase))
 83        rows.append({'seed':s,'observed_projection':proj(y),'baseline_projection':proj(bp),'idea_projection':proj(ip),'baseline_mse':bm,'idea_mse':im})
 84    obs=np.array([r['observed_projection'] for r in rows]); ba=np.array([r['baseline_projection'] for r in rows]); ia=np.array([r['idea_projection'] for r in rows])
 85    be=float(np.mean(np.abs(ba-obs))); ie=float(np.mean(np.abs(ia-obs)))
 86    return {'prediction':'lacunary features should represent the observed q=21 oscillatory projection more accurately than raw coordinates',
 87            'observed_projection_mean':float(obs.mean()),'baseline_projection_mean':float(ba.mean()),'idea_projection_mean':float(ia.mean()),
 88            'baseline_abs_projection_error':be,'idea_abs_projection_error':ie,'confirmed':bool(ie<be),'per_seed':rows}
 89
 90def main():
 91    check=math_check()
 92    # Union parity: all idea settings are included in baseline sweep.
 93    grid=[{'lr':1e-3,'epochs':40},{'lr':3e-3,'epochs':60},{'lr':1e-2,'epochs':40}]
 94    base=sweep_baseline(lambda cfg: factory('baseline',cfg), grid)
 95    best=base['best_cfg']
 96    idea_runs=[evaluate(factory('idea',cfg)) for cfg in grid]
 97    idx=int(np.argmin([r['mean'] for r in idea_runs])); idea=idea_runs[idx]; chosen=grid[idx]
 98    # paired comparison uses same chosen config on baseline and idea, both all 8 seeds
 99    baseline_full=evaluate(factory('baseline',chosen))
100    sig=signature(chosen,chosen)
101    rep=make_report(TRACK,'mlp_tiny',base,idea,{'core_math':check,**sig})
102    rep['baseline_full_at_idea_cfg']=baseline_full
103    rep['idea_config_candidates']=grid; rep['chosen_idea_cfg']=chosen
104    rep['track_justification']='Registered oscillatory elliptic PDE field: the task contains the high-frequency coordinate structure targeted by the embedding.'
105    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
106    print(json.dumps(rep,indent=2))
107if __name__=='__main__': main()