herculesabqp 0.1.2

use ndarray::Array1;

use super::types::Diagnostics;

pub(crate) type DenseVec = Array1<f64>;

// Evaluate the quadratic objective from `x` and cached `Qx`.
pub(crate) fn objective_from_qx(c: &DenseVec, x: &DenseVec, qx: &DenseVec) -> f64 {
    0.5 * x.dot(qx) + c.dot(x)
}

// Build the internal convergence diagnostics from cached first-order information.
pub(crate) fn compute_diagnostics_from_qx(
    qx: &DenseVec,
    c: &DenseVec,
    x: &DenseVec,
    lb: &DenseVec,
    ub: &DenseVec,
    bound_tol: f64,
    dual_certification: bool,
) -> Diagnostics {
    let objective = objective_from_qx(c, x, qx);
    let (kkt_inf, gap, rel_gap) =
        fused_box_qp_diagnostics(qx, c, x, lb, ub, bound_tol, dual_certification);
    let c_scale_inf = norm_inf(c).max(1.0);
    Diagnostics {
        kkt_inf,
        rel_kkt: kkt_inf / c_scale_inf,
        gap,
        rel_gap,
        certified_lower_bound: if dual_certification {
            objective - gap
        } else {
            f64::NAN
        },
    }
}

// Compute primal-gap and box-KKT diagnostics in one scan to avoid rebuilding the gradient twice.
fn fused_box_qp_diagnostics(
    qx: &DenseVec,
    c: &DenseVec,
    x: &DenseVec,
    lb: &DenseVec,
    ub: &DenseVec,
    bound_tol: f64,
    dual_certification: bool,
) -> (f64, f64, f64) {
    let mut kkt_inf: f64 = 0.0;
    let mut gap = if dual_certification { 0.0 } else { f64::NAN };
    for i in 0..x.len() {
        let g_i = qx[i] + c[i];
        let r = if ub[i] <= lb[i] + bound_tol {
            0.0
        } else if x[i] > lb[i] + bound_tol && x[i] < ub[i] - bound_tol {
            g_i
        } else if x[i] <= lb[i] + bound_tol {
            g_i.min(0.0)
        } else {
            g_i.max(0.0)
        };
        kkt_inf = kkt_inf.max(r.abs());
        if dual_certification {
            let alpha = g_i.max(0.0);
            let beta = (-g_i).max(0.0);
            gap += alpha * (x[i] - lb[i]) + beta * (ub[i] - x[i]);
        }
    }
    let rel_gap = if dual_certification {
        gap / objective_from_qx(c, x, qx).abs().max(1.0)
    } else {
        f64::NAN
    };
    (kkt_inf, gap, rel_gap)
}

// Project a dense vector onto elementwise lower and upper bounds.
pub(crate) fn clip_in_place(x: &mut DenseVec, lb: &DenseVec, ub: &DenseVec) {
    for i in 0..x.len() {
        x[i] = x[i].clamp(lb[i], ub[i]);
    }
}

// Compute the infinity norm of a dense vector.
pub(crate) fn norm_inf(x: &DenseVec) -> f64 {
    x.iter().fold(0.0_f64, |best, value| best.max(value.abs()))
}

// Copy a slice into the solver's dense-vector representation.
// Copy one dense workspace buffer into another.
pub(crate) fn copy_into(dst: &mut DenseVec, src: &DenseVec) {
    debug_assert_eq!(dst.len(), src.len());
    for i in 0..dst.len() {
        dst[i] = src[i];
    }
}

// Fill `dst = a - alpha * b`.
pub(crate) fn fill_scaled_sub(dst: &mut DenseVec, a: &DenseVec, alpha: f64, b: &DenseVec) {
    debug_assert_eq!(dst.len(), a.len());
    debug_assert_eq!(dst.len(), b.len());
    for i in 0..dst.len() {
        dst[i] = a[i] - alpha * b[i];
    }
}

// Perform the in-place axpy update `a += alpha * b`.
pub(crate) fn add_assign_scaled(a: &mut DenseVec, alpha: f64, b: &DenseVec) {
    debug_assert_eq!(a.len(), b.len());
    for i in 0..a.len() {
        a[i] += alpha * b[i];
    }
}

// Perform the in-place update `a -= alpha * b`.
pub(crate) fn sub_assign_scaled(a: &mut DenseVec, alpha: f64, b: &DenseVec) {
    debug_assert_eq!(a.len(), b.len());
    for i in 0..a.len() {
        a[i] -= alpha * b[i];
    }
}

// Add one dense vector into another in place.
pub(crate) fn add_assign(a: &mut DenseVec, b: &DenseVec) {
    debug_assert_eq!(a.len(), b.len());
    for i in 0..a.len() {
        a[i] += b[i];
    }
}

// Compute `(a - b)^T (c - d)` without allocating temporaries.
pub(crate) fn dot_diff(a: &DenseVec, b: &DenseVec, c: &DenseVec, d: &DenseVec) -> f64 {
    debug_assert_eq!(a.len(), b.len());
    debug_assert_eq!(a.len(), c.len());
    debug_assert_eq!(a.len(), d.len());
    let mut value = 0.0;
    for i in 0..a.len() {
        value += (a[i] - b[i]) * (c[i] - d[i]);
    }
    value
}

// Fill the accelerated momentum state `x + beta * (x - x_old)`.
pub(crate) fn fill_momentum(dst: &mut DenseVec, x: &DenseVec, x_old: &DenseVec, beta: f64) {
    debug_assert_eq!(dst.len(), x.len());
    debug_assert_eq!(dst.len(), x_old.len());
    for i in 0..dst.len() {
        dst[i] = x[i] + beta * (x[i] - x_old[i]);
    }
}

// Compute the Euclidean norm of a dense vector.
pub(crate) fn l2_norm(x: &DenseVec) -> f64 {
    x.dot(x).sqrt()
}