import json, math, random, sys from pathlib import Path import numpy as np import torch from torch import nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) GRID = [ {'lr': 1e-3, 'weight_decay': 0.0}, {'lr': 3e-3, 'weight_decay': 0.0}, {'lr': 1e-2, 'weight_decay': 0.0}, ] EPOCHS = 18 BATCH = 128 def seed_all(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): try: torch.cuda.manual_seed_all(seed) except Exception: pass class FourierTumbleRNN(nn.Module): """Drop-in recurrent replacement for bench rnn_small. The official dynamics track supplies flattened (8,3) sequences. Each of 32 two-dimensional channels has its own softmax circular distribution; its first Fourier coefficient induces exact damped-rotation dynamics. """ def __init__(self, input_dim=3, hidden=64, bins=24): super().__init__() assert hidden % 2 == 0 self.hidden = hidden self.channels = hidden // 2 self.bins = bins self.inp = nn.Linear(input_dim, hidden) self.bias_in = nn.Linear(input_dim, hidden, bias=False) self.head = nn.Linear(hidden, 1) self.logits = nn.Parameter(torch.zeros(self.channels, bins)) self.log_alpha = nn.Parameter(torch.full((self.channels,), math.log(0.25))) self.register_buffer('angles', torch.linspace(-math.pi, math.pi, bins + 1)[:-1]) def induced(self): q = torch.softmax(self.logits, dim=-1) pi = (q * torch.exp(1j * self.angles)).sum(dim=-1) alpha = torch.nn.functional.softplus(self.log_alpha) + 1e-5 gamma = alpha * (1.0 - pi.real) omega = alpha * pi.imag rho = torch.exp(-gamma) c, s = torch.cos(omega), torch.sin(omega) A = torch.zeros(self.channels, 2, 2, device=self.logits.device) A[:, 0, 0] = rho * c A[:, 0, 1] = -rho * s A[:, 1, 0] = rho * s A[:, 1, 1] = rho * c return A, pi, gamma, omega, rho def forward(self, x, return_states=False): seq = x.view(x.shape[0], -1, 3) h = torch.zeros(x.shape[0], self.hidden, device=x.device) states = [] A, _, _, _, _ = self.induced() for t in range(seq.shape[1]): u = self.inp(seq[:, t]) + self.bias_in(seq[:, t]) hp = h.view(x.shape[0], self.channels, 2) hp = torch.einsum('cij,bcj->bci', A, hp) h = (hp + u.view(x.shape[0], self.channels, 2)).reshape(x.shape[0], self.hidden) states.append(h) out = self.head(h) if return_states: return out, torch.stack(states, dim=1) return out def run(kind, cfg, seed, capture=False): seed_all(seed) ds = get_dataset('dynamics', seed=seed, n_train=400, n_test=200) if kind == 'baseline': model = make_model('rnn_small', ds['input_shape'], ds['out_dim']) else: model = FourierTumbleRNN(input_dim=3, hidden=64, bins=24) model, metric, hist = train_model( model, ds, epochs=EPOCHS, lr=float(cfg['lr']), batch=BATCH, weight_decay=float(cfg.get('weight_decay', 0.0)), log=lambda *a, **k: None) if model is None or metric is None: return (float('nan'), None, ds) if capture else float('nan') return (float(metric), model, ds) if capture else float(metric) def sanity_check(): seed_all(2703) m = FourierTumbleRNN() A, pi, gamma, omega, rho = m.induced() p = pi.detach().cpu().numpy() n = torch.linalg.matrix_norm(A, ord=2, dim=(1, 2)).detach().cpu().numpy() return { 'predictions': {'abs_Pi_le_1': True, 'Jacobian_norm_le_1': True, 'strict_contraction_when_gamma_positive': True}, 'max_abs_Pi': float(np.max(np.abs(p))), 'max_jacobian_norm': float(np.max(n)), 'min_gamma': float(gamma.min().detach().cpu()), 'all_abs_Pi_le_1': bool(np.max(np.abs(p)) <= 1.0 + 1e-6), 'all_jacobian_norms_le_1': bool(np.max(n) <= 1.0 + 1e-6), } def mechanism_signature(cfg): rows = [] for seed in SEEDS: metric, model, ds = run('idea', cfg, seed, capture=True) if model is None: continue model.eval() with torch.no_grad(): A, pi, gamma, omega, rho = model.induced() # Empirical homogeneous response of the trained transition: use # random states and measure ||A h||/||h|| over many samples. h = torch.randn(512, model.channels, 2, device=A.device) nxt = torch.einsum('cij,bcj->bci', A, h) ratios = torch.linalg.vector_norm(nxt, dim=-1) / torch.linalg.vector_norm(h, dim=-1).clamp_min(1e-8) observed = float(ratios.max().cpu()) predicted = float(torch.linalg.matrix_norm(A, ord=2, dim=(1, 2)).max().cpu()) abs_pi = float(torch.abs(pi).max().cpu()) gmin = float(gamma.min().cpu()) rows.append({'seed': seed, 'metric': metric, 'predicted_norm': predicted, 'observed_max_ratio': observed, 'max_abs_pi': abs_pi, 'min_gamma': gmin}) max_obs = max(r['observed_max_ratio'] for r in rows) max_pred = max(r['predicted_norm'] for r in rows) max_pi = max(r['max_abs_pi'] for r in rows) return { 'prediction': 'trained Fourier transition has ||rho R||_2=rho<=1 and |Pi_1|<=1', 'trained_model_samples': len(rows), 'predicted_max_norm': max_pred, 'observed_max_homogeneous_ratio': max_obs, 'observed_max_abs_pi': max_pi, 'within_tolerance': bool(max_obs <= 1.0001 and max_pi <= 1.0001), 'confirmed': bool(max_obs <= 1.0001 and max_pi <= 1.0001), 'per_seed': rows, } def main(): math_sanity = sanity_check() baseline = sweep_baseline( lambda cfg: (lambda seed: run('baseline', cfg, seed)), GRID, seeds=SWEEP_SEEDS) idea_trials = [] for cfg in GRID: result = evaluate(lambda seed, cfg=cfg: run('idea', cfg, seed), seeds=SEEDS) idea_trials.append({'cfg': cfg, 'result': result}) best_idea = min(idea_trials, key=lambda z: z['result']['mean']) idea_cfg, idea_result = best_idea['cfg'], best_idea['result'] signature = mechanism_signature(idea_cfg) report = make_report( 'dynamics', 'rnn_small', baseline, idea_result, extra={ 'track_match': 'stability/control and recurrent memory -> official dynamics track', 'math_sanity': math_sanity, 'idea_sweep': idea_trials, 'selected_idea_cfg': idea_cfg, 'mechanism_signature': signature, }) Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()