1#![allow(clippy::needless_range_loop)]
21
22use std::ops::{Index, IndexMut};
23
24use crate::error::{all_finite, NumericError, NumericResult};
25
26const U: f64 = f64::EPSILON / 2.0;
27
28#[derive(Debug, Clone, PartialEq)]
30pub struct Matrix {
31 rows: usize,
32 cols: usize,
33 data: Vec<f64>,
34}
35
36impl Matrix {
37 pub fn zeros(rows: usize, cols: usize) -> Self {
39 Self {
40 rows,
41 cols,
42 data: vec![0.0; rows * cols],
43 }
44 }
45
46 pub fn identity(n: usize) -> Self {
48 let mut m = Self::zeros(n, n);
49 for i in 0..n {
50 m[(i, i)] = 1.0;
51 }
52 m
53 }
54
55 pub fn from_fn(rows: usize, cols: usize, mut f: impl FnMut(usize, usize) -> f64) -> Self {
57 let mut m = Self::zeros(rows, cols);
58 for i in 0..rows {
59 for j in 0..cols {
60 m[(i, j)] = f(i, j);
61 }
62 }
63 m
64 }
65
66 pub fn from_rows(rows: &[&[f64]]) -> NumericResult<Self> {
72 let cols = rows.first().map_or(0, |r| r.len());
73 let mut data = Vec::with_capacity(rows.len() * cols);
74 for r in rows {
75 if r.len() != cols {
76 return Err(NumericError::DimensionMismatch {
77 name: "matrix row",
78 expected: cols,
79 found: r.len(),
80 });
81 }
82 data.extend_from_slice(r);
83 }
84 Ok(Self {
85 rows: rows.len(),
86 cols,
87 data,
88 })
89 }
90
91 pub fn rows(&self) -> usize {
93 self.rows
94 }
95
96 pub fn cols(&self) -> usize {
98 self.cols
99 }
100
101 pub fn row(&self, i: usize) -> &[f64] {
103 &self.data[i * self.cols..(i + 1) * self.cols]
104 }
105
106 pub fn transpose(&self) -> Self {
108 Self::from_fn(self.cols, self.rows, |i, j| self[(j, i)])
109 }
110
111 pub fn mul_vec(&self, x: &[f64]) -> NumericResult<Vec<f64>> {
117 if x.len() != self.cols {
118 return Err(NumericError::DimensionMismatch {
119 name: "vector",
120 expected: self.cols,
121 found: x.len(),
122 });
123 }
124 Ok((0..self.rows)
125 .map(|i| self.row(i).iter().zip(x).map(|(a, b)| a * b).sum())
126 .collect())
127 }
128
129 pub fn norm1(&self) -> f64 {
131 (0..self.cols)
132 .map(|j| (0..self.rows).map(|i| self[(i, j)].abs()).sum::<f64>())
133 .fold(0.0, f64::max)
134 }
135
136 fn require_finite(&self, name: &'static str) -> NumericResult<()> {
137 all_finite(&self.data, name)
138 }
139
140 fn require_square(&self, name: &'static str) -> NumericResult<()> {
141 if self.rows == self.cols {
142 Ok(())
143 } else {
144 Err(NumericError::DimensionMismatch {
145 name,
146 expected: self.rows,
147 found: self.cols,
148 })
149 }
150 }
151}
152
153impl Index<(usize, usize)> for Matrix {
154 type Output = f64;
155 fn index(&self, (i, j): (usize, usize)) -> &f64 {
156 assert!(i < self.rows && j < self.cols, "matrix index out of range");
157 &self.data[i * self.cols + j]
158 }
159}
160
161impl IndexMut<(usize, usize)> for Matrix {
162 fn index_mut(&mut self, (i, j): (usize, usize)) -> &mut f64 {
163 assert!(i < self.rows && j < self.cols, "matrix index out of range");
164 &mut self.data[i * self.cols + j]
165 }
166}
167
168#[derive(Debug, Clone, PartialEq)]
170#[must_use]
171pub struct LinearSolution {
172 pub x: Vec<f64>,
174 pub residual_norm: f64,
176 pub relative_error_estimate: f64,
180}
181
182#[derive(Debug, Clone, PartialEq)]
184#[must_use]
185pub struct LeastSquaresSolution {
186 pub x: Vec<f64>,
188 pub residual_norm: f64,
190 pub relative_error_estimate: f64,
194}
195
196fn inverse_norm1_estimate(
199 n: usize,
200 mut solve: impl FnMut(&mut [f64]),
201 mut solve_transpose: impl FnMut(&mut [f64]),
202) -> f64 {
203 if n == 0 {
204 return 0.0;
205 }
206 let mut x = vec![1.0 / n as f64; n];
207 let mut estimate = 0.0f64;
208 let mut last_j = usize::MAX;
209 for _ in 0..5 {
210 let mut y = x.clone();
211 solve(&mut y);
212 estimate = estimate.max(y.iter().map(|v| v.abs()).sum());
213 let mut z: Vec<f64> = y
214 .iter()
215 .map(|v| if *v >= 0.0 { 1.0 } else { -1.0 })
216 .collect();
217 solve_transpose(&mut z);
218 let (j, zmax) = z
219 .iter()
220 .enumerate()
221 .map(|(i, v)| (i, v.abs()))
222 .fold((0, -1.0), |best, c| if c.1 > best.1 { c } else { best });
223 let ztx: f64 = z.iter().zip(&x).map(|(a, b)| a * b).sum();
224 if zmax <= ztx || j == last_j {
225 break;
226 }
227 last_j = j;
228 x = vec![0.0; n];
229 x[j] = 1.0;
230 }
231 let mut alt: Vec<f64> = (0..n)
233 .map(|i| {
234 let s = if i % 2 == 0 { 1.0 } else { -1.0 };
235 s * (1.0 + i as f64 / (n.max(2) - 1) as f64)
236 })
237 .collect();
238 solve(&mut alt);
239 let alt_estimate = 2.0 * alt.iter().map(|v| v.abs()).sum::<f64>() / (3.0 * n as f64);
240 estimate.max(alt_estimate)
241}
242
243fn forward_error_estimate(a: &Matrix, x: &[f64], b: &[f64], condition: f64) -> (f64, f64) {
244 let n = a.cols;
245 let mut r_norm = 0.0;
246 let mut scale = 0.0;
247 for (i, bi) in b.iter().enumerate() {
248 let row = a.row(i);
249 let ax: f64 = row.iter().zip(x).map(|(p, q)| p * q).sum();
250 let mag: f64 = row.iter().zip(x).map(|(p, q)| (p * q).abs()).sum::<f64>() + bi.abs();
251 r_norm += (bi - ax).abs();
252 scale += mag;
253 }
254 let x_norm: f64 = x.iter().map(|v| v.abs()).sum();
255 let b_norm: f64 = b.iter().map(|v| v.abs()).sum();
256 let denominator = a.norm1() * x_norm + b_norm;
257 if denominator == 0.0 {
258 return (r_norm, 0.0);
259 }
260 let k = (n + 1) as f64;
263 let residual_rounding = k * U / (1.0 - k * U) * scale;
264 let eta = (r_norm + residual_rounding) / denominator;
265 let kappa_eta = condition * eta;
266 let estimate = if kappa_eta < 0.5 {
267 2.0 * kappa_eta / (1.0 - kappa_eta)
268 } else {
269 f64::INFINITY
270 };
271 (r_norm, estimate)
272}
273
274#[derive(Debug, Clone, PartialEq)]
276pub struct Lu {
277 a: Matrix,
278 lu: Matrix,
279 perm: Vec<usize>,
280 condition: f64,
281}
282
283impl Lu {
284 pub fn new(a: &Matrix) -> NumericResult<Self> {
292 a.require_square("system matrix")?;
293 a.require_finite("system matrix")?;
294 let n = a.rows;
295 let mut lu = a.clone();
296 let mut perm: Vec<usize> = (0..n).collect();
297 for k in 0..n {
298 let p = (k..n)
299 .max_by(|&i, &j| lu[(i, k)].abs().total_cmp(&lu[(j, k)].abs()))
300 .unwrap_or(k);
301 if lu[(p, k)] == 0.0 {
302 return Err(NumericError::Singular {
303 name: "system matrix",
304 });
305 }
306 if p != k {
307 for j in 0..n {
308 lu.data.swap(k * n + j, p * n + j);
309 }
310 perm.swap(k, p);
311 }
312 let pivot = lu[(k, k)];
313 for i in (k + 1)..n {
314 let factor = lu[(i, k)] / pivot;
315 lu[(i, k)] = factor;
316 if factor != 0.0 {
317 for j in (k + 1)..n {
318 let v = lu[(k, j)];
319 lu[(i, j)] -= factor * v;
320 }
321 }
322 }
323 }
324 let mut out = Self {
325 a: a.clone(),
326 lu,
327 perm,
328 condition: 0.0,
329 };
330 let inverse_norm = inverse_norm1_estimate(
331 n,
332 |v| out.solve_in_place(v),
333 |v| out.solve_transpose_in_place(v),
334 );
335 out.condition = a.norm1() * inverse_norm;
336 if !out.condition.is_finite() || out.condition * f64::EPSILON >= 1.0 {
337 return Err(NumericError::Singular {
338 name: "system matrix",
339 });
340 }
341 Ok(out)
342 }
343
344 pub fn condition_estimate(&self) -> f64 {
347 self.condition
348 }
349
350 pub fn solve(&self, b: &[f64]) -> NumericResult<LinearSolution> {
356 check_rhs(b, self.a.rows)?;
357 let mut x = b.to_vec();
358 self.solve_in_place(&mut x);
359 let (residual_norm, relative_error_estimate) =
360 forward_error_estimate(&self.a, &x, b, self.condition);
361 Ok(LinearSolution {
362 x,
363 residual_norm,
364 relative_error_estimate,
365 })
366 }
367
368 fn solve_in_place(&self, b: &mut [f64]) {
369 let n = self.a.rows;
370 let mut y: Vec<f64> = self.perm.iter().map(|&p| b[p]).collect();
371 for i in 0..n {
372 let mut s = y[i];
373 for j in 0..i {
374 s -= self.lu[(i, j)] * y[j];
375 }
376 y[i] = s;
377 }
378 for i in (0..n).rev() {
379 let mut s = y[i];
380 for j in (i + 1)..n {
381 s -= self.lu[(i, j)] * y[j];
382 }
383 y[i] = s / self.lu[(i, i)];
384 }
385 b.copy_from_slice(&y);
386 }
387
388 fn solve_transpose_in_place(&self, c: &mut [f64]) {
389 let n = self.a.rows;
390 let mut w = c.to_vec();
392 for i in 0..n {
393 let mut s = w[i];
394 for j in 0..i {
395 s -= self.lu[(j, i)] * w[j];
396 }
397 w[i] = s / self.lu[(i, i)];
398 }
399 for i in (0..n).rev() {
400 let mut s = w[i];
401 for j in (i + 1)..n {
402 s -= self.lu[(j, i)] * w[j];
403 }
404 w[i] = s;
405 }
406 for (i, &p) in self.perm.iter().enumerate() {
407 c[p] = w[i];
408 }
409 }
410}
411
412fn check_rhs(b: &[f64], expected: usize) -> NumericResult<()> {
413 if b.len() != expected {
414 return Err(NumericError::DimensionMismatch {
415 name: "right-hand side",
416 expected,
417 found: b.len(),
418 });
419 }
420 all_finite(b, "right-hand side")
421}
422
423#[derive(Debug, Clone, PartialEq)]
425pub struct Cholesky {
426 a: Matrix,
427 l: Matrix,
428 condition: f64,
429}
430
431impl Cholesky {
432 pub fn new(a: &Matrix) -> NumericResult<Self> {
444 a.require_square("symmetric matrix")?;
445 a.require_finite("symmetric matrix")?;
446 let n = a.rows;
447 let max_abs = a.data.iter().fold(0.0f64, |m, v| m.max(v.abs()));
448 let symmetry_tolerance = n as f64 * f64::EPSILON * max_abs;
449 for i in 0..n {
450 for j in 0..i {
451 if (a[(i, j)] - a[(j, i)]).abs() > symmetry_tolerance {
452 return Err(NumericError::NotSymmetric { row: i, column: j });
453 }
454 }
455 }
456 let mut l = Matrix::zeros(n, n);
457 for j in 0..n {
458 let mut d = a[(j, j)];
459 for k in 0..j {
460 d -= l[(j, k)] * l[(j, k)];
461 }
462 if d <= 0.0 || !d.is_finite() {
463 return Err(NumericError::NotPositiveDefinite { pivot: j });
464 }
465 let ljj = d.sqrt();
466 l[(j, j)] = ljj;
467 for i in (j + 1)..n {
468 let mut s = a[(i, j)];
469 for k in 0..j {
470 s -= l[(i, k)] * l[(j, k)];
471 }
472 l[(i, j)] = s / ljj;
473 }
474 }
475 let mut out = Self {
476 a: a.clone(),
477 l,
478 condition: 0.0,
479 };
480 let inverse_norm =
481 inverse_norm1_estimate(n, |v| out.solve_in_place(v), |v| out.solve_in_place(v));
482 out.condition = a.norm1() * inverse_norm;
483 if !out.condition.is_finite() || out.condition * f64::EPSILON >= 1.0 {
484 return Err(NumericError::Singular {
485 name: "symmetric matrix",
486 });
487 }
488 Ok(out)
489 }
490
491 pub fn condition_estimate(&self) -> f64 {
493 self.condition
494 }
495
496 pub fn factor(&self) -> &Matrix {
498 &self.l
499 }
500
501 pub fn solve(&self, b: &[f64]) -> NumericResult<LinearSolution> {
507 check_rhs(b, self.a.rows)?;
508 let mut x = b.to_vec();
509 self.solve_in_place(&mut x);
510 let (residual_norm, relative_error_estimate) =
511 forward_error_estimate(&self.a, &x, b, self.condition);
512 Ok(LinearSolution {
513 x,
514 residual_norm,
515 relative_error_estimate,
516 })
517 }
518
519 fn solve_in_place(&self, b: &mut [f64]) {
520 let n = self.a.rows;
521 for i in 0..n {
522 let mut s = b[i];
523 for k in 0..i {
524 s -= self.l[(i, k)] * b[k];
525 }
526 b[i] = s / self.l[(i, i)];
527 }
528 for i in (0..n).rev() {
529 let mut s = b[i];
530 for k in (i + 1)..n {
531 s -= self.l[(k, i)] * b[k];
532 }
533 b[i] = s / self.l[(i, i)];
534 }
535 }
536}
537
538#[derive(Debug, Clone, PartialEq)]
540pub struct Qr {
541 a: Matrix,
542 r: Matrix,
544 reflectors: Vec<(Vec<f64>, f64)>,
546 perm: Vec<usize>,
547 rank: usize,
548 condition: f64,
549}
550
551impl Qr {
552 pub fn new(a: &Matrix) -> NumericResult<Self> {
559 a.require_finite("least-squares matrix")?;
560 let (m, n) = (a.rows, a.cols);
561 let mut r = a.clone();
562 let mut perm: Vec<usize> = (0..n).collect();
563 let mut reflectors = Vec::new();
564 let steps = m.min(n);
565 let mut rank = steps;
566 let mut first_pivot = 0.0;
567 for k in 0..steps {
568 let norms: Vec<f64> = (k..n)
570 .map(|j| (k..m).map(|i| r[(i, j)] * r[(i, j)]).sum::<f64>().sqrt())
571 .collect();
572 let (offset, &best) = norms
573 .iter()
574 .enumerate()
575 .max_by(|x, y| x.1.total_cmp(y.1))
576 .expect("k < n");
577 if k == 0 {
578 first_pivot = best;
579 }
580 if best <= (m.max(n) as f64) * f64::EPSILON * first_pivot || best == 0.0 {
581 rank = k;
582 break;
583 }
584 let p = k + offset;
585 if p != k {
586 for i in 0..m {
587 r.data.swap(i * n + k, i * n + p);
588 }
589 perm.swap(k, p);
590 }
591 let x0 = r[(k, k)];
592 let alpha = if x0 >= 0.0 { -best } else { best };
593 let mut v: Vec<f64> = (k..m).map(|i| r[(i, k)]).collect();
594 v[0] -= alpha;
595 let vtv: f64 = v.iter().map(|t| t * t).sum();
596 let beta = if vtv == 0.0 { 0.0 } else { 2.0 / vtv };
597 for j in k..n {
598 let dot: f64 = (k..m).map(|i| v[i - k] * r[(i, j)]).sum();
599 let s = beta * dot;
600 for i in k..m {
601 r[(i, j)] -= s * v[i - k];
602 }
603 }
604 r[(k, k)] = alpha;
605 for i in (k + 1)..m {
606 r[(i, k)] = 0.0;
607 }
608 reflectors.push((v, beta));
609 }
610 let mut out = Self {
611 a: a.clone(),
612 r,
613 reflectors,
614 perm,
615 rank,
616 condition: 0.0,
617 };
618 let k = out.rank;
619 let inverse_norm =
620 inverse_norm1_estimate(k, |v| out.solve_r(v), |v| out.solve_r_transpose(v));
621 let r_norm = (0..k)
622 .map(|j| (0..=j).map(|i| out.r[(i, j)].abs()).sum::<f64>())
623 .fold(0.0, f64::max);
624 out.condition = r_norm * inverse_norm;
625 Ok(out)
626 }
627
628 pub fn rank(&self) -> usize {
630 self.rank
631 }
632
633 pub fn condition_estimate(&self) -> f64 {
637 self.condition
638 }
639
640 pub fn solve_least_squares(&self, b: &[f64]) -> NumericResult<LeastSquaresSolution> {
648 let (m, n) = (self.a.rows, self.a.cols);
649 check_rhs(b, m)?;
650 if self.rank < n {
651 return Err(NumericError::RankDeficient {
652 name: "least-squares matrix",
653 rank: self.rank,
654 required: n,
655 });
656 }
657 let mut c = b.to_vec();
658 self.apply_qt(&mut c);
659 let mut z = c[..n].to_vec();
660 self.solve_r(&mut z);
661 let mut x = vec![0.0; n];
662 for (k, &p) in self.perm.iter().enumerate() {
663 x[p] = z[k];
664 }
665 let ax = self.a.mul_vec(&x)?;
666 let residual_norm = ax
667 .iter()
668 .zip(b)
669 .map(|(p, q)| (p - q) * (p - q))
670 .sum::<f64>()
671 .sqrt();
672 let a_norm = self.a.data.iter().map(|v| v * v).sum::<f64>().sqrt();
673 let x_norm = x.iter().map(|v| v * v).sum::<f64>().sqrt();
674 let kappa = self.condition;
675 let relative_error_estimate = if x_norm == 0.0 {
676 0.0
677 } else {
678 (m * n) as f64 * U * (2.0 * kappa + kappa * kappa * residual_norm / (a_norm * x_norm))
679 };
680 Ok(LeastSquaresSolution {
681 x,
682 residual_norm,
683 relative_error_estimate,
684 })
685 }
686
687 fn apply_qt(&self, c: &mut [f64]) {
689 for (k, (v, beta)) in self.reflectors.iter().enumerate() {
690 reflect(&mut c[k..], v, *beta);
691 }
692 }
693
694 fn apply_q(&self, c: &mut [f64]) {
696 for (k, (v, beta)) in self.reflectors.iter().enumerate().rev() {
697 reflect(&mut c[k..], v, *beta);
698 }
699 }
700
701 fn solve_r(&self, c: &mut [f64]) {
703 let k = self.rank;
704 for i in (0..k).rev() {
705 let mut s = c[i];
706 for j in (i + 1)..k {
707 s -= self.r[(i, j)] * c[j];
708 }
709 c[i] = s / self.r[(i, i)];
710 }
711 }
712
713 fn solve_r_transpose(&self, c: &mut [f64]) {
715 let k = self.rank;
716 for i in 0..k {
717 let mut s = c[i];
718 for j in 0..i {
719 s -= self.r[(j, i)] * c[j];
720 }
721 c[i] = s / self.r[(i, i)];
722 }
723 }
724}
725
726fn reflect(c: &mut [f64], v: &[f64], beta: f64) {
727 let dot: f64 = v.iter().zip(c.iter()).map(|(a, b)| a * b).sum();
728 let s = beta * dot;
729 for (ci, vi) in c.iter_mut().zip(v) {
730 *ci -= s * vi;
731 }
732}
733
734pub fn least_squares(a: &Matrix, b: &[f64]) -> NumericResult<LeastSquaresSolution> {
742 Qr::new(a)?.solve_least_squares(b)
743}
744
745#[derive(Debug, Clone, PartialEq)]
747#[must_use]
748pub struct ConstrainedSolution {
749 pub x: Vec<f64>,
751 pub residual_norm: f64,
753 pub constraint_residual_norm: f64,
755 pub relative_error_estimate: f64,
758}
759
760pub fn constrained_least_squares(
772 a: &Matrix,
773 b: &[f64],
774 c: &Matrix,
775 d: &[f64],
776) -> NumericResult<ConstrainedSolution> {
777 let n = a.cols;
778 if c.cols != n {
779 return Err(NumericError::DimensionMismatch {
780 name: "constraint matrix columns",
781 expected: n,
782 found: c.cols,
783 });
784 }
785 check_rhs(b, a.rows)?;
786 if d.len() != c.rows {
787 return Err(NumericError::DimensionMismatch {
788 name: "constraint right-hand side",
789 expected: c.rows,
790 found: d.len(),
791 });
792 }
793 all_finite(d, "constraint right-hand side")?;
794 a.require_finite("least-squares matrix")?;
795 c.require_finite("constraint matrix")?;
796 let p = c.rows;
797
798 let ct = Qr::new(&c.transpose())?;
800 if ct.rank < p {
801 return Err(NumericError::RankDeficient {
802 name: "constraint matrix",
803 rank: ct.rank,
804 required: p,
805 });
806 }
807 let mut y1: Vec<f64> = ct.perm.iter().map(|&i| d[i]).collect();
809 ct.solve_r_transpose(&mut y1);
810
811 let mut aq = Matrix::zeros(a.rows, n);
813 for i in 0..a.rows {
814 let mut row = a.row(i).to_vec();
815 ct.apply_qt(&mut row);
816 aq.data[i * n..(i + 1) * n].copy_from_slice(&row);
817 }
818 let mut rhs = b.to_vec();
819 for (i, value) in rhs.iter_mut().enumerate() {
820 let row = aq.row(i);
821 *value -= row[..p].iter().zip(&y1).map(|(s, t)| s * t).sum::<f64>();
822 }
823 let mut y = y1.clone();
824 let mut reduced_estimate = 0.0;
825 if p < n {
826 let a2 = Matrix::from_fn(a.rows, n - p, |i, j| aq[(i, p + j)]);
827 let qr = Qr::new(&a2)?;
828 if qr.rank < n - p {
829 return Err(NumericError::RankDeficient {
830 name: "least-squares matrix on the constraints' null space",
831 rank: qr.rank,
832 required: n - p,
833 });
834 }
835 let reduced = qr.solve_least_squares(&rhs)?;
836 reduced_estimate = reduced.relative_error_estimate;
837 y.extend_from_slice(&reduced.x);
838 }
839 let mut x = y;
840 ct.apply_q(&mut x);
841
842 let residual = |m: &Matrix, rhs: &[f64]| -> NumericResult<f64> {
843 Ok(m.mul_vec(&x)?
844 .iter()
845 .zip(rhs)
846 .map(|(s, t)| (s - t) * (s - t))
847 .sum::<f64>()
848 .sqrt())
849 };
850 let residual_norm = residual(a, b)?;
851 let constraint_residual_norm = residual(c, d)?;
852 let constraint_estimate = (n * p) as f64 * U * 2.0 * ct.condition;
853 Ok(ConstrainedSolution {
854 x,
855 residual_norm,
856 constraint_residual_norm,
857 relative_error_estimate: constraint_estimate + reduced_estimate,
858 })
859}