import os, sys, json, math, random import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, train_model, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) NMOD, H, DELAY = 8, 16, 2 SIGMA = 0.42 def seed_all(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def symmetric_ring(n=NMOD): w = np.zeros((n, n), dtype=np.float32) for i in range(n): w[i, (i - 1) % n] = 0.5 w[i, (i + 1) % n] = 0.5 return w def directed_heterogeneous(n=NMOD, seed=17): rng = np.random.default_rng(seed) return rng.dirichlet(0.22 * np.ones(n), size=n).astype(np.float32) def transverse_radius(w): vals = np.linalg.eigvals(w) j = int(np.argmin(np.abs(vals - 1.0))) return float(np.max(np.abs(np.delete(vals, j)))) def companion_radius(a, alpha, delay): c = np.zeros((delay + 1, delay + 1), dtype=complex if np.iscomplexobj(alpha) else float) c[0, 0] = a c[0, -1] = alpha c[1:, :-1] = np.eye(delay) return float(np.max(np.abs(np.linalg.eigvals(c)))) def sanity_check(): # Constant-Jacobian modal recurrence: empirical exponent should equal log spectral radius. a, alpha, d = 0.72, 0.15, 2 c = np.zeros((d + 1, d + 1)) c[0, 0], c[0, -1] = a, alpha c[1:, :-1] = np.eye(d) rng = np.random.default_rng(0) v = rng.normal(size=d + 1); v /= np.linalg.norm(v) logs = [] for _ in range(2500): v = c @ v z = np.linalg.norm(v) logs.append(math.log(max(z, 1e-300))) v /= max(z, 1e-300) empirical = float(np.mean(logs[300:])) predicted = math.log(companion_radius(a, alpha, d)) return {'predicted_log_rho': predicted, 'measured_exponent': empirical, 'absolute_error': abs(predicted - empirical), 'passed': abs(predicted - empirical) < 1e-5} class CoupledRNN(nn.Module): def __init__(self, coupling, sigma=SIGMA, delay=DELAY, hidden=H): super().__init__() self.hidden = hidden self.n = coupling.shape[0] self.sigma = sigma self.delay = delay self.inp = nn.Linear(3, hidden) self.cell = nn.GRUCell(hidden, hidden) self.head = nn.Linear(hidden, 1) self.register_buffer('coupling', torch.tensor(coupling)) self.last_disagreement = None def forward(self, x): b = x.shape[0] seq = x.view(b, -1, 3) hs = torch.zeros(b, self.n, self.hidden, device=x.device, dtype=x.dtype) history = [hs.clone() for _ in range(self.delay + 1)] disagreements = [] for t in range(seq.shape[1]): z = self.inp(seq[:, t]).unsqueeze(1).expand(-1, self.n, -1) local = self.cell(z.reshape(b * self.n, -1), hs.reshape(b * self.n, -1)).view(b, self.n, self.hidden) delayed = history[0] # Row-stochastic coupling preserves the synchronized mode; only topology differs. hs = local + self.sigma * torch.einsum('ij,bjh->bih', self.coupling, delayed) history = history[1:] + [hs.clone()] disagreements.append((hs - hs.mean(dim=1, keepdim=True)).pow(2).mean()) self.last_disagreement = float(torch.stack(disagreements).detach().cpu()[-1]) return self.head(hs.mean(dim=1)) def train_one(track, seed, cfg, topology): seed_all(seed) ds = get_dataset(track, seed, n_train=400, n_test=400) model = CoupledRNN(topology, sigma=cfg['sigma'], delay=cfg['delay']) _, metric, history = train_model(model, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128) return float(metric), model def evaluate_config(cfg, topology, seeds=SEEDS, keep_models=False): vals, models = [], [] for s in seeds: metric, model = train_one('dynamics', s, cfg, topology) vals.append(metric) if keep_models: models.append(model) out = {'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'per_seed': [float(v) for v in vals], 'n': len(vals)} if keep_models: out['models'] = models return out def main(): sanity = sanity_check() base_w = symmetric_ring() idea_w = directed_heterogeneous() # Shared union: every idea lr is also evaluated by baseline; topology is the method knob. grid = [ {'lr': 0.001, 'epochs': 10, 'sigma': 0.42, 'delay': 2}, {'lr': 0.003, 'epochs': 10, 'sigma': 0.42, 'delay': 2}, {'lr': 0.006, 'epochs': 10, 'sigma': 0.42, 'delay': 2}, ] def make_base(cfg): return lambda seed: train_one('dynamics', seed, cfg, base_w)[0] base = sweep_baseline(make_base, grid, seeds=SWEEP_SEEDS) idea_candidates = [base['best_cfg'], grid[0], grid[2]] idea_candidates = list({tuple(sorted(c.items())): c for c in idea_candidates}.values()) idea_runs = [] for cfg in idea_candidates: r = evaluate_config(cfg, idea_w, SEEDS) idea_runs.append({'cfg': cfg, 'result': r}) best = min(idea_runs, key=lambda z: z['result']['mean']) idea = best['result'] # Re-test the trained systems' observed disagreement, not an analytical toy graph. base_full = evaluate_config(base['best_cfg'], base_w, SEEDS, keep_models=True) idea_full = evaluate_config(best['cfg'], idea_w, SEEDS, keep_models=True) base_dis = float(np.mean([m.last_disagreement for m in base_full['models']])) idea_dis = float(np.mean([m.last_disagreement for m in idea_full['models']])) sig = { 'predicted_transverse_radius_baseline': transverse_radius(base_w), 'predicted_transverse_radius_idea': transverse_radius(idea_w), 'observed_final_disagreement_baseline': base_dis, 'observed_final_disagreement_idea': idea_dis, 'prediction': 'smaller transverse spectral radius should yield smaller trained hidden disagreement', 'confirmed': bool(idea_dis < base_dis) } # Remove model objects before JSON serialization and use the canonical report schema. base_clean = dict(base); base_clean['full'] = base_full.copy(); base_clean['full'].pop('models', None) report = make_report('dynamics', 'rnn_small', base_clean, idea, { 'math_sanity': sanity, 'trained_model_signature': sig, 'topology': {'baseline': 'symmetric reciprocal ring', 'idea': 'directed heterogeneous row-stochastic'} }) report['idea_sweep'] = [{'cfg': z['cfg'], 'mean': z['result']['mean']} for z in idea_runs] report['math_sanity'] = sanity with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()