1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5# Graph-certified switching SSM MVP.  The one-node graph is path-complete
  6# when it contains one self-edge for every mode; multiple-node certificates
  7# are supported by the loss implementation below.
  8
  9def edge_loss(a, rho, p=1.0, tau=0.02):
 10    """Mean softplus edge penalty for V(z)=p z^2, evaluated at z=1."""
 11    # softplus argument is the normalized violation; this is the exact
 12    # expression apart from the harmless tau smoothing.
 13    e = p * (np.asarray(a) ** 2 - rho)
 14    return np.logaddexp(0.0, e / tau) * tau
 15
 16def core_verification():
 17    rho = 0.81
 18    boundary = math.sqrt(rho)
 19    # Prediction 1: violation switches at |a|=sqrt(rho).
 20    amps = np.linspace(0.0, 1.2, 1201)
 21    hard_violation = amps**2 > rho
 22    transition = amps[np.flatnonzero(hard_violation)[0]]
 23    # Prediction 2: arbitrary switching norm is product; worst sequence is max amp^T.
 24    modes = np.array([0.55, 0.72, 0.89, 1.03])
 25    T = 80
 26    predicted_log_growth = T * math.log(np.max(np.abs(modes)))
 27    rng = np.random.default_rng(7)
 28    seq = rng.integers(0, len(modes), size=T)
 29    observed_log_norm = np.sum(np.log(np.abs(modes[seq])))
 30    worst_log_norm = np.sum(np.log(np.max(np.abs(modes))) * np.ones(T))
 31    # Prediction 3: unsmoothed violation is exactly quadratic in amplitude excess.
 32    excess = np.array([0.01, 0.03, 0.07, 0.12])
 33    exact = (boundary + excess)**2 - rho
 34    slopes = np.diff(exact) / np.diff(excess)
 35    predicted_slope_at_boundary = 2 * boundary
 36    # Confirm direct rollout agrees with product formula.
 37    z = 1.0
 38    for a in modes[np.argmax(np.abs(modes))] * np.ones(T):
 39        z *= a
 40    return {
 41        "rho": rho, "predicted_boundary": boundary,
 42        "observed_boundary_grid": float(transition),
 43        "boundary_abs_error": float(abs(transition-boundary)),
 44        "modes": modes.tolist(), "T": T,
 45        "predicted_worst_log_norm": float(predicted_log_growth),
 46        "observed_worst_log_norm": float(math.log(abs(z))),
 47        "random_sequence_log_norm": float(observed_log_norm),
 48        "quadratic_excess": excess.tolist(), "exact_violation": exact.tolist(),
 49        "local_slopes": slopes.tolist(),
 50        "predicted_boundary_slope": predicted_slope_at_boundary,
 51        "slope_relative_error_last": float(abs(slopes[-1]-predicted_slope_at_boundary)/predicted_slope_at_boundary),
 52        "all_predictions_confirmed": bool(abs(transition-boundary)<0.002 and abs(math.log(abs(z))-predicted_log_growth)<1e-9 and abs(slopes[-1]-predicted_slope_at_boundary)/predicted_slope_at_boundary<0.15)
 53    }
 54
 55def train_tiny(seed, regularized, steps=450):
 56    # Torch is optional; numpy fallback still produces a useful comparison.
 57    try:
 58        import torch
 59        torch.manual_seed(seed)
 60        device = "cuda" if torch.cuda.is_available() else "cpu"
 61        try:
 62            torch.tensor([0.], device=device)
 63        except Exception:
 64            device = "cpu"
 65        dtype = torch.float32
 66        # A switching scalar recurrence learns a discounted accumulator.
 67        M, batch, T = 4, 64, 20
 68        rho, tau, lam = .81, .03, .25
 69        raw_a = torch.nn.Parameter(torch.tensor([0.15, 0.25, 0.35, 0.45], device=device))
 70        raw_b = torch.nn.Parameter(torch.zeros(M, device=device))
 71        out_w = torch.nn.Parameter(torch.tensor(1., device=device))
 72        opt = torch.optim.Adam([raw_a, raw_b, out_w], lr=.025)
 73        rng = np.random.default_rng(seed)
 74        for _ in range(steps):
 75            x = torch.tensor(rng.normal(size=(batch,T)), dtype=dtype, device=device)
 76            # random modes are used during training and all modes in penalty
 77            modes = torch.tensor(rng.integers(0,M,size=(batch,T)), dtype=torch.long, device=device)
 78            z = torch.zeros(batch, device=device)
 79            target = torch.zeros(batch, device=device)
 80            discount = 1.
 81            for t in range(T):
 82                target = target + discount*x[:,t]
 83                discount *= .8
 84                ai, bi = raw_a[modes[:,t]], raw_b[modes[:,t]]
 85                z = ai*z + bi*x[:,t]
 86            pred = out_w*z
 87            task = ((pred-target)**2).mean()
 88            # V(z)=z^2, one graph node, every mode self-edge. Use sampled z.
 89            zs = torch.randn(128, device=device)
 90            ai = raw_a[None,:]
 91            violations = ai**2 * zs[:,None]**2 - rho*zs[:,None]**2
 92            cert = torch.nn.functional.softplus(violations/tau).mean()*tau
 93            loss = task + (lam*cert if regularized else 0.)
 94            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_([raw_a,raw_b,out_w], 10.); opt.step()
 95        with torch.no_grad():
 96            a = raw_a.detach().cpu().numpy()
 97            # Worst arbitrary switching rollout, no inputs, initialized at one.
 98            worst = float(1000*np.log(max(np.max(np.abs(a)), 1e-30)))
 99            violation_rate = float(np.mean(a*a > rho))
100            final_task = float(task.detach().cpu())
101        return {"task_mse":final_task, "max_abs_a":float(np.max(np.abs(a))), "edge_violation_rate":violation_rate, "norm_after_1000_worst":worst, "device":device}
102    except Exception as exc:
103        return {"error": repr(exc)}
104
105def main():
106    result = {"verification": core_verification(), "mini_experiment": {
107        "baseline_unregularized": train_tiny(11, False),
108        "graph_certified": train_tiny(11, True)}}
109    Path("results.json").write_text(json.dumps(result, indent=2))
110    print(json.dumps(result, indent=2))
111
112if __name__ == "__main__": main()