1import sys, json, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import make_model, train_model
  7from bench.protocol import sweep_baseline, evaluate, make_report
  8from phase_track import N, EDGES, get_dataset
  9
 10SEEDS = tuple(range(8))
 11# Union is shared by baseline and idea; the baseline is swept at all idea rates.
 12LR_GRID = [1e-3, 3e-3, 1e-2]
 13EPOCHS = 24
 14BATCH = 64
 15
 16def spectral_align(x):
 17    """Align each sample's 8 two-dimensional views using its noisy phases."""
 18    a = np.asarray(x, dtype=np.float32).copy()
 19    out = a[:, :2*N].reshape(-1, N, 2).copy()
 20    ph = a[:, 2*N:]
 21    aligned = np.empty_like(out)
 22    for k in range(len(a)):
 23        C = np.zeros((N, N), dtype=np.complex64)
 24        for e, (i, j) in enumerate(EDGES):
 25            z = np.exp(1j * ph[k, e])
 26            C[i, j] = z
 27            C[j, i] = np.conj(z)
 28        v = np.ones(N, dtype=np.complex64) / np.sqrt(N)
 29        for _ in range(15):
 30            q = C @ v
 31            nq = np.linalg.norm(q)
 32            if nq > 1e-8: v = q / nq
 33        theta = np.angle(v)
 34        theta -= theta[0]
 35        # R(-theta) removes the estimated view rotation.
 36        for i in range(N):
 37            c, s = np.cos(theta[i]), np.sin(theta[i])
 38            aligned[k, i] = (np.array([[c, s], [-s, c]], dtype=np.float32) @ out[k, i])
 39    return np.concatenate([aligned.reshape(len(a), -1), ph], axis=1).astype(np.float32)
 40
 41def ds(seed, aligned):
 42    d = get_dataset(seed, 400, 200)
 43    for k in ('xtr', 'ytr', 'xte', 'yte'):
 44        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
 45    if aligned:
 46        d['xtr'] = torch.from_numpy(spectral_align(d['xtr'].numpy()))
 47        d['xte'] = torch.from_numpy(spectral_align(d['xte'].numpy()))
 48    return d
 49
 50def set_seed(seed):
 51    np.random.seed(seed); torch.manual_seed(seed)
 52    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 53
 54def train_one(seed, lr, aligned):
 55    set_seed(seed + (10000 if aligned else 0))
 56    d = ds(seed, aligned)
 57    model = make_model('mlp_tiny', tuple(d['xtr'].shape[1:]), 1)
 58    _, metric, hist = train_model(model, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
 59    return float(metric)
 60
 61def make_side(aligned):
 62    return lambda cfg: (lambda seed: train_one(seed, cfg['lr'], aligned))
 63
 64def mechanism_signature():
 65    # NN-scale check: compare model predictions on identical raw vs aligned views,
 66    # and use the trained idea model's test MSE to assess the predicted stabilization.
 67    rows=[]
 68    for seed in SEEDS:
 69        set_seed(seed + 20000)
 70        raw=ds(seed, False); ali=ds(seed, True)
 71        m=make_model('mlp_tiny', tuple(raw['xtr'].shape[1:]), 1)
 72        m,_,_=train_model(m, ali, epochs=EPOCHS, lr=3e-3, batch=BATCH, log=lambda *_: None)
 73        m = m.cpu(); m.eval()
 74        with torch.no_grad():
 75            pr=m(ali['xte'].cpu()).numpy().ravel()
 76        y=ali['yte'].numpy().ravel()
 77        # observed task prediction error and residual variation, measured from trained NN
 78        rows.append((float(np.mean((pr-y)**2)), float(np.std(pr-y))))
 79    mse=float(np.mean([r[0] for r in rows])); resid=float(np.mean([r[1] for r in rows]))
 80    # Prediction: synchronization should reduce downstream error versus random/raw input.
 81    rawm=float(np.mean([train_one(s,3e-3,False) for s in SEEDS]))
 82    return {'prediction':'spectral alignment reduces downstream test MSE and residual spread',
 83            'predicted_baseline_mse':rawm, 'observed_aligned_mse':mse,
 84            'observed_aligned_residual_std':resid,
 85            'confirmed': bool(mse < rawm)}
 86
 87def main():
 88    t=time.time()
 89    base=sweep_baseline(make_side(False), [{'lr':x} for x in LR_GRID], seeds=(0,1,2,3))
 90    # Same three settings on idea side, including best baseline setting and neighbors.
 91    idea_cfgs=[{'lr':x} for x in LR_GRID]
 92    idea_trials=[]
 93    for cfg in idea_cfgs:
 94        r=evaluate(make_side(True)(cfg), seeds=SEEDS)
 95        idea_trials.append({'cfg':cfg,'result':r})
 96    best=min(idea_trials, key=lambda z:z['result']['mean'])
 97    sig=mechanism_signature()
 98    report=make_report('phase_synchronization_regression','mlp_tiny',base,best['result'],{
 99        'custom_track':{'name':'phase_synchronization_regression','file':'phase_track.py','domain':'geometry/graph synchronization'},
100        'idea_sweep':idea_trials,'mechanism_signature':sig,
101        'runtime_sec':time.time()-t})
102    Path('bench_report.json').write_text(json.dumps(report,indent=2))
103    print(json.dumps(report,indent=2))
104if __name__=='__main__': main()