sundial-core 0.1.0

//! CPU f64 reference implementation of restarted PDHG — the executable
//! statement of this project's math conventions. The GPU engine mirrors this
//! loop exactly, so this file is the arbiter when the two disagree.
use crate::kkt::{self, KktResiduals};
use crate::problem::*;
use crate::scale;
use web_time::Instant;

pub fn power_iteration_norm(a: &CsrMatrix, at: &CsrMatrix, iters: usize, seed: u64) -> f64 {
    crate::linop::power_iteration_norm_op(&crate::linop::CsrOp { a, at }, iters, seed)
}

fn unscale(scaling: Option<&scale::Scaling>, x: &[f64], is_col: bool) -> Vec<f64> {
    match scaling {
        Some(s) if is_col => s.unscale_x(x),
        Some(s) => s.unscale_y(x),
        None => x.to_vec(),
    }
}

struct State {
    x: Vec<f64>,
    y: Vec<f64>,
    x_avg: Vec<f64>,
    y_avg: Vec<f64>,
    avg_count: u64,
}

impl State {
    /// Restart: optionally adopt the running average as the current iterate,
    /// then reset the average to the current iterate.
    fn restart_from(&mut self, from_avg: bool) {
        if from_avg {
            self.x.copy_from_slice(&self.x_avg);
            self.y.copy_from_slice(&self.y_avg);
        }
        self.x_avg.copy_from_slice(&self.x);
        self.y_avg.copy_from_slice(&self.y);
        self.avg_count = 1;
    }
}

pub fn solve(
    p: &LpProblem,
    opts: &SolveOptions,
    progress: &mut dyn FnMut(ProgressEvent),
) -> Solution {
    let (sp, s) = scale::ruiz_pc(p, 10);
    let norm_a = power_iteration_norm(&sp.a, &sp.at, 100, opts.seed);
    solve_view(&sp.view(), &p.view(), Some(&s), norm_a, opts, progress)
}

pub fn solve_op<O: crate::linop::LinOp>(
    p: &crate::problem::OpProblem<O>,
    opts: &SolveOptions,
    progress: &mut dyn FnMut(ProgressEvent),
) -> Solution {
    // Matrix-free problems solve UNSCALED: the transport operator is an
    // all-ones incidence structure, so it is already balanced.
    let norm_a = crate::linop::op_norm2(&p.op, opts.seed);
    let v = p.view();
    solve_view(&v, &v, None, norm_a, opts, progress)
}

pub fn solve_view(
    iterate: &LpView,
    original: &LpView,
    scaling: Option<&scale::Scaling>,
    norm_a: f64,
    opts: &SolveOptions,
    progress: &mut dyn FnMut(ProgressEvent),
) -> Solution {
    assert!(
        opts.check_every > 0,
        "SolveOptions::check_every must be > 0"
    );
    let start = Instant::now();
    let (m, n) = (iterate.op.n_rows(), iterate.op.n_cols());

    // Primal-weight balancing (ω) runs ONLY on the unscaled path (`scaling ==
    // None`, the matrix-free operator entry). The explicit path arrives
    // Ruiz+PC-equilibrated — that scaling already sets the primal/dual step
    // balance, so applying PDLP's ω = ‖c‖/‖q‖ on top double-corrects it. On the
    // unscaled transport problems ω is the only step balancing and is what
    // converges the 1M gate. Mirrors the GPU `solve_core` `primal_weight` gate.
    let mut omega = if scaling.is_none() {
        let (q_it, c_it) = kkt::denominators_view(iterate);
        crate::weight::initial_primal_weight(q_it, c_it)
    } else {
        1.0
    };
    let mut tau = 0.9 / (norm_a * omega);
    let mut sigma = 0.9 * omega / norm_a;

    let mut st = State {
        x: vec![0.0; n],
        y: vec![0.0; m],
        x_avg: vec![0.0; n],
        y_avg: vec![0.0; m],
        avg_count: 1,
    };
    // start feasible w.r.t. boxes: clamp 0 into [l_v, u_v]
    for j in 0..n {
        st.x[j] = 0.0f64.clamp(iterate.col_lower[j], iterate.col_upper[j]);
    }
    st.x_avg.copy_from_slice(&st.x);

    let mut aty = vec![0.0; n];
    let mut axt = vec![0.0; m];
    let mut x_new = vec![0.0; n];
    let mut x_tilde = vec![0.0; n];

    // Movement-based ω (experimental) measures ‖Δx‖/‖Δy‖ between consecutive
    // RESTART points, in iterate space (where τ/σ act). Explicit path only —
    // the matrix-free path keeps the residual-ratio rule that converges the 1M
    // gate. Allocated only when armed, so the default path is untouched.
    let track_movement = opts.movement_weight && scaling.is_some();
    let mut x_prev_restart = if track_movement {
        st.x.clone()
    } else {
        Vec::new()
    };
    let mut y_prev_restart = if track_movement {
        st.y.clone()
    } else {
        Vec::new()
    };

    let mut mu_last_restart = f64::INFINITY;
    let mut iters_since_restart: u64 = 0;
    let mut restarts: u32 = 0;
    // Divergence-detection state (see farkas.rs constants): candidate norms
    // and residuals at the LAST restart, plus growth streak counters.
    let mut y_norm_prev = 0.0f64;
    let mut x_norm_prev = 0.0f64;
    let mut relp_prev = f64::INFINITY;
    let mut reld_prev = f64::INFINITY;
    let mut infeas_streak: u32 = 0;
    let mut unbound_streak: u32 = 0;
    let mut status = SolveStatus::IterationLimit;
    let mut iter: u64 = 0;
    let mut last_check_time = Instant::now();
    let mut last_check_iter: u64 = 0;

    while iter < opts.max_iters {
        // one PDHG iteration (see Math conventions)
        iterate.op.apply_t(&st.y, &mut aty);
        for j in 0..n {
            let v = st.x[j] - tau * (iterate.c[j] + aty[j]);
            let xn = v.clamp(iterate.col_lower[j], iterate.col_upper[j]);
            x_new[j] = xn;
            x_tilde[j] = 2.0 * xn - st.x[j];
        }
        iterate.op.apply(&x_tilde, &mut axt);
        for (i, y_i) in st.y.iter_mut().enumerate() {
            let v = *y_i + sigma * axt[i];
            *y_i = v - sigma * (v / sigma).clamp(iterate.row_lower[i], iterate.row_upper[i]);
        }
        std::mem::swap(&mut st.x, &mut x_new);
        iter += 1;
        iters_since_restart += 1;

        // incremental running average
        st.avg_count += 1;
        let w = 1.0 / st.avg_count as f64;
        for j in 0..n {
            st.x_avg[j] += w * (st.x[j] - st.x_avg[j]);
        }
        for i in 0..m {
            st.y_avg[i] += w * (st.y[i] - st.y_avg[i]);
        }

        if iter.is_multiple_of(opts.check_every as u64) {
            if st.x.iter().any(|v| !v.is_finite()) || st.y.iter().any(|v| !v.is_finite()) {
                status = SolveStatus::NumericalBreakdown;
                break;
            }
            // IMPORTANT: residuals for termination/restart are evaluated on the
            // ORIGINAL problem (scaled-space residuals passing tol does NOT
            // imply the real ones do). Unscale candidates first.
            let xc_u = unscale(scaling, &st.x, true);
            let yc_u = unscale(scaling, &st.y, false);
            let xa_u = unscale(scaling, &st.x_avg, true);
            let ya_u = unscale(scaling, &st.y_avg, false);
            let r_cur = kkt::residuals_view(original, &xc_u, &yc_u);
            let r_avg = kkt::residuals_view(original, &xa_u, &ya_u);
            let (mu_cand, cand_is_avg) = if r_avg.mu() < r_cur.mu() {
                (r_avg.mu(), true)
            } else {
                (r_cur.mu(), false)
            };

            let now = Instant::now();
            let ms_per_iter = now.duration_since(last_check_time).as_secs_f64() * 1000.0
                / (iter - last_check_iter).max(1) as f64;
            last_check_time = now;
            last_check_iter = iter;
            progress(ProgressEvent {
                iter,
                rel_primal: r_cur.rel_primal,
                rel_dual: r_cur.rel_dual,
                rel_gap: r_cur.rel_gap,
                ms_per_iter,
            });

            if mu_cand <= opts.tol {
                if cand_is_avg {
                    st.restart_from(true);
                }
                status = SolveStatus::Optimal;
                break;
            }
            // restart rule
            if mu_cand <= 0.5 * mu_last_restart || iters_since_restart >= 4096 {
                st.restart_from(cand_is_avg);
                if scaling.is_none() {
                    let r_cand = if cand_is_avg { &r_avg } else { &r_cur };
                    omega = crate::weight::update_primal_weight(
                        omega,
                        r_cand.rel_primal,
                        r_cand.rel_dual,
                    );
                    tau = 0.9 / (norm_a * omega);
                    sigma = 0.9 * omega / norm_a;
                } else if track_movement {
                    let l2 = |a: &[f64], b: &[f64]| -> f64 {
                        a.iter()
                            .zip(b)
                            .map(|(u, v)| (u - v) * (u - v))
                            .sum::<f64>()
                            .sqrt()
                    };
                    let dx = l2(&st.x, &x_prev_restart);
                    let dy = l2(&st.y, &y_prev_restart);
                    omega = crate::weight::update_primal_weight_movement(omega, dx, dy);
                    tau = 0.9 / (norm_a * omega);
                    sigma = 0.9 * omega / norm_a;
                    x_prev_restart.copy_from_slice(&st.x);
                    y_prev_restart.copy_from_slice(&st.y);
                }
                // ---- divergence detection (restart cadence only) ----
                // Reuse the unscaled CANDIDATE rather than re-unscaling st.x/st.y:
                // restart_from() above has already copied the candidate into st,
                // so unscale(st.x) here is by construction xc_u (cand_is_avg =
                // false, st untouched) or xa_u (true, st.x ← st.x_avg). Selecting
                // by cand_is_avg keeps this bit-identical while dropping two
                // redundant unscales per live-streak check.
                let r_cand = if cand_is_avg { &r_avg } else { &r_cur };
                let (x_u, y_u) = if cand_is_avg {
                    (&xa_u, &ya_u)
                } else {
                    (&xc_u, &yc_u)
                };
                let y_norm = y_u.iter().fold(0.0f64, |a, &v| a.max(v.abs()));
                let x_norm = x_u.iter().fold(0.0f64, |a, &v| a.max(v.abs()));
                if y_norm >= crate::farkas::GROWTH * y_norm_prev
                    && r_cand.rel_primal > crate::farkas::STALL * relp_prev
                {
                    infeas_streak += 1;
                } else {
                    infeas_streak = 0;
                }
                if x_norm >= crate::farkas::GROWTH * x_norm_prev
                    && r_cand.rel_dual > crate::farkas::STALL * reld_prev
                {
                    unbound_streak += 1;
                } else {
                    unbound_streak = 0;
                }
                y_norm_prev = y_norm;
                x_norm_prev = x_norm;
                relp_prev = r_cand.rel_primal;
                reld_prev = r_cand.rel_dual;
                if infeas_streak >= crate::farkas::STREAK_K {
                    if crate::farkas::verify_infeasible(original, y_u).is_some() {
                        status = SolveStatus::Infeasible;
                        break;
                    }
                    infeas_streak = 0; // failed verification: don't re-hammer
                }
                if unbound_streak >= crate::farkas::STREAK_K {
                    if crate::farkas::verify_unbounded(original, x_u).is_some() {
                        status = SolveStatus::Unbounded;
                        break;
                    }
                    unbound_streak = 0;
                }
                mu_last_restart = mu_cand;
                iters_since_restart = 0;
                restarts += 1;
            }
            if let Some(limit) = opts.time_limit_ms {
                if start.elapsed().as_secs_f64() * 1000.0 > limit {
                    status = SolveStatus::TimeLimit;
                    break;
                }
            }
        }
    }

    // unscale and record the authoritative f64 verification on the ORIGINAL problem
    //
    // No dual sign-projection here, unlike gpu/engine.rs — and that asymmetry is
    // deliberate, not an oversight. The GPU projects because f32 iteration leaves
    // wrong-sign noise on the duals of rows with an open bound; this path's dual
    // prox step (`v − σ·clamp(v/σ, l, u)` above) lands in the sign cone exactly
    // in f64 — open above ⇒ clamp ≥ v/σ ⇒ y ≤ 0, mirrored below — and Ruiz
    // unscaling multiplies by positive factors, preserving sign. A `project_dual`
    // call here would be provably dead code. Pinned by the sign-cone tests in
    // tests/farkas.rs, which fail if the prox step ever stops guaranteeing it.
    let x = unscale(scaling, &st.x, true);
    let y = unscale(scaling, &st.y, false);
    let verified: KktResiduals = kkt::residuals_view(original, &x, &y);
    assert!(
        status != SolveStatus::Optimal || verified.mu() <= opts.tol,
        "honesty violation: Optimal status with verified mu {} > tol {}",
        verified.mu(),
        opts.tol
    );
    let primal_obj = verified.primal_obj;
    Solution {
        x,
        y,
        primal_obj,
        status,
        stats: SolveStats {
            iterations: iter,
            restarts,
            solve_ms: start.elapsed().as_secs_f64() * 1000.0,
            verified,
        },
    }
}