1import json
  2from pathlib import Path
  3import numpy as np
  4
  5
  6def cg(A, b, tol=1e-11, max_iter=1000):
  7    x = np.zeros_like(b, dtype=float)
  8    r = b - A(x)
  9    p = r.copy()
 10    rr = float(r @ r)
 11    r0 = np.sqrt(rr)
 12    if r0 == 0:
 13        return x, 0, 0.0
 14    for k in range(1, max_iter + 1):
 15        Ap = A(p)
 16        den = float(p @ Ap)
 17        if den <= 0 or not np.isfinite(den):
 18            raise RuntimeError("CG operator was not positive definite")
 19        alpha = rr / den
 20        x += alpha * p
 21        r -= alpha * Ap
 22        new_rr = float(r @ r)
 23        if np.sqrt(new_rr) <= tol * max(1.0, r0):
 24            return x, k, np.sqrt(new_rr)
 25        p = r + (new_rr / rr) * p
 26        rr = new_rr
 27    return x, max_iter, np.linalg.norm(r)
 28
 29
 30def rel(a, b):
 31    return float(np.linalg.norm(a - b) / max(1e-15, np.linalg.norm(b)))
 32
 33
 34def main():
 35    rng = np.random.default_rng(2796)
 36    q, _ = np.linalg.qr(rng.normal(size=(4, 4)))
 37    eig = np.array([0.0, 0.1, 1.0, 10.0])
 38    H = q @ np.diag(eig) @ q.T
 39    # B is in Range(H), so the pseudoinverse derivative is finite.
 40    B = q @ np.array([0.0, 1.0, -0.7, 0.4])
 41    v_true = np.linalg.pinv(H) @ B
 42    eps_values = np.logspace(-1, -8, 8)
 43
 44    stable = []
 45    for eps in eps_values:
 46        A = H + eps * np.eye(4)
 47        v, it, residual = cg(lambda z: A @ z, B)
 48        stable.append({
 49            "eps": float(eps), "relative_error": rel(v, v_true),
 50            "iterations": int(it), "residual": float(residual),
 51        })
 52
 53    # Prediction 1: stable-range Tikhonov derivative converges to -H^+B.
 54    # Prediction 2: for each positive eigencomponent, bias is exactly
 55    # eps/(lambda+eps), hence asymptotically linear in eps.
 56    component_errors = []
 57    for eps in eps_values:
 58        v = np.linalg.solve(H + eps * np.eye(4), B)
 59        coeff = q.T @ v
 60        true_coeff = q.T @ v_true
 61        component_errors.append({
 62            "eps": float(eps),
 63            "lambda_0_error": float(abs(coeff[0] - true_coeff[0])),
 64            "lambda_0p1_relative_error": float(abs(coeff[1] - true_coeff[1]) / abs(true_coeff[1])),
 65            "predicted_lambda_0p1_relative_error": float(eps / (0.1 + eps)),
 66        })
 67
 68    # Prediction 3: if B has a nullspace component, the regularized solve
 69    # diverges as (B_null)/eps. We sweep eps and estimate the log-log slope.
 70    B_bad = B + 0.8 * q[:, 0]
 71    unstable = []
 72    bad_norms = []
 73    for eps in eps_values:
 74        v = np.linalg.solve(H + eps * np.eye(4), B_bad)
 75        null_coeff = float(q[:, 0] @ v)
 76        bad_norms.append(np.linalg.norm(v))
 77        unstable.append({"eps": float(eps), "norm": float(np.linalg.norm(v)), "null_component": null_coeff,
 78                         "eps_times_null_component": float(eps * null_coeff)})
 79    slope = float(np.polyfit(np.log(eps_values[-5:]), np.log(bad_norms[-5:]), 1)[0])
 80
 81    # Tiny bilevel quadratic: f(x,theta)=1/2 x^T Hx - theta B^T x.
 82    # The minimum-norm inner solution is x*=H^+B theta. Outer loss is
 83    # 1/2||x-c||^2. Exact hypergradient is (x-c)^T H^+B.
 84    c = q @ np.array([0.3, -0.2, 0.4, 0.1])
 85    direction = v_true
 86    theta0 = 2.0
 87    exact_x = direction * theta0
 88    exact_loss = 0.5 * np.sum((exact_x - c) ** 2)
 89    exact_grad = float((exact_x - c) @ direction)
 90
 91    def damped_grad(theta, eps):
 92        x = np.linalg.solve(H + eps * np.eye(4), B) * theta
 93        dx = np.linalg.solve(H + eps * np.eye(4), B)
 94        return float((x - c) @ dx), float(0.5 * np.sum((x - c) ** 2))
 95
 96    # Compare fixed damping and continuation (eps halves toward a floor).
 97    fixed_eps = 0.1
 98    theta_fixed = theta0
 99    theta_cont = theta0
100    fixed_trace, cont_trace = [], []
101    lr = 0.015
102    for step in range(40):
103        gf, lf = damped_grad(theta_fixed, fixed_eps)
104        gc, lc = damped_grad(theta_cont, max(1e-7, 0.1 * 0.7 ** step))
105        theta_fixed -= lr * gf
106        theta_cont -= lr * gc
107        fixed_trace.append(lf)
108        cont_trace.append(lc)
109
110    out = {
111        "seed": 2796,
112        "eigenvalues": eig.tolist(),
113        "predictions": {
114            "stable_range_limit": "relative error -> 0 for B in Range(H)",
115            "positive_eigen_bias": "relative bias on lambda=0.1 component = eps/(0.1+eps)",
116            "nullspace_violation": "norm scales as eps^-1 when B has null component",
117        },
118        "stable_sweep": stable,
119        "component_sweep": component_errors,
120        "unstable_sweep": unstable,
121        "unstable_loglog_slope": slope,
122        "bilevel": {
123            "exact_minimum_norm_loss_at_theta0": exact_loss,
124            "exact_hypergradient_at_theta0": exact_grad,
125            "fixed_eps": fixed_eps,
126            "fixed_final_loss": fixed_trace[-1],
127            "continuation_final_loss": cont_trace[-1],
128            "fixed_final_theta": theta_fixed,
129            "continuation_final_theta": theta_cont,
130            "fixed_loss_trace": fixed_trace,
131            "continuation_loss_trace": cont_trace,
132        },
133    }
134    Path("results.json").write_text(json.dumps(out, indent=2))
135    print(json.dumps({
136        "stable_last_relative_error": stable[-1]["relative_error"],
137        "lambda01_observed_vs_predicted": [component_errors[-1]["lambda_0p1_relative_error"], component_errors[-1]["predicted_lambda_0p1_relative_error"]],
138        "unstable_loglog_slope": slope,
139        "exact_grad": exact_grad,
140        "fixed_final_loss": fixed_trace[-1],
141        "continuation_final_loss": cont_trace[-1],
142    }, indent=2))
143
144
145if __name__ == "__main__":
146    main()