1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  9
 10SEEDS = tuple(range(8))
 11# Shared union: every lr and radius is evaluated by both methods.
 12GRID = [
 13    {'lr': 0.0015, 'radius': 0.05},
 14    {'lr': 0.0030, 'radius': 0.05},
 15    {'lr': 0.0060, 'radius': 0.05},
 16]
 17EPOCHS = 12
 18BATCH = 128
 19WEIGHT_DECAY = 0.0
 20
 21
 22def seed_all(seed):
 23    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 24    if torch.cuda.is_available():
 25        torch.cuda.manual_seed_all(seed)
 26
 27
 28def project_exact(g, R):
 29    n = torch.linalg.vector_norm(g)
 30    return g * torch.minimum(torch.ones_like(n), torch.as_tensor(R, device=g.device) / (n + 1e-12))
 31
 32
 33def project_smooth(g, R):
 34    # Differentiable radial approximation in the proposal.
 35    n = torch.linalg.vector_norm(g)
 36    return g / torch.sqrt(1.0 + (n / R) ** 2)
 37
 38
 39def train_limited(seed, cfg, smooth, collect=False):
 40    seed_all(seed)
 41    ds = get_dataset('dynamics', seed, n_train=400, n_test=160)
 42    model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 43    # Keep the intervention loop otherwise identical to bench.train_model's Adam path.
 44    try:
 45        device = 'cuda' if torch.cuda.is_available() else 'cpu'
 46        model = model.to(device)
 47        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 48        opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=WEIGHT_DECAY)
 49        lossf = nn.MSELoss()
 50        norms, clipped = [], []
 51        for _ in range(EPOCHS):
 52            model.train(); perm = torch.randperm(len(x), device=device)
 53            for i in range(0, len(x), BATCH):
 54                idx = perm[i:i+BATCH]
 55                loss = lossf(model(x[idx]), y[idx]); opt.zero_grad(); loss.backward()
 56                grad = torch.cat([p.grad.reshape(-1) for p in model.parameters() if p.grad is not None])
 57                gn = float(torch.linalg.vector_norm(grad).detach().cpu()); norms.append(gn)
 58                newg = project_smooth(grad, cfg['radius']) if smooth else project_exact(grad, cfg['radius'])
 59                scale = torch.linalg.vector_norm(newg) / (torch.linalg.vector_norm(grad) + 1e-30)
 60                for p in model.parameters():
 61                    if p.grad is not None: p.grad.mul_(scale)
 62                clipped.append(float(torch.linalg.vector_norm(newg).detach().cpu()))
 63                opt.step()
 64        model.eval()
 65        with torch.no_grad():
 66            pred = model(ds['xte'].to(device)); metric = float(((pred-ds['yte'].to(device))**2).mean().cpu())
 67        if collect:
 68            return metric, {'grad_norm_mean': float(np.mean(norms)), 'grad_norm_p95': float(np.percentile(norms,95)),
 69                            'applied_norm_mean': float(np.mean(clipped)), 'applied_norm_p95': float(np.percentile(clipped,95)),
 70                            'fraction_over_radius': float(np.mean(np.asarray(norms)>cfg['radius']))}
 71        return metric
 72    except RuntimeError:
 73        # Explicit CPU fallback, rebuilding the model and reproducing the same seed.
 74        if device == 'cuda':
 75            seed_all(seed)
 76            ds = get_dataset('dynamics', seed, n_train=400, n_test=160)
 77            model = make_model('rnn_small', ds['input_shape'], ds['out_dim']).cpu()
 78            old = torch.cuda.is_available
 79            torch.cuda.is_available = lambda: False
 80            try: return train_limited(seed, cfg, smooth, collect)
 81            finally: torch.cuda.is_available = old
 82        raise
 83
 84
 85def make_fn(smooth, cfg):
 86    return lambda seed: train_limited(seed, cfg, smooth)
 87
 88
 89def main():
 90    # Core math check first: exact boundedness and smooth local gain are measured numerically.
 91    torch.manual_seed(11)
 92    checks = []
 93    for n in [0.01, 0.05, 0.2, 1.0]:
 94        v = torch.randn(257, dtype=torch.float64); v *= n / torch.linalg.vector_norm(v)
 95        ex = project_exact(v, 0.05); sm = project_smooth(v, 0.05)
 96        checks.append({'input_norm': n, 'exact_norm': float(torch.linalg.vector_norm(ex)),
 97                       'smooth_norm': float(torch.linalg.vector_norm(sm)),
 98                       'exact_bound_ok': bool(torch.linalg.vector_norm(ex) <= 0.05 + 1e-12)})
 99    # Baseline sweep over the full shared union, then final full-seed evaluation.
100    base = sweep_baseline(lambda cfg: make_fn(False, cfg), GRID, seeds=(0,1,2,3))
101    idea_candidates = []
102    for cfg in GRID:
103        r = evaluate(make_fn(True, cfg), seeds=SEEDS)
104        idea_candidates.append((r, cfg))
105    idea, idea_cfg = min(idea_candidates, key=lambda z: z[0]['mean'])
106    # Train one paired model at selected setting for behavior signature.
107    bmet, bsig = train_limited(0, idea_cfg, False, True)
108    imet, isig = train_limited(0, idea_cfg, True, True)
109    signature = {
110        'prediction': 'smooth clipping can apply a different local gain and create larger gradients than exact projection near the boundary',
111        'observed': {'cfg': idea_cfg, 'baseline_seed0_metric': bmet, 'idea_seed0_metric': imet,
112                     'baseline': bsig, 'idea': isig,
113                     'smooth_applied_norm_exceeds_radius': bool(isig['applied_norm_p95'] > idea_cfg['radius'])},
114        'confirmed': bool(isig['applied_norm_p95'] > idea_cfg['radius'] and bsig['applied_norm_p95'] <= idea_cfg['radius'] + 1e-7)
115    }
116    report = make_report('dynamics', 'rnn_small', base, idea, extra=signature)
117    report['idea_grid'] = [{'cfg': c, 'full': r} for r,c in idea_candidates]
118    report['math_sanity'] = checks
119    report['protocol'] = {'epochs': EPOCHS, 'batch': BATCH, 'paired_seeds': list(SEEDS), 'track_reason': 'dynamics matches optimizer stability and bifurcation structure'}
120    Path('bench_report.json').write_text(json.dumps(report, indent=2))
121    print(json.dumps(report, indent=2))
122
123if __name__ == '__main__': main()