Orbital-Stable Dancing RNN / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 2841
  7np.random.seed(SEED)
  8random.seed(SEED)
  9torch.manual_seed(SEED)
 10
 11# Two-node bidirectional instance of the paper's phase Jacobian.
 12def jacobian(phi, kappa=1.0):
 13    c = math.cos(phi)
 14    return kappa * np.array([[-c, c], [c, -c]], dtype=float)
 15
 16
 17def math_checks():
 18    # Prediction 1: global phase is neutral: J @ 1 = 0.
 19    shifts = np.linspace(0, math.pi, 25)
 20    neutral_err = max(float(np.linalg.norm(jacobian(p) @ np.ones(2))) for p in shifts)
 21
 22    # Prediction 2: transverse boundary is phi=pi/2 and lambda=-2*kappa*cos(phi).
 23    grid = np.linspace(0, math.pi, 10001)
 24    lambdas = -2.0 * np.cos(grid)
 25    boundary = float(grid[np.argmin(np.abs(lambdas))])
 26
 27    # Prediction 3: finite-Euler perturbation slope agrees with the transverse eigenvalue.
 28    decay_rows = []
 29    dt, steps, eps = 0.002, 1500, 1e-5
 30    for phi in [0.0, math.pi/6, math.pi/3, 0.70*math.pi/2,
 31                math.pi/2, 0.60*math.pi]:
 32        kappa = 0.8
 33        J = jacobian(phi, kappa)
 34        d = np.array([eps, -eps], dtype=float)
 35        norms = []
 36        for _ in range(steps):
 37            norms.append(np.linalg.norm(d))
 38            d = d + dt * J.dot(d)
 39        fit_n = np.arange(100, 1100)
 40        slope = float(np.polyfit(fit_n*dt,
 41                                 np.log(np.maximum(np.asarray(norms)[fit_n], 1e-30)), 1)[0])
 42        predicted = -2.0*kappa*math.cos(phi)
 43        decay_rows.append({
 44            "phi": float(phi),
 45            "predicted_lambda": float(predicted),
 46            "measured_slope": slope,
 47            "stable_predicted": bool(predicted < 0),
 48            "stable_measured": bool(slope < -1e-5),
 49            "abs_error": abs(slope-predicted),
 50        })
 51
 52    # Sweep boundary directly and record signs.
 53    boundary_sweep = []
 54    for phi in np.linspace(0, math.pi, 9):
 55        eig = np.linalg.eigvals(jacobian(phi, 1.0))
 56        transverse = float(sorted(eig)[0])
 57        boundary_sweep.append({"phi_over_pi": float(phi/math.pi),
 58                              "transverse_eigenvalue": transverse})
 59    return {
 60        "neutral_mode_max_norm": neutral_err,
 61        "predicted_boundary_phi": math.pi/2,
 62        "measured_boundary_phi": boundary,
 63        "boundary_error": abs(boundary-math.pi/2),
 64        "decay_scaling": decay_rows,
 65        "boundary_sweep": boundary_sweep,
 66    }
 67
 68
 69class PhaseRNN(nn.Module):
 70    def __init__(self, hidden=16):
 71        super().__init__()
 72        self.hidden = hidden
 73        self.inp = nn.Linear(1, hidden, bias=False)
 74        self.omega = nn.Parameter(torch.randn(hidden)*0.15)
 75        # Fixed, positive directed ring with a skip edge; stable synchronized offsets.
 76        A = torch.zeros(hidden, hidden)
 77        for i in range(hidden):
 78            A[i, (i+1) % hidden] = 0.35
 79            A[i, (i+2) % hidden] = 0.15
 80        self.register_buffer("A", A)
 81        self.readout = nn.Linear(2*hidden, 2)
 82
 83    def forward(self, x):
 84        b, t, _ = x.shape
 85        theta = torch.zeros(b, self.hidden, device=x.device)
 86        outs = []
 87        h = 0.15
 88        kappa = 0.8
 89        for n in range(t):
 90            drive = self.inp(x[:, n])
 91            diff = theta[:, None, :] - theta[:, :, None]  # theta_j-theta_i
 92            coupling = torch.einsum("bij,ij->bi", torch.sin(diff), self.A)
 93            theta = theta + h*(self.omega + 0.08*drive + kappa*coupling)
 94            outs.append(torch.cat([torch.cos(theta), torch.sin(theta)], dim=1))
 95        return self.readout(outs[-1])
 96
 97
 98class TanhRNN(nn.Module):
 99    def __init__(self, hidden=16):
100        super().__init__()
101        self.cell = nn.RNNCell(1, hidden, nonlinearity="tanh")
102        self.readout = nn.Linear(hidden, 2)
103        self.hidden = hidden
104
105    def forward(self, x):
106        b, t, _ = x.shape
107        h = torch.zeros(b, self.hidden, device=x.device)
108        for n in range(t):
109            h = self.cell(x[:, n], h)
110        return self.readout(h)
111
112
113def train_model(model, train_x, train_y, test_x, test_y, device, steps=220):
114    model.to(device)
115    opt = torch.optim.Adam(model.parameters(), lr=0.01)
116    loss_fn = nn.CrossEntropyLoss()
117    model.train()
118    for step in range(steps):
119        idx = torch.randint(0, len(train_x), (64,), device=device)
120        loss = loss_fn(model(train_x[idx]), train_y[idx])
121        opt.zero_grad(); loss.backward(); opt.step()
122    model.eval()
123    with torch.no_grad():
124        pred = model(test_x).argmax(1)
125        acc = float((pred == test_y).float().mean().cpu())
126        test_loss = float(loss_fn(model(test_x), test_y).cpu())
127    return {"accuracy": acc, "loss": test_loss, "steps": steps}
128
129
130def mini_experiment():
131    # Simple sequential sine-frequency classification: task has genuine temporal memory.
132    rng = np.random.default_rng(SEED)
133    def make(n):
134        xs, ys = [], []
135        for _ in range(n):
136            label = int(rng.integers(0, 2))
137            freq = 0.22 if label == 0 else 0.39
138            phase = rng.uniform(0, 2*math.pi)
139            t = np.arange(24)
140            seq = np.sin(freq*t + phase) + 0.18*rng.normal(size=24)
141            xs.append(seq[:, None]); ys.append(label)
142        return torch.tensor(np.asarray(xs), dtype=torch.float32), torch.tensor(ys)
143    train_x, train_y = make(512); test_x, test_y = make(256)
144    try:
145        device = "cuda" if torch.cuda.is_available() else "cpu"
146        torch.set_num_threads(4)
147        a = train_model(PhaseRNN(), train_x.to(device), train_y.to(device),
148                        test_x.to(device), test_y.to(device), device)
149        torch.manual_seed(SEED)
150        b = train_model(TanhRNN(), train_x.to(device), train_y.to(device),
151                        test_x.to(device), test_y.to(device), device)
152    except Exception as exc:
153        device = "cpu"
154        torch.manual_seed(SEED)
155        a = train_model(PhaseRNN(), train_x, train_y, test_x, test_y, device)
156        torch.manual_seed(SEED)
157        b = train_model(TanhRNN(), train_x, train_y, test_x, test_y, device)
158        return {"device": device, "phase_rnn": a, "tanh_rnn": b, "fallback_error": repr(exc)}
159    return {"device": device, "phase_rnn": a, "tanh_rnn": b}
160
161
162if __name__ == "__main__":
163    result = {"math": math_checks(), "mini_experiment": mini_experiment()}
164    with open("results.json", "w") as f:
165        json.dump(result, f, indent=2)
166    print(json.dumps(result, indent=2))