1import sys, json, random
  2import numpy as np
  3import torch
  4from torch import nn
  5from scipy.stats import wasserstein_distance
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, sweep_baseline, evaluate, make_report
  8
  9SEED = 461
 10LR_GRID = [1.5e-3, 3e-3, 6e-3]
 11EPOCHS = 12
 12BATCH = 128
 13K = 8
 14
 15def seed_all(seed):
 16    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 17    if torch.cuda.is_available():
 18        try: torch.cuda.manual_seed_all(seed)
 19        except Exception: pass
 20
 21class PoolNet(nn.Module):
 22    """Matched end-to-end systems; only the set pooling map differs."""
 23    def __init__(self, kind, k=K):
 24        super().__init__()
 25        self.kind = kind
 26        self.k = k
 27        # A shared-size output head makes the comparison parameter-matched.
 28        self.head = nn.Sequential(nn.Linear(k, 32), nn.Tanh(), nn.Linear(32, 1))
 29        if kind == 'idea':
 30            # scalar probes with |slope|<=1; subtracting phi(0) enforces phi(theta)=0
 31            self.raw_slope = nn.Parameter(torch.randn(k) * 0.15)
 32            self.bias = nn.Parameter(torch.zeros(k))
 33        elif kind != 'baseline':
 34            raise ValueError(kind)
 35
 36    def probe_values(self, x):
 37        slope = torch.tanh(self.raw_slope)
 38        return torch.tanh(x.unsqueeze(-1) * slope + self.bias) - torch.tanh(self.bias)
 39
 40    def pooled(self, x):
 41        if self.kind == 'baseline':
 42            z = x.mean(dim=1, keepdim=True).expand(-1, self.k)
 43        else:
 44            z = self.probe_values(x).mean(dim=1)
 45        return z
 46
 47    def forward(self, x):
 48        return self.head(self.pooled(x))
 49
 50def train_one(kind, seed, lr, retain=False):
 51    seed_all(seed)
 52    d = get_dataset('sequence', seed, n_train=400, n_test=400)
 53    model = PoolNet(kind)
 54    net, metric, hist = train_model(model, d, epochs=EPOCHS, lr=lr,
 55                                    batch=BATCH, weight_decay=0.0)
 56    if net is None: raise RuntimeError('benchmark training failed')
 57    return float(metric), net, d
 58
 59def fn(kind, lr):
 60    return lambda seed: train_one(kind, seed, lr)[0]
 61
 62def mechanism_signature(net, d, n_pairs=100):
 63    net.eval()
 64    dev = next(net.parameters()).device
 65    x = d['xte'].to(dev)
 66    with torch.no_grad():
 67        vals = net.probe_values(x).mean(dim=1).detach().cpu().numpy()
 68    # For equal-weight 1D empirical measures, W1 is an independent task-free
 69    # geometric reference; D_hat is computed from the trained probes.
 70    rng = np.random.default_rng(9001)
 71    dh, w1 = [], []
 72    xx = d['xte'].cpu().numpy()
 73    for _ in range(n_pairs):
 74        a, b = rng.integers(0, len(xx), 2)
 75        # Re-evaluate individual probe pooled values, not model predictions.
 76        va, vb = vals[a], vals[b]
 77        dh.append(float(np.max(np.abs(va - vb))))
 78        w1.append(float(wasserstein_distance(xx[a], xx[b])))
 79    dh, w1 = np.asarray(dh), np.asarray(w1)
 80    corr = float(np.corrcoef(dh, w1)[0,1]) if np.std(dh)>1e-12 and np.std(w1)>1e-12 else 0.0
 81    ratio = float(np.mean(dh / np.maximum(w1, 1e-8)))
 82    violations = float(np.mean(dh > w1 + 1e-5))
 83    # Prediction: a valid 1-Lipschitz dual estimate should not exceed W1.
 84    return {'quantity': 'trained probe D_hat versus independent 1D Wasserstein-1',
 85            'predicted': 'D_hat <= W1 for every pair; positive geometric association',
 86            'observed_mean_D_hat': float(np.mean(dh)),
 87            'observed_mean_W1': float(np.mean(w1)),
 88            'mean_ratio_D_hat_over_W1': ratio,
 89            'pearson_correlation': corr,
 90            'violation_fraction': violations,
 91            'confirmed': bool(violations <= 0.05 and corr > 0.2)}
 92
 93def main():
 94    # Baseline sweep and idea sweep use identical learning-rate union and seeds.
 95    grid = [{'lr': lr, 'epochs': EPOCHS, 'batch': BATCH} for lr in LR_GRID]
 96    base = sweep_baseline(lambda cfg: fn('baseline', cfg['lr']), grid)
 97    idea_sweep = []
 98    for cfg in grid:
 99        r = evaluate(fn('idea', cfg['lr']))
100        idea_sweep.append({'cfg': cfg, 'mean': r['mean'], 'std': r['std'], 'per_seed': r['per_seed']})
101    best_i = min(idea_sweep, key=lambda q: q['mean'])['cfg']
102    idea_full = evaluate(fn('idea', best_i['lr']))
103    # Re-train one paired model for the behavioral signature.
104    _, sig_net, sig_d = train_one('idea', 0, best_i['lr'])
105    sig = mechanism_signature(sig_net, sig_d)
106    report = make_report('sequence', 'transformer_tiny_pooling_mvp', base, idea_full,
107                         {'mechanism_signature': sig,
108                          'idea_sweep': idea_sweep,
109                          'selection_seeds': [0,1,2,3],
110                          'epochs': EPOCHS, 'batch': BATCH,
111                          'parameter_counts': {'baseline': sum(p.numel() for p in PoolNet('baseline').parameters()),
112                                               'idea': sum(p.numel() for p in PoolNet('idea').parameters())},
113                          'structural_match': 'sequence-level window forecasting with token-set pooling'})
114    report['idea']['selected_cfg'] = best_i
115    with open('bench_report.json','w') as f: json.dump(report, f, indent=2)
116    print(json.dumps(report, indent=2))
117
118if __name__ == '__main__': main()