pounce-algorithm 0.10.0

//! Warm-start iterate initializer — port of
//! `IpWarmStartIterateInitializer.{hpp,cpp}`. Used when a previous
//! solve has left a trial point that should be reused.
//!
//! There are two callers we serve:
//!
//! * **A full primal-dual warm restart** installed via
//!   `Application::set_warm_start_iterate` and consumed by the next
//!   `optimize_tnlp` (e.g. the debugger `resolve` re-solve): `data.curr`
//!   already carries the previous solve's iterate, so we keep it, clamp
//!   multipliers, and optionally override `mu`.
//! * **First solves from `OptimizeTNLP`** that opt into
//!   `warm_start_init_point=yes` to forward user-supplied
//!   primal/dual seeds via `TNLP::get_starting_point`. Here
//!   `data.curr` carries only dim metadata (uninitialized vectors);
//!   we pull seeds from the NLP, push primals/slacks into the bound
//!   interior with warm-start `bound_push`/`bound_frac`, and then
//!   apply the same multiplier clamps.
//!
//! Wired options today: `bound_push`, `bound_frac`,
//! `slack_bound_push`, `slack_bound_frac`, `mult_bound_push`,
//! `mult_init_max`, `target_mu`. `mult_bound_push` floors the four
//! bound-multiplier blocks (mirroring upstream's `ElementWiseMax`
//! with `warm_start_mult_bound_push`): a user-seeded `z = 0` would
//! otherwise start the barrier on its boundary. The remaining knobs
//! (`entire_iterate`, `same_structure`) are parsed for Ipopt option
//! compatibility but not yet consumed — they require the
//! `GetWarmStartIterate` TNLP surface, which pounce does not expose.

use crate::alg_builder::WarmStartOptions;
use crate::init::default::push_x_into_interior;
use crate::init::r#trait::IterateInitializer;
use crate::ipopt_cq::IpoptCqHandle;
use crate::ipopt_data::IpoptDataHandle;
use crate::ipopt_nlp::IpoptNlp;
use crate::iterates_vector::IteratesVector;
use crate::kkt::aug_system_solver::AugSystemSolver;
use pounce_linalg::Vector;
use pounce_linalg::compound_vector::CompoundVector;
use pounce_linalg::dense_vector::{DenseVector, DenseVectorSpace};
use std::cell::RefCell;
use std::rc::Rc;

pub struct WarmStartIterateInitializer {
    opts: WarmStartOptions,
}

impl WarmStartIterateInitializer {
    pub fn new() -> Self {
        Self {
            opts: WarmStartOptions::default(),
        }
    }

    pub fn with_options(opts: WarmStartOptions) -> Self {
        Self { opts }
    }
}

impl Default for WarmStartIterateInitializer {
    fn default() -> Self {
        Self::new()
    }
}

impl IterateInitializer for WarmStartIterateInitializer {
    fn set_initial_iterates(
        &mut self,
        data: &IpoptDataHandle,
        _cq: &IpoptCqHandle,
        nlp: &Rc<RefCell<dyn IpoptNlp>>,
        _aug_solver: &mut dyn AugSystemSolver,
    ) -> bool {
        // Two entry points share this initializer: the re-optimize path
        // (curr.x carries values from the prior solve) and the first
        // OptimizeTNLP call that opted into warm_start_init_point=yes
        // (curr.x is the application's placeholder seed — allocated but
        // never written). Detect the latter and rebuild `curr` from the
        // NLP's get_starting_x/y/z hooks before clamping.
        let needs_seed_from_nlp = {
            let borrow = data.borrow();
            match borrow.curr.as_ref() {
                None => return false,
                Some(c) => !is_initialized(&c.x),
            }
        };

        if needs_seed_from_nlp && !seed_from_nlp(data, nlp, &self.opts) {
            return false;
        }

        {
            // Rebuild `curr` with clamped multipliers. Components are
            // shared via `Rc` with previous solves, so we make fresh
            // copies before mutating to avoid clobbering downstream
            // borrowers. Bound multipliers are additionally floored at
            // `mult_bound_push` (upstream `warm_start_mult_bound_push`):
            // the barrier needs them strictly positive, and a carried-in
            // 0 (e.g. an inactive bound in the previous solution) would
            // otherwise start on the boundary. This block runs even
            // with both clamps disabled (cap = inf, floor = 0; the
            // floor still clamps a negative z/v to 0) because it also
            // resolves NaN seeds: NaN in a user-supplied multiplier
            // means "unseeded", and takes `bound_mult_init_val` for
            // bound multipliers, or 0 for equality multipliers. That 0
            // is the warm path's existing unseeded value (what
            // `seed_from_nlp` produced already), NOT the cold path's
            // least-squares estimate; routing NaN duals through the
            // least-squares calculator is a possible refinement.
            let mut borrow = data.borrow_mut();
            let curr = borrow.curr.as_ref().unwrap();
            let cap = if self.opts.mult_init_max > 0.0 {
                self.opts.mult_init_max
            } else {
                f64::INFINITY
            };
            let z_floor = self.opts.mult_bound_push.max(0.0);
            let z_nan = self.opts.bound_mult_init_val;
            let new_curr = IteratesVector::new(
                Rc::clone(&curr.x),
                Rc::clone(&curr.s),
                clone_clamped(&curr.y_c, -cap, cap, 0.0),
                clone_clamped(&curr.y_d, -cap, cap, 0.0),
                clone_clamped(&curr.z_l, z_floor, cap, z_nan),
                clone_clamped(&curr.z_u, z_floor, cap, z_nan),
                clone_clamped(&curr.v_l, z_floor, cap, z_nan),
                clone_clamped(&curr.v_u, z_floor, cap, z_nan),
            );
            borrow.set_curr(new_curr);
        }

        if self.opts.target_mu > 0.0 {
            data.borrow_mut().curr_mu = self.opts.target_mu;
        }

        true
    }
}

/// Pull a fresh starting iterate from the NLP (which routes to
/// `TNLP::get_starting_point` with `init_x` / `init_lambda` /
/// `init_z` all true), push the primals and slacks into the bound
/// interior using warm-start-specific `bound_push`/`bound_frac`, and
/// install the result on `data.curr`. Mirrors steps 1-4 of
/// `DefaultIterateInitializer::set_initial_iterates`, but with
/// upstream's warm-start option block governing the push.
fn seed_from_nlp(
    data: &IpoptDataHandle,
    nlp: &Rc<RefCell<dyn IpoptNlp>>,
    opts: &WarmStartOptions,
) -> bool {
    if !nlp.borrow_mut().prepare_warm_start() {
        return false;
    }
    let (n_x, n_s, n_yc, n_yd, n_zl, n_zu, n_vl, n_vu) = {
        let borrow = data.borrow();
        let c = borrow.curr.as_ref().unwrap();
        (
            c.x.dim(),
            c.s.dim(),
            c.y_c.dim(),
            c.y_d.dim(),
            c.z_l.dim(),
            c.z_u.dim(),
            c.v_l.dim(),
            c.v_u.dim(),
        )
    };

    let mut x = DenseVectorSpace::new(n_x).make_new_dense();
    nlp.borrow_mut().get_starting_x(&mut x);
    {
        let nlp_ref = nlp.borrow();
        push_x_into_interior(
            &mut x,
            &*nlp_ref.px_l(),
            nlp_ref.x_l(),
            &*nlp_ref.px_u(),
            nlp_ref.x_u(),
            opts.bound_push,
            opts.bound_frac,
        );
    }

    let mut s = DenseVectorSpace::new(n_s).make_new_dense();
    nlp.borrow_mut().eval_d(&x, &mut s);
    {
        let nlp_ref = nlp.borrow();
        push_x_into_interior(
            &mut s,
            &*nlp_ref.pd_l(),
            nlp_ref.d_l(),
            &*nlp_ref.pd_u(),
            nlp_ref.d_u(),
            opts.slack_bound_push,
            opts.slack_bound_frac,
        );
    }

    let mut y_c = DenseVectorSpace::new(n_yc).make_new_dense();
    let mut y_d = DenseVectorSpace::new(n_yd).make_new_dense();
    y_c.set(0.0);
    y_d.set(0.0);
    nlp.borrow_mut().get_starting_y(&mut y_c, &mut y_d);

    let mut z_l = DenseVectorSpace::new(n_zl).make_new_dense();
    let mut z_u = DenseVectorSpace::new(n_zu).make_new_dense();
    let mut v_l = DenseVectorSpace::new(n_vl).make_new_dense();
    let mut v_u = DenseVectorSpace::new(n_vu).make_new_dense();
    z_l.set(0.0);
    z_u.set(0.0);
    v_l.set(0.0);
    v_u.set(0.0);
    nlp.borrow_mut()
        .get_starting_z(&mut z_l, &mut z_u, &mut v_l, &mut v_u);
    nlp.borrow_mut().finish_warm_start();

    let iv = IteratesVector::new(
        Rc::new(x),
        Rc::new(s),
        Rc::new(y_c),
        Rc::new(y_d),
        Rc::new(z_l),
        Rc::new(z_u),
        Rc::new(v_l),
        Rc::new(v_u),
    );
    data.borrow_mut().set_curr(iv);
    true
}

fn is_initialized(v: &Rc<dyn Vector>) -> bool {
    if v.dim() == 0 {
        return true;
    }
    v.as_any()
        .downcast_ref::<DenseVector>()
        .map(|d| d.is_initialized())
        .unwrap_or(true)
}

/// Replace every NaN entry of `v` with `fill`, in place.
///
/// NaN in a user-supplied multiplier seed means "unseeded" (see the
/// `Problem.solve` contract), and has to be resolved before the
/// clamps: `element_wise_min`/`element_wise_max` would propagate it
/// into the iterate, poisoning the solve.
///
/// Both `Vector` storage layouts are handled. A dense block is
/// scanned directly; a compound block recurses into its components,
/// so the contract holds wherever the iterate's multiplier blocks
/// live — the seed path (`seed_from_nlp`) always builds dense
/// vectors, but the re-optimize path reuses whatever the previous
/// solve's spaces produced, and a debug-only guard would be compiled
/// out of exactly the release builds that ship.
fn resolve_nan_seeds(v: &mut dyn Vector, fill: f64) {
    // Type-test before taking the mutable borrow: `if let Some(d) =
    // v.as_any_mut()… else` would keep that borrow live across the
    // else arm.
    if v.as_any().is::<DenseVector>() {
        let d = v.as_any_mut().downcast_mut::<DenseVector>().unwrap();
        for e in d.values_mut() {
            if e.is_nan() {
                *e = fill;
            }
        }
    } else if v.as_any().is::<CompoundVector>() {
        let c = v.as_any_mut().downcast_mut::<CompoundVector>().unwrap();
        for i in 0..c.n_comps() {
            resolve_nan_seeds(c.comp_mut(i), fill);
        }
    } else {
        // `DenseVector` and `CompoundVector` are the only `Vector`
        // implementations; a third one must be handled here, or NaN
        // rides the clamps into the iterate as a silent poison.
        debug_assert!(false, "resolve_nan_seeds: unhandled Vector implementation");
    }
}

/// Clone `v` into a fresh owned vector and clamp every entry to
/// `[lo, hi]` componentwise. Empty vectors short-circuit. Vectors that
/// were never written to (the application's placeholder seed iterates
/// before any solve ran) collapse to a zero-initialized vector — `0`
/// is inside every well-formed warm-start clamp range, so this matches
/// upstream's behavior when a multiplier block has no carry-over
/// value.
fn clone_clamped(v: &Rc<dyn Vector>, lo: f64, hi: f64, nan_fill: f64) -> Rc<dyn Vector> {
    let n = v.dim();
    if n == 0 {
        return Rc::clone(v);
    }
    let mut out = v.make_new();
    let initialized = v
        .as_any()
        .downcast_ref::<DenseVector>()
        .map(|d| d.is_initialized())
        .unwrap_or(true);
    if initialized {
        out.copy(&**v);
        // NaN marks an unseeded entry; resolve it before the clamps
        // (element-wise min/max would just propagate it)
        resolve_nan_seeds(&mut *out, nan_fill);
    } else {
        out.set(0.0);
    }
    let mut cap_hi = v.make_new();
    cap_hi.set(hi);
    out.element_wise_min(&*cap_hi);
    let mut cap_lo = v.make_new();
    cap_lo.set(lo);
    out.element_wise_max(&*cap_lo);
    Rc::from(out)
}

#[cfg(test)]
mod tests_nan_seed {
    use super::*;
    use pounce_linalg::compound_vector::CompoundVectorSpace;
    use pounce_linalg::dense_vector::DenseVectorSpace;

    #[test]
    fn nan_entries_take_the_fill_before_clamping() {
        let space = DenseVectorSpace::new(3);
        let mut d = space.make_new_dense();
        d.values_mut().copy_from_slice(&[0.5, f64::NAN, 2e7]);
        let v: Rc<dyn Vector> = Rc::from(d);
        let out = clone_clamped(&v, 1e-3, 1e6, 7.0);
        let out = out.as_any().downcast_ref::<DenseVector>().unwrap();
        assert_eq!(out.values()[0], 0.5);
        assert_eq!(out.values()[1], 7.0); // unseeded -> fill
        assert_eq!(out.values()[2], 1e6); // then the cap applies
    }

    /// The re-optimize path reuses the previous solve's vector spaces,
    /// which are compound for a blocked NLP. NaN has to resolve there
    /// too: a debug-only guard is compiled out of the release builds
    /// that ship, so an unresolved NaN would ride the clamps into the
    /// iterate and poison the solve.
    #[test]
    fn nan_resolves_inside_a_compound_vector() {
        let inner = DenseVectorSpace::new(2);
        let space = CompoundVectorSpace::new(2, 4);
        for icomp in 0..2 {
            let inner = Rc::clone(&inner);
            space.set_comp(icomp, 2, move || {
                let mut d = inner.make_new_dense();
                d.set(0.0);
                Box::new(d)
            });
        }
        let mut cv = CompoundVector::new(Rc::clone(&space));
        for (icomp, vals) in [[0.5, f64::NAN], [f64::NAN, 2e7]].into_iter().enumerate() {
            let c = cv.comp_mut(icomp as pounce_common::types::Index);
            let d = c.as_any_mut().downcast_mut::<DenseVector>().unwrap();
            d.values_mut().copy_from_slice(&vals);
        }

        let v: Rc<dyn Vector> = Rc::from(cv);
        let out = clone_clamped(&v, 1e-3, 1e6, 7.0);

        let out = out.as_any().downcast_ref::<CompoundVector>().unwrap();
        let flat: Vec<f64> = (0..out.n_comps())
            .flat_map(|i| {
                out.comp(i)
                    .as_any()
                    .downcast_ref::<DenseVector>()
                    .unwrap()
                    .values()
                    .to_vec()
            })
            .collect();
        assert_eq!(flat[0], 0.5);
        assert_eq!(flat[1], 7.0); // unseeded -> fill, not NaN
        assert_eq!(flat[2], 7.0);
        assert_eq!(flat[3], 1e6); // then the cap applies
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use pounce_linalg::dense_vector::DenseVectorSpace;

    fn dense(n: i32, fill: f64) -> Rc<dyn Vector> {
        let space = DenseVectorSpace::new(n);
        let mut v = space.make_new_dense();
        v.set(fill);
        Rc::new(v)
    }

    #[test]
    fn clamps_multipliers_to_cap() {
        let v = dense(3, 1e10);
        let out = clone_clamped(&v, 0.0, 1e6, 0.0);
        assert_eq!(out.amax(), 1e6);
        let v2 = dense(3, -1e10);
        let out2 = clone_clamped(&v2, -1e6, 1e6, 0.0);
        assert_eq!(out2.amax(), 1e6);
    }

    #[test]
    fn clamps_bound_mults_nonneg() {
        let v = dense(3, -5.0);
        let out = clone_clamped(&v, 0.0, 1e6, 0.0);
        assert_eq!(out.amax(), 0.0);
    }

    #[test]
    fn empty_vector_short_circuits() {
        let v = dense(0, 0.0);
        let out = clone_clamped(&v, 0.0, 1.0, 0.0);
        assert_eq!(out.dim(), 0);
    }

    #[test]
    fn in_range_values_pass_through_untouched() {
        let v = dense(3, 0.5);
        let out = clone_clamped(&v, 0.0, 1.0, 0.0);
        assert!((out.max() - 0.5).abs() < 1e-15);
        assert!((out.min() - 0.5).abs() < 1e-15);
    }

    #[test]
    fn mult_bound_push_floors_zero_bound_multipliers() {
        // A carried-in z = 0 (inactive bound in the previous solution)
        // must be floored at warm_start_mult_bound_push, matching
        // upstream's ElementWiseMax — the barrier needs z > 0.
        let v = dense(3, 0.0);
        let out = clone_clamped(&v, 1e-3, 1e6, 0.0);
        assert!((out.min() - 1e-3).abs() < 1e-18);
        // Values already above the floor pass through.
        let v2 = dense(3, 0.7);
        let out2 = clone_clamped(&v2, 1e-3, 1e6, 0.0);
        assert!((out2.max() - 0.7).abs() < 1e-15);
    }

    #[test]
    fn uninitialized_source_collapses_to_zero() {
        // Application's placeholder seed iterate: vector allocated but
        // never written. `clone_clamped` must fall back to zero instead
        // of tripping the dense-vector "must be initialized" assert.
        let space = DenseVectorSpace::new(4);
        let v: Rc<dyn Vector> = Rc::new(space.make_new_dense());
        let out = clone_clamped(&v, 0.0, 1e6, 0.0);
        assert_eq!(out.amax(), 0.0);
    }
}