1import json
  2from pathlib import Path
  3import numpy as np
  4
  5
  6def schur_monitor(H, G, lam=1e-3, beta=1.0):
  7    """Return R=G.T (H+lam I)^-1 G and the trust update operator."""
  8    H = np.asarray(H, dtype=float)
  9    G = np.asarray(G, dtype=float)
 10    A = H + lam * np.eye(H.shape[0])
 11    X = np.linalg.solve(A, G)
 12    R = 0.5 * (G.T @ X + X.T @ G)
 13    # eigh is stable and makes tiny numerical asymmetry harmless
 14    ev, V = np.linalg.eigh(R)
 15    ev = np.maximum(ev, 0.0)
 16    R_psd = (V * ev) @ V.T
 17    trust = np.linalg.inv(np.eye(G.shape[1]) + beta * R_psd)
 18    return R_psd, trust
 19
 20
 21def reduced_energy(sigma, u, H, G, a):
 22    # E is affine in u before relaxation: .5 sigma'Hsigma + sigma'Gu + a'u
 23    return .5 * sigma @ H @ sigma + sigma @ G @ u + a @ u
 24
 25
 26def verify_identity(rng):
 27    k, m = 4, 3
 28    A = rng.normal(size=(k, k))
 29    H = A.T @ A + 0.7 * np.eye(k)
 30    G = rng.normal(size=(k, m))
 31    a = rng.normal(size=m)
 32    HinvG = np.linalg.solve(H, G)
 33    R = G.T @ HinvG
 34    # Relax sigma*(u)=-H^-1 G u and finite-difference the reduced energy.
 35    def reduced(u):
 36        sigma = -np.linalg.solve(H, G @ u)
 37        return reduced_energy(sigma, u, H, G, a)
 38    u0 = rng.normal(size=m)
 39    eps = 2e-4
 40    numerical = np.zeros((m, m))
 41    for i in range(m):
 42        ei = np.eye(m)[i]
 43        for j in range(m):
 44            ej = np.eye(m)[j]
 45            numerical[i, j] = (reduced(u0+eps*ei+eps*ej)-reduced(u0+eps*ei-eps*ej)
 46                - reduced(u0-eps*ei+eps*ej)+reduced(u0-eps*ei-eps*ej))/(4*eps**2)
 47    target = -R
 48    eig_R = np.linalg.eigvalsh(R)
 49    return {
 50        "max_abs_hessian_error": float(np.max(np.abs(numerical-target))),
 51        "R_eigenvalues": eig_R.tolist(),
 52        "min_R_eigenvalue": float(eig_R.min()),
 53        "identity_pass": bool(np.max(np.abs(numerical-target)) < 2e-6 and eig_R.min() >= -1e-10),
 54    }
 55
 56
 57def run_control(rng, steps=80):
 58    # Two mechanism amplitudes. R has one very interactive direction.
 59    H = np.diag([0.25, 1.0])
 60    G = np.array([[2.7, 2.7], [0.15, -0.15]])
 61    R, trust = schur_monitor(H, G, lam=0.02, beta=1.0)
 62    # Smooth bounded target loss, with a deliberately aggressive hyper-step.
 63    target = np.array([1.0, -1.0])
 64    Q = np.diag([1.0, 1.0])
 65    eta = 0.72
 66    rows = {"baseline": [], "schur": []}
 67    noise = 0.015 * rng.normal(size=(steps, 2))
 68    for name, P in [("baseline", np.eye(2)), ("schur", trust)]:
 69        u = np.array([0.0, 0.0])
 70        for t in range(steps):
 71            # Hypergradient of a quadratic validation proxy, plus reproducible mild noise.
 72            h = Q @ (u-target) + noise[t]
 73            delta = -eta * P @ h
 74            u = np.clip(u + delta, -3.0, 3.0)
 75            loss = 0.5 * (u-target) @ Q @ (u-target)
 76            rows[name].append({"step": t+1, "loss": float(loss), "u_norm": float(np.linalg.norm(u)),
 77                               "delta_norm": float(np.linalg.norm(delta))})
 78    def summarize(x):
 79        losses = np.array([z["loss"] for z in x])
 80        return {"final_loss": float(losses[-1]), "best_loss": float(losses.min()),
 81                "loss_spikes": int(np.sum(losses[1:] > 1.2 * losses[:-1])),
 82                "mean_last10": float(losses[-10:].mean()),
 83                "max_u_norm": float(max(z["u_norm"] for z in x))}
 84    return {"H": H.tolist(), "G": G.tolist(), "R": R.tolist(),
 85            "R_eigenvalues": np.linalg.eigvalsh(R).tolist(),
 86            "trust_matrix": trust.tolist(), "eta": eta,
 87            "summary": {k: summarize(v) for k,v in rows.items()}, "trace": rows}
 88
 89
 90def main():
 91    rng = np.random.default_rng(491)
 92    verification = verify_identity(rng)
 93    experiment = run_control(rng)
 94    out = {"verification": verification, "experiment": experiment}
 95    Path("results.json").write_text(json.dumps(out, indent=2))
 96    print(json.dumps({"verification": verification, "summary": experiment["summary"],
 97                      "R_eigenvalues": experiment["R_eigenvalues"]}, indent=2))
 98
 99if __name__ == "__main__":
100    main()