use std::cell::{Ref, RefCell};
use std::rc::Rc;
use pounce_algorithm::application::IpoptApplication;
use pounce_common::types::{Index, Number};
use pounce_nlp::TNLP;
use pounce_nlp::return_codes::ApplicationReturnStatus;
use pounce_sens_core::converged::ConvergedSolve;
use crate::PdSensBacksolver;
use crate::activity::{ActivityReport, ReducedActivityReport, ReducedRowActivityReport};
use crate::backsolver::SensBacksolver;
use crate::boundcheck::{PathOperator, PathSegment, RefineStop, WeakBound};
use crate::corrector::CorrectorReport;
use crate::vec_util::dense_to_vec;
pub use pounce_sens_core::converged::{BARRIER_SIGN, SolverError};
pub struct ConvergedState {
pub status: ApplicationReturnStatus,
pub x: Vec<Number>,
pub obj_val: Number,
pub bound_relax_factor: Number,
pub exact_hessian: bool,
backsolver: PdSensBacksolver,
}
impl ConvergedState {
pub fn block_dims(&self) -> [usize; 8] {
self.backsolver.block_dims()
}
pub fn kkt_dim(&self) -> usize {
self.backsolver.dim()
}
}
pub struct Solver {
app: IpoptApplication,
tnlp: Rc<RefCell<dyn TNLP>>,
state: Rc<RefCell<Option<ConvergedState>>>,
}
impl Solver {
pub fn new(app: IpoptApplication, tnlp: Rc<RefCell<dyn TNLP>>) -> Self {
Self {
app,
tnlp,
state: Rc::new(RefCell::new(None)),
}
}
pub fn app(&self) -> &IpoptApplication {
&self.app
}
pub fn app_mut(&mut self) -> &mut IpoptApplication {
&mut self.app
}
pub fn solve(&mut self) -> ApplicationReturnStatus {
self.state.borrow_mut().take();
let brf = self
.app
.options()
.get_numeric_value("bound_relax_factor", "")
.map(|(v, _)| v)
.expect("bound_relax_factor is a registered core option");
let exact_hessian = self
.app
.options()
.get_string_value("hessian_approximation", "")
.map(|(v, _)| v == "exact")
.expect("hessian_approximation is a registered core option");
let attempts: Rc<RefCell<Vec<(u64, ConvergedState)>>> = Rc::new(RefCell::new(Vec::new()));
let state_cb = Rc::clone(&attempts);
self.app
.set_on_converged(Box::new(move |data, cq, nlp, pd| {
let (curr, id) = {
let d = data.borrow();
match (d.curr.clone(), d.on_converged_id) {
(Some(c), Some(id)) => (c, id),
_ => return,
}
};
let backsolver = match PdSensBacksolver::new(data, cq, nlp, Rc::clone(&pd)) {
Ok(b) => b,
Err(e) => {
eprintln!("pounce: Solver could not capture the KKT factor: {e}");
return;
}
};
let mut x = dense_to_vec(&*curr.x);
if let Some(d) = backsolver.variable_scaling() {
debug_assert_eq!(x.len(), d.len());
for (xi, &di) in x.iter_mut().zip(d.iter()) {
*xi /= di;
}
}
let obj_val = cq.borrow_mut().curr_f();
state_cb.borrow_mut().push((
id,
ConvergedState {
status: ApplicationReturnStatus::InternalError,
x,
obj_val,
bound_relax_factor: brf,
exact_hessian,
backsolver,
},
));
}));
let status = crate::optimize_tnlp_for_sensitivity(&mut self.app, Rc::clone(&self.tnlp));
let reported = self.app.reported_on_converged_id();
let held = attempts
.borrow_mut()
.drain(..)
.find(|(id, _)| Some(*id) == reported)
.map(|(_, mut s)| {
s.status = status;
s
});
*self.state.borrow_mut() = held;
status
}
pub fn converged(&self) -> Option<Ref<'_, ConvergedState>> {
let r = self.state.borrow();
r.as_ref()?;
Some(Ref::map(r, |o| {
o.as_ref()
.unwrap_or_else(|| unreachable!("checked is_some above"))
}))
}
pub fn kkt_dim(&self) -> Option<usize> {
self.converged().map(|c| c.kkt_dim())
}
pub fn block_dims(&self) -> Option<[usize; 8]> {
self.converged().map(|c| c.block_dims())
}
pub fn with_converged<R>(
&self,
f: impl FnOnce(&ConvergedSolve<'_, PdSensBacksolver>) -> Result<R, SolverError>,
) -> Result<R, SolverError> {
let state = self.state.borrow();
let state = state.as_ref().ok_or(SolverError::NotConverged)?;
f(&ConvergedSolve::new(
&state.backsolver,
state.bound_relax_factor,
state.exact_hessian,
))
}
pub fn classify_activity(&self) -> Result<ActivityReport, SolverError> {
self.with_converged(|c| c.classify_activity())
}
pub fn reduced_activity(
&self,
user_vars: &[usize],
) -> Result<ReducedActivityReport, SolverError> {
self.with_converged(|c| c.reduced_activity(user_vars))
}
pub fn reduced_row_activity(
&self,
user_rows: &[usize],
) -> Result<ReducedRowActivityReport, SolverError> {
self.with_converged(|c| c.reduced_row_activity(user_rows))
}
pub fn row_normal(&self, user_row: usize) -> Result<Vec<Number>, SolverError> {
self.with_converged(|c| c.row_normal(user_row))
}
pub fn hessian_vec(&self, v: &[Number]) -> Result<Vec<Number>, SolverError> {
self.with_converged(|c| c.hessian_vec(v))
}
pub fn kkt_solve(&self, rhs: &[Number], lhs: &mut [Number]) -> Result<(), SolverError> {
self.with_converged(|c| c.kkt_solve(rhs, lhs))
}
pub fn kkt_solve_scaled(&self, rhs: &[Number], lhs: &mut [Number]) -> Result<(), SolverError> {
let state = self.state.borrow();
let state = state.as_ref().ok_or(SolverError::NotConverged)?;
let total = state.backsolver.dim();
if rhs.len() != total {
return Err(SolverError::BadShape {
what: "rhs",
got: rhs.len(),
expected: total,
});
}
if lhs.len() != total {
return Err(SolverError::BadShape {
what: "lhs",
got: lhs.len(),
expected: total,
});
}
if state.backsolver.solve_scaled_space(rhs, lhs) {
Ok(())
} else {
Err(SolverError::BacksolveFailed)
}
}
pub fn kkt_solve_many(
&self,
rhs_flat: &[Number],
lhs_flat: &mut [Number],
n_rhs: usize,
) -> Result<(), SolverError> {
self.with_converged(|c| c.kkt_solve_many(rhs_flat, lhs_flat, n_rhs))
}
pub fn kkt_solve_many_scaled(
&self,
rhs_flat: &[Number],
lhs_flat: &mut [Number],
n_rhs: usize,
) -> Result<(), SolverError> {
let state = self.state.borrow();
let state = state.as_ref().ok_or(SolverError::NotConverged)?;
let total = state.backsolver.dim();
let expected = n_rhs * total;
if rhs_flat.len() != expected {
return Err(SolverError::BadShape {
what: "rhs",
got: rhs_flat.len(),
expected,
});
}
if lhs_flat.len() != expected {
return Err(SolverError::BadShape {
what: "lhs",
got: lhs_flat.len(),
expected,
});
}
if state
.backsolver
.solve_many_scaled_space(rhs_flat, lhs_flat, n_rhs)
{
Ok(())
} else {
Err(SolverError::BacksolveFailed)
}
}
pub fn parametric_step(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
) -> Result<Vec<Number>, SolverError> {
check_deltas(pin_constraint_indices, deltas)?;
self.with_converged(|c| c.parametric_step(pin_constraint_indices, deltas))
}
pub fn parametric_step_bounded(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
max_iter: usize,
bound_eps: Option<Number>,
) -> Result<(Vec<Number>, Vec<Index>, RefineStop), SolverError> {
self.with_converged(|c| {
c.parametric_step_bounded(pin_constraint_indices, deltas, max_iter, bound_eps)
})
}
pub fn parametric_step_path(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
max_iter: usize,
) -> Result<(Vec<Number>, Vec<PathSegment>), SolverError> {
self.with_converged(|c| c.parametric_step_path(pin_constraint_indices, deltas, max_iter))
}
pub fn parametric_step_path_decided(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
max_iter: usize,
held_var_rows: &[Index],
) -> Result<(Vec<Number>, Vec<PathSegment>), SolverError> {
self.with_converged(|c| {
c.parametric_step_path_decided(pin_constraint_indices, deltas, max_iter, held_var_rows)
})
}
pub fn correct_step(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
step: &[Number],
max_iter: usize,
) -> Result<(Vec<Number>, CorrectorReport), SolverError> {
self.with_converged(|c| c.correct_step(pin_constraint_indices, deltas, step, max_iter))
}
pub fn parametric_step_bounded_decided(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
max_iter: usize,
held_var_rows: &[Index],
bound_eps: Option<Number>,
) -> Result<(Vec<Number>, Vec<Index>, RefineStop), SolverError> {
self.with_converged(|c| {
c.parametric_step_bounded_decided(
pin_constraint_indices,
deltas,
max_iter,
held_var_rows,
bound_eps,
)
})
}
pub fn path_direction_decided(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
released_bound_rows: &[Index],
held_primal_rows: &[Index],
operator: PathOperator,
) -> Result<(Vec<Number>, Vec<Number>), SolverError> {
self.with_converged(|c| {
c.path_direction_decided(
pin_constraint_indices,
deltas,
released_bound_rows,
held_primal_rows,
operator,
)
})
}
pub fn parametric_step_release_all(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
) -> Result<(Vec<Number>, usize), SolverError> {
self.with_converged(|c| c.parametric_step_release_all(pin_constraint_indices, deltas))
}
pub fn parametric_step_directional(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
max_iter: usize,
) -> Result<(Vec<Number>, Vec<usize>, usize), SolverError> {
self.with_converged(|c| {
c.parametric_step_directional(pin_constraint_indices, deltas, max_iter)
})
}
pub fn weakly_active_bounds(&self) -> Result<Vec<WeakBound>, SolverError> {
self.with_converged(|c| c.weakly_active_bounds())
}
pub fn parametric_step_full(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
) -> Result<Vec<Number>, SolverError> {
check_deltas(pin_constraint_indices, deltas)?;
self.with_converged(|c| c.parametric_step_full(pin_constraint_indices, deltas))
}
pub fn parametric_rhs_full(
&self,
pin_constraint_indices: &[Index],
deltas: &[Number],
) -> Result<Vec<Number>, SolverError> {
self.with_converged(|c| c.parametric_rhs_full(pin_constraint_indices, deltas))
}
pub fn g_multiplier_rows(
&self,
g_indices: &[Index],
) -> Result<Vec<Option<Index>>, SolverError> {
self.with_converged(|c| c.g_multiplier_rows(g_indices))
}
pub fn d_multiplier_rows(
&self,
g_indices: &[Index],
) -> Result<Vec<Option<Index>>, SolverError> {
self.with_converged(|c| c.d_multiplier_rows(g_indices))
}
pub fn d_slack_rows(&self, g_indices: &[Index]) -> Result<Vec<Option<Index>>, SolverError> {
self.with_converged(|c| c.d_slack_rows(g_indices))
}
pub fn x_primal_rows(&self, x_indices: &[Index]) -> Result<Vec<Option<Index>>, SolverError> {
self.with_converged(|c| c.x_primal_rows(x_indices))
}
pub fn bound_multiplier_rows(
&self,
x_indices: &[Index],
) -> Result<Vec<(Option<Index>, Option<Index>)>, SolverError> {
self.with_converged(|c| c.bound_multiplier_rows(x_indices))
}
pub fn n_full_x(&self) -> Result<usize, SolverError> {
self.with_converged(|c| c.n_full_x())
}
pub fn n_full_g(&self) -> Result<usize, SolverError> {
self.with_converged(|c| c.n_full_g())
}
pub fn compute_reduced_hessian(
&self,
pin_constraint_indices: &[Index],
obj_scal: Number,
) -> Result<Vec<Number>, SolverError> {
self.with_converged(|c| c.compute_reduced_hessian(pin_constraint_indices, obj_scal))
}
pub fn compute_reduced_hessian_eigen(
&self,
pin_constraint_indices: &[Index],
obj_scal: Number,
) -> Result<(Vec<Number>, Vec<Number>, Vec<Number>), SolverError> {
self.with_converged(|c| c.compute_reduced_hessian_eigen(pin_constraint_indices, obj_scal))
}
pub fn compute_reduced_hessian_scaled(
&self,
pin_constraint_indices: &[Index],
obj_scal: Number,
) -> Result<Vec<Number>, SolverError> {
self.with_converged(|c| c.compute_reduced_hessian_scaled(pin_constraint_indices, obj_scal))
}
pub fn nlp_scaling(
&self,
) -> Result<(Number, Option<Vec<Number>>, Option<Vec<Number>>), SolverError> {
self.with_converged(|c| c.nlp_scaling())
}
pub fn variable_scaling(&self) -> Result<Option<Vec<Number>>, SolverError> {
let state = self.state.borrow();
let state = state.as_ref().ok_or(SolverError::NotConverged)?;
Ok(state.backsolver.variable_scaling_full().map(|d| d.to_vec()))
}
pub fn kkt_perturbations(&self) -> Result<[Number; 4], SolverError> {
let state = self.state.borrow();
let state = state.as_ref().ok_or(SolverError::NotConverged)?;
Ok(state.backsolver.kkt_perturbations())
}
pub fn pin_g_scaling(
&self,
pin_constraint_indices: &[Index],
) -> Result<Vec<Number>, SolverError> {
self.with_converged(|c| c.pin_g_scaling(pin_constraint_indices))
}
}
fn check_deltas(pins: &[Index], deltas: &[Number]) -> Result<(), SolverError> {
if pins.len() != deltas.len() {
return Err(SolverError::BadShape {
what: "deltas",
got: deltas.len(),
expected: pins.len(),
});
}
Ok(())
}