1import os, sys, json, math, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, sweep_baseline, make_report
  8
  9SEEDS = tuple(range(8))
 10SWEEP_SEEDS = (0, 1, 2, 3)
 11NMOD, H, DELAY = 8, 16, 2
 12SIGMA = 0.42
 13
 14
 15def seed_all(seed):
 16    random.seed(seed)
 17    np.random.seed(seed)
 18    torch.manual_seed(seed)
 19    if torch.cuda.is_available():
 20        torch.cuda.manual_seed_all(seed)
 21
 22
 23def symmetric_ring(n=NMOD):
 24    w = np.zeros((n, n), dtype=np.float32)
 25    for i in range(n):
 26        w[i, (i - 1) % n] = 0.5
 27        w[i, (i + 1) % n] = 0.5
 28    return w
 29
 30
 31def directed_heterogeneous(n=NMOD, seed=17):
 32    rng = np.random.default_rng(seed)
 33    return rng.dirichlet(0.22 * np.ones(n), size=n).astype(np.float32)
 34
 35
 36def transverse_radius(w):
 37    vals = np.linalg.eigvals(w)
 38    j = int(np.argmin(np.abs(vals - 1.0)))
 39    return float(np.max(np.abs(np.delete(vals, j))))
 40
 41
 42def companion_radius(a, alpha, delay):
 43    c = np.zeros((delay + 1, delay + 1), dtype=complex if np.iscomplexobj(alpha) else float)
 44    c[0, 0] = a
 45    c[0, -1] = alpha
 46    c[1:, :-1] = np.eye(delay)
 47    return float(np.max(np.abs(np.linalg.eigvals(c))))
 48
 49
 50def sanity_check():
 51    # Constant-Jacobian modal recurrence: empirical exponent should equal log spectral radius.
 52    a, alpha, d = 0.72, 0.15, 2
 53    c = np.zeros((d + 1, d + 1))
 54    c[0, 0], c[0, -1] = a, alpha
 55    c[1:, :-1] = np.eye(d)
 56    rng = np.random.default_rng(0)
 57    v = rng.normal(size=d + 1); v /= np.linalg.norm(v)
 58    logs = []
 59    for _ in range(2500):
 60        v = c @ v
 61        z = np.linalg.norm(v)
 62        logs.append(math.log(max(z, 1e-300)))
 63        v /= max(z, 1e-300)
 64    empirical = float(np.mean(logs[300:]))
 65    predicted = math.log(companion_radius(a, alpha, d))
 66    return {'predicted_log_rho': predicted, 'measured_exponent': empirical,
 67            'absolute_error': abs(predicted - empirical),
 68            'passed': abs(predicted - empirical) < 1e-5}
 69
 70
 71class CoupledRNN(nn.Module):
 72    def __init__(self, coupling, sigma=SIGMA, delay=DELAY, hidden=H):
 73        super().__init__()
 74        self.hidden = hidden
 75        self.n = coupling.shape[0]
 76        self.sigma = sigma
 77        self.delay = delay
 78        self.inp = nn.Linear(3, hidden)
 79        self.cell = nn.GRUCell(hidden, hidden)
 80        self.head = nn.Linear(hidden, 1)
 81        self.register_buffer('coupling', torch.tensor(coupling))
 82        self.last_disagreement = None
 83
 84    def forward(self, x):
 85        b = x.shape[0]
 86        seq = x.view(b, -1, 3)
 87        hs = torch.zeros(b, self.n, self.hidden, device=x.device, dtype=x.dtype)
 88        history = [hs.clone() for _ in range(self.delay + 1)]
 89        disagreements = []
 90        for t in range(seq.shape[1]):
 91            z = self.inp(seq[:, t]).unsqueeze(1).expand(-1, self.n, -1)
 92            local = self.cell(z.reshape(b * self.n, -1), hs.reshape(b * self.n, -1)).view(b, self.n, self.hidden)
 93            delayed = history[0]
 94            # Row-stochastic coupling preserves the synchronized mode; only topology differs.
 95            hs = local + self.sigma * torch.einsum('ij,bjh->bih', self.coupling, delayed)
 96            history = history[1:] + [hs.clone()]
 97            disagreements.append((hs - hs.mean(dim=1, keepdim=True)).pow(2).mean())
 98        self.last_disagreement = float(torch.stack(disagreements).detach().cpu()[-1])
 99        return self.head(hs.mean(dim=1))
100
101
102def train_one(track, seed, cfg, topology):
103    seed_all(seed)
104    ds = get_dataset(track, seed, n_train=400, n_test=400)
105    model = CoupledRNN(topology, sigma=cfg['sigma'], delay=cfg['delay'])
106    _, metric, history = train_model(model, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128)
107    return float(metric), model
108
109
110def evaluate_config(cfg, topology, seeds=SEEDS, keep_models=False):
111    vals, models = [], []
112    for s in seeds:
113        metric, model = train_one('dynamics', s, cfg, topology)
114        vals.append(metric)
115        if keep_models: models.append(model)
116    out = {'mean': float(np.mean(vals)), 'std': float(np.std(vals)),
117           'per_seed': [float(v) for v in vals], 'n': len(vals)}
118    if keep_models: out['models'] = models
119    return out
120
121
122def main():
123    sanity = sanity_check()
124    base_w = symmetric_ring()
125    idea_w = directed_heterogeneous()
126    # Shared union: every idea lr is also evaluated by baseline; topology is the method knob.
127    grid = [
128        {'lr': 0.001, 'epochs': 10, 'sigma': 0.42, 'delay': 2},
129        {'lr': 0.003, 'epochs': 10, 'sigma': 0.42, 'delay': 2},
130        {'lr': 0.006, 'epochs': 10, 'sigma': 0.42, 'delay': 2},
131    ]
132    def make_base(cfg):
133        return lambda seed: train_one('dynamics', seed, cfg, base_w)[0]
134    base = sweep_baseline(make_base, grid, seeds=SWEEP_SEEDS)
135    idea_candidates = [base['best_cfg'], grid[0], grid[2]]
136    idea_candidates = list({tuple(sorted(c.items())): c for c in idea_candidates}.values())
137    idea_runs = []
138    for cfg in idea_candidates:
139        r = evaluate_config(cfg, idea_w, SEEDS)
140        idea_runs.append({'cfg': cfg, 'result': r})
141    best = min(idea_runs, key=lambda z: z['result']['mean'])
142    idea = best['result']
143    # Re-test the trained systems' observed disagreement, not an analytical toy graph.
144    base_full = evaluate_config(base['best_cfg'], base_w, SEEDS, keep_models=True)
145    idea_full = evaluate_config(best['cfg'], idea_w, SEEDS, keep_models=True)
146    base_dis = float(np.mean([m.last_disagreement for m in base_full['models']]))
147    idea_dis = float(np.mean([m.last_disagreement for m in idea_full['models']]))
148    sig = {
149        'predicted_transverse_radius_baseline': transverse_radius(base_w),
150        'predicted_transverse_radius_idea': transverse_radius(idea_w),
151        'observed_final_disagreement_baseline': base_dis,
152        'observed_final_disagreement_idea': idea_dis,
153        'prediction': 'smaller transverse spectral radius should yield smaller trained hidden disagreement',
154        'confirmed': bool(idea_dis < base_dis)
155    }
156    # Remove model objects before JSON serialization and use the canonical report schema.
157    base_clean = dict(base); base_clean['full'] = base_full.copy(); base_clean['full'].pop('models', None)
158    report = make_report('dynamics', 'rnn_small', base_clean, idea, {
159        'math_sanity': sanity, 'trained_model_signature': sig,
160        'topology': {'baseline': 'symmetric reciprocal ring', 'idea': 'directed heterogeneous row-stochastic'}
161    })
162    report['idea_sweep'] = [{'cfg': z['cfg'], 'mean': z['result']['mean']} for z in idea_runs]
163    report['math_sanity'] = sanity
164    with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2)
165    print(json.dumps(report, indent=2))
166
167
168if __name__ == '__main__':
169    main()