Skip to main content

axiolid_numeric/
linalg.rs

1//! Dense linear algebra for small and medium systems.
2//!
3//! - [`Lu`]: Gaussian elimination with partial pivoting, for square
4//!   systems (OCCT `math_Gauss`).
5//! - [`Cholesky`]: `L L^T` for symmetric positive definite systems, such as
6//!   normal equations and fairing energies (OCCT `math_Crout` covers the
7//!   symmetric case there).
8//! - [`Qr`]: Householder QR with column pivoting, for linear least squares
9//!   and rank detection (OCCT `math_Householder`, LAPACK `xGEQP3`).
10//! - [`constrained_least_squares`]: least squares subject to linear
11//!   equality constraints, by the null-space method (Golub and Van Loan,
12//!   *Matrix Computations*, 6.2).
13//!
14//! Every solve reports a condition estimate (Hager's 1-norm estimator as
15//! refined by Higham, LAPACK `xLACON`) and a relative forward error estimate
16//! derived from it. Singular and rank-deficient systems are refused by name
17//! rather than solved into infinities.
18
19// Triangular solves read clearest with the textbook indices.
20#![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/// A dense row-major matrix of `f64`.
29#[derive(Debug, Clone, PartialEq)]
30pub struct Matrix {
31    rows: usize,
32    cols: usize,
33    data: Vec<f64>,
34}
35
36impl Matrix {
37    /// A `rows x cols` matrix of zeros.
38    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    /// The `n x n` identity.
47    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    /// A matrix whose entry `(i, j)` is `f(i, j)`.
56    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    /// A matrix from rows of equal length.
67    ///
68    /// # Errors
69    ///
70    /// Refuses rows of different lengths.
71    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    /// Number of rows.
92    pub fn rows(&self) -> usize {
93        self.rows
94    }
95
96    /// Number of columns.
97    pub fn cols(&self) -> usize {
98        self.cols
99    }
100
101    /// Row `i` as a slice.
102    pub fn row(&self, i: usize) -> &[f64] {
103        &self.data[i * self.cols..(i + 1) * self.cols]
104    }
105
106    /// The transpose.
107    pub fn transpose(&self) -> Self {
108        Self::from_fn(self.cols, self.rows, |i, j| self[(j, i)])
109    }
110
111    /// `A x`.
112    ///
113    /// # Errors
114    ///
115    /// Refuses `x` whose length is not the column count.
116    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    /// The 1-norm: the largest column sum of absolute values.
130    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/// The solution of a square linear system.
169#[derive(Debug, Clone, PartialEq)]
170#[must_use]
171pub struct LinearSolution {
172    /// The solution vector.
173    pub x: Vec<f64>,
174    /// `||b - A x||_1`, computed in `f64`.
175    pub residual_norm: f64,
176    /// Estimate of `||x - x_exact||_1 / ||x_exact||_1`: twice the condition
177    /// estimate times the normwise backward error, the latter enlarged by
178    /// a bound on the rounding of the computed residual itself.
179    pub relative_error_estimate: f64,
180}
181
182/// The solution of a linear least-squares problem.
183#[derive(Debug, Clone, PartialEq)]
184#[must_use]
185pub struct LeastSquaresSolution {
186    /// The minimiser of `||A x - b||_2`.
187    pub x: Vec<f64>,
188    /// `||A x - b||_2` at the solution.
189    pub residual_norm: f64,
190    /// Estimate of `||x - x_exact|| / ||x_exact||` from the standard
191    /// first-order perturbation bound for Householder least squares,
192    /// `m n u (2 kappa + kappa^2 ||r|| / (||A|| ||x||))`.
193    pub relative_error_estimate: f64,
194}
195
196/// Estimate `||M^-1||_1` given solves with `M` and `M^T` (Hager 1984, with
197/// Higham's 1988 alternative vector).
198fn 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    // Higham's alternating vector catches cases the power-like loop misses.
232    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    // The computed residual carries its own rounding, up to
261    // gamma(n + 1) * sum(|A||x| + |b|).
262    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/// `P A = L U` with partial pivoting.
275#[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    /// Factor a square matrix.
285    ///
286    /// # Errors
287    ///
288    /// Refuses a non-square or non-finite matrix, and a matrix singular to
289    /// working precision (a zero pivot, or a condition estimate above
290    /// `1 / eps`).
291    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    /// Estimate of the 1-norm condition number `||A||_1 ||A^-1||_1`. It is
345    /// a lower bound, in practice within a factor of 3.
346    pub fn condition_estimate(&self) -> f64 {
347        self.condition
348    }
349
350    /// Solve `A x = b`.
351    ///
352    /// # Errors
353    ///
354    /// Refuses `b` of the wrong length or with non-finite entries.
355    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        // A^T = U^T L^T P: solve U^T w = c, L^T v = w, then z[perm[i]] = v[i].
391        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/// `A = L L^T` for a symmetric positive definite matrix.
424#[derive(Debug, Clone, PartialEq)]
425pub struct Cholesky {
426    a: Matrix,
427    l: Matrix,
428    condition: f64,
429}
430
431impl Cholesky {
432    /// Factor a symmetric positive definite matrix.
433    ///
434    /// Symmetry is checked to `n eps max|a_ij|`; only the lower triangle is
435    /// used.
436    ///
437    /// # Errors
438    ///
439    /// Refuses a non-square or non-finite matrix, an asymmetric one
440    /// ([`NumericError::NotSymmetric`]), one with a non-positive pivot
441    /// ([`NumericError::NotPositiveDefinite`]), and one singular to working
442    /// precision.
443    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    /// Estimate of the 1-norm condition number.
492    pub fn condition_estimate(&self) -> f64 {
493        self.condition
494    }
495
496    /// The factor `L` (lower triangular, positive diagonal).
497    pub fn factor(&self) -> &Matrix {
498        &self.l
499    }
500
501    /// Solve `A x = b`.
502    ///
503    /// # Errors
504    ///
505    /// Refuses `b` of the wrong length or with non-finite entries.
506    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/// Householder QR with column pivoting: `A P = Q R`.
539#[derive(Debug, Clone, PartialEq)]
540pub struct Qr {
541    a: Matrix,
542    /// Upper triangle holds `R`; the rest is scratch.
543    r: Matrix,
544    /// Householder vectors `v_k` (length `m - k`) and their `beta_k`.
545    reflectors: Vec<(Vec<f64>, f64)>,
546    perm: Vec<usize>,
547    rank: usize,
548    condition: f64,
549}
550
551impl Qr {
552    /// Factor an `m x n` matrix and determine its numerical rank: a pivot
553    /// counts when it exceeds `max(m, n) eps |r_00|`.
554    ///
555    /// # Errors
556    ///
557    /// Refuses a non-finite matrix.
558    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            // Pivot: the remaining column of largest norm.
569            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    /// Numerical rank.
629    pub fn rank(&self) -> usize {
630        self.rank
631    }
632
633    /// Estimate of the 1-norm condition number of the leading `rank x rank`
634    /// block of `R`, which equals `A`'s 2-norm condition number up to a
635    /// factor of `n`.
636    pub fn condition_estimate(&self) -> f64 {
637        self.condition
638    }
639
640    /// Minimise `||A x - b||_2`.
641    ///
642    /// # Errors
643    ///
644    /// Refuses `b` of the wrong length or with non-finite entries, and a
645    /// matrix whose rank is below its column count
646    /// ([`NumericError::RankDeficient`]): the minimiser is then not unique.
647    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    /// `c <- Q^T c` (length `m`).
688    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    /// `c <- Q c` (length `m`).
695    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    /// Solve `R_11 z = c` for the leading `rank` entries.
702    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    /// Solve `R_11^T z = c`.
714    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
734/// Minimise `||A x - b||_2`: a convenience for [`Qr::new`] followed by
735/// [`Qr::solve_least_squares`]. Factor once with [`Qr`] when several
736/// right-hand sides share `A` (one per coordinate when fitting points).
737///
738/// # Errors
739///
740/// As [`Qr::new`] and [`Qr::solve_least_squares`].
741pub fn least_squares(a: &Matrix, b: &[f64]) -> NumericResult<LeastSquaresSolution> {
742    Qr::new(a)?.solve_least_squares(b)
743}
744
745/// The solution of an equality-constrained least-squares problem.
746#[derive(Debug, Clone, PartialEq)]
747#[must_use]
748pub struct ConstrainedSolution {
749    /// The minimiser of `||A x - b||_2` subject to `C x = d`.
750    pub x: Vec<f64>,
751    /// `||A x - b||_2` at the solution.
752    pub residual_norm: f64,
753    /// `||C x - d||_2` at the solution: zero up to rounding.
754    pub constraint_residual_norm: f64,
755    /// Sum of the relative error estimates of the two stages (constraint
756    /// solve and reduced least squares).
757    pub relative_error_estimate: f64,
758}
759
760/// Minimise `||A x - b||_2` subject to `C x = d` by the null-space method.
761///
762/// `C` (`p x n`) must have full row rank `p <= n`, and `A` must have full
763/// column rank on the null space of `C`, so the minimiser is unique. Use
764/// this to fit with interpolated points or prescribed end tangents.
765///
766/// # Errors
767///
768/// Refuses mismatched sizes, non-finite input, a rank-deficient constraint
769/// matrix (redundant or conflicting constraints) and an objective that does
770/// not determine the solution on the constraints' null space.
771pub 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    // C^T Pi = Q [R; 0].
799    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    // R^T y1 = Pi^T d.
808    let mut y1: Vec<f64> = ct.perm.iter().map(|&i| d[i]).collect();
809    ct.solve_r_transpose(&mut y1);
810
811    // A Q, row by row: (A Q)_i = (Q^T a_i^T)^T.
812    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}