1use crate::kkt::aug_system_solver::{AugSysCoeffs, AugSysRhs, AugSysSol, AugSystemSolver};
30use pounce_common::diagnostics::{DiagCategory, DiagnosticsState};
31use pounce_common::timing::TimingStatistics;
32use pounce_common::types::{Index, Number};
33use pounce_linalg::Vector;
34use pounce_linalg::compound_vector::CompoundVector;
35use pounce_linalg::dense_vector::DenseVector;
36use pounce_linalg::diag_matrix::DiagMatrix;
37use pounce_linalg::triplet::{GenTMatrix, SymTMatrix};
38use pounce_linsol::{ESymSolverStatus, FactorPattern, SymLinearSolver, TSymLinearSolver};
39use std::ops::Range;
40use std::rc::Rc;
41
42pub struct StdAugSystemSolver {
44 linsol: TSymLinearSolver,
45
46 initialized: bool,
48 struct_sig: Option<(usize, usize, usize, Index, Index, Index)>,
57 n_x: Index,
58 n_s: Index,
59 n_c: Index,
60 n_d: Index,
61 dim: Index,
63
64 irn: Vec<Index>,
66 jcn: Vec<Index>,
68 vals: Vec<Number>,
70
71 w_range: Range<usize>,
73 dx_range: Range<usize>,
74 ds_range: Range<usize>,
75 jc_range: Range<usize>,
76 dc_range: Range<usize>,
77 jd_range: Range<usize>,
78 minus_i_range: Range<usize>,
79 dd_range: Range<usize>,
80
81 last_neg_evals: Index,
82 last_status: Option<ESymSolverStatus>,
83
84 have_factor: bool,
88
89 timing: Option<Rc<TimingStatistics>>,
93
94 diagnostics: Option<Rc<DiagnosticsState>>,
100
101 legacy_dump_path: std::cell::OnceCell<Option<String>>,
109
110 inertia_sink:
115 Option<std::sync::Arc<std::sync::Mutex<pounce_linsol::summary::LinearSolverSummary>>>,
116}
117
118impl std::fmt::Debug for StdAugSystemSolver {
119 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120 f.debug_struct("StdAugSystemSolver")
121 .field("dim", &self.dim)
122 .field("nnz", &self.vals.len())
123 .field("initialized", &self.initialized)
124 .field("last_neg_evals", &self.last_neg_evals)
125 .field("last_status", &self.last_status)
126 .finish_non_exhaustive()
127 }
128}
129
130impl StdAugSystemSolver {
131 pub fn new(linsol: TSymLinearSolver) -> Self {
133 Self {
134 linsol,
135 initialized: false,
136 struct_sig: None,
137 legacy_dump_path: std::cell::OnceCell::new(),
138 n_x: 0,
139 n_s: 0,
140 n_c: 0,
141 n_d: 0,
142 dim: 0,
143 irn: Vec::new(),
144 jcn: Vec::new(),
145 vals: Vec::new(),
146 w_range: 0..0,
147 dx_range: 0..0,
148 ds_range: 0..0,
149 jc_range: 0..0,
150 dc_range: 0..0,
151 jd_range: 0..0,
152 minus_i_range: 0..0,
153 dd_range: 0..0,
154 last_neg_evals: 0,
155 last_status: None,
156 have_factor: false,
157 timing: None,
158 diagnostics: None,
159 inertia_sink: None,
160 }
161 }
162
163 pub fn with_inertia_sink(
166 mut self,
167 sink: std::sync::Arc<std::sync::Mutex<pounce_linsol::summary::LinearSolverSummary>>,
168 ) -> Self {
169 self.inertia_sink = Some(sink);
170 self
171 }
172
173 fn build_structure(&mut self, coeffs: &AugSysCoeffs<'_>) -> ESymSolverStatus {
174 let n_x = coeffs.j_c.n_cols();
175 let n_c = coeffs.j_c.n_rows();
176 let n_d = coeffs.j_d.n_rows();
177 debug_assert_eq!(coeffs.j_d.n_cols(), n_x);
178 let n_s = n_d;
179
180 let w_nnz = match coeffs.w {
181 None => 0_usize,
182 Some(w) => w_nonzeros(w),
183 };
184 let jc_nnz = gen_t_downcast(coeffs.j_c).nonzeros() as usize;
185 let jd_nnz = gen_t_downcast(coeffs.j_d).nonzeros() as usize;
186
187 let total = w_nnz
188 + (n_x as usize) + (n_s as usize) + jc_nnz
191 + (n_c as usize) + jd_nnz
193 + (n_s as usize) + (n_d as usize); self.irn = Vec::with_capacity(total);
197 self.jcn = Vec::with_capacity(total);
198 self.vals = vec![0.0; total];
199
200 let w_start = self.irn.len();
202 if let Some(w) = coeffs.w {
203 if let Some(t) = w.as_any().downcast_ref::<SymTMatrix>() {
204 self.irn.extend_from_slice(t.irows());
205 self.jcn.extend_from_slice(t.jcols());
206 } else if let Some(dm) = w.as_any().downcast_ref::<DiagMatrix>() {
207 let n = w_diag_dim(dm);
210 for i in 0..n {
211 self.irn.push(i + 1);
212 self.jcn.push(i + 1);
213 }
214 } else {
215 unreachable!("StdAugSystemSolver: W must be a SymTMatrix or DiagMatrix in v1.0")
216 }
217 }
218 self.w_range = w_start..self.irn.len();
219
220 let dx_start = self.irn.len();
222 for i in 0..n_x {
223 self.irn.push(i + 1);
224 self.jcn.push(i + 1);
225 }
226 self.dx_range = dx_start..self.irn.len();
227
228 let ds_start = self.irn.len();
230 for i in 0..n_s {
231 let r = n_x + i + 1;
232 self.irn.push(r);
233 self.jcn.push(r);
234 }
235 self.ds_range = ds_start..self.irn.len();
236
237 let jc_start = self.irn.len();
239 let j_c = gen_t_downcast(coeffs.j_c);
240 let row_off_c = n_x + n_s;
241 for (&i, &j) in j_c.irows().iter().zip(j_c.jcols().iter()) {
242 self.irn.push(row_off_c + i);
245 self.jcn.push(j);
246 }
247 self.jc_range = jc_start..self.irn.len();
248
249 let dc_start = self.irn.len();
251 for i in 0..n_c {
252 let r = n_x + n_s + i + 1;
253 self.irn.push(r);
254 self.jcn.push(r);
255 }
256 self.dc_range = dc_start..self.irn.len();
257
258 let jd_start = self.irn.len();
260 let j_d = gen_t_downcast(coeffs.j_d);
261 let row_off_d = n_x + n_s + n_c;
262 for (&i, &j) in j_d.irows().iter().zip(j_d.jcols().iter()) {
263 self.irn.push(row_off_d + i);
264 self.jcn.push(j);
265 }
266 self.jd_range = jd_start..self.irn.len();
267
268 let mi_start = self.irn.len();
270 for i in 0..n_s {
271 self.irn.push(n_x + n_s + n_c + i + 1);
272 self.jcn.push(n_x + i + 1);
273 }
274 self.minus_i_range = mi_start..self.irn.len();
275
276 let dd_start = self.irn.len();
278 for i in 0..n_d {
279 let r = n_x + n_s + n_c + i + 1;
280 self.irn.push(r);
281 self.jcn.push(r);
282 }
283 self.dd_range = dd_start..self.irn.len();
284
285 debug_assert_eq!(self.irn.len(), total);
286 debug_assert_eq!(self.jcn.len(), total);
287
288 self.n_x = n_x;
289 self.n_s = n_s;
290 self.n_c = n_c;
291 self.n_d = n_d;
292 self.dim = n_x + n_s + n_c + n_d;
293
294 let status = self
295 .linsol
296 .initialize_structure(self.dim, &self.irn, &self.jcn);
297 if status == ESymSolverStatus::Success {
298 self.initialized = true;
299 }
300 status
301 }
302
303 fn refill_values(&mut self, coeffs: &AugSysCoeffs<'_>) {
304 if !self.w_range.is_empty() {
306 let Some(w_dyn) = coeffs.w else {
307 unreachable!("structure pinned with W; W cannot be None now")
308 };
309 let dst = &mut self.vals[self.w_range.clone()];
310 if let Some(t) = w_dyn.as_any().downcast_ref::<SymTMatrix>() {
311 for (d, &v) in dst.iter_mut().zip(t.values().iter()) {
312 *d = coeffs.w_factor * v;
313 }
314 } else if let Some(dm) = w_dyn.as_any().downcast_ref::<DiagMatrix>() {
315 let diag = w_diag_values(dm);
316 for (d, &v) in dst.iter_mut().zip(diag.iter()) {
317 *d = coeffs.w_factor * v;
318 }
319 } else {
320 unreachable!("StdAugSystemSolver: W must be a SymTMatrix or DiagMatrix in v1.0")
321 }
322 }
323 fill_diag(
325 &mut self.vals[self.dx_range.clone()],
326 coeffs.d_x,
327 coeffs.delta_x,
328 1.0,
329 );
330 fill_diag(
332 &mut self.vals[self.ds_range.clone()],
333 coeffs.d_s,
334 coeffs.delta_s,
335 1.0,
336 );
337 {
339 let j_c = gen_t_downcast(coeffs.j_c);
340 self.vals[self.jc_range.clone()].copy_from_slice(j_c.values());
341 }
342 fill_diag(
344 &mut self.vals[self.dc_range.clone()],
345 coeffs.d_c,
346 coeffs.delta_c,
347 -1.0,
348 );
349 {
351 let j_d = gen_t_downcast(coeffs.j_d);
352 self.vals[self.jd_range.clone()].copy_from_slice(j_d.values());
353 }
354 for v in &mut self.vals[self.minus_i_range.clone()] {
356 *v = -1.0;
357 }
358 fill_diag(
360 &mut self.vals[self.dd_range.clone()],
361 coeffs.d_d,
362 coeffs.delta_d,
363 -1.0,
364 );
365 }
366
367 pub(crate) fn assemble(&mut self, coeffs: &AugSysCoeffs<'_>) -> ESymSolverStatus {
376 let sig = {
377 let w_nnz = coeffs.w.map(w_nonzeros).unwrap_or(0);
378 let jc_nnz = gen_t_downcast(coeffs.j_c).nonzeros() as usize;
379 let jd_nnz = gen_t_downcast(coeffs.j_d).nonzeros() as usize;
380 (
381 w_nnz,
382 jc_nnz,
383 jd_nnz,
384 coeffs.j_c.n_cols(),
385 coeffs.j_c.n_rows(),
386 coeffs.j_d.n_rows(),
387 )
388 };
389 if !self.initialized || self.struct_sig != Some(sig) {
390 let timing = self.timing.clone();
402 let _sym_guard = timing
403 .as_deref()
404 .map(|t| t.linear_system_symbolic_factorization.guard());
405 let s = self.build_structure(coeffs);
406 if s != ESymSolverStatus::Success {
407 self.last_status = Some(s);
408 return s;
409 }
410 self.struct_sig = Some(sig);
411 }
412 self.refill_values(coeffs);
413 ESymSolverStatus::Success
414 }
415
416 pub(crate) fn assembled_dim(&self) -> Index {
418 self.dim
419 }
420 pub(crate) fn assembled_triplet(&self) -> (&[Index], &[Index], &[Number]) {
423 (&self.irn, &self.jcn, &self.vals)
424 }
425
426 fn legacy_dump_path(&self) -> Option<&String> {
429 self.legacy_dump_path
430 .get_or_init(|| std::env::var("POUNCE_DUMP_KKT").ok())
431 .as_ref()
432 }
433
434 fn kkt_dump_active(&self) -> bool {
438 self.diagnostics
439 .as_deref()
440 .is_some_and(|d| d.want(DiagCategory::Kkt))
441 || self.legacy_dump_path().is_some()
442 }
443
444 pub(crate) fn pack_rhs(&self, rhs: &AugSysRhs<'_>, packed: &mut [Number]) {
445 let n_x = self.n_x as usize;
446 let n_s = self.n_s as usize;
447 let n_c = self.n_c as usize;
448 let n_d = self.n_d as usize;
449 copy_vec(rhs.rhs_x, &mut packed[..n_x]);
450 copy_vec(rhs.rhs_s, &mut packed[n_x..n_x + n_s]);
451 copy_vec(rhs.rhs_c, &mut packed[n_x + n_s..n_x + n_s + n_c]);
452 copy_vec(
453 rhs.rhs_d,
454 &mut packed[n_x + n_s + n_c..n_x + n_s + n_c + n_d],
455 );
456 }
457
458 pub(crate) fn unpack_sol(&self, packed: &[Number], sol: &mut AugSysSol<'_>) {
459 let n_x = self.n_x as usize;
460 let n_s = self.n_s as usize;
461 let n_c = self.n_c as usize;
462 let n_d = self.n_d as usize;
463 write_vec(sol.sol_x, &packed[..n_x]);
464 write_vec(sol.sol_s, &packed[n_x..n_x + n_s]);
465 write_vec(sol.sol_c, &packed[n_x + n_s..n_x + n_s + n_c]);
466 write_vec(sol.sol_d, &packed[n_x + n_s + n_c..n_x + n_s + n_c + n_d]);
467 }
468}
469
470impl AugSystemSolver for StdAugSystemSolver {
471 fn provides_inertia(&self) -> bool {
472 self.linsol.provides_inertia()
473 }
474
475 fn singularity_certified(&self) -> bool {
476 self.linsol.singularity_certified()
477 }
478
479 fn number_of_neg_evals(&self) -> Index {
480 self.last_neg_evals
481 }
482
483 fn system_dim(&self) -> Index {
484 self.dim
485 }
486
487 fn kkt_triplets(&self) -> Option<(Index, Vec<Index>, Vec<Index>, Vec<Number>)> {
488 if self.irn.is_empty() {
489 return None;
490 }
491 Some((
492 self.dim,
493 self.irn.clone(),
494 self.jcn.clone(),
495 self.vals.clone(),
496 ))
497 }
498
499 fn l_factor(&self, want_values: bool) -> Option<FactorPattern> {
500 self.linsol.factor_pattern(want_values)
501 }
502
503 fn increase_quality(&mut self) -> bool {
504 self.have_factor = false;
508 self.linsol.increase_quality()
509 }
510
511 fn last_solve_status(&self) -> ESymSolverStatus {
512 self.last_status.unwrap_or(ESymSolverStatus::FatalError)
513 }
514
515 fn solve(
516 &mut self,
517 coeffs: &AugSysCoeffs<'_>,
518 rhs: &AugSysRhs<'_>,
519 sol: &mut AugSysSol<'_>,
520 check_neg_evals: bool,
521 num_neg_evals: Index,
522 ) -> ESymSolverStatus {
523 let s = self.assemble(coeffs);
531 if s != ESymSolverStatus::Success {
532 return s;
533 }
534
535 let mut packed = vec![0.0; self.dim as usize];
536 self.pack_rhs(rhs, &mut packed);
537
538 let dump_rhs = packed.clone();
539
540 let _factor_guard = self
544 .timing
545 .as_deref()
546 .map(|t| t.linear_system_factorization.guard());
547 let factors_before = self
548 .inertia_sink
549 .as_ref()
550 .and_then(|s| s.lock().ok().map(|g| g.n_factors));
551 let status = self.linsol.multi_solve(
552 &self.vals,
553 true,
554 1,
555 &mut packed,
556 check_neg_evals,
557 num_neg_evals,
558 );
559 drop(_factor_guard);
560 self.last_status = Some(status);
561 if let Some(sink) = &self.inertia_sink
567 && coeffs.delta_x == 0.0
568 && coeffs.delta_s == 0.0
569 && coeffs.delta_c == 0.0
570 && coeffs.delta_d == 0.0
571 && let Ok(mut g) = sink.lock()
572 && factors_before.is_some_and(|b| g.n_factors > b)
573 {
574 g.last_inertia_unregularized = g.last_inertia;
575 }
576 if self.linsol.provides_inertia()
588 && matches!(
589 status,
590 ESymSolverStatus::Success
591 | ESymSolverStatus::WrongInertia
592 | ESymSolverStatus::Singular
593 )
594 {
595 self.last_neg_evals = self.linsol.number_of_neg_evals();
596 }
597 if status == ESymSolverStatus::Success {
598 self.unpack_sol(&packed, sol);
599 self.have_factor = true;
600 }
601
602 if let Some(diag) = self.diagnostics.clone() {
607 if diag.want(DiagCategory::Kkt) {
608 let solve_idx = diag.next_solve_index();
609 let filename = format!("kkt_solve_{solve_idx:03}.jsonl");
610 let variant = diag.config.kkt_variant;
617 let factor_pattern =
618 if status == ESymSolverStatus::Success && variant.wants_l_pattern() {
619 self.linsol.factor_pattern(variant.wants_l_values())
620 } else {
621 None
622 };
623 if let Some(mut w) = diag.open_writer(&filename) {
624 let _ = write_kkt_record(
625 &mut w,
626 self.dim,
627 &self.irn,
628 &self.jcn,
629 &self.vals,
630 &dump_rhs,
631 &packed,
632 check_neg_evals,
633 num_neg_evals,
634 status,
635 self.last_neg_evals,
636 factor_pattern.as_ref(),
637 );
638 }
639 }
640 }
641 if let Some(path) = self.legacy_dump_path().cloned() {
642 use std::sync::atomic::{AtomicBool, Ordering};
643 static WARNED: AtomicBool = AtomicBool::new(false);
644 if !WARNED.swap(true, Ordering::SeqCst) {
645 tracing::warn!(target: "pounce::linsol",
646 "warning: POUNCE_DUMP_KKT is deprecated; prefer `--dump kkt:<iter-spec>` (see pounce --help)"
647 );
648 }
649 dump_kkt(
650 &path,
651 self.dim,
652 &self.irn,
653 &self.jcn,
654 &self.vals,
655 &dump_rhs,
656 &packed,
657 check_neg_evals,
658 num_neg_evals,
659 status,
660 self.last_neg_evals,
661 );
662 }
663
664 status
665 }
666
667 fn resolve(
668 &mut self,
669 coeffs: &AugSysCoeffs<'_>,
670 rhs: &AugSysRhs<'_>,
671 sol: &mut AugSysSol<'_>,
672 ) -> ESymSolverStatus {
673 if !self.have_factor {
680 return self.solve(coeffs, rhs, sol, false, 0);
681 }
682
683 let mut packed = vec![0.0; self.dim as usize];
684 self.pack_rhs(rhs, &mut packed);
685
686 let _back_guard = self
689 .timing
690 .as_deref()
691 .map(|t| t.linear_system_back_solve.guard());
692 let status = self
693 .linsol
694 .multi_solve(&self.vals, false, 1, &mut packed, false, 0);
695 drop(_back_guard);
696 self.last_status = Some(status);
697 if status == ESymSolverStatus::Success {
698 self.unpack_sol(&packed, sol);
699 }
700 status
701 }
702
703 fn set_diagnostics(&mut self, diag: Rc<DiagnosticsState>) {
704 self.diagnostics = Some(diag);
705 }
706
707 fn set_slack_scaling(&mut self, nx: Index, s_scale: &[Number]) {
708 self.linsol.set_slack_scaling(nx, s_scale);
709 }
710
711 fn set_timing_stats(&mut self, timing: Rc<TimingStatistics>) {
712 self.timing = Some(timing);
713 }
714
715 fn multi_solve_matches_single_solve(&self, nrhs: usize) -> bool {
716 self.linsol.multi_solve_matches_single_solve(nrhs)
717 }
718
719 fn try_solve_many_flat(
720 &mut self,
721 coeffs: &AugSysCoeffs<'_>,
722 packed_rhs: &mut [Number],
723 nrhs: usize,
724 check_neg_evals: bool,
725 num_neg_evals: Index,
726 ) -> Option<ESymSolverStatus> {
727 if self.kkt_dump_active() {
731 return None;
732 }
733 let s = self.assemble(coeffs);
734 if s != ESymSolverStatus::Success {
735 self.last_status = Some(s);
736 return Some(s);
737 }
738 if packed_rhs.len() != (self.dim as usize) * nrhs {
744 return None;
745 }
746
747 let _factor_guard = self
754 .timing
755 .as_deref()
756 .map(|t| t.linear_system_factorization.guard());
757 let status = self.linsol.multi_solve(
758 &self.vals,
759 true,
760 nrhs as Index,
761 packed_rhs,
762 check_neg_evals,
763 num_neg_evals,
764 );
765 drop(_factor_guard);
766 self.last_status = Some(status);
767 if self.linsol.provides_inertia()
770 && matches!(
771 status,
772 ESymSolverStatus::Success
773 | ESymSolverStatus::WrongInertia
774 | ESymSolverStatus::Singular
775 )
776 {
777 self.last_neg_evals = self.linsol.number_of_neg_evals();
778 }
779 if status == ESymSolverStatus::Success {
780 self.have_factor = true;
781 }
782 Some(status)
783 }
784
785 fn try_resolve_many_flat(
786 &mut self,
787 _coeffs: &AugSysCoeffs<'_>,
788 packed_rhs: &mut [Number],
789 nrhs: usize,
790 ) -> Option<ESymSolverStatus> {
791 if !self.have_factor {
796 return None;
797 }
798 if packed_rhs.len() != (self.dim as usize) * nrhs {
799 return Some(ESymSolverStatus::FatalError);
800 }
801 let _back_guard = self
802 .timing
803 .as_deref()
804 .map(|t| t.linear_system_back_solve.guard());
805 let status =
806 self.linsol
807 .multi_solve(&self.vals, false, nrhs as Index, packed_rhs, false, 0);
808 drop(_back_guard);
809 self.last_status = Some(status);
810 Some(status)
811 }
812}
813
814#[allow(clippy::too_many_arguments)]
817fn write_kkt_record(
822 w: &mut dyn std::io::Write,
823 dim: Index,
824 irn: &[Index],
825 jcn: &[Index],
826 vals: &[Number],
827 rhs: &[Number],
828 sol: &[Number],
829 check_neg_evals: bool,
830 num_neg_evals: Index,
831 status: ESymSolverStatus,
832 last_neg_evals: Index,
833 factor_pattern: Option<&FactorPattern>,
834) -> std::io::Result<()> {
835 use std::fmt::Write as _;
836
837 let mut line = String::with_capacity(64 * vals.len());
838 line.push('{');
839 let _ = write!(line, "\"n\":{dim},");
840 let _ = write!(line, "\"check_neg_evals\":{check_neg_evals},");
841 let _ = write!(line, "\"num_neg_evals_expected\":{num_neg_evals},");
842 let _ = write!(line, "\"num_neg_evals_actual\":{last_neg_evals},");
843 let _ = write!(line, "\"status\":\"{status:?}\",");
844
845 line.push_str("\"irn\":[");
846 for (i, v) in irn.iter().enumerate() {
847 if i > 0 {
848 line.push(',');
849 }
850 let _ = write!(line, "{v}");
851 }
852 line.push_str("],\"jcn\":[");
853 for (i, v) in jcn.iter().enumerate() {
854 if i > 0 {
855 line.push(',');
856 }
857 let _ = write!(line, "{v}");
858 }
859 line.push_str("],\"vals\":[");
860 for (i, v) in vals.iter().enumerate() {
861 if i > 0 {
862 line.push(',');
863 }
864 let _ = write!(line, "{v:.17e}");
865 }
866 line.push_str("],\"rhs\":[");
867 for (i, v) in rhs.iter().enumerate() {
868 if i > 0 {
869 line.push(',');
870 }
871 let _ = write!(line, "{v:.17e}");
872 }
873 line.push_str("],\"sol\":[");
874 for (i, v) in sol.iter().enumerate() {
875 if i > 0 {
876 line.push(',');
877 }
878 let _ = write!(line, "{v:.17e}");
879 }
880 line.push(']');
881
882 if let Some(fp) = factor_pattern {
887 line.push_str(",\"L_irn\":[");
888 for (i, v) in fp.l_irn.iter().enumerate() {
889 if i > 0 {
890 line.push(',');
891 }
892 let _ = write!(line, "{v}");
893 }
894 line.push_str("],\"L_jcn\":[");
895 for (i, v) in fp.l_jcn.iter().enumerate() {
896 if i > 0 {
897 line.push(',');
898 }
899 let _ = write!(line, "{v}");
900 }
901 line.push_str("],\"perm\":[");
902 for (i, v) in fp.perm.iter().enumerate() {
903 if i > 0 {
904 line.push(',');
905 }
906 let _ = write!(line, "{v}");
907 }
908 line.push(']');
909 if let Some(vals) = fp.l_vals.as_ref() {
910 line.push_str(",\"L_vals\":[");
911 for (i, v) in vals.iter().enumerate() {
912 if i > 0 {
913 line.push(',');
914 }
915 let _ = write!(line, "{v:.17e}");
916 }
917 line.push(']');
918 }
919 }
920
921 line.push_str("}\n");
922
923 w.write_all(line.as_bytes())
924}
925
926fn dump_kkt(
927 path: &str,
928 dim: Index,
929 irn: &[Index],
930 jcn: &[Index],
931 vals: &[Number],
932 rhs: &[Number],
933 sol: &[Number],
934 check_neg_evals: bool,
935 num_neg_evals: Index,
936 status: ESymSolverStatus,
937 last_neg_evals: Index,
938) {
939 if let Ok(mut f) = std::fs::OpenOptions::new()
940 .create(true)
941 .append(true)
942 .open(path)
943 {
944 let _ = write_kkt_record(
945 &mut f,
946 dim,
947 irn,
948 jcn,
949 vals,
950 rhs,
951 sol,
952 check_neg_evals,
953 num_neg_evals,
954 status,
955 last_neg_evals,
956 None, );
958 }
959}
960
961fn w_nonzeros(w: &dyn pounce_linalg::SymMatrix) -> usize {
966 if let Some(t) = w.as_any().downcast_ref::<SymTMatrix>() {
967 t.nonzeros() as usize
968 } else if let Some(dm) = w.as_any().downcast_ref::<DiagMatrix>() {
969 w_diag_dim(dm) as usize
970 } else {
971 unreachable!("StdAugSystemSolver: W must be a SymTMatrix or DiagMatrix in v1.0")
972 }
973}
974
975fn w_diag_dim(dm: &DiagMatrix) -> Index {
976 dm.get_diag()
977 .expect("DiagMatrix W has no diagonal set")
978 .dim()
979}
980
981fn w_diag_values(dm: &DiagMatrix) -> Vec<Number> {
982 let diag = dm.get_diag().expect("DiagMatrix W has no diagonal set");
983 diag.as_any()
984 .downcast_ref::<DenseVector>()
985 .expect("StdAugSystemSolver: DiagMatrix W diagonal must be DenseVector in v1.0")
986 .expanded_values()
987}
988
989fn gen_t_downcast(m: &dyn pounce_linalg::Matrix) -> &GenTMatrix {
990 let Some(t) = m.as_any().downcast_ref::<GenTMatrix>() else {
991 unreachable!("StdAugSystemSolver: J_c / J_d must be GenTMatrix in v1.0")
992 };
993 t
994}
995
996fn flat_read(v: &dyn Vector) -> Vec<Number> {
1002 if let Some(dv) = v.as_any().downcast_ref::<DenseVector>() {
1003 return dv.expanded_values();
1004 }
1005 if let Some(cv) = v.as_any().downcast_ref::<CompoundVector>() {
1006 let mut out = Vec::with_capacity(cv.dim() as usize);
1007 for k in 0..cv.n_comps() {
1008 let blk = cv.comp(k);
1009 let dblk = blk
1010 .as_any()
1011 .downcast_ref::<DenseVector>()
1012 .expect("StdAugSystemSolver: CompoundVector blocks must be DenseVectors");
1013 out.extend_from_slice(&dblk.expanded_values());
1014 }
1015 return out;
1016 }
1017 unreachable!(
1018 "StdAugSystemSolver: D_*/rhs/sol must be DenseVector or CompoundVector of DenseVectors in v1.0"
1019 )
1020}
1021
1022fn flat_write(dst: &mut dyn Vector, src: &[Number]) {
1024 if let Some(dv) = dst.as_any_mut().downcast_mut::<DenseVector>() {
1025 dv.set_values(src);
1026 return;
1027 }
1028 if let Some(cv) = dst.as_any_mut().downcast_mut::<CompoundVector>() {
1029 let mut off = 0usize;
1030 for k in 0..cv.n_comps() {
1031 let blk = cv.comp_mut(k);
1032 let dim = blk.dim() as usize;
1033 let dblk = blk
1034 .as_any_mut()
1035 .downcast_mut::<DenseVector>()
1036 .expect("StdAugSystemSolver: CompoundVector blocks must be DenseVectors");
1037 dblk.set_values(&src[off..off + dim]);
1038 off += dim;
1039 }
1040 return;
1041 }
1042 unreachable!(
1043 "StdAugSystemSolver: sol must be DenseVector or CompoundVector of DenseVectors in v1.0"
1044 )
1045}
1046
1047fn fill_diag(dst: &mut [Number], d: Option<&dyn Vector>, delta: Number, sign: Number) {
1050 match d {
1051 None => {
1052 for v in dst.iter_mut() {
1053 *v = sign * delta;
1054 }
1055 }
1056 Some(d) => {
1057 let xs = flat_read(d);
1058 debug_assert_eq!(xs.len(), dst.len());
1059 for (out, &x) in dst.iter_mut().zip(xs.iter()) {
1060 *out = sign * (x + delta);
1061 }
1062 }
1063 }
1064}
1065
1066fn copy_vec(src: &dyn Vector, dst: &mut [Number]) {
1067 let xs = flat_read(src);
1068 debug_assert_eq!(xs.len(), dst.len());
1069 dst.copy_from_slice(&xs);
1070}
1071
1072fn write_vec(dst: &mut dyn Vector, src: &[Number]) {
1073 flat_write(dst, src);
1074}
1075
1076#[cfg(test)]
1077mod tests {
1078 use super::*;
1079 use pounce_common::types::{Index, Number};
1080 use pounce_linalg::dense_vector::DenseVectorSpace;
1081 use pounce_linalg::triplet::{GenTMatrixSpace, SymTMatrixSpace};
1082 use pounce_linsol::EMatrixFormat;
1083 use pounce_linsol::sparse_sym_iface::SparseSymLinearSolverInterface;
1084
1085 struct DenseMock {
1088 dim: Index,
1089 nz: Index,
1090 a: Vec<Number>,
1091 last_factor: Vec<Number>, neg_evals: Index,
1093 }
1094
1095 impl DenseMock {
1096 fn new() -> Self {
1097 Self {
1098 dim: 0,
1099 nz: 0,
1100 a: Vec::new(),
1101 last_factor: Vec::new(),
1102 neg_evals: 0,
1103 }
1104 }
1105 }
1106
1107 impl SparseSymLinearSolverInterface for DenseMock {
1108 fn initialize_structure(
1109 &mut self,
1110 dim: Index,
1111 nz: Index,
1112 _ia: &[Index],
1113 _ja: &[Index],
1114 ) -> ESymSolverStatus {
1115 self.dim = dim;
1116 self.nz = nz;
1117 self.a = vec![0.0; nz as usize];
1118 ESymSolverStatus::Success
1119 }
1120 fn values_array_mut(&mut self) -> &mut [Number] {
1121 &mut self.a
1122 }
1123 fn multi_solve(
1124 &mut self,
1125 new_matrix: bool,
1126 ia: &[Index],
1127 ja: &[Index],
1128 nrhs: Index,
1129 rhs_vals: &mut [Number],
1130 _check: bool,
1131 _nev: Index,
1132 ) -> ESymSolverStatus {
1133 let n = self.dim as usize;
1134 if new_matrix {
1135 let mut dense = vec![0.0; n * n];
1138 for k in 0..self.nz as usize {
1139 let i = (ia[k] - 1) as usize;
1140 let j = (ja[k] - 1) as usize;
1141 dense[i * n + j] += self.a[k];
1142 if i != j {
1143 dense[j * n + i] += self.a[k];
1144 }
1145 }
1146 self.last_factor = dense;
1147 }
1148 for col in 0..nrhs as usize {
1150 let mut a = self.last_factor.clone();
1151 let b = &mut rhs_vals[col * n..col * n + n];
1152 let mut neg = 0_i32;
1153 for k in 0..n {
1154 let mut piv = k;
1156 let mut piv_abs = a[k * n + k].abs();
1157 for r in (k + 1)..n {
1158 let av = a[r * n + k].abs();
1159 if av > piv_abs {
1160 piv_abs = av;
1161 piv = r;
1162 }
1163 }
1164 if piv != k {
1165 for c in 0..n {
1166 a.swap(k * n + c, piv * n + c);
1167 }
1168 b.swap(k, piv);
1169 }
1170 let p = a[k * n + k];
1171 if p.abs() < 1e-14 {
1172 return ESymSolverStatus::Singular;
1173 }
1174 if p < 0.0 {
1175 neg += 1;
1176 }
1177 for r in (k + 1)..n {
1178 let f = a[r * n + k] / p;
1179 for c in k..n {
1180 a[r * n + c] -= f * a[k * n + c];
1181 }
1182 b[r] -= f * b[k];
1183 }
1184 }
1185 for k in (0..n).rev() {
1187 let mut s = b[k];
1188 for c in (k + 1)..n {
1189 s -= a[k * n + c] * b[c];
1190 }
1191 b[k] = s / a[k * n + k];
1192 }
1193 self.neg_evals = neg;
1194 }
1195 ESymSolverStatus::Success
1196 }
1197 fn number_of_neg_evals(&self) -> Index {
1198 self.neg_evals
1199 }
1200 fn increase_quality(&mut self) -> bool {
1201 false
1202 }
1203 fn provides_inertia(&self) -> bool {
1204 true
1205 }
1206 fn matrix_format(&self) -> EMatrixFormat {
1207 EMatrixFormat::TripletFormat
1208 }
1209 }
1210
1211 #[test]
1223 fn solves_5x5_kkt_through_dense_mock() {
1224 let w_space = SymTMatrixSpace::new(2, vec![1, 2], vec![1, 2]);
1226 let mut w = SymTMatrix::new(w_space);
1227 w.set_values(&[2.0, 3.0]);
1228
1229 let jc_space = GenTMatrixSpace::new(1, 2, vec![1, 1], vec![1, 2]);
1231 let mut j_c = GenTMatrix::new(jc_space);
1232 j_c.set_values(&[1.0, 1.0]);
1233
1234 let jd_space = GenTMatrixSpace::new(1, 2, vec![1], vec![1]);
1236 let mut j_d = GenTMatrix::new(jd_space);
1237 j_d.set_values(&[1.0]);
1238
1239 let s_space = DenseVectorSpace::new(1);
1241 let mut d_s = s_space.make_new_dense();
1242 d_s.set_values(&[1.0]);
1243
1244 let xs = DenseVectorSpace::new(2);
1252 let mut rx = xs.make_new_dense();
1253 rx.set_values(&[4.0, 4.0]);
1254 let mut rs = s_space.make_new_dense();
1255 rs.set_values(&[0.0]);
1256 let cs = DenseVectorSpace::new(1);
1257 let mut rc = cs.make_new_dense();
1258 rc.set_values(&[2.0]);
1259 let ds_space = DenseVectorSpace::new(1);
1260 let mut rd = ds_space.make_new_dense();
1261 rd.set_values(&[0.0]);
1262
1263 let mut sx = xs.make_new_dense();
1264 let mut ss = s_space.make_new_dense();
1265 let mut sc = cs.make_new_dense();
1266 let mut sd = ds_space.make_new_dense();
1267
1268 let linsol = TSymLinearSolver::new(Box::new(DenseMock::new()), None, false);
1269 let mut solver = StdAugSystemSolver::new(linsol);
1270
1271 let coeffs = AugSysCoeffs {
1272 w: Some(&w),
1273 w_factor: 1.0,
1274 d_x: None,
1275 delta_x: 0.0,
1276 d_s: Some(&d_s),
1277 delta_s: 0.0,
1278 j_c: &j_c,
1279 d_c: None,
1280 delta_c: 0.0,
1281 j_d: &j_d,
1282 d_d: None,
1283 delta_d: 0.0,
1284 };
1285 let rhs = AugSysRhs {
1286 rhs_x: &rx,
1287 rhs_s: &rs,
1288 rhs_c: &rc,
1289 rhs_d: &rd,
1290 };
1291 let mut sol = AugSysSol {
1292 sol_x: &mut sx,
1293 sol_s: &mut ss,
1294 sol_c: &mut sc,
1295 sol_d: &mut sd,
1296 };
1297 let status = solver.solve(&coeffs, &rhs, &mut sol, false, 0);
1298 assert_eq!(status, ESymSolverStatus::Success);
1299
1300 for v in sx.values() {
1301 assert!((v - 1.0).abs() < 1e-10, "sol_x = {v}");
1302 }
1303 for v in ss.values() {
1304 assert!((v - 1.0).abs() < 1e-10, "sol_s = {v}");
1305 }
1306 for v in sc.values() {
1307 assert!((v - 1.0).abs() < 1e-10, "sol_c = {v}");
1308 }
1309 for v in sd.values() {
1310 assert!((v - 1.0).abs() < 1e-10, "sol_d = {v}");
1311 }
1312 }
1313
1314 #[test]
1324 fn lowrank_smw_matches_dense_w_on_constrained_system() {
1325 use crate::kkt::low_rank_aug_system_solver::LowRankAugSystemSolver;
1326 use pounce_linalg::diag_matrix::DiagMatrix;
1327 use pounce_linalg::low_rank_update_sym_matrix::LowRankUpdateSymMatrixSpace;
1328 use pounce_linalg::multi_vector_matrix::MultiVectorMatrixSpace;
1329
1330 let n = 4usize;
1331 let sigma = 2.0;
1332 let vcols = [
1337 vec![0.6, 0.1, -0.2, 0.3],
1338 vec![0.2, 0.5, 0.1, -0.1],
1339 vec![-0.1, 0.2, 0.4, 0.2],
1340 vec![0.3, -0.2, 0.1, 0.4],
1341 vec![0.15, 0.25, -0.3, 0.1],
1342 vec![-0.2, 0.1, 0.2, 0.35],
1343 ];
1344 let ucols = [
1345 vec![0.3, -0.1, 0.2, 0.1],
1346 vec![0.1, 0.3, -0.2, 0.2],
1347 vec![0.2, 0.1, 0.1, -0.3],
1348 vec![-0.1, 0.2, 0.15, 0.1],
1349 vec![0.25, -0.15, 0.1, 0.2],
1350 vec![0.1, 0.2, -0.25, 0.15],
1351 ];
1352 let mut wfull = vec![0.0_f64; n * n];
1354 for i in 0..n {
1355 wfull[i * n + i] = sigma;
1356 }
1357 for c in vcols.iter() {
1358 for i in 0..n {
1359 for j in 0..n {
1360 wfull[i * n + j] += c[i] * c[j];
1361 }
1362 }
1363 }
1364 for c in ucols.iter() {
1365 for i in 0..n {
1366 for j in 0..n {
1367 wfull[i * n + j] -= c[i] * c[j];
1368 }
1369 }
1370 }
1371
1372 let make_jc = || {
1374 let sp = GenTMatrixSpace::new(1, 4, vec![1, 1, 1, 1], vec![1, 2, 3, 4]);
1375 let mut m = GenTMatrix::new(sp);
1376 m.set_values(&[1.0, 1.0, 1.0, 1.0]);
1377 m
1378 };
1379 let make_jd = || {
1383 let sp = GenTMatrixSpace::new(1, 4, vec![1, 1], vec![1, 3]);
1384 let mut m = GenTMatrix::new(sp);
1385 m.set_values(&[1.0, 1.0]);
1386 m
1387 };
1388
1389 let xs = DenseVectorSpace::new(4);
1390 let cs = DenseVectorSpace::new(1);
1391 let mk = |sp: &Rc<DenseVectorSpace>, vals: &[Number]| {
1392 let mut d = sp.make_new_dense();
1393 d.set_values(vals);
1394 d
1395 };
1396
1397 let solve_with = |w: &dyn pounce_linalg::SymMatrix,
1398 aug: &mut dyn AugSystemSolver|
1399 -> (Vec<Number>, Vec<Number>) {
1400 let j_c = make_jc();
1401 let j_d = make_jd();
1402 let rx = mk(&xs, &[1.0, 2.0, -1.0, 0.5]);
1403 let rs = mk(&cs, &[0.4]);
1404 let rc = mk(&cs, &[3.0]);
1405 let rd = mk(&cs, &[0.7]);
1406 let mut sx = mk(&xs, &[0.0, 0.0, 0.0, 0.0]);
1407 let mut ss = mk(&cs, &[0.0]);
1408 let mut sc = mk(&cs, &[0.0]);
1409 let mut sd = mk(&cs, &[0.0]);
1410 let d_s = mk(&cs, &[1.5]);
1411 let coeffs = AugSysCoeffs {
1412 w: Some(w),
1413 w_factor: 1.0,
1414 d_x: None,
1415 delta_x: 0.0,
1416 d_s: Some(&d_s),
1417 delta_s: 0.0,
1418 j_c: &j_c,
1419 d_c: None,
1420 delta_c: 0.0,
1421 j_d: &j_d,
1422 d_d: None,
1423 delta_d: 0.0,
1424 };
1425 let rhs = AugSysRhs {
1426 rhs_x: &rx,
1427 rhs_s: &rs,
1428 rhs_c: &rc,
1429 rhs_d: &rd,
1430 };
1431 let mut sol = AugSysSol {
1432 sol_x: &mut sx,
1433 sol_s: &mut ss,
1434 sol_c: &mut sc,
1435 sol_d: &mut sd,
1436 };
1437 let status = aug.solve(&coeffs, &rhs, &mut sol, false, 1);
1438 assert_eq!(status, ESymSolverStatus::Success);
1439 (sx.expanded_values(), sc.expanded_values())
1440 };
1441
1442 let mut wi = Vec::new();
1444 let mut wj = Vec::new();
1445 let mut wv = Vec::new();
1446 for i in 0..n {
1447 for j in 0..=i {
1448 wi.push(i as Index + 1);
1449 wj.push(j as Index + 1);
1450 wv.push(wfull[i * n + j]);
1451 }
1452 }
1453 let w_space = SymTMatrixSpace::new(4, wi, wj);
1454 let mut w_dense = SymTMatrix::new(w_space);
1455 w_dense.set_values(&wv);
1456 let mut std_solver = StdAugSystemSolver::new(TSymLinearSolver::new(
1457 Box::new(pounce_feral::FeralSolverInterface::new()),
1458 None,
1459 false,
1460 ));
1461 let (ref_x, ref_c) = solve_with(&w_dense, &mut std_solver);
1462
1463 let lr_space = LowRankUpdateSymMatrixSpace::new(4, None, false);
1465 let mut lr = lr_space.make_new_low_rank();
1466 let mut diag = xs.make_new_dense();
1467 diag.set_values(&[sigma; 4]);
1468 lr.set_diag(Rc::new(diag) as Rc<dyn Vector>);
1469 let build_mvm = |cols: &[Vec<Number>]| {
1470 let sp = MultiVectorMatrixSpace::new(cols.len() as Index, Rc::clone(&xs));
1471 let mut mvm = sp.make_new_multi_vector();
1472 for (k, c) in cols.iter().enumerate() {
1473 let mut cv = xs.make_new_dense();
1474 cv.set_values(c);
1475 mvm.set_vector(k as Index, Rc::new(cv) as Rc<dyn Vector>);
1476 }
1477 mvm
1478 };
1479 lr.set_v(Rc::new(build_mvm(&vcols)));
1480 lr.set_u(Rc::new(build_mvm(&ucols)));
1481 let _ = DiagMatrix::new(4); let mut lr_solver =
1484 LowRankAugSystemSolver::new(Box::new(StdAugSystemSolver::new(TSymLinearSolver::new(
1485 Box::new(pounce_feral::FeralSolverInterface::new()),
1486 None,
1487 false,
1488 ))));
1489 let (lr_x, lr_c) = solve_with(&lr, &mut lr_solver);
1490
1491 for (a, b) in ref_x.iter().zip(lr_x.iter()) {
1492 assert!((a - b).abs() < 1e-9, "sol_x mismatch: dense={a} smw={b}");
1493 }
1494 for (a, b) in ref_c.iter().zip(lr_c.iter()) {
1495 assert!((a - b).abs() < 1e-9, "sol_c mismatch: dense={a} smw={b}");
1496 }
1497 }
1498}