1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED = 3155
  8
  9def seed_all(seed=SEED):
 10    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 11    if torch.cuda.is_available():
 12        torch.cuda.manual_seed_all(seed)
 13
 14
 15def resolvent_peak_np(W, nfreq=512):
 16    n = W.shape[0]
 17    grid = np.linspace(0.0, np.pi, nfreq)
 18    vals = []
 19    eye = np.eye(n)
 20    for th in grid:
 21        z = np.exp(1j * th)
 22        vals.append(np.linalg.svd(np.linalg.inv(z * eye - W), compute_uv=False)[0])
 23    vals = np.asarray(vals)
 24    i = int(np.argmax(vals))
 25    return float(vals[i]), float(grid[i]), vals
 26
 27
 28def toy_verification():
 29    # Both matrices have the same stable eigenvalues, but W_nonnormal has
 30    # strong off-diagonal coupling and therefore a much larger resolvent peak.
 31    W_normal = np.diag([0.72, 0.55]).astype(float)
 32    W_nonnormal = np.array([[0.72, 8.0], [0.0, 0.55]], dtype=float)
 33    pn, wn, curve_n = resolvent_peak_np(W_normal)
 34    pnn, wnn, curve_nn = resolvent_peak_np(W_nonnormal)
 35
 36    # Discrete-time transient response ||W^k||, another direct manifestation
 37    # of nonnormal amplification despite spectral radius below one.
 38    transient = []
 39    P = np.eye(2)
 40    for k in range(31):
 41        transient.append(float(np.linalg.svd(P, compute_uv=False)[0]))
 42        P = P @ W_nonnormal
 43    transient = np.asarray(transient)
 44    return {
 45        "normal_eigenvalues": np.linalg.eigvals(W_normal).tolist(),
 46        "nonnormal_eigenvalues": np.linalg.eigvals(W_nonnormal).tolist(),
 47        "normal_rho": float(max(abs(np.linalg.eigvals(W_normal)))),
 48        "nonnormal_rho": float(max(abs(np.linalg.eigvals(W_nonnormal)))),
 49        "normal_resolvent_peak": pn,
 50        "nonnormal_resolvent_peak": pnn,
 51        "normal_peak_theta": wn,
 52        "nonnormal_peak_theta": wnn,
 53        "resolvent_amplification_ratio": pnn / pn,
 54        "max_transient_gain": float(transient.max()),
 55        "transient_peak_step": int(transient.argmax()),
 56        "transient_curve": transient.tolist(),
 57    }
 58
 59
 60class TinyRNN(nn.Module):
 61    def __init__(self, hidden=8):
 62        super().__init__()
 63        self.hidden = hidden
 64        self.inp = nn.Linear(1, hidden)
 65        self.rec = nn.Linear(hidden, hidden, bias=False)
 66        self.out = nn.Linear(hidden, 1)
 67        nn.init.orthogonal_(self.rec.weight, gain=0.85)
 68
 69    def forward(self, x, return_states=False):
 70        b, t, _ = x.shape
 71        h = torch.zeros(b, self.hidden, device=x.device)
 72        states = []
 73        for k in range(t):
 74            h = torch.tanh(self.inp(x[:, k]) + self.rec(h))
 75            states.append(h)
 76        y = self.out(h)
 77        return (y, states) if return_states else y
 78
 79
 80def frequency_penalty(model, theta_grid, tau=0.15):
 81    # Exact Jacobian for the linearized recurrent map at h=0 and input u=0.
 82    # For tanh, derivative is identity there: J=rec.weight, B=inp.weight[:,0],
 83    # C=out.weight, D=0. We use the discrete transfer formula.
 84    J = model.rec.weight
 85    B = model.inp.weight[:, :1]
 86    C = model.out.weight
 87    I = torch.eye(J.shape[0], device=J.device, dtype=J.dtype)
 88    vals = []
 89    for theta in theta_grid:
 90        z = torch.complex(torch.cos(theta), torch.sin(theta))
 91        # Solve (zI-J)v=B in complex arithmetic.
 92        M = z * I.to(torch.complex64) - J.to(torch.complex64)
 93        v = torch.linalg.solve(M, B.to(torch.complex64))
 94        g = C.to(torch.complex64) @ v
 95        vals.append(torch.abs(g).reshape(()))
 96    vals = torch.stack(vals)
 97    return tau * torch.logsumexp(vals / tau, dim=0), vals.max().detach(), vals
 98
 99
100def make_batch(batch, length, device):
101    x = torch.randn(batch, length, 1, device=device)
102    # Sequence-to-one delayed sum task; it encourages memory without requiring
103    # a large model or dataset.
104    y = x.sum(dim=1)
105    return x, y
106
107
108def train_one(kind, device, steps=500):
109    seed_all(SEED + (0 if kind == "baseline" else 1))
110    model = TinyRNN().to(device)
111    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
112    theta = torch.linspace(0, math.pi, 24, device=device)
113    losses, penalties, peaks = [], [], []
114    for step in range(steps):
115        x, y = make_batch(64, 12, device)
116        pred = model(x)
117        task = ((pred - y) ** 2).mean()
118        if kind == "frequency":
119            reg, peak, _ = frequency_penalty(model, theta)
120            loss = task + 0.002 * reg
121            penalties.append(float(reg.detach().cpu()))
122        else:
123            peak = torch.tensor(float("nan"), device=device)
124            loss = task
125        opt.zero_grad(set_to_none=True)
126        loss.backward()
127        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
128        opt.step()
129        losses.append(float(task.detach().cpu())); peaks.append(float(peak.detach().cpu()))
130    with torch.no_grad():
131        _, _, gridvals = frequency_penalty(model, theta)
132        final_peak = float(gridvals.max().cpu())
133        final_task = float(np.mean(losses[-50:]))
134        W = model.rec.weight.detach().cpu().numpy()
135        rho = float(max(abs(np.linalg.eigvals(W))))
136    return {"final_task_loss": final_task, "final_resolvent_peak": final_peak,
137            "spectral_radius": rho, "loss_start": float(np.mean(losses[:20])),
138            "loss_curve": losses, "penalty_curve": penalties}
139
140
141def main():
142    seed_all()
143    # CUDA is allowed but failure must gracefully fall back to CPU.
144    device = "cuda" if torch.cuda.is_available() else "cpu"
145    try:
146        if device == "cuda":
147            torch.zeros(1, device=device)
148    except Exception:
149        device = "cpu"
150    result = {"seed": SEED, "device": device, "toy_verification": toy_verification()}
151    result["baseline"] = train_one("baseline", device)
152    result["frequency"] = train_one("frequency", device)
153    result["observed_peak_reduction"] = (result["baseline"]["final_resolvent_peak"] - result["frequency"]["final_resolvent_peak"]) / result["baseline"]["final_resolvent_peak"]
154    result["observed_task_change"] = result["frequency"]["final_task_loss"] - result["baseline"]["final_task_loss"]
155    Path("results.json").write_text(json.dumps(result, indent=2))
156    print(json.dumps({k: v for k, v in result.items() if k not in ("baseline", "frequency")}, indent=2))
157    print(json.dumps({"baseline": {k:v for k,v in result["baseline"].items() if not k.endswith("curve")}, "frequency": {k:v for k,v in result["frequency"].items() if not k.endswith("curve")}}, indent=2))
158
159if __name__ == "__main__":
160    main()