1import sys, json, time, numpy as np, torch
 2sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 3from bench import make_model, train_model, sweep_baseline, make_report
 4
 5# Local custom track: observations are Fourier measurements of two separated sources.
 6# Target is source coordinates (fixed two-source instance per sample); this is the
 7# closest valid supervised neural benchmark for the continuous localization mechanism.
 8META={'name':'fourier_source_localization','domain':'continuous_localization','description':'Noisy Fourier measurements of two separated planar sources; regress their coordinates.'}
 9def get_dataset(seed,n_train,n_test):
10    def make(n, ss):
11        rng=np.random.default_rng(ss); k=np.array([2,3,5,7,11,13,17,19],float)
12        def one():
13            x=np.sort(rng.uniform(.12,.88,2)); y=rng.uniform(.12,.88,2)
14            # fixed Fourier feature bank, complex coefficients encoded real/imag
15            ph=np.outer(k,x)+np.outer(k*0.73,y)
16            a=(np.exp(1j*ph[:,0])+np.exp(1j*ph[:,1]))/np.sqrt(2)
17            z=np.concatenate([a.real,a.imag])+rng.normal(0,.035,16)
18            return z, np.array([x[0],y[0],x[1],y[1]])
19        q=[one() for _ in range(n)]; return np.float32([a for a,b in q]),np.float32([b for a,b in q])
20    xtr,ytr=make(n_train,seed); xte,yte=make(n_test,seed+5000)
21    return {'xtr':torch.from_numpy(xtr),'ytr':torch.from_numpy(ytr),'xte':torch.from_numpy(xte),'yte':torch.from_numpy(yte),'task':'regression','metric':'mse','out_dim':4,'input_shape':(xtr.shape[1],)}
22
23def train_eval(seed, lr, idea, epochs=12):
24    d=get_dataset(seed,400,200); torch.manual_seed(seed+991)
25    net=make_model('mlp_tiny',d['input_shape'],4)
26    # Same architecture and training path. Idea is a differentiable soft coarse-grid
27    # proposal/refinement layer applied to the model's coordinate hypotheses.
28    if idea:
29        class Refined(torch.nn.Module):
30            def __init__(self, base): super().__init__(); self.base=base
31            def forward(self,x):
32                raw=self.base(x).view(-1,2,2).sigmoid()
33                # calibrated fixed-step contraction toward the nearest coarse proposal;
34                # this is the coarse-to-fine inductive bias, not an oracle.
35                grid=torch.tensor([.25,.5,.75],dtype=x.dtype,device=x.device)
36                g=torch.stack(torch.meshgrid(grid,grid,indexing='ij'),-1).reshape(-1,2)
37                dist=((raw[:,:,None,:]-g[None,None,:,:])**2).sum(-1)
38                w=torch.softmax(-30*dist,-1)
39                prop=(w[:,:,:,None]*g[None,None,:,:]).sum(2)
40                return (raw + 0.35*(prop-raw)).reshape(-1,4)
41        net=Refined(net)
42    net, metric, hist=train_model(net,d,epochs=epochs,lr=lr)
43    dev=next(net.parameters()).device
44    with torch.no_grad():
45        pred=net(d['xte'].to(dev)).cpu().numpy(); obs=d['yte'].numpy()
46    mse=float(np.mean((pred-obs)**2))
47    # trained-model mechanism signature: measured proposal displacement and residual.
48    raw=net.base(d['xte'].to(dev)).view(-1,2,2).sigmoid().detach().cpu().numpy() if idea else pred.reshape(-1,2,2)
49    final=pred.reshape(-1,2,2); displacement=float(np.mean(np.linalg.norm(final-raw,axis=-1)))
50    return mse, {'mse':mse,'proposal_displacement':displacement,'n':len(obs)}
51
52def main():
53    # shared union of baseline/idea learning rates; baseline sweep includes all.
54    grid=[{'lr':1e-3,'epochs':12},{'lr':3e-3,'epochs':12},{'lr':1e-2,'epochs':12}]
55    seeds=list(range(8)); base_sweep=[]
56    for cfg in grid:
57        vals=[train_eval(s,cfg['lr'],False,cfg['epochs'])[0] for s in seeds[:4]]
58        base_sweep.append({'config':cfg,'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':4})
59    best=min(base_sweep,key=lambda z:z['mean'])['config']
60    base=[]; idea=[]; sig=[]
61    for s in seeds:
62        b,_=train_eval(s,best['lr'],False,best['epochs']); i,sg=train_eval(s,best['lr'],True,best['epochs'])
63        base.append(b); idea.append(i); sig.append(sg)
64    # nearby idea settings are run, but only shared/best setting is reported to keep
65    # the paired comparison transparent and budget equal.
66    nearby={str(lr):[train_eval(s,lr,True,12)[0] for s in seeds[:4]] for lr in [1e-3,1e-2]}
67    delta=np.asarray(idea)-np.asarray(base)
68    rng=np.random.default_rng(0); signs=rng.choice([-1,1],(20000,8)); p=float(np.mean(np.abs((signs*delta).mean(1))>=abs(delta.mean())))
69    report={'track':'custom_fourier_source_localization','model':'mlp_tiny','baseline':{'best_cfg':best,'sweep':base_sweep,'full':{'mean':float(np.mean(base)),'std':float(np.std(base)),'per_seed':base,'n':8}},'idea':{'mean':float(np.mean(idea)),'std':float(np.std(idea)),'per_seed':idea,'n':8,'nearby':nearby},'paired_delta_mean':float(delta.mean()),'permutation_pvalue':p,'mechanism_signature':{'predicted':'coarse proposal contracts raw coordinates with fixed bandwidth-independent shrinkage','observed_mean_proposal_displacement':float(np.mean([x['proposal_displacement'] for x in sig])),'observed_mse':float(np.mean(idea)),'confirmed':False},'custom_track':{'name':'fourier_source_localization','file':'bench_experiment.py','domain':'continuous_localization'}}
70    print(json.dumps(report,indent=2))
71if __name__=='__main__': main()