benchmark_qr.ipynb

solvers.py

import numpy as np
from ortools.linear_solver import pywraplp
from scipy.optimize import linprog
from sklearn.linear_model import QuantileRegressor

from statsmodels.regression.quantile_regression import QuantReg

# %%

def solve_quantile_regression_sm(
    X, y, tau
):
    m = QuantReg(y, X)
    return m.fit(q=tau).params



def solve_quantile_regression_scikit(
    X,
    y,
    tau,
):
    """
    Solves the quantile regression problem using Scikit-learn's QuantileRegressor.
    """
    model = QuantileRegressor(quantile=tau, alpha=0)
    model.fit(X, y)
    return model.coef_


# %%
######################################################################
# solving Qreg using google's ORTools library
# https://github.com/google/or-tools
# C++ w python bindings
def solve_quantile_regression_ortools(
    X,
    y,
    tau,
    solver_id,
):
    """
    Solves the quantile regression LP using Google OR-Tools.

    Args:
        X (np.ndarray): Design matrix.
        y (np.ndarray): Target vector.
        tau (float): Quantile.
        solver_id (str): The solver to use ('GLOP' or 'PDLP').

    Returns:
        np.ndarray: The estimated quantile regression coefficients.
    """
    n_samples, n_features = X.shape
    solver = pywraplp.Solver.CreateSolver(solver_id)
    if not solver:
        raise RuntimeError(f"Solver {solver_id} not available.")

    # 1. Define variables: beta = beta_plus - beta_minus, and residuals = u - v
    beta_plus = [
        solver.NumVar(0, solver.infinity(), f"bp_{j}") for j in range(n_features)
    ]
    beta_minus = [
        solver.NumVar(0, solver.infinity(), f"bm_{j}") for j in range(n_features)
    ]
    u = [solver.NumVar(0, solver.infinity(), f"u_{i}") for i in range(n_samples)]
    v = [solver.NumVar(0, solver.infinity(), f"v_{i}") for i in range(n_samples)]

    # 2. Define Objective Function: min(tau * sum(u) + (1-tau) * sum(v))
    objective = solver.Objective()
    for i in range(n_samples):
        objective.SetCoefficient(u[i], tau)
        objective.SetCoefficient(v[i], (1 - tau))
    objective.SetMinimization()

    # 3. Define Constraints: X*beta_plus - X*beta_minus + u - v = y
    for i in range(n_samples):
        constraint = solver.Constraint(
            y[i], y[i], f"ct_{i}"
        )  # Lower and upper bound are the same for equality
        for j in range(n_features):
            constraint.SetCoefficient(beta_plus[j], X[i, j])
            constraint.SetCoefficient(beta_minus[j], -X[i, j])
        constraint.SetCoefficient(u[i], 1)
        constraint.SetCoefficient(v[i], -1)

    # 4. Solve the LP
    status = solver.Solve()

    # 5. Extract results
    if status == pywraplp.Solver.OPTIMAL:
        bp_sol = np.array([var.solution_value() for var in beta_plus])
        bm_sol = np.array([var.solution_value() for var in beta_minus])
        return bp_sol - bm_sol
    else:
        raise RuntimeError(f"OR-Tools solver {solver_id} failed with status: {status}")


# %%


# -----------------------------------------------------------------------------
# Method 1: Standard Linear Programming using SciPy
# -----------------------------------------------------------------------------
def solve_quantile_regression_lp(
    X,
    y,
    tau=0.5,
):
    """
    Solves quantile regression using a standard LP solver.
    """
    n_samples, n_features = X.shape

    # Solution Variables: [beta_plus, beta_minus, u, v]
    c = np.concatenate(
        [
            np.zeros(2 * n_features),
            tau * np.ones(n_samples),
            (1 - tau) * np.ones(n_samples),
        ]
    )
    # A
    A_eq = np.hstack([X, -X, np.eye(n_samples), -np.eye(n_samples)])
    b_eq = y
    # A c = b

    bounds = [(0, None) for _ in range(len(c))]

    result = linprog(c, A_eq=A_eq, b_eq=b_eq, bounds=bounds, method="highs")

    if result.success:
        beta_plus = result.x[:n_features]
        beta_minus = result.x[n_features : 2 * n_features]
        return beta_plus - beta_minus
    else:
        raise RuntimeError("LP solver failed: " + result.message)


# %%
######################################################################
# -----------------------------------------------------------------------------
# Method 3: Frisch-Newton Interior Point Method (as implemented in pyfixest)
# -----------------------------------------------------------------------------
######################################################################
def _duality_gap(x, z, s, w):
    return x @ z + s @ w


def _bound(v: np.ndarray, dv: np.ndarray, backoff: float):
    mask = dv < 0
    if not mask.any():
        return 1.0
    alpha_max = (-v[mask] / dv[mask]).min()
    return min(backoff * alpha_max, 1.0)


def _step_length(a: tuple, b: tuple, backoff: float) -> float:
    x, dx = a
    s, ds = b
    return min(_bound(x, dx, backoff), _bound(s, ds, backoff))


def cold_start(A: np.ndarray, c: np.ndarray, q: float):
    n = A.shape[1]
    x = np.full(n, 1.0 - q)
    s = np.full_like(x, q)
    d_plus = np.maximum(c, 0.0)
    d_minus = np.maximum(-c, 0.0)
    U = x @ d_plus + s @ d_minus
    mu0 = max(1, U / n)
    alpha = (n * mu0 - U) / (np.sum(1 / x) + np.sum(1 / s))
    eps = 1e-8
    z = np.maximum(d_plus, eps) + alpha / x
    w = np.maximum(d_minus, eps) + alpha / s
    y = np.linalg.solve(A @ A.T, A @ (c - z + w))
    return x, s, z, w, y


def frisch_newton_solver(
    A: np.ndarray,
    b: np.ndarray,
    c: np.ndarray,
    u: np.ndarray,
    q: float,
    tol: float,
    max_iter: int,
    backoff: float = 0.9995,
):
    m, n = A.shape
    c, b, u = c.ravel(), b.ravel(), u.ravel()

    x, s, z, w, y = cold_start(A=A, c=c, q=q)

    mu_curr = _duality_gap(x=x, z=z, s=s, w=w)

    for _it in range(max_iter):
        if mu_curr < tol:
            break

        r1_tilde = c - A.T @ y
        r2_tilde = b - A @ x

        Qinv = 1.0 / (z / x + w / s)
        M = A @ (Qinv[:, None] * A.T)
        work_m = r2_tilde + A @ (Qinv * r1_tilde)
        dy_aff = np.linalg.solve(M, work_m)

        dx_aff = Qinv * (A.T @ dy_aff - r1_tilde)
        ds_aff = -dx_aff
        dz_aff = -z - (z / x) * dx_aff
        dw_aff = -w - (w / s) * ds_aff

        alpha_p_aff = _step_length(a=(x, dx_aff), b=(s, ds_aff), backoff=backoff)
        alpha_d_aff = _step_length(a=(z, dz_aff), b=(w, dw_aff), backoff=backoff)

        x_pred = x + alpha_p_aff * dx_aff
        s_pred = s + alpha_p_aff * ds_aff
        y_pred = y + alpha_d_aff * dy_aff
        z_pred = z + alpha_d_aff * dz_aff
        w_pred = w + alpha_d_aff * dw_aff

        mu_aff = _duality_gap(x=x_pred, z=z_pred, s=s_pred, w=w_pred)

        sigma = (mu_aff / mu_curr) ** 2
        mu_targ = sigma * mu_curr / n

        r1_hat = (
            mu_targ * (1 / s - 1 / x) + (dx_aff * dz_aff) / x - (ds_aff * dw_aff) / s
        )
        work_m = A @ (Qinv * r1_hat)
        dy_cor = np.linalg.solve(M, work_m)
        dx_cor = Qinv * (A.T @ dy_cor - r1_hat)
        ds_cor = -dx_cor
        dz_cor = -(z / x) * dx_cor + (mu_targ - dx_aff * dz_aff) / x
        dw_cor = -(w / s) * ds_cor + (mu_targ - ds_aff * dw_aff) / s

        alpha_p_cor = _step_length(
            a=(x_pred, dx_cor), b=(s_pred, ds_cor), backoff=backoff
        )
        alpha_d_cor = _step_length(
            a=(z_pred, dz_cor), b=(w_pred, dw_cor), backoff=backoff
        )

        x = x_pred + alpha_p_cor * dx_cor
        s = s_pred + alpha_p_cor * ds_cor
        y = y_pred + alpha_d_cor * dy_cor
        z = z_pred + alpha_d_cor * dz_cor
        w = w_pred + alpha_d_cor * dw_cor

        mu_curr = _duality_gap(x=x, z=z, s=s, w=w)

    return -y

# %%