Detailed Diagnostics and Implementation

import os
for variable in ["OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS",
                 "VECLIB_MAXIMUM_THREADS", "RAYON_NUM_THREADS"]:
    os.environ[variable] = "1"
os.environ["OMP_CANCELLATION"] = "TRUE"
os.environ["OMP_PROC_BIND"] = "TRUE"

from pathlib import Path
from time import perf_counter
import json
import hashlib
import zipfile
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.lines import Line2D
import platform
import importlib.metadata as metadata
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.special import expit
from scipy.linalg import solve_triangular
from scipy.optimize import minimize, least_squares, linprog
from threadpoolctl import threadpool_limits, threadpool_info
from galahad import tru, arc, qpb
import pyensmallen as pe
import crabbymetrics as cm
from IPython.display import display, Markdown

threadpool_limits(1)
Path("figures").mkdir(exist_ok=True)
FAMILIES = ["OLS", "Ridge", "Logit", "Poisson", "IV-GMM", "NLS", "SCM"]
MAXIT, FW_MAXIT, SECONDS = 1000, 20000, 10.0
SMOOTH = ["GALAHAD TRU", "GALAHAD ARC", "pyensmallen L-BFGS",
          "CrabbyMetrics L-BFGS", "CrabbyMetrics BFGS"]
VERSIONS = {name: metadata.version(name) for name in
            ["galahad-optrove", "pyensmallen", "crabbymetrics", "numpy", "scipy"]}
environment = {
    "versions": VERSIONS, "python": platform.python_version(),
    "platform": platform.platform(), "machine": platform.machine(),
    "cpu": "Apple M4", "memory_bytes": 17179869184,
    "galahad_source_commit": "4275b98e4fdf9af1141d302f04ef944bcda23d32",
    "threads": [{k: v for k, v in info.items() if k != "filepath"}
                for info in threadpool_info()],
    "max_iterations": MAXIT, "fw_max_iterations": FW_MAXIT,
    "callback_time_limit_seconds": SECONDS,
}
Path("environment.json").write_text(json.dumps(environment, indent=2))
WOLFE = [0.9, 0.1]
REFERENCE = ["Ref Newton", "Ref BHHH", "Ref DFP", "Ref BFGS"]
LIKELIHOODS = ["OLS", "Ridge", "Logit", "Poisson", "NLS"]
NATIVE = "CrabbyMetrics native"
NAMES = SMOOTH + REFERENCE + ["GALAHAD QPB", "pyensmallen Frank-Wolfe", NATIVE]
COLORS = dict(zip(NAMES, ["#00838f", "#5bb6ad", "#ba5a18", "#334f96", "#9463a5",
                          "#263238", "#c34242", "#ae8400", "#467a37", "#00838f", "#ba5a18", "#764e20"]))
plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False,
                     "savefig.dpi": 180, "axes.titleweight": "semibold"})
environment.update(reference_wolfe_c2=WOLFE, reference_eigenvalue_floor=1e-5,
                   diagnostic_traces="untimed warm-up; timed fits have logging disabled")
Path("environment.json").write_text(json.dumps(environment, indent=2))

def make_problem(family, size, level, seed):
    group = 0 if family in ["OLS", "Ridge"] else FAMILIES.index(family)
    rng = np.random.default_rng(20261021 + 100 * group + 10 * size + seed)
    n, p = [(2000, 8), (20000, 24)][size]
    if family == "SCM":
        n, p = [(100, 12), (500, 50)][size]
        factors, loadings = rng.normal(size=(n, 3)), rng.normal(size=(3, p))
        X = factors @ loadings / np.sqrt(3) + [1, .1, .01][level] * rng.normal(size=(n, p))
        beta = np.r_[.2, .3, .5, np.zeros(p - 3)]
        y = X @ beta + .05 * rng.normal(size=n)
        return dict(family=family, X=X, y=y, beta=beta, x0=np.full(p, 1 / p),
                    penalty=np.zeros(p), Z=None)
    features = rng.normal(size=(n, p - 1))
    if family in ["OLS", "Ridge"]:
        Q, _ = np.linalg.qr(features - features.mean(axis=0))
        V, _ = np.linalg.qr(rng.normal(size=(p - 1, p - 1)))
        singular = np.geomspace(1, 1 / [1, 100, 10000][level], p - 1)
        features = np.sqrt(n) * (Q * singular) @ V.T
    X = np.column_stack([np.ones(n), features])
    beta = rng.normal(size=p)
    beta[1:] *= .4 / np.linalg.norm(beta[1:])
    beta[0] = .2
    Z = None
    if family == "Logit":
        beta[1:] *= [1, 4, 12][level] / .4
        y = rng.binomial(1, expit(X @ beta)).astype(float)
    elif family == "Poisson":
        scales = np.r_[1, np.geomspace(1, [1, 100, 10000][level], p - 1)]
        X *= scales
        beta /= scales
        y = rng.poisson(np.exp(X @ beta)).astype(float)
    elif family == "IV-GMM":
        Z = np.column_stack([np.ones(n), rng.normal(size=(n, p + 3))])
        v = rng.normal(size=(n, p - 1))
        strength = np.ones(p - 1)
        strength[0] = [1, .1, .01][level]
        X[:, 1:] = Z[:, 1:p] * strength + v
        y = X @ beta + .7 * v[:, 0] + np.sqrt(1 - .7**2) * rng.normal(size=n)
    elif family == "NLS":
        y = np.exp(X @ beta) + .1 * rng.normal(size=n)
    else:
        y = X @ beta + .25 * rng.normal(size=n)
    penalty = np.r_[0., np.full(p - 1, .01 if family == "Ridge" else 0.)]
    x0 = [0, 4, -6][level] * beta if family == "NLS" else np.zeros(p)
    return dict(family=family, X=np.ascontiguousarray(X), y=y, beta=beta,
                x0=x0, penalty=penalty, Z=Z)

class Criterion:
    def __init__(self, problem, trace=False, origin=None):
        self.d = problem
        self.kind, self.X, self.y = problem["family"], problem["X"], problem["y"]
        self.n, self.p = self.X.shape
        self.penalty = problem["penalty"]
        self.last = self.value = self.gradient = None
        self.nfg = self.nh = self.nopg = 0
        self.trace = [] if trace else None
        self.origin = perf_counter() if origin is None else origin
        self.scales = np.sqrt(np.mean(self.X**2, axis=0)) if trace else None
        self.deadline = float("inf")
        self.H = None
        if self.kind == "IV-GMM":
            Z = problem["Z"]
            L = np.linalg.cholesky(Z.T @ Z / self.n)
            self.A = solve_triangular(L, Z.T @ self.X / self.n, lower=True)
            self.b = solve_triangular(L, Z.T @ self.y / self.n, lower=True)
        elif self.kind in ["OLS", "Ridge", "SCM"]:
            self.A, self.b = self.X / np.sqrt(self.n), self.y / np.sqrt(self.n)
        if self.kind in ["OLS", "Ridge", "IV-GMM", "SCM"]:
            self.H = self.A.T @ self.A + np.diag(self.penalty)
            self.linear = self.A.T @ self.b
            self.constant = .5 * self.b @ self.b

    @np.errstate(over="ignore", invalid="ignore")
    def fg(self, x):
        if perf_counter() > self.deadline:
            raise TimeoutError("10-second callback budget reached")
        if self.last is None or not np.array_equal(x, self.last):
            self.nfg += 1
            if self.H is not None:
                g = self.H @ x - self.linear
                f = .5 * x @ self.H @ x - self.linear @ x + self.constant
            else:
                eta = self.X @ x
                if self.kind == "Logit":
                    mu = expit(eta)
                    f, residual = np.mean(np.logaddexp(0, eta) - self.y * eta), mu - self.y
                elif self.kind == "Poisson":
                    mu = np.exp(eta)
                    f, residual = np.mean(mu - self.y * eta), mu - self.y
                else:
                    mu = np.exp(eta)
                    f, residual = .5 * np.mean((mu - self.y)**2), mu * (mu - self.y)
                g = self.X.T @ residual / self.n
            self.last, self.value, self.gradient = x.copy(), float(f), g
            if self.trace is not None:
                self.trace.append(dict(evaluation=self.nfg, elapsed=perf_counter() - self.origin,
                                       objective=float(f), scaled_gradient=float(np.max(np.abs(g / self.scales))),
                                       fw_gap=float(g @ x - g.min()) if self.kind == "SCM" else np.nan,
                                       hessian_evaluations=self.nh, opg_evaluations=self.nopg))
        return self.value, self.gradient

    @np.errstate(over="ignore", invalid="ignore")
    def hessian(self, x):
        self.nh += 1
        if self.H is not None:
            return self.H
        eta = self.X @ x
        if self.kind == "Logit":
            mu = expit(eta)
            weights = mu * (1 - mu)
        else:
            mu = np.exp(eta)
            weights = mu if self.kind == "Poisson" else mu * (2 * mu - self.y)
        return self.X.T @ (weights[:, None] * self.X) / self.n

    def fun(self, x):
        return self.fg(x)[0]

    def grad(self, x):
        return self.fg(x)[1]

    def ensmallen(self, x, gradient):
        f, g = self.fg(x)
        gradient[:] = g
        return f

    @np.errstate(over="ignore", invalid="ignore")
    def opg(self, x):
        self.nopg += 1
        eta = self.X @ x
        if self.kind == "Logit":
            scores = (expit(eta) - self.y)[:, None] * self.X
            variance = 1.
        elif self.kind == "Poisson":
            scores = (np.exp(eta) - self.y)[:, None] * self.X
            variance = 1.
        else:
            mu = np.exp(eta) if self.kind == "NLS" else eta
            derivative = mu if self.kind == "NLS" else np.ones(self.n)
            scores = (derivative * (mu - self.y))[:, None] * self.X
            variance = .1**2 if self.kind == "NLS" else .25**2
        return scores.T @ scores / (self.n * variance) + np.diag(self.penalty)

def strong_wolfe(x, f, g, direction, c2, criterion):
    slope0, calls = float(g @ direction), 0
    assert slope0 < 0

    def evaluate(alpha):
        nonlocal calls
        calls += 1
        ft, gt = criterion.fg(x + alpha * direction)
        return ft, gt, float(gt @ direction)

    def zoom(lo, hi, flo):
        while calls < 40:
            alpha = (lo + hi) / 2
            ft, gt, slope = evaluate(alpha)
            if not np.isfinite(ft + slope) or ft > f + 1e-4 * alpha * slope0 or ft >= flo:
                hi = alpha
            else:
                if abs(slope) <= -c2 * slope0:
                    return alpha, ft, gt
                if slope * (hi - lo) >= 0:
                    hi = lo
                lo, flo = alpha, ft
            if abs(hi - lo) < 1e-16:
                break
        raise RuntimeError("line search failed")

    previous, previous_f, alpha = 0., f, 1.
    while calls < 40:
        ft, gt, slope = evaluate(alpha)
        if (not np.isfinite(ft + slope) or ft > f + 1e-4 * alpha * slope0
                or (calls > 1 and ft >= previous_f)):
            return zoom(previous, alpha, previous_f)
        if abs(slope) <= -c2 * slope0:
            return alpha, ft, gt
        if slope >= 0:
            return zoom(alpha, previous, ft)
        previous, previous_f, alpha = alpha, ft, 2 * alpha
    raise RuntimeError("line search failed")


def inverse_update(M, step, change, method):
    sy = float(step @ change)
    if method == "Ref BFGS":
        A = np.eye(len(step)) - np.outer(step, change) / sy
        updated = A @ M @ A.T + np.outer(step, step) / sy
    else:
        My = M @ change
        updated = M + np.outer(step, step) / sy - np.outer(My, My) / (change @ My)
    return (updated + updated.T) / 2


def solve_reference(d, method, c2, tracing=False):
    start = perf_counter()
    c = Criterion(d, trace=tracing, origin=start)
    prepared = perf_counter()
    c.deadline = start + SECONDS
    x, M = d["x0"].copy(), np.eye(c.p)
    iterations = clipped = skipped = resets = 0
    error, status = "", "iteration limit"
    try:
        f, g = c.fg(x)
        for _ in range(MAXIT):
            if np.max(np.abs(g)) <= 1e-8:
                status = "gradient tolerance"
                break
            if method in ["Ref Newton", "Ref BHHH"]:
                C = c.hessian(x) if method == "Ref Newton" else c.opg(x)
                eigenvalues, Q = np.linalg.eigh(C)
                clipped += int(eigenvalues.min() < 1e-5)
                direction = -Q @ ((Q.T @ g) / np.maximum(eigenvalues, 1e-5))
            else:
                direction = -M @ g
            if not np.all(np.isfinite(direction)) or g @ direction >= 0:
                direction, M = -g, np.eye(c.p)
                resets += 1
            try:
                alpha, fn, gn = strong_wolfe(x, f, g, direction, c2, c)
            except RuntimeError:
                status = "line search failed"
                break
            assert fn <= f + 1e-4 * alpha * float(g @ direction)
            assert abs(float(gn @ direction)) <= c2 * abs(float(g @ direction))
            step, change = alpha * direction, gn - g
            if method in ["Ref DFP", "Ref BFGS"]:
                if step @ change > 1e-12 * np.linalg.norm(step) * np.linalg.norm(change):
                    M = inverse_update(M, step, change, method)
                else:
                    skipped += 1
            x, f, g = x + step, fn, gn
            iterations += 1
        if np.max(np.abs(g)) <= 1e-8:
            status = "gradient tolerance"
    except Exception as exc:
        error, status = str(exc), "exception"
    return x, dict(seconds=perf_counter() - start, preparation_seconds=prepared - start,
                   status=status, iterations=iterations, error=error,
                   fg_evaluations=c.nfg, hessian_evaluations=c.nh, opg_evaluations=c.nopg,
                   clipped=clipped, skipped=skipped, resets=resets), c.trace

def solve_shared(d, name, tracing=False):
    start = perf_counter()
    c = Criterion(d, trace=tracing, origin=start)
    prepared = perf_counter()
    c.deadline = start + SECONDS
    x = d["x0"].copy()
    status, iterations, error = "returned", np.nan, ""
    module = None
    try:
        if name in ["GALAHAD TRU", "GALAHAD ARC"]:
            module = tru if name.endswith("TRU") else arc
            options = module.initialize()
            options.update(print_level=0, maxit=MAXIT, stop_g_absolute=1e-8,
                           stop_g_relative=0., norm=-1, subproblem_direct=True,
                           clock_time_limit=SECONDS)
            lower = np.tril_indices(c.p)
            ne = len(lower[0])
            module.load(c.p, "dense", ne, None, None, None, options=options)
            x, _ = module.solve(c.p, ne, x, c.fun, c.grad,
                                lambda x: c.hessian(x)[lower])
            info = module.information()
            status, iterations = str(info["status"]), info["iter"]
        elif name == "pyensmallen L-BFGS":
            solver = pe.L_BFGS()
            solver.numBasis, solver.maxIterations = 7, MAXIT
            solver.minGradientNorm, solver.factr = 1e-8, 1e-15
            x = solver.optimize(c.ensmallen, x)
        elif name.startswith("CrabbyMetrics"):
            method = cm.Optimizers.minimize_lbfgs if name.endswith("L-BFGS") else cm.Optimizers.minimize_bfgs
            result = method(c.fun, x, c.grad, max_iterations=MAXIT, tolerance=1e-12)
            x, status, iterations = result["x"], result["message"], result["nit"]
        elif name == "GALAHAD QPB":
            module = qpb
            options = module.initialize()
            options.update(print_level=0, maxit=MAXIT, stop_p=1e-10, stop_d=1e-10, stop_c=1e-10)
            lower = np.tril_indices(c.p)
            ne = len(lower[0])
            module.load(c.p, 1, "dense", ne, None, None, None,
                        "dense", c.p, None, None, None, options)
            x, *_ = module.solve_qp(c.p, 1, c.constant, -c.linear, ne, c.H[lower],
                                   c.p, np.ones(c.p), np.ones(1), np.ones(1),
                                   np.zeros(c.p), np.ones(c.p), x, np.zeros(1), np.zeros(c.p))
            info = module.information()
            status, iterations = str(info["status"]), info["iter"]
        else:
            solver = pe.SimplexFrankWolfe(maxIterations=FW_MAXIT, tolerance=1e-10)
            x = solver.optimize(c.ensmallen, x)
    except Exception as exc:
        error, status = str(exc), "exception"
    finally:
        if module is not None:
            module.terminate()
    return x, dict(seconds=perf_counter() - start, preparation_seconds=prepared - start,
                   status=status, iterations=iterations, error=error,
                   fg_evaluations=c.nfg, hessian_evaluations=c.nh, opg_evaluations=c.nopg), c.trace

def solve_native(d):
    start = perf_counter()
    family, X, y = d["family"], d["X"], d["y"]
    constructors = {
        "OLS": lambda: cm.OLS(), "Ridge": lambda: cm.Ridge(penalty=len(y) * .01),
        "Logit": lambda: cm.Logit(max_iterations=MAXIT, gradient_tolerance=1e-8),
        "Poisson": lambda: cm.Poisson(max_iterations=MAXIT, tolerance=1e-10),
        "IV-GMM": lambda: cm.TwoSLS(), "SCM": lambda: cm.SyntheticControl(max_iterations=MAXIT),
    }
    try:
        model = constructors[family]()
        if family == "SCM":
            model.fit(X, y)
        elif family == "IV-GMM":
            model.fit(X[:, 1:], np.empty((len(y), 0)), d["Z"][:, 1:], y)
        else:
            model.fit(X[:, 1:], y.astype(np.int32) if family == "Logit" else y)
        elapsed = perf_counter() - start
        summary = model.summary()
        x = summary["weights"] if family == "SCM" else np.r_[summary["intercept"], summary["coef"]]
        return x, dict(seconds=elapsed, preparation_seconds=np.nan,
                       iterations=summary.get("iterations", np.nan), status=str(summary.get("converged", "direct solve")),
                       error="", fg_evaluations=np.nan, hessian_evaluations=np.nan, opg_evaluations=np.nan), None
    except Exception as exc:
        return d["x0"].copy(), dict(seconds=perf_counter() - start, preparation_seconds=np.nan,
                                    iterations=np.nan, status="exception", error=str(exc),
                                    fg_evaluations=np.nan, hessian_evaluations=np.nan, opg_evaluations=np.nan), None

def reference_solution(d):
    c = Criterion(d)
    scales = np.sqrt(np.mean(d["X"]**2, axis=0))
    if d["family"] in ["OLS", "Ridge", "IV-GMM"]:
        A = np.vstack([c.A, np.diag(np.sqrt(d["penalty"]))])
        b = np.r_[c.b, np.zeros(c.p)]
        x = np.linalg.lstsq(A, b, rcond=None)[0]
    elif d["family"] == "SCM":
        result = minimize(c.fun, d["x0"], jac=c.grad, method="SLSQP",
                          bounds=[(0, 1)] * c.p,
                          constraints={"type": "eq", "fun": lambda x: x.sum() - 1,
                                       "jac": lambda x: np.ones(c.p)},
                          options={"ftol": 1e-14, "maxiter": 3000})
        x = result.x
    elif d["family"] == "NLS":
        X, y = d["X"], d["y"]
        fits = [least_squares(lambda x: (np.exp(X @ x) - y) / np.sqrt(c.n), start,
                             jac=lambda x: np.exp(X @ x)[:, None] * X / np.sqrt(c.n),
                             ftol=1e-12, xtol=1e-12, gtol=1e-12, max_nfev=2000)
                for start in [np.zeros(c.p), d["beta"], d["x0"]]]
        x = min(fits, key=lambda fit: np.dot(fit.fun, fit.fun)).x
    else:
        result = minimize(lambda z: c.fun(z / scales), np.zeros(c.p),
                          jac=lambda z: c.grad(z / scales) / scales,
                          hess=lambda z: c.hessian(z / scales) / np.outer(scales, scales),
                          method="trust-exact", options={"gtol": 1e-11, "maxiter": 1000})
        x = result.x / scales
    return x

@np.errstate(over="ignore", invalid="ignore")
def evaluate_solution(d, x, ref):
    c = Criterion(d)
    f, g = c.fg(x)
    fref = c.fun(ref)
    # Score quadratic fits from residuals, independently of the Gram-form callback.
    if c.H is not None:
        residual, reference_residual = c.A @ x - c.b, c.A @ ref - c.b
        f = .5 * (residual @ residual + np.dot(c.penalty, x**2))
        fref = .5 * (reference_residual @ reference_residual + np.dot(c.penalty, ref**2))
        g = c.A.T @ residual + c.penalty * x
    scale = np.sqrt(np.mean(d["X"]**2, axis=0))
    score = np.max(np.abs(g / scale))
    gap = (f - fref) / max(1, abs(fref))
    feasibility = max(0., -float(x.min()), abs(float(x.sum()) - 1)) if c.kind == "SCM" else 0.
    fw_gap = float(g @ x - g.min()) if c.kind == "SCM" else np.nan
    passed = (fw_gap <= 1e-7 and feasibility <= 1e-8) if c.kind == "SCM" else (gap <= 1e-8 and score <= 1e-6)
    passed = bool(passed and np.isfinite(f) and np.isfinite(score) and np.all(np.isfinite(x)))
    eta, eta_ref = d["X"] @ x, d["X"] @ ref
    transform = expit if c.kind == "Logit" else np.exp if c.kind in ["Poisson", "NLS"] else lambda a: a
    pred, predref = transform(eta), transform(eta_ref)
    return dict(objective=f, objective_gap=gap, scaled_gradient=score,
                fw_gap=fw_gap, feasibility=feasibility, accepted=passed,
                coefficient_error=np.linalg.norm(x - ref) / max(1, np.linalg.norm(ref)),
                prediction_error=np.sqrt(np.mean((pred - predref)**2)) / max(1, np.sqrt(np.mean(predref**2))))

derivative_checks = []
for family in FAMILIES:
    d = make_problem(family, 0, 0, 0)
    c = Criterion(d)
    x = d["beta"].copy()
    step = 1e-5
    directions = np.eye(len(x))
    numeric_g = np.array([(c.fun(x + step * v) - c.fun(x - step * v)) / (2 * step) for v in directions])
    numeric_h = np.column_stack([(c.grad(x + step * v) - c.grad(x - step * v)) / (2 * step) for v in directions])
    g, H = c.grad(x), c.hessian(x)
    ge = np.max(np.abs(g - numeric_g)) / max(1, np.max(np.abs(g)))
    he = np.max(np.abs(H - numeric_h)) / max(1, np.max(np.abs(H)))
    assert ge < 1e-6 and he < 1e-6, (family, ge, he)
    ref = reference_solution(d)
    assert evaluate_solution(d, ref, ref)["accepted"], family
    derivative_checks.append(dict(family=family, gradient_error=ge, hessian_error=he))
display(pd.DataFrame(derivative_checks).style.hide(axis="index").format(
    {"gradient_error": "{:.2e}", "hessian_error": "{:.2e}"}))

# Verify likelihood-score curvature against finite differences of individual losses.
opg_checks = []
for family in LIKELIHOODS:
    d = make_problem(family, 0, 0, 0)
    d = dict(d, X=d["X"][:32], y=d["y"][:32])
    x = d["beta"]
    variance = .1**2 if family == "NLS" else .25**2 if family in ["OLS", "Ridge"] else 1.
    def individual_losses(b):
        eta = d["X"] @ b
        if family == "Logit":
            return np.logaddexp(0, eta) - d["y"] * eta
        if family == "Poisson":
            return np.exp(eta) - d["y"] * eta
        mu = np.exp(eta) if family == "NLS" else eta
        return .5 * (mu - d["y"])**2 / variance
    scores = np.column_stack([(individual_losses(x + 1e-5 * direction) -
                                individual_losses(x - 1e-5 * direction)) / 2e-5
                               for direction in np.eye(len(x))])
    numeric = variance * scores.T @ scores / len(d["y"]) + np.diag(d["penalty"])
    analytic = Criterion(d).opg(x)
    relative = np.max(np.abs(numeric - analytic)) / max(1, np.max(np.abs(analytic)))
    assert relative < 1e-6, (family, relative)
    opg_checks.append(dict(family=family, relative_error=float(relative)))
for method in ["Ref DFP", "Ref BFGS"]:
    M = np.diag([1., 2., 3.])
    step, change = np.array([1., -.5, .2]), np.array([2., -.2, .7])
    updated = inverse_update(M, step, change, method)
    assert np.allclose(updated @ change, step) and np.linalg.eigvalsh(updated).min() > 0

# Full-rank logit designs have finite MLEs if no complete/quasi-separating direction exists.
# Feasibility of S v >= 0 and 1' S v >= 1 would exhibit such a direction.
for size in range(2):
    for level in range(3):
        for seed in range(2):
            d = make_problem("Logit", size, level, seed)
            S = (2 * d["y"] - 1)[:, None] * d["X"]
            check = linprog(np.zeros(S.shape[1]), A_ub=np.vstack([-S, -S.sum(axis=0)]),
                            b_ub=np.r_[np.zeros(len(S)), -1.],
                            bounds=[(None, None)] * S.shape[1], method="highs")
            assert np.linalg.matrix_rank(S) == S.shape[1] and check.status == 2
display(Markdown("**Derivative checks pass for all seven families. All 12 logit samples "
                 "also pass a linear-program check excluding complete or quasi-separation.**"))

def methods_for(family):
    if family == "SCM":
        return [("GALAHAD QPB", "native"), ("pyensmallen Frank-Wolfe", "native"), (NATIVE, "native")]
    methods = [(name, "native") for name in SMOOTH]
    for c2 in WOLFE:
        methods += [(name, f"Wolfe {c2}") for name in REFERENCE
                    if name != "Ref BHHH" or family in LIKELIHOODS]
    if family != "NLS":
        methods += [(NATIVE, "native")]
    return methods

rows, diagnostic_rows, problem_rows, trace_rows, solutions = [], [], [], [], {}
expected_pairs = sum(12 * len(methods_for(family)) for family in FAMILIES)
assert expected_pairs == 1008
for family in FAMILIES:
    for size in range(2):
        for level in range(3):
            for seed in range(2):
                d = make_problem(family, size, level, seed)
                case = f"{family}-{size}-{level}-{seed}"
                ref = reference_solution(d)
                reference_check = evaluate_solution(d, ref, ref)
                assert reference_check["accepted"], (case, reference_check)
                c = Criterion(d)
                H, H0 = c.hessian(ref), c.hessian(d["x0"])
                eigen = np.linalg.eigvalsh(H)
                assert eigen[0] > 0, case
                digest = hashlib.sha256()
                for a in [d["X"], d["y"], d["beta"], d["x0"], d["penalty"], d["Z"]]:
                    if a is not None:
                        digest.update(np.ascontiguousarray(a).tobytes())
                problem_rows.append(dict(case=case, family=family, size=size, level=level, seed=seed,
                                         n=len(d["y"]), p=len(ref), condition=np.linalg.cond(H),
                                         min_eigenvalue=eigen[0], initial_min_eigenvalue=np.linalg.eigvalsh(H0)[0],
                                         reference_objective=c.fun(ref), reference_score=reference_check["scaled_gradient"],
                                         reference_fw_gap=reference_check["fw_gap"], data_sha256=digest.hexdigest()))
                solutions[case + "_reference"] = ref
                methods = methods_for(family)
                for repetition in [-1, 0, 1, 2]:
                    rng = np.random.default_rng(7800 + 100 * size + 10 * level + seed + 1000 * (repetition + 1))
                    for index in rng.permutation(len(methods)):
                        method, setting = methods[index]
                        diagnostic = repetition == -1
                        if method.startswith("Ref "):
                            x, result, history = solve_reference(d, method, float(setting.split()[-1]), diagnostic)
                        elif method == NATIVE:
                            x, result, history = solve_native(d)
                        else:
                            x, result, history = solve_shared(d, method, diagnostic)
                        quality = evaluate_solution(d, x, ref)
                        if result["error"]:
                            quality["accepted"] = False
                        fit_id = f"{case}|{method}|{setting}|{repetition}"
                        row = dict(fit_id=fit_id, case=case, family=family, size=size, level=level, seed=seed,
                                   method=method, setting=setting, repetition=repetition, **result, **quality)
                        solutions[fit_id] = np.asarray(x).copy()
                        if diagnostic:
                            diagnostic_rows.append(row)
                            if history:
                                trace_rows.extend([dict(fit_id=fit_id, case=case, family=family,
                                                        method=method, setting=setting, **point) for point in history])
                        else:
                            rows.append(row)
                    pd.DataFrame(rows).to_csv("results.csv", index=False)
                pd.DataFrame(diagnostic_rows).to_csv("diagnostics.csv", index=False)
                pd.DataFrame(trace_rows).to_csv("traces.csv.gz", index=False, compression="gzip")
                print(f"Completed {case}: {len(rows)}/3024 timed fits", flush=True)
results, problems = pd.DataFrame(rows), pd.DataFrame(problem_rows)
diagnostics, traces = pd.DataFrame(diagnostic_rows), pd.DataFrame(trace_rows)
assert len(results) == 3024 and len(diagnostics) == expected_pairs
problems.to_csv("problems.csv", index=False)
np.savez_compressed("solutions.npz", **solutions)

def independent_score(d, x, ref):
    X, y, family = d["X"], d["y"], d["family"]
    n = len(y)
    eta, eta_ref = X @ x, X @ ref
    if family in ["OLS", "Ridge", "SCM"]:
        r, rr = eta - y, eta_ref - y
        f = .5 * (r @ r / n + np.dot(d["penalty"], x**2))
        fref = .5 * (rr @ rr / n + np.dot(d["penalty"], ref**2))
        g = X.T @ r / n + d["penalty"] * x
    elif family == "IV-GMM":
        Z = d["Z"]
        moment, ref_moment = Z.T @ (eta - y) / n, Z.T @ (eta_ref - y) / n
        weighted = np.linalg.solve(Z.T @ Z / n, moment)
        f = .5 * moment @ weighted
        fref = .5 * ref_moment @ np.linalg.solve(Z.T @ Z / n, ref_moment)
        g = X.T @ Z @ weighted / n
    else:
        if family == "Logit":
            mu, muref = expit(eta), expit(eta_ref)
            f, fref = np.mean(np.logaddexp(0, eta) - y * eta), np.mean(np.logaddexp(0, eta_ref) - y * eta_ref)
            residual = mu - y
        elif family == "Poisson":
            mu, muref = np.exp(eta), np.exp(eta_ref)
            f, fref = np.mean(mu - y * eta), np.mean(muref - y * eta_ref)
            residual = mu - y
        else:
            mu, muref = np.exp(eta), np.exp(eta_ref)
            f, fref = .5 * np.mean((mu - y)**2), .5 * np.mean((muref - y)**2)
            residual = mu * (mu - y)
        g = X.T @ residual / n
    scaled = float(np.max(np.abs(g) / np.sqrt(np.mean(X**2, axis=0))))
    gap = float((f - fref) / max(1, abs(fref)))
    feasible = max(0., -float(x.min()), abs(float(x.sum()) - 1)) if family == "SCM" else 0.
    fw_gap = float(g @ x - g.min()) if family == "SCM" else np.nan
    passed = (fw_gap <= 1e-7 and feasible <= 1e-8) if family == "SCM" else (gap <= 1e-8 and scaled <= 1e-6)
    return dict(objective=float(f), scaled_gradient=scaled, objective_gap=gap,
                fw_gap=fw_gap, feasibility=feasible,
                accepted=bool(passed and np.isfinite(f) and np.isfinite(scaled) and np.all(np.isfinite(x))))


def recovery(d, x, ref):
    scale = np.sqrt(np.mean(d["X"]**2, axis=0)) if d["family"] != "SCM" else np.ones(len(x))
    transform = expit if d["family"] == "Logit" else np.exp if d["family"] in ["Poisson", "NLS"] else lambda a: a
    pred, truth, predref = [transform(d["X"] @ b) for b in [x, d["beta"], ref]]
    return dict(parameter_rmse=float(np.sqrt(np.mean((scale * (x - d["beta"]))**2))),
                raw_parameter_rmse=float(np.sqrt(np.mean((x - d["beta"])**2))),
                prediction_rmse=float(np.sqrt(np.mean((pred - truth)**2))),
                numerical_parameter_rmse=float(np.sqrt(np.mean((scale * (x - ref))**2))),
                numerical_prediction_rmse=float(np.sqrt(np.mean((pred - predref)**2))))

audit_rows, recovery_rows, reference_recovery = [], [], []
for problem in problems.to_dict("records"):
    case = problem["case"]
    d = make_problem(problem["family"], problem["size"], problem["level"], problem["seed"])
    ref = solutions[case + "_reference"]
    reference_recovery.append(dict(case=case, family=d["family"], size=problem["size"],
                                   level=problem["level"], seed=problem["seed"], **recovery(d, ref, ref)))
    for frame in [results[results.case == case], diagnostics[diagnostics.case == case]]:
        for row in frame.to_dict("records"):
            x = solutions[row["fit_id"]]
            scored = independent_score(d, x, ref)
            scored["accepted"] = scored["accepted"] and not row["error"]
            assert scored["accepted"] == row["accepted"], (row["fit_id"], scored, row)
            differences = {key: abs(scored[key] - row[key]) for key in ["objective", "scaled_gradient", "objective_gap"]}
            assert differences["objective"] <= 1e-9 * max(1, abs(scored["objective"])), (row["fit_id"], differences)
            assert differences["scaled_gradient"] <= 1e-8 * max(1, scored["scaled_gradient"]), (row["fit_id"], differences)
            audit_rows.append(dict(fit_id=row["fit_id"], repetition=row["repetition"], accepted=scored["accepted"], **differences))
            recovery_rows.append(dict(fit_id=row["fit_id"], case=case, family=d["family"], method=row["method"],
                                      setting=row["setting"], repetition=row["repetition"], level=row["level"],
                                      size=row["size"], seed=row["seed"], accepted=row["accepted"], **recovery(d, x, ref)))
audit = pd.DataFrame(audit_rows)
recoveries = pd.DataFrame(recovery_rows)
reference_recoveries = pd.DataFrame(reference_recovery)
audit.to_csv("independent-audit.csv", index=False)
recoveries.to_csv("recovery.csv", index=False)
reference_recoveries.to_csv("reference-recovery.csv", index=False)
matched_controls = []
for size in range(2):
    for seed in range(2):
        for level in range(3):
            a, b = make_problem("OLS", size, level, seed), make_problem("Ridge", size, level, seed)
            assert np.array_equal(a["X"], b["X"]) and np.array_equal(a["y"], b["y"])
        poisson = [make_problem("Poisson", size, level, seed) for level in range(3)]
        nls = [make_problem("NLS", size, level, seed) for level in range(3)]
        for level in [1, 2]:
            assert np.array_equal(poisson[0]["y"], poisson[level]["y"])
            assert np.allclose(poisson[0]["X"] @ poisson[0]["beta"], poisson[level]["X"] @ poisson[level]["beta"])
            assert np.array_equal(nls[0]["X"], nls[level]["X"]) and np.array_equal(nls[0]["y"], nls[level]["y"])
        matched_controls.append(dict(size=size, seed=seed, passed=True))
validation = dict(instances=len(problems), timed_fits=len(results), diagnostic_fits=len(diagnostics),
                  independent_rescores=len(audit), all_acceptance_decisions_agree=True,
                  matched_controls=matched_controls, derivative_checks=derivative_checks, opg_checks=opg_checks,
                  maximum_objective_difference=float(audit.objective.max()),
                  maximum_scaled_gradient_difference=float(audit.scaled_gradient.max()))
Path("validation.json").write_text(json.dumps(validation, indent=2))

paired = results.groupby(["case", "family", "size", "level", "seed", "method", "setting"], as_index=False).agg(
    accepted=("accepted", "all"), seconds=("seconds", "median"),
    minimum_seconds=("seconds", "min"), maximum_seconds=("seconds", "max"),
    fg_evaluations=("fg_evaluations", "median"), hessian_evaluations=("hessian_evaluations", "median"),
    opg_evaluations=("opg_evaluations", "median"), objective_gap=("objective_gap", "max"),
    scaled_gradient=("scaled_gradient", "max"), fw_gap=("fw_gap", "max"),
    coefficient_error=("coefficient_error", "max"), prediction_error=("prediction_error", "max"))
paired["accepted_ms"] = (1000 * paired.seconds).where(paired.accepted)
summary = paired.groupby(["family", "method", "setting"], as_index=False).agg(
    solved=("accepted", "sum"), attempted=("accepted", "size"),
    median_accepted_ms=("accepted_ms", "median"), worst_gradient=("scaled_gradient", "max"))
paired.to_csv("paired-results.csv", index=False)
summary.to_csv("summary.csv", index=False)
smooth = paired[paired.family != "SCM"]
coverage = smooth.groupby(["method", "setting"], as_index=False).agg(
    solved=("accepted", "sum"), attempted=("accepted", "size"), median_accepted_ms=("accepted_ms", "median"))
display(Markdown(f"**{len(problems)} instances; {len(paired):,} solver–instance–setting combinations; "
                 f"{len(results):,} freshly timed fits and {len(diagnostics):,} diagnostic warm-ups.** "
                 f"{int(results.accepted.sum()):,}/{len(results):,} timed fits pass. An instance is accepted "
                 "for a method and setting only when all three measured repeats pass."))
display(coverage.style.hide(axis="index").format({"median_accepted_ms": "{:.3f}"}))