1use pounce_common::types::{Index, Number};
9use pounce_nlp::tnlp::Linearity;
10
11#[derive(Debug, Clone, Copy)]
15pub struct ProbeView<'a> {
16 pub n_vars: usize,
18 pub m_rows: usize,
20 pub jac_irow: &'a [Index],
22 pub jac_jcol: &'a [Index],
23 pub jac_values: Option<&'a [Number]>,
28 pub g_l: &'a [Number],
29 pub g_u: &'a [Number],
30 pub linearity: Option<&'a [Linearity]>,
33 pub one_based: bool,
35 pub eq_tol: Number,
38 pub excluded_vars: Option<&'a [bool]>,
42 pub excluded_rows: Option<&'a [bool]>,
46}
47
48#[derive(Debug, Clone, Default)]
80pub struct EqualityIncidence {
81 pub n_vars: usize,
83 pub eq_row_inner_idx: Vec<usize>,
86 pub adj_ptr: Vec<usize>,
88 pub vars: Vec<usize>,
90}
91
92impl EqualityIncidence {
93 pub fn from_probe(p: &ProbeView<'_>) -> Self {
95 let mut eq_row_inner_idx: Vec<usize> = Vec::new();
98 let mut inner_to_eq: Vec<Option<usize>> = vec![None; p.m_rows];
99 for (i, slot) in inner_to_eq.iter_mut().enumerate() {
100 if let Some(mask) = p.excluded_rows {
101 if mask[i] {
102 continue;
103 }
104 }
105 if (p.g_u[i] - p.g_l[i]).abs() <= p.eq_tol {
106 *slot = Some(eq_row_inner_idx.len());
107 eq_row_inner_idx.push(i);
108 }
109 }
110 let n_eq = eq_row_inner_idx.len();
111
112 let mut per_row: Vec<Vec<usize>> = vec![Vec::new(); n_eq];
114 let nnz = p.jac_irow.len();
115 for k in 0..nnz {
116 if let Some(vals) = p.jac_values {
118 if vals[k] == 0.0 {
119 continue;
120 }
121 }
122 let i = if p.one_based {
123 (p.jac_irow[k] as isize - 1) as usize
124 } else {
125 p.jac_irow[k] as usize
126 };
127 if i >= p.m_rows {
128 continue;
129 }
130 let Some(eq_k) = inner_to_eq[i] else { continue };
131 let j = if p.one_based {
132 (p.jac_jcol[k] as isize - 1) as usize
133 } else {
134 p.jac_jcol[k] as usize
135 };
136 if j >= p.n_vars {
137 continue;
138 }
139 if let Some(mask) = p.excluded_vars {
141 if mask[j] {
142 continue;
143 }
144 }
145 per_row[eq_k].push(j);
146 }
147
148 let mut adj_ptr: Vec<usize> = Vec::with_capacity(n_eq + 1);
150 let mut vars: Vec<usize> = Vec::new();
151 adj_ptr.push(0);
152 for row in per_row.iter_mut() {
153 row.sort_unstable();
154 row.dedup();
155 vars.extend_from_slice(row);
156 adj_ptr.push(vars.len());
157 }
158
159 Self {
160 n_vars: p.n_vars,
161 eq_row_inner_idx,
162 adj_ptr,
163 vars,
164 }
165 }
166
167 pub fn n_eq_rows(&self) -> usize {
169 self.eq_row_inner_idx.len()
170 }
171
172 pub fn neighbors(&self, k: usize) -> &[usize] {
175 let lo = self.adj_ptr[k];
176 let hi = self.adj_ptr[k + 1];
177 &self.vars[lo..hi]
178 }
179}
180
181#[derive(Debug, Clone, Default)]
188pub struct InequalityIncidence {
189 pub n_vars: usize,
191 pub ineq_row_inner_idx: Vec<usize>,
193 pub adj_ptr: Vec<usize>,
195 pub vars: Vec<usize>,
197 pub col_to_rows_ptr: Vec<usize>,
200 pub col_to_rows: Vec<usize>,
201}
202
203impl InequalityIncidence {
204 pub fn from_probe(p: &ProbeView<'_>) -> Self {
208 let mut ineq_row_inner_idx: Vec<usize> = Vec::new();
212 let mut inner_to_ineq: Vec<Option<usize>> = vec![None; p.m_rows];
213 for (i, slot) in inner_to_ineq.iter_mut().enumerate() {
214 if let Some(mask) = p.excluded_rows {
215 if mask[i] {
216 continue;
217 }
218 }
219 if (p.g_u[i] - p.g_l[i]).abs() > p.eq_tol {
220 *slot = Some(ineq_row_inner_idx.len());
221 ineq_row_inner_idx.push(i);
222 }
223 }
224 let n_ineq = ineq_row_inner_idx.len();
225
226 let mut per_row: Vec<Vec<usize>> = vec![Vec::new(); n_ineq];
228 let nnz = p.jac_irow.len();
229 for k in 0..nnz {
230 if let Some(vals) = p.jac_values {
231 if vals[k] == 0.0 {
232 continue;
233 }
234 }
235 let i = if p.one_based {
236 (p.jac_irow[k] as isize - 1) as usize
237 } else {
238 p.jac_irow[k] as usize
239 };
240 if i >= p.m_rows {
241 continue;
242 }
243 let Some(ineq_k) = inner_to_ineq[i] else {
244 continue;
245 };
246 let j = if p.one_based {
247 (p.jac_jcol[k] as isize - 1) as usize
248 } else {
249 p.jac_jcol[k] as usize
250 };
251 if j >= p.n_vars {
252 continue;
253 }
254 if let Some(mask) = p.excluded_vars {
256 if mask[j] {
257 continue;
258 }
259 }
260 per_row[ineq_k].push(j);
261 }
262
263 let mut adj_ptr: Vec<usize> = Vec::with_capacity(n_ineq + 1);
265 let mut vars: Vec<usize> = Vec::new();
266 adj_ptr.push(0);
267 for row in per_row.iter_mut() {
268 row.sort_unstable();
269 row.dedup();
270 vars.extend_from_slice(row);
271 adj_ptr.push(vars.len());
272 }
273
274 let mut col_to_rows_ptr = vec![0usize; p.n_vars + 1];
277 for &v in &vars {
278 col_to_rows_ptr[v + 1] += 1;
279 }
280 for i in 1..=p.n_vars {
281 col_to_rows_ptr[i] += col_to_rows_ptr[i - 1];
282 }
283 let mut col_to_rows = vec![0usize; col_to_rows_ptr[p.n_vars]];
284 let mut cursor = col_to_rows_ptr[..p.n_vars].to_vec();
285 for (ineq_k, row) in per_row.iter().enumerate() {
286 for &v in row {
287 col_to_rows[cursor[v]] = ineq_k;
288 cursor[v] += 1;
289 }
290 }
291
292 Self {
293 n_vars: p.n_vars,
294 ineq_row_inner_idx,
295 adj_ptr,
296 vars,
297 col_to_rows_ptr,
298 col_to_rows,
299 }
300 }
301
302 pub fn n_ineq_rows(&self) -> usize {
304 self.ineq_row_inner_idx.len()
305 }
306
307 pub fn neighbors(&self, k: usize) -> &[usize] {
309 let lo = self.adj_ptr[k];
310 let hi = self.adj_ptr[k + 1];
311 &self.vars[lo..hi]
312 }
313
314 pub fn rows_for_var(&self, j: usize) -> &[usize] {
318 let lo = self.col_to_rows_ptr[j];
319 let hi = self.col_to_rows_ptr[j + 1];
320 &self.col_to_rows[lo..hi]
321 }
322
323 pub fn var_in_inequality(&self, j: usize) -> bool {
325 !self.rows_for_var(j).is_empty()
326 }
327}
328
329#[cfg(test)]
330mod tests {
331 use super::*;
332
333 fn probe<'a>(
334 n_vars: usize,
335 m_rows: usize,
336 irow: &'a [Index],
337 jcol: &'a [Index],
338 g_l: &'a [Number],
339 g_u: &'a [Number],
340 ) -> ProbeView<'a> {
341 ProbeView {
342 n_vars,
343 m_rows,
344 jac_irow: irow,
345 jac_jcol: jcol,
346 jac_values: None,
347 g_l,
348 g_u,
349 linearity: None,
350 one_based: false,
351 eq_tol: 1e-12,
352 excluded_vars: None,
353 excluded_rows: None,
354 }
355 }
356
357 #[test]
358 fn incidence_empty_problem_is_empty() {
359 let p = probe(0, 0, &[], &[], &[], &[]);
360 let inc = EqualityIncidence::from_probe(&p);
361 assert_eq!(inc.n_eq_rows(), 0);
362 assert_eq!(inc.vars.len(), 0);
363 assert_eq!(inc.adj_ptr, vec![0]);
364 }
365
366 #[test]
367 fn incidence_filters_inequalities() {
368 let p = probe(2, 2, &[0, 0, 1], &[0, 1, 0], &[1.0, 0.0], &[1.0, 5.0]);
370 let inc = EqualityIncidence::from_probe(&p);
371 assert_eq!(inc.n_eq_rows(), 1);
372 assert_eq!(inc.eq_row_inner_idx, vec![0]);
373 assert_eq!(inc.neighbors(0), &[0, 1]);
374 }
375
376 #[test]
377 fn incidence_dedupes_and_sorts_columns() {
378 let p = probe(2, 1, &[0, 0, 0, 0], &[1, 1, 0, 1], &[2.5], &[2.5]);
381 let inc = EqualityIncidence::from_probe(&p);
382 assert_eq!(inc.neighbors(0), &[0, 1]);
383 }
384
385 #[test]
386 fn incidence_respects_fortran_indexing() {
387 let mut p = probe(
388 2,
389 1,
390 &[1, 1], &[1, 2], &[0.0],
393 &[0.0],
394 );
395 p.one_based = true;
396 let inc = EqualityIncidence::from_probe(&p);
397 assert_eq!(inc.n_eq_rows(), 1);
398 assert_eq!(inc.neighbors(0), &[0, 1]);
399 }
400
401 #[test]
402 fn incidence_drops_structural_zeros_when_values_provided() {
403 let vals = [3.5, 0.0];
406 let p = ProbeView {
407 n_vars: 2,
408 m_rows: 1,
409 jac_irow: &[0, 0],
410 jac_jcol: &[0, 1],
411 jac_values: Some(&vals),
412 g_l: &[1.0],
413 g_u: &[1.0],
414 linearity: None,
415 one_based: false,
416 eq_tol: 1e-12,
417 excluded_vars: None,
418 excluded_rows: None,
419 };
420 let inc = EqualityIncidence::from_probe(&p);
421 assert_eq!(inc.neighbors(0), &[0]);
422 }
423
424 #[test]
425 fn inequality_incidence_filters_equalities() {
426 let p = probe(2, 2, &[0, 0, 1], &[0, 1, 0], &[1.0, 0.0], &[1.0, 5.0]);
429 let ineq = InequalityIncidence::from_probe(&p);
430 assert_eq!(ineq.n_ineq_rows(), 1);
431 assert_eq!(ineq.ineq_row_inner_idx, vec![1]);
432 assert_eq!(ineq.neighbors(0), &[0]);
433 assert!(ineq.var_in_inequality(0));
434 assert!(!ineq.var_in_inequality(1));
435 }
436
437 #[test]
438 fn inequality_incidence_range_row() {
439 let p = probe(2, 1, &[0, 0], &[0, 1], &[-2.0], &[5.0]);
441 let ineq = InequalityIncidence::from_probe(&p);
442 assert_eq!(ineq.n_ineq_rows(), 1);
443 assert_eq!(ineq.neighbors(0), &[0, 1]);
444 assert_eq!(ineq.rows_for_var(0), &[0]);
446 assert_eq!(ineq.rows_for_var(1), &[0]);
447 }
448
449 #[test]
450 fn inequality_incidence_one_sided() {
451 let p = probe(1, 1, &[0], &[0], &[-1e19], &[5.0]);
454 let ineq = InequalityIncidence::from_probe(&p);
455 assert_eq!(ineq.n_ineq_rows(), 1);
456 assert_eq!(ineq.neighbors(0), &[0]);
457 }
458}