repro_bundle/code/eval_analytic.py

"""Evaluation for the analytic experiments.

Produces the evidence for:
  * Claim 2  -- learned projection: secant residual + Newton alignment, on vs off.
  * Claim 4  -- performance profiles over the 6-function x 5-dim suite, and
                Newton-direction alignment vs L-BFGS on quadratics.
  * Claim 5  -- (buffer L, wall-clock recorded by the trainer).

L-SR1 runs batched (dimension-invariant); baselines run per problem.
"""
from __future__ import annotations

import argparse
import json

import numpy as np
import torch

from lsr1 import (LSR1Config, LSR1Optimizer, Quadratic, Rastrigin, Rosenbrock,
                  make_quadratics, init_points, rollout)
from baselines import (run_adam, run_lbfgs, run_adahessian, iters_to_target,
                       performance_profile, auc)


def load_model(ckpt_path, device="cpu", buffer_L=None, K=None):
    ck = torch.load(ckpt_path, map_location=device, weights_only=False)
    cfg = LSR1Config(**ck["cfg"])
    if buffer_L is not None:
        cfg.buffer_L = buffer_L         # larger test-time buffer (Table 6)
    if K is not None:
        cfg.K = K
    opt = LSR1Optimizer(cfg).to(device)
    opt.load_state_dict(ck["state_dict"])
    opt.eval()
    return opt, cfg


def f_star(problem, x0):
    if isinstance(problem, Quadratic):
        xs = problem.optimum()
        return problem.value(xs).detach()
    # Rosenbrock min 0 at ones, Rastrigin min 0 at zero
    return torch.zeros(x0.shape[0], device=x0.device)


# ---------------------------------------------------------------------------
# Newton-direction alignment on quadratics (Claim 2 & 4)
# ---------------------------------------------------------------------------
def newton_alignment(opt, cfg, N, steps=20, n_problems=64, cond=1000.0,
                     device="cpu", seed=1):
    gen = torch.Generator(device=device).manual_seed(seed)
    prob = make_quadratics(n_problems, N, cond=cond, generator=gen,
                           device=device)
    x0 = init_points(n_problems, N, "quadratic", generator=gen, device=device)
    with torch.enable_grad():
        out = rollout(opt, prob, x0, K=steps, create_graph=False, record=True)
    traj = out["traj"]           # [steps+1, B, N]
    dirs = out["dirs"]           # [steps, B, N]  (actual update = -step)
    cos = []
    for k in range(steps):
        xk = traj[k]
        nd = prob.newton_dir(xk)                    # Newton *step* direction
        d = dirs[k]
        c = torch.nn.functional.cosine_similarity(d, nd, dim=-1)
        cos.append(float(c.mean()))
    return cos


def lbfgs_newton_alignment(N, steps=20, n_problems=64, cond=1000.0,
                           device="cpu", seed=1):
    gen = torch.Generator(device=device).manual_seed(seed)
    prob = make_quadratics(n_problems, N, cond=cond, generator=gen,
                           device=device)
    x0 = init_points(n_problems, N, "quadratic", generator=gen, device=device)
    cos_per_iter = [[] for _ in range(steps)]
    for i in range(n_problems):
        pi = Quadratic(prob.H[i:i+1], prob.b[i:i+1])
        x = x0[i:i+1].clone().detach().requires_grad_(True)
        opt = torch.optim.LBFGS([x], lr=1.0, max_iter=1, history_size=10,
                                line_search_fn="strong_wolfe")
        for k in range(steps):
            x_prev = x.detach().clone()
            nd = pi.newton_dir(x_prev)
            def closure():
                opt.zero_grad(); loss = pi.value(x).sum(); loss.backward()
                return loss
            opt.step(closure)
            d = (x.detach() - x_prev)
            c = torch.nn.functional.cosine_similarity(d, nd, dim=-1)
            cos_per_iter[k].append(float(c.mean()))
    return [float(np.mean(c)) for c in cos_per_iter]


# ---------------------------------------------------------------------------
# Performance profiles over the 6-function x 5-dim suite (Claim 4)
# ---------------------------------------------------------------------------
def make_problem(family, N, cond, B, gen, device):
    if family.startswith("quad"):
        prob = make_quadratics(B, N, cond=cond, diagonal=True, generator=gen,
                               device=device)
        x0 = init_points(B, N, "quadratic", generator=gen, device=device)
    elif family == "rosenbrock":
        prob = Rosenbrock()
        x0 = init_points(B, N, "rosenbrock", generator=gen, device=device)
    elif family == "rastrigin":
        prob = Rastrigin()
        x0 = init_points(B, N, "rastrigin", generator=gen, device=device)
    return prob, x0


def lsr1_costs(opt, prob, x0, steps, rtol):
    with torch.enable_grad():
        out = rollout(opt, prob, x0, K=steps, create_graph=False, record=False)
    f_traj = out["f_traj"].detach().cpu().numpy()    # [steps, B]
    f0 = prob.value(x0).detach().cpu().numpy()
    fs = f_star(prob, x0).cpu().numpy()
    full = np.concatenate([f0[None, :], f_traj], axis=0)  # [steps+1, B]
    return np.array([iters_to_target(full[:, i], fs[i], f0[i], rtol)
                     for i in range(full.shape[1])])


def baseline_costs(runner, prob, x0, steps, rtol, **kw):
    costs = []
    B = x0.shape[0]
    for i in range(B):
        pi = (Quadratic(prob.H[i:i+1], prob.b[i:i+1])
              if isinstance(prob, Quadratic) else prob)
        xi = x0[i:i+1]
        traj = runner(pi, xi, steps, **kw)
        f0 = float(pi.value(xi))
        fs = float(f_star(pi, xi)[0]) if isinstance(pi, Quadratic) else 0.0
        costs.append(iters_to_target(traj, fs, f0, rtol))
    return np.array(costs)


def performance_suite(models, dims=(50, 100, 250, 500, 1000),
                      quad_conds=(1.0, 100.0, 1000.0, 10000.0),
                      instances=5, steps=60, rtol=1e-2, test_L=64,
                      device="cpu", seed=7, include_hard=True):
    """models: dict family-> (ckpt path). Returns cost dict + profile + per-type."""
    gen = torch.Generator(device=device).manual_seed(seed)
    problems = []
    for c in quad_conds:
        for N in dims:
            problems.append(("quadratic", N, c))
    if include_hard:
        for fam in ("rosenbrock", "rastrigin"):
            for N in dims:
                problems.append((fam, N, None))

    loaded = {}
    for fam, path in models.items():
        loaded[fam] = load_model(path, device=device, buffer_L=test_L)

    solver_costs = {"L-SR1": [], "L-BFGS": [], "Adam": [], "AdaHessian": []}
    per_type = []
    for (family, N, cond) in problems:
        prob, x0 = make_problem(family, N, cond, instances, gen, device)
        key = "quadratic" if family.startswith("quad") else family
        opt, cfg = loaded[key]
        c_lsr1 = lsr1_costs(opt, prob, x0, steps, rtol)
        c_lbfgs = baseline_costs(run_lbfgs, prob, x0, steps, rtol)
        c_adam = baseline_costs(run_adam, prob, x0, steps, rtol, lr=1e-2)
        c_adah = baseline_costs(run_adahessian, prob, x0, steps, rtol, lr=0.15)
        for s, c in [("L-SR1", c_lsr1), ("L-BFGS", c_lbfgs),
                     ("Adam", c_adam), ("AdaHessian", c_adah)]:
            solver_costs[s].append(c)
        per_type.append(dict(family=family, N=N, cond=cond,
                             solved={s: int(np.isfinite(v).sum())
                                     for s, v in [("L-SR1", c_lsr1),
                                                  ("L-BFGS", c_lbfgs),
                                                  ("Adam", c_adam),
                                                  ("AdaHessian", c_adah)]},
                             median={s: (float(np.median(v[np.isfinite(v)]))
                                         if np.isfinite(v).any() else None)
                                     for s, v in [("L-SR1", c_lsr1),
                                                  ("L-BFGS", c_lbfgs),
                                                  ("Adam", c_adam),
                                                  ("AdaHessian", c_adah)]}))
    cost_matrix = {s: np.concatenate(v) for s, v in solver_costs.items()}
    taus, profiles = performance_profile(cost_matrix)
    aucs = {s: auc(taus, profiles[s]) for s in profiles}
    return dict(
        taus=taus.tolist(),
        profiles={s: profiles[s].tolist() for s in profiles},
        aucs=aucs,
        per_type=per_type,
        problems=[dict(family=f, N=N, cond=c) for (f, N, c) in problems],
    )


def convergence_curves(opt, N=10, steps=30, n=200, cond=1000.0, device="cpu",
                       seed=3):
    """Median f-gap vs iteration for L-SR1 vs L-BFGS/Adam on quadratics."""
    gen = torch.Generator(device=device).manual_seed(seed)
    prob = make_quadratics(n, N, cond=cond, generator=gen, device=device)
    x0 = init_points(n, N, "quadratic", generator=gen, device=device)
    fs = f_star(prob, x0).cpu().numpy()
    with torch.enable_grad():
        out = rollout(opt, prob, x0, K=steps, create_graph=False, record=False)
    f_lsr1 = out["f_traj"].detach().cpu().numpy()
    f0 = prob.value(x0).detach().cpu().numpy()
    lsr1 = np.median(np.concatenate([f0[None], f_lsr1], 0) - fs[None], axis=1)
    lb, ad = [], []
    for i in range(n):
        pi = Quadratic(prob.H[i:i+1], prob.b[i:i+1])
        lb.append(run_lbfgs(pi, x0[i:i+1], steps) - fs[i])
        ad.append(run_adam(pi, x0[i:i+1], steps, lr=1e-2) - fs[i])
    return dict(iters=list(range(steps + 1)),
                lsr1=lsr1.tolist(),
                lbfgs=np.median(np.stack(lb), 0).tolist(),
                adam=np.median(np.stack(ad), 0).tolist())


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--quad2", default="outputs/quad2_proj1.pt")
    ap.add_argument("--quad2-noproj", default="outputs/quad2_proj0.pt")
    ap.add_argument("--quad100", default="outputs/quad100_proj1.pt")
    ap.add_argument("--rosen100", default="outputs/rosen100_proj1.pt")
    ap.add_argument("--rastr100", default="outputs/rastr100_proj1.pt")
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--out", default="outputs/eval.json")
    ap.add_argument("--quick", type=int, default=0)
    args = ap.parse_args()
    dev = args.device
    res = {}

    # --- Claim 2: projection ablation via Newton alignment + secant ---------
    opt_on, cfg_on = load_model(args.quad2, dev)
    opt_off, cfg_off = load_model(args.quad2_noproj, dev)
    steps = 20
    res["newton_alignment"] = dict(
        lsr1_proj=newton_alignment(opt_on, cfg_on, N=10, steps=steps, device=dev),
        lsr1_noproj=newton_alignment(opt_off, cfg_off, N=10, steps=steps, device=dev),
        lbfgs=lbfgs_newton_alignment(N=10, steps=steps, device=dev),
        N=10, cond=1000.0,
    )
    res["convergence"] = convergence_curves(opt_on, N=10, device=dev)

    # --- Claim 4: performance profiles --------------------------------------
    models = {"quadratic": args.quad100, "rosenbrock": args.rosen100,
              "rastrigin": args.rastr100}
    dims = (50, 100) if args.quick else (50, 100, 250, 500, 1000)
    res["performance"] = performance_suite(models, dims=dims, device=dev,
                                           instances=3 if args.quick else 5,
                                           steps=40 if args.quick else 60)
    with open(args.out, "w") as f:
        json.dump(res, f, indent=2)
    print("wrote", args.out)
    print("AUCs:", res["performance"]["aucs"])
    na = res["newton_alignment"]
    print("Newton align (final iter): proj=%.3f noproj=%.3f lbfgs=%.3f" %
          (na["lsr1_proj"][-1], na["lsr1_noproj"][-1], na["lbfgs"][-1]))