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 fn dv_from_slice(&self, space: &Rc<DenseVectorSpace>, s: &[Number]) -> DenseVector {
123 let mut dv = space.make_new_dense();
124 dv.set_values(s);
125 dv
126 }
127}
128
129impl SqpProblemSpec for IpoptNlpAdapter {
130 fn n(&self) -> usize {
131 self.n
132 }
133 fn m(&self) -> usize {
134 self.m_c + self.m_d
135 }
136
137 fn x_init(&self) -> Vec<Number> {
138 self.x_init.clone()
139 }
140
141 fn variable_bounds(&self) -> (Vec<Number>, Vec<Number>) {
142 (self.x_l.clone(), self.x_u.clone())
143 }
144
145 fn constraint_bounds(&self) -> (Vec<Number>, Vec<Number>) {
146 let mut bl = vec![0.0; self.m_c];
147 bl.extend_from_slice(&self.d_l);
148 let mut bu = vec![0.0; self.m_c];
149 bu.extend_from_slice(&self.d_u);
150 (bl, bu)
151 }
152
153 fn eval_f(&mut self, x: &[Number]) -> Number {
154 let x_dv = self.dv_from_slice(&self.x_space, x);
155 let mut nlp = self.nlp.borrow_mut();
156 nlp.eval_f(&x_dv)
157 }
158
159 fn eval_grad_f(&mut self, x: &[Number]) -> Vec<Number> {
160 let x_dv = self.dv_from_slice(&self.x_space, x);
161 let mut g = self.x_space.make_new_dense();
162 {
163 let mut nlp = self.nlp.borrow_mut();
164 nlp.eval_grad_f(&x_dv, &mut g);
165 }
166 g.expanded_values()
167 }
168
169 fn eval_c(&mut self, x: &[Number]) -> Vec<Number> {
170 let x_dv = self.dv_from_slice(&self.x_space, x);
171 let mut combined = Vec::with_capacity(self.m_c + self.m_d);
172 if self.m_c > 0 {
173 let mut c_out = self.c_space.make_new_dense();
174 {
175 let mut nlp = self.nlp.borrow_mut();
176 nlp.eval_c(&x_dv, &mut c_out);
177 }
178 combined.extend(c_out.expanded_values());
179 }
180 if self.m_d > 0 {
181 let mut d_out = self.d_space.make_new_dense();
182 {
183 let mut nlp = self.nlp.borrow_mut();
184 nlp.eval_d(&x_dv, &mut d_out);
185 }
186 combined.extend(d_out.expanded_values());
187 }
188 combined
189 }
190
191 fn eval_jac_c(&mut self, x: &[Number]) -> Triplet {
192 let x_dv = self.dv_from_slice(&self.x_space, x);
193 let mut irow = Vec::new();
194 let mut jcol = Vec::new();
195 let mut vals = Vec::new();
196
197 if self.m_c > 0 {
198 let jac_c = {
199 let mut nlp = self.nlp.borrow_mut();
200 nlp.eval_jac_c(&x_dv)
201 };
202 let t = gen_t_downcast(&*jac_c);
203 irow.extend_from_slice(t.irows());
204 jcol.extend_from_slice(t.jcols());
205 vals.extend_from_slice(t.values());
206 }
207
208 if self.m_d > 0 {
209 let jac_d = {
210 let mut nlp = self.nlp.borrow_mut();
211 nlp.eval_jac_d(&x_dv)
212 };
213 let t = gen_t_downcast(&*jac_d);
214 let shift = self.m_c as pounce_common::Index;
215 irow.extend(t.irows().iter().map(|&r| r + shift));
216 jcol.extend_from_slice(t.jcols());
217 vals.extend_from_slice(t.values());
218 }
219
220 Triplet {
221 n_rows: self.m_c + self.m_d,
222 n_cols: self.n,
223 irow,
224 jcol,
225 vals,
226 }
227 }
228
229 fn eval_hess_lag(&mut self, x: &[Number], lambda_g: &[Number]) -> Triplet {
230 let x_dv = self.dv_from_slice(&self.x_space, x);
231 let y_c_dv = self.dv_from_slice(&self.c_space, &lambda_g[..self.m_c]);
232 let y_d_dv = self.dv_from_slice(&self.d_space, &lambda_g[self.m_c..]);
233
234 let h = {
235 let mut nlp = self.nlp.borrow_mut();
236 nlp.eval_h(&x_dv, 1.0, &y_c_dv, &y_d_dv)
237 };
238 let t = sym_t_downcast(&*h);
239 Triplet {
240 n_rows: self.n,
241 n_cols: self.n,
242 irow: t.irows().to_vec(),
243 jcol: t.jcols().to_vec(),
244 vals: t.values().to_vec(),
245 }
246 }
247}
248
249fn vec_from_dyn(v: &dyn pounce_linalg::Vector) -> Vec<Number> {
250 let dv = v
251 .as_any()
252 .downcast_ref::<DenseVector>()
253 .expect("IpoptNlp bound accessors must return DenseVector");
254 dv.expanded_values()
255}
256
257fn scatter_bound(
262 expansion: &dyn pounce_linalg::Matrix,
263 small: &[Number],
264 n_large: usize,
265 fill: Number,
266) -> Vec<Number> {
267 let em = expansion
268 .as_any()
269 .downcast_ref::<ExpansionMatrix>()
270 .expect("px_l / px_u / pd_l / pd_u must be ExpansionMatrix");
271 let exp_pos = em.expanded_pos_indices();
272 debug_assert_eq!(small.len(), exp_pos.len());
273 let mut out = vec![fill; n_large];
274 for (i, &pos) in exp_pos.iter().enumerate() {
275 out[pos as usize] = small[i];
276 }
277 out
278}
279
280fn gen_t_downcast(m: &dyn pounce_linalg::Matrix) -> &GenTMatrix {
281 m.as_any()
282 .downcast_ref::<GenTMatrix>()
283 .expect("IpoptNlp::eval_jac_* must return GenTMatrix")
284}
285
286fn sym_t_downcast(m: &dyn pounce_linalg::matrix::SymMatrix) -> &SymTMatrix {
287 m.as_any()
288 .downcast_ref::<SymTMatrix>()
289 .expect("IpoptNlp::eval_h must return SymTMatrix")
290}