1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import make_model, train_model, evaluate, sweep_baseline, make_report
  8from bloch_track import get_dataset
  9
 10K = 2.0 * math.pi
 11SEEDS = tuple(range(8))
 12SWEEP_SEEDS = tuple(range(4))
 13
 14class GaugeModel(nn.Module):
 15    """Same MLP as baseline, with canonical q and exact output gauge."""
 16    def __init__(self, base):
 17        super().__init__(); self.base = base
 18    def forward(self, z):
 19        c, x, q = z[:, 0], z[:, 1], z[:, 2]
 20        m = torch.round(q / K)
 21        q0 = q - m * K
 22        raw = self.base(torch.stack([c, x, q0], dim=1))
 23        phase = torch.exp(-1j * m * K * x)
 24        out = torch.complex(raw[:, 0], raw[:, 1]) * phase
 25        return torch.stack([out.real, out.imag], dim=1)
 26
 27def seed_all(seed):
 28    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 29    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 30
 31def make_system(kind, ds):
 32    base = make_model('mlp_tiny', ds['input_shape'], 2)
 33    return GaugeModel(base) if kind == 'idea' else base
 34
 35def run_one(kind, cfg, seed):
 36    seed_all(seed)
 37    ds = get_dataset(seed, 400, 400)
 38    ds['xtr'] = torch.as_tensor(ds['xtr'], dtype=torch.float32); ds['ytr'] = torch.as_tensor(ds['ytr'], dtype=torch.float32)
 39    ds['xte'] = torch.as_tensor(ds['xte'], dtype=torch.float32); ds['yte'] = torch.as_tensor(ds['yte'], dtype=torch.float32)
 40    ds['input_shape'] = tuple(ds['xtr'].shape[1:]); ds['out_dim'] = 2
 41    net = make_system(kind, ds)
 42    trained, metric, hist = train_model(net, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128, weight_decay=cfg['weight_decay'], log=lambda *_: None)
 43    return float(metric) if metric is not None else float('nan')
 44
 45def model_for_signature(kind, cfg, seed):
 46    seed_all(seed); ds = get_dataset(seed, 400, 400)
 47    ds['xtr'] = torch.as_tensor(ds['xtr'], dtype=torch.float32); ds['ytr'] = torch.as_tensor(ds['ytr'], dtype=torch.float32)
 48    ds['xte'] = torch.as_tensor(ds['xte'], dtype=torch.float32); ds['yte'] = torch.as_tensor(ds['yte'], dtype=torch.float32)
 49    ds['input_shape'] = tuple(ds['xtr'].shape[1:]); ds['out_dim'] = 2
 50    net = make_system(kind, ds)
 51    net, _, _ = train_model(net, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128, weight_decay=cfg['weight_decay'], log=lambda *_: None)
 52    return net, ds
 53
 54def signature(cfg):
 55    rows = {'baseline': [], 'idea': []}
 56    # Query the trained systems at q0 and q0+K, then compare observed shift with exact phase prediction.
 57    rng = np.random.default_rng(991)
 58    c = torch.tensor(rng.uniform(-.9, .9, 96), dtype=torch.float32)
 59    x = torch.tensor(rng.uniform(0, 1, 96), dtype=torch.float32)
 60    q0 = torch.tensor(rng.uniform(-.8*math.pi, .8*math.pi, 96), dtype=torch.float32)
 61    for kind in ('baseline', 'idea'):
 62        net, _ = model_for_signature(kind, cfg, 0)
 63        net.eval()
 64        dev = next(net.parameters()).device
 65        with torch.no_grad():
 66            probe0 = torch.stack([c, x, q0], 1).to(dev)
 67            probe1 = torch.stack([c, x, q0 + K], 1).to(dev)
 68            p0 = net(probe0)
 69            p1 = net(probe1)
 70            cdev, xdev, qdev = c.to(dev), x.to(dev), q0.to(dev)
 71        z0 = torch.complex(p0[:,0], p0[:,1]); z1 = torch.complex(p1[:,0], p1[:,1])
 72        expected = z0 * torch.exp(-1j*K*xdev)
 73        err = torch.sqrt(torch.mean(torch.abs(z1-expected)**2)).item()
 74        scale = torch.sqrt(torch.mean(torch.abs(z0)**2)).item()
 75        rows[kind] = {'shift_m': 1, 'observed_gauge_rmse': err, 'relative_rmse': err/(scale+1e-8)}
 76    # Quantitative stage-1 prediction: exact covariance should be near zero for the wrapped model,
 77    # while the raw model should have a nonzero alias discrepancy.
 78    confirmed = rows['idea']['relative_rmse'] < 0.08 and rows['idea']['relative_rmse'] < rows['baseline']['relative_rmse']
 79    return {'prediction': 'canonicalized model has much smaller trained-model reciprocal-shift gauge discrepancy', 'rows': rows, 'confirmed': bool(confirmed)}
 80
 81def main():
 82    grid = [
 83        {'lr': 0.001, 'epochs': 25, 'weight_decay': 0.0},
 84        {'lr': 0.003, 'epochs': 25, 'weight_decay': 0.0},
 85        {'lr': 0.006, 'epochs': 25, 'weight_decay': 0.0},
 86    ]
 87    base = sweep_baseline(lambda cfg: lambda s: run_one('baseline', cfg, s), grid, seeds=SWEEP_SEEDS)
 88    idea_sweep = []
 89    for cfg in grid:
 90        r = evaluate(lambda s, cfg=cfg: run_one('idea', cfg, s), seeds=SWEEP_SEEDS)
 91        idea_sweep.append({'cfg': cfg, 'mean': r['mean']})
 92    best_idea_cfg = min(idea_sweep, key=lambda x: x['mean'])['cfg']
 93    idea_full = evaluate(lambda s: run_one('idea', best_idea_cfg, s), seeds=SEEDS)
 94    rep = make_report('bloch_gauge_pde', 'mlp_tiny',
 95                      {'best_cfg': base['best_cfg'], 'sweep': base['sweep'], 'full': base['full']},
 96                      dict(idea_full, best_cfg=best_idea_cfg, sweep=idea_sweep),
 97                      extra=signature(best_idea_cfg))
 98    rep['custom_track'] = {'name': 'bloch_gauge_pde', 'file': 'bloch_track.py', 'domain': 'pde'}
 99    Path('bench_report.json').write_text(json.dumps(rep, indent=2))
100    print(json.dumps(rep, indent=2))
101if __name__ == '__main__': main()