1import json, math, random, sys
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import 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))
 11SWEEP_SEEDS = (0, 1, 2, 3)
 12GRID = [
 13    {'lr': 1e-3, 'weight_decay': 0.0},
 14    {'lr': 3e-3, 'weight_decay': 0.0},
 15    {'lr': 1e-2, 'weight_decay': 0.0},
 16]
 17EPOCHS = 18
 18BATCH = 128
 19
 20
 21def seed_all(seed):
 22    random.seed(seed)
 23    np.random.seed(seed)
 24    torch.manual_seed(seed)
 25    if torch.cuda.is_available():
 26        try:
 27            torch.cuda.manual_seed_all(seed)
 28        except Exception:
 29            pass
 30
 31
 32class FourierTumbleRNN(nn.Module):
 33    """Drop-in recurrent replacement for bench rnn_small.
 34
 35    The official dynamics track supplies flattened (8,3) sequences. Each of
 36    32 two-dimensional channels has its own softmax circular distribution;
 37    its first Fourier coefficient induces exact damped-rotation dynamics.
 38    """
 39    def __init__(self, input_dim=3, hidden=64, bins=24):
 40        super().__init__()
 41        assert hidden % 2 == 0
 42        self.hidden = hidden
 43        self.channels = hidden // 2
 44        self.bins = bins
 45        self.inp = nn.Linear(input_dim, hidden)
 46        self.bias_in = nn.Linear(input_dim, hidden, bias=False)
 47        self.head = nn.Linear(hidden, 1)
 48        self.logits = nn.Parameter(torch.zeros(self.channels, bins))
 49        self.log_alpha = nn.Parameter(torch.full((self.channels,), math.log(0.25)))
 50        self.register_buffer('angles', torch.linspace(-math.pi, math.pi, bins + 1)[:-1])
 51
 52    def induced(self):
 53        q = torch.softmax(self.logits, dim=-1)
 54        pi = (q * torch.exp(1j * self.angles)).sum(dim=-1)
 55        alpha = torch.nn.functional.softplus(self.log_alpha) + 1e-5
 56        gamma = alpha * (1.0 - pi.real)
 57        omega = alpha * pi.imag
 58        rho = torch.exp(-gamma)
 59        c, s = torch.cos(omega), torch.sin(omega)
 60        A = torch.zeros(self.channels, 2, 2, device=self.logits.device)
 61        A[:, 0, 0] = rho * c
 62        A[:, 0, 1] = -rho * s
 63        A[:, 1, 0] = rho * s
 64        A[:, 1, 1] = rho * c
 65        return A, pi, gamma, omega, rho
 66
 67    def forward(self, x, return_states=False):
 68        seq = x.view(x.shape[0], -1, 3)
 69        h = torch.zeros(x.shape[0], self.hidden, device=x.device)
 70        states = []
 71        A, _, _, _, _ = self.induced()
 72        for t in range(seq.shape[1]):
 73            u = self.inp(seq[:, t]) + self.bias_in(seq[:, t])
 74            hp = h.view(x.shape[0], self.channels, 2)
 75            hp = torch.einsum('cij,bcj->bci', A, hp)
 76            h = (hp + u.view(x.shape[0], self.channels, 2)).reshape(x.shape[0], self.hidden)
 77            states.append(h)
 78        out = self.head(h)
 79        if return_states:
 80            return out, torch.stack(states, dim=1)
 81        return out
 82
 83
 84def run(kind, cfg, seed, capture=False):
 85    seed_all(seed)
 86    ds = get_dataset('dynamics', seed=seed, n_train=400, n_test=200)
 87    if kind == 'baseline':
 88        model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 89    else:
 90        model = FourierTumbleRNN(input_dim=3, hidden=64, bins=24)
 91    model, metric, hist = train_model(
 92        model, ds, epochs=EPOCHS, lr=float(cfg['lr']), batch=BATCH,
 93        weight_decay=float(cfg.get('weight_decay', 0.0)), log=lambda *a, **k: None)
 94    if model is None or metric is None:
 95        return (float('nan'), None, ds) if capture else float('nan')
 96    return (float(metric), model, ds) if capture else float(metric)
 97
 98
 99def sanity_check():
100    seed_all(2703)
101    m = FourierTumbleRNN()
102    A, pi, gamma, omega, rho = m.induced()
103    p = pi.detach().cpu().numpy()
104    n = torch.linalg.matrix_norm(A, ord=2, dim=(1, 2)).detach().cpu().numpy()
105    return {
106        'predictions': {'abs_Pi_le_1': True, 'Jacobian_norm_le_1': True,
107                        'strict_contraction_when_gamma_positive': True},
108        'max_abs_Pi': float(np.max(np.abs(p))),
109        'max_jacobian_norm': float(np.max(n)),
110        'min_gamma': float(gamma.min().detach().cpu()),
111        'all_abs_Pi_le_1': bool(np.max(np.abs(p)) <= 1.0 + 1e-6),
112        'all_jacobian_norms_le_1': bool(np.max(n) <= 1.0 + 1e-6),
113    }
114
115
116def mechanism_signature(cfg):
117    rows = []
118    for seed in SEEDS:
119        metric, model, ds = run('idea', cfg, seed, capture=True)
120        if model is None:
121            continue
122        model.eval()
123        with torch.no_grad():
124            A, pi, gamma, omega, rho = model.induced()
125            # Empirical homogeneous response of the trained transition: use
126            # random states and measure ||A h||/||h|| over many samples.
127            h = torch.randn(512, model.channels, 2, device=A.device)
128            nxt = torch.einsum('cij,bcj->bci', A, h)
129            ratios = torch.linalg.vector_norm(nxt, dim=-1) / torch.linalg.vector_norm(h, dim=-1).clamp_min(1e-8)
130            observed = float(ratios.max().cpu())
131            predicted = float(torch.linalg.matrix_norm(A, ord=2, dim=(1, 2)).max().cpu())
132            abs_pi = float(torch.abs(pi).max().cpu())
133            gmin = float(gamma.min().cpu())
134        rows.append({'seed': seed, 'metric': metric,
135                     'predicted_norm': predicted, 'observed_max_ratio': observed,
136                     'max_abs_pi': abs_pi, 'min_gamma': gmin})
137    max_obs = max(r['observed_max_ratio'] for r in rows)
138    max_pred = max(r['predicted_norm'] for r in rows)
139    max_pi = max(r['max_abs_pi'] for r in rows)
140    return {
141        'prediction': 'trained Fourier transition has ||rho R||_2=rho<=1 and |Pi_1|<=1',
142        'trained_model_samples': len(rows),
143        'predicted_max_norm': max_pred,
144        'observed_max_homogeneous_ratio': max_obs,
145        'observed_max_abs_pi': max_pi,
146        'within_tolerance': bool(max_obs <= 1.0001 and max_pi <= 1.0001),
147        'confirmed': bool(max_obs <= 1.0001 and max_pi <= 1.0001),
148        'per_seed': rows,
149    }
150
151
152def main():
153    math_sanity = sanity_check()
154    baseline = sweep_baseline(
155        lambda cfg: (lambda seed: run('baseline', cfg, seed)),
156        GRID, seeds=SWEEP_SEEDS)
157    idea_trials = []
158    for cfg in GRID:
159        result = evaluate(lambda seed, cfg=cfg: run('idea', cfg, seed), seeds=SEEDS)
160        idea_trials.append({'cfg': cfg, 'result': result})
161    best_idea = min(idea_trials, key=lambda z: z['result']['mean'])
162    idea_cfg, idea_result = best_idea['cfg'], best_idea['result']
163    signature = mechanism_signature(idea_cfg)
164    report = make_report(
165        'dynamics', 'rnn_small', baseline, idea_result,
166        extra={
167            'track_match': 'stability/control and recurrent memory -> official dynamics track',
168            'math_sanity': math_sanity,
169            'idea_sweep': idea_trials,
170            'selected_idea_cfg': idea_cfg,
171            'mechanism_signature': signature,
172        })
173    Path('bench_report.json').write_text(json.dumps(report, indent=2))
174    print(json.dumps(report, indent=2))
175
176
177if __name__ == '__main__':
178    main()