use std::time::Instant;
use ndarray_linalg::SolveH;
use crate::matrix::QuadraticMatrix;
use super::problem::BoxQPProblem;
use super::types::SolverOptions;
use super::utils::{DenseVec, add_assign_scaled, copy_into, l2_norm, sub_assign_scaled};
struct ReducedPolishSystem {
free: Vec<usize>,
qff: QuadraticMatrix,
rhs_base: DenseVec,
x_free_base: DenseVec,
diagonal: DenseVec,
}
pub(crate) fn improved_proxqp_like_polish_box_qp(
problem: &BoxQPProblem,
x: &DenseVec,
qx: &DenseVec,
options: &SolverOptions,
) -> (DenseVec, DenseVec, f64) {
let start = Instant::now();
let mut x_best = x.clone();
let mut qx_best = qx.clone();
let mut f_best = problem.objective_from_qx(&x_best, &qx_best);
let mut grad = DenseVec::zeros(problem.c.len());
let mut prev_free: Option<Vec<usize>> = None;
let mut prev_active: Option<Vec<usize>> = None;
let mut qx_trial = DenseVec::zeros(problem.c.len());
for _pass_idx in 0..options.polish.num_passes {
let (free, active) = identify_active_set_from_gradient(
problem,
&x_best,
options.polish.bound_tol,
options.polish.grad_tol,
&mut grad,
);
let same_partition_as_last =
prev_free.as_ref() == Some(&free) && prev_active.as_ref() == Some(&active);
let Some(system) = build_reduced_system(problem, &x_best, &free, &active) else {
break;
};
let mut pass_accepted = false;
let mut accepted_rho = 0.0;
let mut accepted_tau = 0.0;
let mut rhs_work = system.rhs_base.clone();
let f_before_pass = f_best;
for &rho in &options.polish.rho_sequence {
let (x_trial, solve_success) =
reduced_box_qp_candidate(problem, &x_best, &system, rho, &mut rhs_work, options);
if !solve_success {
continue;
}
problem.matvec_into(&x_trial, &mut qx_trial);
let (x_ls, qx_ls, f_ls, accepted, tau_used) = line_search_accept_polish(
problem,
&x_best,
&qx_best,
&x_trial,
&qx_trial,
f_best,
options.polish.max_backtracks,
);
if accepted {
x_best = x_ls;
qx_best = qx_ls;
f_best = f_ls;
pass_accepted = true;
accepted_rho = rho;
accepted_tau = tau_used;
break;
}
}
if !pass_accepted {
break;
}
let improve_tol = 1e-12 * f_before_pass.abs().max(1.0);
if f_before_pass - f_best <= improve_tol {
break;
}
if same_partition_as_last
&& accepted_rho == 0.0
&& (accepted_tau - 1.0).abs() <= f64::EPSILON
{
break;
}
prev_free = Some(free);
prev_active = Some(active);
}
(x_best, qx_best, start.elapsed().as_secs_f64())
}
fn identify_active_set_from_gradient(
problem: &BoxQPProblem,
x: &DenseVec,
bound_tol: f64,
grad_tol: f64,
grad: &mut DenseVec,
) -> (Vec<usize>, Vec<usize>) {
problem.gradient_into(x, grad);
let mut free = Vec::new();
let mut active = Vec::new();
for i in 0..x.len() {
let active_lb = x[i] <= problem.lb[i] + bound_tol && grad[i] >= -grad_tol;
let active_ub = x[i] >= problem.ub[i] - bound_tol && grad[i] <= grad_tol;
if active_lb || active_ub {
active.push(i);
} else {
free.push(i);
}
}
(free, active)
}
fn build_reduced_system(
problem: &BoxQPProblem,
x: &DenseVec,
free: &[usize],
active: &[usize],
) -> Option<ReducedPolishSystem> {
if free.is_empty() {
return None;
}
let qff = problem.principal(free, free);
let qfa = problem.principal(free, active);
let x_active = DenseVec::from_iter(active.iter().map(|&idx| x[idx]));
let x_free_base = DenseVec::from_iter(free.iter().map(|&idx| x[idx]));
let c_free = DenseVec::from_iter(free.iter().map(|&idx| problem.c[idx]));
let mut rhs = DenseVec::zeros(free.len());
qfa.matvec_into(&x_active, &mut rhs);
for i in 0..rhs.len() {
rhs[i] = -rhs[i] - c_free[i];
}
let diagonal = qff.diagonal();
Some(ReducedPolishSystem {
free: free.to_vec(),
qff,
rhs_base: rhs,
x_free_base,
diagonal,
})
}
fn reduced_box_qp_candidate(
problem: &BoxQPProblem,
x: &DenseVec,
system: &ReducedPolishSystem,
rho: f64,
rhs: &mut DenseVec,
options: &SolverOptions,
) -> (DenseVec, bool) {
copy_into(rhs, &system.rhs_base);
if rho > 0.0 {
add_assign_scaled(rhs, rho, &system.x_free_base);
}
let (sol, success, _method) = solve_reduced_spd_system(
&system.qff,
&system.diagonal,
rhs,
rho,
options.polish.prefer_direct,
options.polish.direct_max_n,
options.polish.cg_rtol,
options.polish.cg_maxiter,
options.polish.use_preconditioner,
);
if !success {
return (x.clone(), false);
}
let mut x_trial = x.clone();
for (k, &idx) in system.free.iter().enumerate() {
x_trial[idx] = sol[k];
}
problem.clip_in_place(&mut x_trial);
(x_trial, true)
}
fn solve_reduced_spd_system(
a: &QuadraticMatrix,
diagonal: &DenseVec,
rhs: &DenseVec,
rho: f64,
prefer_direct: bool,
direct_max_n: usize,
cg_rtol: f64,
cg_maxiter: usize,
use_preconditioner: bool,
) -> (Vec<f64>, bool, &'static str) {
if let QuadraticMatrix::Dense(m) = a {
if prefer_direct && m.nrows() <= direct_max_n {
let mut shifted = m.clone();
if rho > 0.0 {
for i in 0..shifted.nrows() {
shifted[[i, i]] += rho;
}
}
if let Ok(sol) = shifted.solveh_into(rhs.clone()) {
return (sol.iter().copied().collect(), true, "dense_cholesky");
}
}
}
let mut x = DenseVec::zeros(rhs.len());
let mut r = rhs.clone();
let mut z = DenseVec::zeros(rhs.len());
if use_preconditioner {
apply_diagonal_preconditioner_into(diagonal, rho, &r, &mut z);
} else {
copy_into(&mut z, &r);
}
let mut p = z.clone();
let mut ap = DenseVec::zeros(rhs.len());
let mut rz_old = r.dot(&z);
let rhs_norm = l2_norm(rhs).max(1.0);
for _ in 0..cg_maxiter {
shifted_matvec_into(a, rho, &p, &mut ap);
let alpha = rz_old / p.dot(&ap).max(1e-30);
add_assign_scaled(&mut x, alpha, &p);
sub_assign_scaled(&mut r, alpha, &ap);
if l2_norm(&r) <= cg_rtol * rhs_norm {
let method = match a {
QuadraticMatrix::Dense(_) => "cg",
QuadraticMatrix::Sparse(_) => "sparse_cg",
};
return (x.iter().copied().collect(), true, method);
}
if use_preconditioner {
apply_diagonal_preconditioner_into(diagonal, rho, &r, &mut z);
} else {
copy_into(&mut z, &r);
}
let rz_new = r.dot(&z);
let beta = rz_new / rz_old.max(1e-30);
for i in 0..p.len() {
p[i] = z[i] + beta * p[i];
}
rz_old = rz_new;
}
let method = match a {
QuadraticMatrix::Dense(_) => "cg",
QuadraticMatrix::Sparse(_) => "sparse_cg",
};
(x.iter().copied().collect(), false, method)
}
fn shifted_matvec_into(a: &QuadraticMatrix, rho: f64, x: &DenseVec, out: &mut DenseVec) {
a.matvec_into(x, out);
if rho > 0.0 {
add_assign_scaled(out, rho, x);
}
}
fn apply_diagonal_preconditioner_into(diag: &DenseVec, rho: f64, r: &DenseVec, out: &mut DenseVec) {
for i in 0..diag.len() {
out[i] = r[i] / (diag[i] + rho).abs().max(1e-12);
}
}
fn line_search_accept_polish(
problem: &BoxQPProblem,
x: &DenseVec,
qx: &DenseVec,
x_trial: &DenseVec,
qx_trial: &DenseVec,
f_x: f64,
max_backtracks: usize,
) -> (DenseVec, DenseVec, f64, bool, f64) {
let mut tau = 1.0;
let mut best_x = x.clone();
let mut best_qx = qx.clone();
let mut best_f = f_x;
let mut x_ls = DenseVec::zeros(x.len());
let mut qx_ls = DenseVec::zeros(x.len());
for _ in 0..=max_backtracks {
for i in 0..x.len() {
x_ls[i] = x[i] + tau * (x_trial[i] - x[i]);
qx_ls[i] = qx[i] + tau * (qx_trial[i] - qx[i]);
}
let f_ls = problem.objective_from_qx(&x_ls, &qx_ls);
if f_ls <= f_x {
return (x_ls, qx_ls, f_ls, true, tau);
}
if f_ls < best_f {
best_x = x_ls.clone();
best_qx = qx_ls.clone();
best_f = f_ls;
}
tau *= 0.5;
}
(best_x, best_qx, best_f, best_f <= f_x, 0.0)
}