1import json, math, os
  2import numpy as np
  3import torch
  4
  5SEED = 2244
  6np.random.seed(SEED)
  7torch.manual_seed(SEED)
  8torch.set_default_dtype(torch.float64)
  9
 10
 11def device():
 12    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
 13
 14
 15def omega(rho, curvature, potential=None):
 16    # A differentiable toy variational functional.  The quartic term makes
 17    # negative curvature bounded below, while rho=0 exposes its Hessian exactly.
 18    if potential is None:
 19        potential = torch.zeros_like(rho)
 20    return 0.5 * (curvature * rho.square()).sum() + 0.025 * rho.pow(4).sum() + (potential * rho).sum()
 21
 22
 23def hvp(rho, curvature, v, potential=None, create_graph=False):
 24    # Keep the graph when the HVP is used inside a trainable regularizer.
 25    rho_in = rho.detach().clone().requires_grad_(True)
 26    g = torch.autograd.grad(omega(rho_in, curvature, potential), rho_in, create_graph=True)[0]
 27    hv = torch.autograd.grad((g * v).sum(), rho_in, create_graph=create_graph)[0]
 28    return hv if create_graph else hv.detach()
 29
 30
 31def direct_hessian(rho, curvature, potential=None):
 32    rho = rho.detach().clone().requires_grad_(True)
 33    return torch.autograd.functional.hessian(lambda x: omega(x, curvature, potential), rho).detach()
 34
 35
 36def hess_regularizer(rho, curvature, nvec=16, eps=1e-3, kappa=2.0):
 37    # Stochastic Rayleigh estimator, exactly in the form proposed in the idea.
 38    vals = []
 39    for _ in range(nvec):
 40        v = torch.randn_like(rho)
 41        v = v / v.norm()
 42        vals.append((v * hvp(rho, curvature, v, create_graph=True)).sum())
 43    q = torch.stack(vals)
 44    neg = torch.relu(-q + eps).square().mean()
 45    high = torch.relu(q - kappa).square().mean()
 46    return neg + high, q.detach()
 47
 48
 49def run(device_name):
 50    dev = torch.device(device_name)
 51    n = 8
 52    # 1. Core autodiff claim: HVP equals dense Hessian times v.
 53    c = torch.linspace(0.15, 1.4, n, device=dev)
 54    r = torch.randn(n, device=dev) * 0.2
 55    v = torch.randn(n, device=dev)
 56    H = direct_hessian(r, c)
 57    hv_err = (hvp(r, c, v) - H @ v).abs().max().item()
 58    eigs = torch.linalg.eigvalsh(H).cpu().numpy()
 59
 60    # 2. Susceptibility sweep. For quadratic Omega, drho/dV=-H^-1 exactly.
 61    gammas = np.array([0.02, 0.04, 0.08, 0.16, 0.32, 0.64])
 62    suscept = []
 63    predicted = []
 64    for g in gammas:
 65        # use a fixed unit perturbation so the norm is analytically 1/gamma
 66        H_g = torch.eye(n, device=dev) * float(g)
 67        dv = torch.zeros(n, device=dev); dv[0] = 1.0
 68        dr = torch.linalg.solve(H_g, -dv)
 69        suscept.append(dr.norm().item()); predicted.append(1.0 / g)
 70    slope = float(np.polyfit(np.log(gammas), np.log(suscept), 1)[0])
 71    rel_err = float(np.max(np.abs(np.array(suscept)-predicted) / np.array(predicted)))
 72
 73    # 3. Zero crossing: explicit Euler response iteration is stable iff gamma*lambda < 2.
 74    # For the fixed-point minimization rho <- rho - eta(H rho + dV), eta=1.
 75    # Its amplification factor is |1-gamma|; instability starts at gamma=2.
 76    cross_g = np.array([0.5, 1.5, 2.0, 2.5])
 77    growth = []
 78    for g in cross_g:
 79        x = 1.0
 80        for _ in range(30): x = (1.0-g)*x
 81        growth.append(abs(x))
 82    # Negative curvature is directly visible as negative Rayleigh quotient.
 83    neg_c = torch.tensor(-0.2, device=dev)
 84    neg_rho = torch.zeros(8, device=dev)
 85    reg_neg, q_neg = hess_regularizer(neg_rho, neg_c, nvec=32, eps=1e-3, kappa=2.0)
 86
 87    # 4. Mini experiment: fit an observed curvature target that is spuriously negative.
 88    # Baseline can retain the negative mode; regularized fit rejects it.
 89    target = -0.20
 90    def fit(use_reg, target_value=target, steps=300):
 91        x = torch.tensor(-0.5, device=dev, requires_grad=True)
 92        opt = torch.optim.Adam([x], lr=0.03)
 93        rr = torch.zeros(8, device=dev)
 94        for _ in range(steps):
 95            opt.zero_grad()
 96            loss = (x-target_value).square()
 97            if use_reg:
 98                reg, _ = hess_regularizer(rr, x, nvec=16, eps=1e-3, kappa=2.0)
 99                loss = loss + 8.0 * reg
100            loss.backward(); opt.step()
101        return float(x.detach().cpu()), float((x-target_value)**2)
102    base, base_fit = fit(False)
103    regular, reg_fit = fit(True)
104
105    # Deliberately near-critical positive mode should not be pushed to kappa;
106    # only the negative-curvature branch is active.
107    soft_target = 0.02
108    soft_base, _ = fit(False, soft_target)
109    soft_reg, _ = fit(True, soft_target)
110
111    result = {
112        "device": str(dev),
113        "hvp_max_abs_error": hv_err,
114        "hessian_eigenvalues": eigs.tolist(),
115        "susceptibility_sweep": [{"lambda": float(g), "observed": float(s), "predicted_1_over_lambda": float(p)} for g,s,p in zip(gammas,suscept,predicted)],
116        "loglog_slope_observed": slope,
117        "loglog_slope_predicted": -1.0,
118        "max_relative_scaling_error": rel_err,
119        "euler_zero_crossing_predicted_gamma": 2.0,
120        "euler_growth_sweep": [{"gamma": float(g), "abs_amplification_after_30": float(a)} for g,a in zip(cross_g,growth)],
121        "negative_mode_mean_rayleigh": float(q_neg.mean().cpu()),
122        "negative_mode_regularizer": float(reg_neg.cpu()),
123        "curvature_fit": {"target_negative": target, "baseline_final": base, "regularized_final": regular, "baseline_mse": base_fit, "regularized_mse": reg_fit, "soft_target": soft_target, "soft_baseline_final": soft_base, "soft_regularized_final": soft_reg},
124    }
125    return result
126
127
128def main():
129    try:
130        d = device()
131        out = run(d)
132    except Exception as e:
133        if torch.cuda.is_available():
134            out = run("cpu")
135            out["cuda_error_fallback"] = repr(e)
136        else:
137            raise
138    with open("results.json", "w") as f: json.dump(out, f, indent=2)
139    print(json.dumps(out, indent=2))
140
141if __name__ == "__main__": main()