repro_bundle/code/baselines.py

"""Baseline optimizers and the Dolan-More performance-profile metric.

Baselines: L-BFGS, Adam, AdaHessian (Sec. 5.1.2). Each is run per problem and
returns the loss trajectory, so we can compute iterations-to-target and build
performance profiles. The learned L-SR1 optimizer is evaluated separately (it
runs batched) and its trajectories are fed into the same metric code.
"""
from __future__ import annotations

import numpy as np
import torch


# ---------------------------------------------------------------------------
# Baseline rollouts (single problem, x0: [1, N])
# ---------------------------------------------------------------------------
def _traj(problem, x):
    with torch.no_grad():
        return problem.value(x).item()


def run_adam(problem, x0, steps, lr=1e-2):
    x = x0.clone().detach().requires_grad_(True)
    opt = torch.optim.Adam([x], lr=lr)
    out = [_traj(problem, x)]
    for _ in range(steps):
        opt.zero_grad()
        loss = problem.value(x).sum()
        loss.backward()
        opt.step()
        out.append(_traj(problem, x))
    return np.array(out)


def run_lbfgs(problem, x0, steps, lr=1.0, history=10):
    x = x0.clone().detach().requires_grad_(True)
    opt = torch.optim.LBFGS([x], lr=lr, max_iter=1, history_size=history,
                            line_search_fn="strong_wolfe")
    out = [_traj(problem, x)]
    for _ in range(steps):
        def closure():
            opt.zero_grad()
            loss = problem.value(x).sum()
            loss.backward()
            return loss
        opt.step(closure)
        out.append(_traj(problem, x))
    return np.array(out)


def run_adahessian(problem, x0, steps, lr=0.1, beta1=0.9, beta2=0.999,
                   eps=1e-4, seed=0):
    """Minimal AdaHessian: diagonal Hessian via one-sample Hutchinson estimator."""
    x = x0.clone().detach().requires_grad_(True)
    g_rng = torch.Generator(device=x.device).manual_seed(seed)
    m = torch.zeros_like(x)
    vhat = torch.zeros_like(x)
    out = [_traj(problem, x)]
    for t in range(1, steps + 1):
        loss = problem.value(x).sum()
        grad = torch.autograd.grad(loss, x, create_graph=True)[0]
        z = torch.randint(0, 2, x.shape, generator=g_rng,
                          device=x.device, dtype=x.dtype) * 2 - 1  # Rademacher
        hv = torch.autograd.grad((grad * z).sum(), x, retain_graph=False)[0]
        d = (hv * z).abs()                      # diagonal Hessian estimate
        m = beta1 * m + (1 - beta1) * grad.detach()
        vhat = beta2 * vhat + (1 - beta2) * d.detach() ** 2
        mhat = m / (1 - beta1 ** t)
        vh = (vhat / (1 - beta2 ** t)).sqrt()
        with torch.no_grad():
            x -= lr * mhat / (vh + eps)
        x = x.detach().requires_grad_(True)
        out.append(_traj(problem, x))
    return np.array(out)


BASELINES = {
    "L-BFGS": run_lbfgs,
    "Adam": run_adam,
    "AdaHessian": run_adahessian,
}


# ---------------------------------------------------------------------------
# Performance profile (Dolan & More, 2002)
# ---------------------------------------------------------------------------
def iters_to_target(traj, f_star, f0, rtol=1e-3, atol=1e-8):
    """First iteration index where f - f* <= rtol*(f0 - f*) + atol; inf if never."""
    thresh = rtol * (f0 - f_star) + atol
    hit = np.where((traj - f_star) <= thresh)[0]
    return int(hit[0]) if len(hit) else np.inf


def performance_profile(cost_matrix, taus=None):
    """cost_matrix: dict[solver] -> np.array[n_problems] of costs (inf = fail).

    Returns (taus, {solver: rho(tau)}) with rho_s(tau) = fraction of problems
    whose performance ratio r_{p,s} = cost/min_over_solvers(cost) is <= tau.
    """
    solvers = list(cost_matrix.keys())
    costs = np.stack([cost_matrix[s] for s in solvers])          # [S, P]
    best = np.min(costs, axis=0, keepdims=True)                  # [1, P]
    ratios = np.where(np.isfinite(costs), costs / np.maximum(best, 1e-12),
                      np.inf)
    if taus is None:
        finite = ratios[np.isfinite(ratios)]
        rmax = float(np.max(finite)) if finite.size else 2.0
        taus = np.geomspace(1.0, max(rmax, 2.0), 60)
    profiles = {}
    P = ratios.shape[1]
    for i, s in enumerate(solvers):
        profiles[s] = np.array([(ratios[i] <= t).sum() / P for t in taus])
    return taus, profiles


def auc(taus, rho):
    """Area under the performance profile over log(tau) -- a scalar summary."""
    lt = np.log(taus)
    trap = np.trapezoid if hasattr(np, "trapezoid") else np.trapz  # np>=2 rename
    return float(trap(rho, lt) / (lt[-1] - lt[0]))