1use crate::ipopt_nlp::IpoptNlp;
27use crate::sqp::problem::SqpProblemSpec;
28use crate::sqp::qp_assembly::Triplet;
29use pounce_common::Number;
30use pounce_linalg::dense_vector::{DenseVector, DenseVectorSpace};
31use pounce_linalg::expansion_matrix::ExpansionMatrix;
32use pounce_linalg::triplet::{GenTMatrix, SymTMatrix};
33use std::cell::RefCell;
34use std::rc::Rc;
35
36pub struct IpoptNlpAdapter {
37 nlp: Rc<RefCell<dyn IpoptNlp>>,
38 n: usize,
39 m_c: usize,
40 m_d: usize,
41 x_l: Vec<Number>,
42 x_u: Vec<Number>,
43 d_l: Vec<Number>,
44 d_u: Vec<Number>,
45 x_init: Vec<Number>,
46 x_space: Rc<DenseVectorSpace>,
47 c_space: Rc<DenseVectorSpace>,
48 d_space: Rc<DenseVectorSpace>,
49}
50
51impl IpoptNlpAdapter {
52 pub fn new(nlp: Rc<RefCell<dyn IpoptNlp>>) -> Self {
56 let (n, m_c, m_d) = {
57 let b = nlp.borrow();
58 (b.n() as usize, b.m_eq() as usize, b.m_ineq() as usize)
59 };
60 let x_space = DenseVectorSpace::new(n as i32);
61 let c_space = DenseVectorSpace::new(m_c as i32);
62 let d_space = DenseVectorSpace::new(m_d as i32);
63
64 let (x_l, x_u, d_l, d_u, x_init) = {
71 let mut n_borrow = nlp.borrow_mut();
72 let x_l_small = vec_from_dyn(n_borrow.x_l());
73 let x_u_small = vec_from_dyn(n_borrow.x_u());
74 let d_l_small = if m_d > 0 {
75 vec_from_dyn(n_borrow.d_l())
76 } else {
77 Vec::new()
78 };
79 let d_u_small = if m_d > 0 {
80 vec_from_dyn(n_borrow.d_u())
81 } else {
82 Vec::new()
83 };
84 let px_l = n_borrow.px_l();
85 let px_u = n_borrow.px_u();
86 let pd_l = n_borrow.pd_l();
87 let pd_u = n_borrow.pd_u();
88 let x_l = scatter_bound(&*px_l, &x_l_small, n, Number::NEG_INFINITY);
89 let x_u = scatter_bound(&*px_u, &x_u_small, n, Number::INFINITY);
90 let d_l = if m_d > 0 {
91 scatter_bound(&*pd_l, &d_l_small, m_d, Number::NEG_INFINITY)
92 } else {
93 Vec::new()
94 };
95 let d_u = if m_d > 0 {
96 scatter_bound(&*pd_u, &d_u_small, m_d, Number::INFINITY)
97 } else {
98 Vec::new()
99 };
100 let mut x = x_space.make_new_dense();
101 let _ = n_borrow.get_starting_x(&mut x);
102 let x_init = x.expanded_values();
103 (x_l, x_u, d_l, d_u, x_init)
104 };
105
106 Self {
107 nlp,
108 n,
109 m_c,
110 m_d,
111 x_l,
112 x_u,
113 d_l,
114 d_u,
115 x_init,
116 x_space,
117 c_space,
118 d_space,
119 }
120 }
121
122 pub fn new_with_declared_bounds(nlp: Rc<RefCell<dyn IpoptNlp>>) -> Self {
139 let mut me = Self::new(Rc::clone(&nlp));
140 let (n, m_d) = (me.n, me.m_d);
141 let b = nlp.borrow();
142 if let Some((x_l_small, x_u_small)) = b.declared_x_bounds() {
143 me.x_l = scatter_bound(&*b.px_l(), &x_l_small, n, Number::NEG_INFINITY);
144 me.x_u = scatter_bound(&*b.px_u(), &x_u_small, n, Number::INFINITY);
145 }
146 if m_d > 0 {
147 if let Some((d_l_small, d_u_small)) = b.declared_d_bounds() {
148 me.d_l = scatter_bound(&*b.pd_l(), &d_l_small, m_d, Number::NEG_INFINITY);
149 me.d_u = scatter_bound(&*b.pd_u(), &d_u_small, m_d, Number::INFINITY);
150 }
151 }
152 drop(b);
153 me
154 }
155
156 fn dv_from_slice(&self, space: &Rc<DenseVectorSpace>, s: &[Number]) -> DenseVector {
157 let mut dv = space.make_new_dense();
158 dv.set_values(s);
159 dv
160 }
161}
162
163impl SqpProblemSpec for IpoptNlpAdapter {
164 fn n(&self) -> usize {
165 self.n
166 }
167 fn m(&self) -> usize {
168 self.m_c + self.m_d
169 }
170
171 fn x_init(&self) -> Vec<Number> {
172 self.x_init.clone()
173 }
174
175 fn variable_bounds(&self) -> (Vec<Number>, Vec<Number>) {
176 (self.x_l.clone(), self.x_u.clone())
177 }
178
179 fn constraint_bounds(&self) -> (Vec<Number>, Vec<Number>) {
180 let mut bl = vec![0.0; self.m_c];
181 bl.extend_from_slice(&self.d_l);
182 let mut bu = vec![0.0; self.m_c];
183 bu.extend_from_slice(&self.d_u);
184 (bl, bu)
185 }
186
187 fn eval_f(&mut self, x: &[Number]) -> Number {
188 let x_dv = self.dv_from_slice(&self.x_space, x);
189 let mut nlp = self.nlp.borrow_mut();
190 nlp.eval_f(&x_dv)
191 }
192
193 fn eval_grad_f(&mut self, x: &[Number]) -> Vec<Number> {
194 let x_dv = self.dv_from_slice(&self.x_space, x);
195 let mut g = self.x_space.make_new_dense();
196 {
197 let mut nlp = self.nlp.borrow_mut();
198 nlp.eval_grad_f(&x_dv, &mut g);
199 }
200 g.expanded_values()
201 }
202
203 fn eval_c(&mut self, x: &[Number]) -> Vec<Number> {
204 let x_dv = self.dv_from_slice(&self.x_space, x);
205 let mut combined = Vec::with_capacity(self.m_c + self.m_d);
206 if self.m_c > 0 {
207 let mut c_out = self.c_space.make_new_dense();
208 {
209 let mut nlp = self.nlp.borrow_mut();
210 nlp.eval_c(&x_dv, &mut c_out);
211 }
212 combined.extend(c_out.expanded_values());
213 }
214 if self.m_d > 0 {
215 let mut d_out = self.d_space.make_new_dense();
216 {
217 let mut nlp = self.nlp.borrow_mut();
218 nlp.eval_d(&x_dv, &mut d_out);
219 }
220 combined.extend(d_out.expanded_values());
221 }
222 combined
223 }
224
225 fn eval_jac_c(&mut self, x: &[Number]) -> Triplet {
226 let x_dv = self.dv_from_slice(&self.x_space, x);
227 let mut irow = Vec::new();
228 let mut jcol = Vec::new();
229 let mut vals = Vec::new();
230
231 if self.m_c > 0 {
232 let jac_c = {
233 let mut nlp = self.nlp.borrow_mut();
234 nlp.eval_jac_c(&x_dv)
235 };
236 let t = gen_t_downcast(&*jac_c);
237 irow.extend_from_slice(t.irows());
238 jcol.extend_from_slice(t.jcols());
239 vals.extend_from_slice(t.values());
240 }
241
242 if self.m_d > 0 {
243 let jac_d = {
244 let mut nlp = self.nlp.borrow_mut();
245 nlp.eval_jac_d(&x_dv)
246 };
247 let t = gen_t_downcast(&*jac_d);
248 let shift = self.m_c as pounce_common::Index;
249 irow.extend(t.irows().iter().map(|&r| r + shift));
250 jcol.extend_from_slice(t.jcols());
251 vals.extend_from_slice(t.values());
252 }
253
254 Triplet {
255 n_rows: self.m_c + self.m_d,
256 n_cols: self.n,
257 irow,
258 jcol,
259 vals,
260 }
261 }
262
263 fn eval_hess_lag(&mut self, x: &[Number], lambda_g: &[Number]) -> Triplet {
264 let x_dv = self.dv_from_slice(&self.x_space, x);
265 let y_c_dv = self.dv_from_slice(&self.c_space, &lambda_g[..self.m_c]);
266 let y_d_dv = self.dv_from_slice(&self.d_space, &lambda_g[self.m_c..]);
267
268 let h = {
269 let mut nlp = self.nlp.borrow_mut();
270 nlp.eval_h(&x_dv, 1.0, &y_c_dv, &y_d_dv)
271 };
272 let t = sym_t_downcast(&*h);
273 Triplet {
274 n_rows: self.n,
275 n_cols: self.n,
276 irow: t.irows().to_vec(),
277 jcol: t.jcols().to_vec(),
278 vals: t.values().to_vec(),
279 }
280 }
281}
282
283fn vec_from_dyn(v: &dyn pounce_linalg::Vector) -> Vec<Number> {
284 let dv = v
285 .as_any()
286 .downcast_ref::<DenseVector>()
287 .expect("IpoptNlp bound accessors must return DenseVector");
288 dv.expanded_values()
289}
290
291fn scatter_bound(
296 expansion: &dyn pounce_linalg::Matrix,
297 small: &[Number],
298 n_large: usize,
299 fill: Number,
300) -> Vec<Number> {
301 let em = expansion
302 .as_any()
303 .downcast_ref::<ExpansionMatrix>()
304 .expect("px_l / px_u / pd_l / pd_u must be ExpansionMatrix");
305 let exp_pos = em.expanded_pos_indices();
306 debug_assert_eq!(small.len(), exp_pos.len());
307 let mut out = vec![fill; n_large];
308 for (i, &pos) in exp_pos.iter().enumerate() {
309 out[pos as usize] = small[i];
310 }
311 out
312}
313
314fn gen_t_downcast(m: &dyn pounce_linalg::Matrix) -> &GenTMatrix {
315 m.as_any()
316 .downcast_ref::<GenTMatrix>()
317 .expect("IpoptNlp::eval_jac_* must return GenTMatrix")
318}
319
320fn sym_t_downcast(m: &dyn pounce_linalg::matrix::SymMatrix) -> &SymTMatrix {
321 m.as_any()
322 .downcast_ref::<SymTMatrix>()
323 .expect("IpoptNlp::eval_h must return SymTMatrix")
324}
325
326fn bound_maps(nlp: &Rc<RefCell<dyn IpoptNlp>>) -> (Vec<usize>, Vec<usize>) {
336 let b = nlp.borrow();
337 let idx = |m: &dyn pounce_linalg::Matrix| -> Vec<usize> {
338 m.as_any()
339 .downcast_ref::<ExpansionMatrix>()
340 .expect("px_l / px_u must be ExpansionMatrix")
341 .expanded_pos_indices()
342 .iter()
343 .map(|&p| p as usize)
344 .collect()
345 };
346 let lo = idx(&*b.px_l());
347 let up = idx(&*b.px_u());
348 (lo, up)
349}
350
351pub fn pack_bound_multipliers(
358 nlp: &Rc<RefCell<dyn IpoptNlp>>,
359 z_l: &[Number],
360 z_u: &[Number],
361) -> Vec<Number> {
362 let n = nlp.borrow().n() as usize;
363 let (lo_map, up_map) = bound_maps(nlp);
364 let mut out = vec![0.0; n];
365 for (i, &pos) in lo_map.iter().enumerate() {
366 if let Some(&v) = z_l.get(i) {
367 out[pos] += v;
368 }
369 }
370 for (i, &pos) in up_map.iter().enumerate() {
371 if let Some(&v) = z_u.get(i) {
372 out[pos] -= v;
373 }
374 }
375 out
376}
377
378pub fn split_bound_multipliers(
395 nlp: &Rc<RefCell<dyn IpoptNlp>>,
396 lambda_x: &[Number],
397) -> (Vec<Number>, Vec<Number>) {
398 let (lo_map, up_map) = bound_maps(nlp);
399 let z_l = lo_map
400 .iter()
401 .map(|&pos| lambda_x.get(pos).copied().unwrap_or(0.0).max(0.0))
402 .collect();
403 let z_u = up_map
404 .iter()
405 .map(|&pos| (-lambda_x.get(pos).copied().unwrap_or(0.0)).max(0.0))
406 .collect();
407 (z_l, z_u)
408}
409
410pub fn split_slack_multipliers(
420 nlp: &Rc<RefCell<dyn IpoptNlp>>,
421 y_d: &[Number],
422) -> (Vec<Number>, Vec<Number>) {
423 let b = nlp.borrow();
424 let idx = |m: &dyn pounce_linalg::Matrix| -> Vec<usize> {
425 m.as_any()
426 .downcast_ref::<ExpansionMatrix>()
427 .expect("pd_l / pd_u must be ExpansionMatrix")
428 .expanded_pos_indices()
429 .iter()
430 .map(|&p| p as usize)
431 .collect()
432 };
433 let lo_map = idx(&*b.pd_l());
434 let up_map = idx(&*b.pd_u());
435 let v_l = lo_map
436 .iter()
437 .map(|&pos| (-y_d.get(pos).copied().unwrap_or(0.0)).max(0.0))
438 .collect();
439 let v_u = up_map
440 .iter()
441 .map(|&pos| y_d.get(pos).copied().unwrap_or(0.0).max(0.0))
442 .collect();
443 (v_l, v_u)
444}