pub(crate) fn flatten(rows: &[Vec<f64>]) -> Vec<f64> {
rows.iter().flatten().copied().collect()
}
pub(crate) fn to_rows(a: &[f64], n_rows: usize, n_cols: usize) -> Vec<Vec<f64>> {
(0..n_rows)
.map(|i| a[i * n_cols..(i + 1) * n_cols].to_vec())
.collect()
}
pub(crate) fn dot(a: &[f64], b: &[f64]) -> f64 {
a.iter().zip(b).map(|(x, y)| x * y).sum()
}
pub(crate) fn matvec(a: &[f64], n_rows: usize, n_cols: usize, x: &[f64]) -> Vec<f64> {
(0..n_rows)
.map(|i| dot(&a[i * n_cols..(i + 1) * n_cols], x))
.collect()
}
pub(crate) fn cholesky(a: &[f64], n: usize, rel_tol: f64) -> Option<Vec<f64>> {
let mut l = vec![0.0; n * n];
for j in 0..n {
let mut s = a[j * n + j];
for k in 0..j {
s -= l[j * n + k] * l[j * n + k];
}
if !s.is_finite() || s <= rel_tol * a[j * n + j].abs() {
return None;
}
let d = s.sqrt();
l[j * n + j] = d;
for i in (j + 1)..n {
let mut s = a[i * n + j];
for k in 0..j {
s -= l[i * n + k] * l[j * n + k];
}
l[i * n + j] = s / d;
}
}
Some(l)
}
pub(crate) fn cholesky_solve(l: &[f64], n: usize, b: &[f64]) -> Vec<f64> {
let mut y = vec![0.0; n];
for i in 0..n {
let mut s = b[i];
for k in 0..i {
s -= l[i * n + k] * y[k];
}
y[i] = s / l[i * n + i];
}
for i in (0..n).rev() {
let mut s = y[i];
for k in (i + 1)..n {
s -= l[k * n + i] * y[k];
}
y[i] = s / l[i * n + i];
}
y
}
pub(crate) fn spd_inverse(l: &[f64], n: usize) -> Vec<f64> {
let mut inv = vec![0.0; n * n];
let mut e = vec![0.0; n];
for c in 0..n {
e[c] = 1.0;
let col = cholesky_solve(l, n, &e);
e[c] = 0.0;
for i in 0..n {
inv[i * n + c] = col[i];
}
}
inv
}
pub(crate) fn quadratic_form(l: &[f64], n: usize, x: &[f64]) -> f64 {
dot(x, &cholesky_solve(l, n, x))
}
pub(crate) fn congruence_inverse(l: &[f64], a: &[f64], n: usize) -> Vec<f64> {
let mut y = vec![0.0; n * n];
for c in 0..n {
for i in 0..n {
let mut s = a[i * n + c];
for k in 0..i {
s -= l[i * n + k] * y[k * n + c];
}
y[i * n + c] = s / l[i * n + i];
}
}
let mut w = vec![0.0; n * n];
for r in 0..n {
for i in 0..n {
let mut s = y[r * n + i];
for k in 0..i {
s -= l[i * n + k] * w[r * n + k];
}
w[r * n + i] = s / l[i * n + i];
}
}
for i in 0..n {
for j in (i + 1)..n {
let v = 0.5 * (w[i * n + j] + w[j * n + i]);
w[i * n + j] = v;
w[j * n + i] = v;
}
}
w
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) struct RegSchedule {
pub rounds: usize,
pub initial_rel: f64,
pub growth: f64,
}
pub(crate) fn solve_spd_regularised(
a: &[f64],
n: usize,
b: &[f64],
schedule: &RegSchedule,
) -> Option<Vec<f64>> {
let mut reg = 0.0;
for _ in 0..schedule.rounds {
let mut m = a.to_vec();
if reg > 0.0 {
for i in 0..n {
m[i * n + i] += reg;
}
}
if let Some(l) = cholesky(&m, n, 0.0) {
return Some(cholesky_solve(&l, n, b));
}
let scale = (0..n)
.map(|i| a[i * n + i].abs())
.fold(0.0, f64::max)
.max(1e-300);
reg = if reg == 0.0 {
scale * schedule.initial_rel
} else {
reg * schedule.growth
};
}
None
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) enum EigenTol {
RelativeFrobenius(f64),
Absolute(f64),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum OnExhaust {
Error,
Accept,
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) struct EigenOpts {
pub tol: EigenTol,
pub max_sweeps: usize,
pub on_exhaust: OnExhaust,
}
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct SymEigen {
pub values: Vec<f64>,
pub vectors: Vec<f64>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum DenseError {
NoConvergence { sweeps: usize },
}
impl std::fmt::Display for DenseError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
DenseError::NoConvergence { sweeps } => {
write!(
f,
"the Jacobi sweeps did not converge in {sweeps} iterations"
)
}
}
}
}
pub(crate) fn sym_eigen(a: &[f64], n: usize, opts: &EigenOpts) -> Result<SymEigen, DenseError> {
let mut m = a.to_vec();
let mut v = vec![0.0; n * n];
for i in 0..n {
v[i * n + i] = 1.0;
}
let threshold = match opts.tol {
EigenTol::RelativeFrobenius(rel) => {
let frob: f64 = m.iter().map(|x| x * x).sum::<f64>().sqrt();
rel * frob.max(f64::MIN_POSITIVE)
}
EigenTol::Absolute(eps) => eps,
};
let mut converged = false;
for _sweep in 0..opts.max_sweeps {
let done = match opts.tol {
EigenTol::RelativeFrobenius(_) => {
let off: f64 = (0..n)
.flat_map(|i| (0..n).filter(move |&j| j != i).map(move |j| (i, j)))
.map(|(i, j)| m[i * n + j] * m[i * n + j])
.sum::<f64>()
.sqrt();
off <= threshold
}
EigenTol::Absolute(_) => {
let mut off = 0.0;
for i in 0..n {
for j in (i + 1)..n {
off += m[i * n + j] * m[i * n + j];
}
}
off < threshold
}
};
if done {
converged = true;
break;
}
for p in 0..n {
for q in (p + 1)..n {
let apq = m[p * n + q];
if apq.abs() < 1e-300 {
continue;
}
let app = m[p * n + p];
let aqq = m[q * n + q];
let theta = (aqq - app) / (2.0 * apq);
let t = theta.signum() / (theta.abs() + (theta * theta + 1.0).sqrt());
let t = if theta == 0.0 { 1.0 } else { t };
let c = 1.0 / (t * t + 1.0).sqrt();
let s = t * c;
for k in 0..n {
let mkp = m[k * n + p];
let mkq = m[k * n + q];
m[k * n + p] = c * mkp - s * mkq;
m[k * n + q] = s * mkp + c * mkq;
}
for k in 0..n {
let mpk = m[p * n + k];
let mqk = m[q * n + k];
m[p * n + k] = c * mpk - s * mqk;
m[q * n + k] = s * mpk + c * mqk;
}
for k in 0..n {
let vkp = v[k * n + p];
let vkq = v[k * n + q];
v[k * n + p] = c * vkp - s * vkq;
v[k * n + q] = s * vkp + c * vkq;
}
}
}
}
if !converged && opts.on_exhaust == OnExhaust::Error {
return Err(DenseError::NoConvergence {
sweeps: opts.max_sweeps,
});
}
Ok(SymEigen {
values: (0..n).map(|i| m[i * n + i]).collect(),
vectors: v,
})
}
pub(crate) fn solve_partial_pivot(a: &[f64], n: usize, b: &[f64]) -> Option<Vec<f64>> {
let mut m: Vec<f64> = a.to_vec();
let mut r = b.to_vec();
for c in 0..n {
let p = (c..n).max_by(|&i, &j| {
m[i * n + c]
.abs()
.partial_cmp(&m[j * n + c].abs())
.unwrap_or(std::cmp::Ordering::Equal)
})?;
if m[p * n + c].abs() < 1e-300 {
return None;
}
if p != c {
for k in 0..n {
m.swap(c * n + k, p * n + k);
}
r.swap(c, p);
}
for i in (c + 1)..n {
let f = m[i * n + c] / m[c * n + c];
if f != 0.0 {
for k in c..n {
m[i * n + k] -= f * m[c * n + k];
}
r[i] -= f * r[c];
}
}
}
let mut x = vec![0.0; n];
for i in (0..n).rev() {
let mut sacc = r[i];
for k in (i + 1)..n {
sacc -= m[i * n + k] * x[k];
}
x[i] = sacc / m[i * n + i];
}
Some(x)
}
pub(crate) fn lstsq_householder(
a: &[f64],
n_rows: usize,
n_cols: usize,
b: &[f64],
) -> Option<Vec<f64>> {
let (m, n) = (n_rows, n_cols);
if m < n || b.len() != m || a.len() != m * n {
return None;
}
let mut a = a.to_vec();
let mut b = b.to_vec();
for k in 0..n {
let norm = (k..m)
.map(|i| a[i * n + k] * a[i * n + k])
.sum::<f64>()
.sqrt();
if norm == 0.0 || !norm.is_finite() {
return None;
}
let alpha = if a[k * n + k] > 0.0 { -norm } else { norm };
let mut v: Vec<f64> = (k..m).map(|i| a[i * n + k]).collect();
v[0] -= alpha;
let vnorm2: f64 = v.iter().map(|x| x * x).sum();
if vnorm2 == 0.0 {
continue;
}
let scale = 2.0 / vnorm2;
let w: Vec<f64> = (k..n)
.map(|j| {
v.iter()
.zip(k..m)
.map(|(vi, i)| vi * a[i * n + j])
.sum::<f64>()
})
.collect();
for (vi, i) in v.iter().zip(k..m) {
for (wj, j) in w.iter().zip(k..n) {
a[i * n + j] -= scale * vi * wj;
}
}
let s: f64 = v.iter().zip(k..m).map(|(vi, i)| vi * b[i]).sum();
let factor = scale * s;
for (vi, i) in v.iter().zip(k..m) {
b[i] -= factor * vi;
}
}
let r_max = (0..n).map(|k| a[k * n + k].abs()).fold(0.0_f64, f64::max);
let threshold = r_max * f64::EPSILON * m as f64;
let mut x = vec![0.0; n];
for r in (0..n).rev() {
let diag = a[r * n + r];
if !diag.is_finite() || diag.abs() <= threshold {
return None;
}
let s = b[r] - ((r + 1)..n).map(|c| a[r * n + c] * x[c]).sum::<f64>();
x[r] = s / diag;
}
Some(x)
}
#[cfg(test)]
mod tests {
use super::*;
fn close(actual: f64, expected: f64) -> bool {
(actual - expected).abs() < 1e-12 * expected.abs().max(1.0)
}
fn assert_close(actual: &[f64], expected: &[f64]) {
assert_eq!(actual.len(), expected.len());
for (a, e) in actual.iter().zip(expected) {
assert!(close(*a, *e), "{actual:?} vs {expected:?}");
}
}
const SPD3: [f64; 9] = [4.0, 12.0, -16.0, 12.0, 37.0, -43.0, -16.0, -43.0, 98.0];
const TRIDIAG4: [f64; 16] = [
2.0, -1.0, 0.0, 0.0, -1.0, 2.0, -1.0, 0.0, 0.0, -1.0, 2.0, -1.0, 0.0, 0.0, -1.0, 2.0,
];
const EIGEN_STRICT: EigenOpts = EigenOpts {
tol: EigenTol::RelativeFrobenius(1e-15),
max_sweeps: 100,
on_exhaust: OnExhaust::Error,
};
#[test]
fn cholesky_3x3_is_the_integer_factor() {
let l = cholesky(&SPD3, 3, 0.0).unwrap();
assert_eq!(l, [2.0, 0.0, 0.0, 6.0, 1.0, 0.0, -8.0, 5.0, 3.0]);
let x = cholesky_solve(&l, 3, &[1.0, 2.0, 3.0]);
assert_close(&x, &[343.0 / 12.0, -23.0 / 3.0, 4.0 / 3.0]);
let inv = spd_inverse(&l, 3);
assert_close(&inv[..3], &[1777.0 / 36.0, -122.0 / 9.0, 19.0 / 9.0]);
assert_close(&inv[3..6], &[-122.0 / 9.0, 34.0 / 9.0, -5.0 / 9.0]);
assert_close(&inv[6..], &[19.0 / 9.0, -5.0 / 9.0, 1.0 / 9.0]);
assert!(close(quadratic_form(&l, 3, &[1.0, 2.0, 3.0]), 17.25));
let w = congruence_inverse(&l, &SPD3, 3);
assert_close(&w, &[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]);
}
#[test]
fn non_spd_matrices_return_none() {
assert!(cholesky(&[1.0, 2.0, 2.0, 1.0], 2, 0.0).is_none());
assert!(cholesky(&[1.0, 2.0, 2.0, 4.0], 2, 0.0).is_none());
assert!(cholesky(&[1.0, 0.0, 0.0, f64::NAN], 2, 0.0).is_none());
assert!(cholesky(&[f64::INFINITY, 0.0, 0.0, 1.0], 2, 0.0).is_none());
assert!(cholesky(&[0.0, 0.0, 0.0, 1.0], 2, 0.0).is_none());
}
#[test]
fn relative_pivot_tolerance() {
let a = [1.0, 1.0, 1.0, 1.0 + 1e-13];
assert!(cholesky(&a, 2, 1e-12).is_none());
assert!(cholesky(&a, 2, 1e-14).is_some());
assert!(cholesky(&a, 2, 0.0).is_some());
let tiny: Vec<f64> = SPD3.iter().map(|v| v * 1e-20).collect();
assert!(cholesky(&tiny, 3, 1e-12).is_some());
}
#[test]
fn regularised_solve_recovers_from_a_singular_system() {
let schedule = RegSchedule {
rounds: 6,
initial_rel: 1e-12,
growth: 100.0,
};
let x = solve_spd_regularised(&SPD3, 3, &[1.0, 2.0, 3.0], &schedule).unwrap();
assert_close(&x, &[343.0 / 12.0, -23.0 / 3.0, 4.0 / 3.0]);
let sing = [1.0, 1.0, 1.0, 1.0];
let x = solve_spd_regularised(&sing, 2, &[2.0, 2.0], &schedule).unwrap();
assert!(x.iter().all(|v| v.is_finite()));
assert!((x[0] + x[1] - 2.0).abs() < 1e-6, "{x:?}");
assert!(
solve_spd_regularised(&[-1.0, 0.0, 0.0, -1.0], 2, &[1.0, 1.0], &schedule).is_none()
);
}
#[test]
fn jacobi_eigen_4x4_matches_eigh() {
let e = sym_eigen(&TRIDIAG4, 4, &EIGEN_STRICT).unwrap();
let mut order: Vec<usize> = (0..4).collect();
order.sort_by(|&a, &b| e.values[a].total_cmp(&e.values[b]));
let sorted: Vec<f64> = order.iter().map(|&i| e.values[i]).collect();
assert_close(
&sorted,
&[
0.3819660112501053,
1.3819660112501055,
2.618033988749895,
3.6180339887498945,
],
);
let c = order[0];
let mut v: Vec<f64> = (0..4).map(|i| e.vectors[i * 4 + c]).collect();
if v[0] < 0.0 {
v.iter_mut().for_each(|x| *x = -*x);
}
assert_close(
&v,
&[
0.3717480344601845,
0.6015009550075455,
0.6015009550075455,
0.3717480344601844,
],
);
for c in 0..4 {
let v: Vec<f64> = (0..4).map(|i| e.vectors[i * 4 + c]).collect();
let av = matvec(&TRIDIAG4, 4, 4, &v);
for i in 0..4 {
assert!((av[i] - e.values[c] * v[i]).abs() < 1e-14);
}
}
}
#[test]
fn jacobi_eigen_3x3_matches_eigh() {
let a = [2.0, 1.0, 0.5, 1.0, 3.0, 1.0, 0.5, 1.0, 4.0];
let mut values = sym_eigen(&a, 3, &EIGEN_STRICT).unwrap().values;
values.sort_by(f64::total_cmp);
assert_close(
&values,
&[1.3770948783558645, 2.6824555614434478, 4.940449560200685],
);
}
#[test]
fn eigen_tolerances_and_exhaustion() {
let absolute = EigenOpts {
tol: EigenTol::Absolute(1e-30),
max_sweeps: 100,
on_exhaust: OnExhaust::Accept,
};
let a = sym_eigen(&TRIDIAG4, 4, &absolute).unwrap();
let r = sym_eigen(&TRIDIAG4, 4, &EIGEN_STRICT).unwrap();
let (mut av, mut rv) = (a.values.clone(), r.values.clone());
av.sort_by(f64::total_cmp);
rv.sort_by(f64::total_cmp);
assert_close(&av, &rv);
let tiny: Vec<f64> = TRIDIAG4.iter().map(|v| v * 1e-16).collect();
assert_eq!(
sym_eigen(&tiny, 4, &absolute).unwrap().values,
vec![2e-16; 4]
);
let mut scaled = sym_eigen(&tiny, 4, &EIGEN_STRICT).unwrap().values;
scaled.sort_by(f64::total_cmp);
let expected: Vec<f64> = rv.iter().map(|v| v * 1e-16).collect();
assert_close(&scaled, &expected);
let never = EigenOpts {
tol: EigenTol::Absolute(0.0),
max_sweeps: 100,
on_exhaust: OnExhaust::Error,
};
assert_eq!(
sym_eigen(&TRIDIAG4, 4, &never),
Err(DenseError::NoConvergence { sweeps: 100 })
);
let mut accepted = sym_eigen(
&TRIDIAG4,
4,
&EigenOpts {
on_exhaust: OnExhaust::Accept,
..never
},
)
.unwrap()
.values;
accepted.sort_by(f64::total_cmp);
assert_close(&accepted, &rv);
let none = EigenOpts {
max_sweeps: 0,
..EIGEN_STRICT
};
assert_eq!(
sym_eigen(&TRIDIAG4, 4, &none),
Err(DenseError::NoConvergence { sweeps: 0 })
);
let accept_none = EigenOpts {
on_exhaust: OnExhaust::Accept,
..none
};
assert_eq!(
sym_eigen(&TRIDIAG4, 4, &accept_none).unwrap().values,
vec![2.0; 4]
);
let d = [3.0, 0.0, 0.0, 1.0];
assert!(sym_eigen(&d, 2, &none).is_err());
let one = EigenOpts {
max_sweeps: 1,
..EIGEN_STRICT
};
assert_eq!(sym_eigen(&d, 2, &one).unwrap().values, vec![3.0, 1.0]);
let e = sym_eigen(
&d,
2,
&EigenOpts {
max_sweeps: 1,
on_exhaust: OnExhaust::Error,
..absolute
},
)
.unwrap();
assert_eq!(e.vectors, vec![1.0, 0.0, 0.0, 1.0]);
}
#[test]
fn partial_pivot_solves_and_detects_singularity() {
let x = solve_partial_pivot(&[0.0, 1.0, 1.0, 0.0], 2, &[3.0, 4.0]).unwrap();
assert_close(&x, &[4.0, 3.0]);
let x = solve_partial_pivot(&SPD3, 3, &[1.0, 2.0, 3.0]).unwrap();
assert_close(&x, &[343.0 / 12.0, -23.0 / 3.0, 4.0 / 3.0]);
assert!(solve_partial_pivot(&[1.0, 2.0, 2.0, 4.0], 2, &[1.0, 2.0]).is_none());
assert!(solve_partial_pivot(&[0.0, 0.0, 0.0, 0.0], 2, &[1.0, 2.0]).is_none());
assert_eq!(solve_partial_pivot(&[], 0, &[]), Some(vec![]));
}
#[test]
fn householder_least_squares_overdetermined() {
let a = [1.0, 0.0, 1.0, 1.0, 1.0, 2.0, 1.0, 3.0];
let x = lstsq_householder(&a, 4, 2, &[1.0, 3.0, 2.0, 5.0]).unwrap();
assert_close(&x, &[1.1, 1.1]);
let x = lstsq_householder(&SPD3, 3, 3, &[1.0, 2.0, 3.0]).unwrap();
assert!((x[0] - 343.0 / 12.0).abs() < 1e-9, "{x:?}");
assert!(
lstsq_householder(&[1.0, 2.0, 2.0, 4.0, 3.0, 6.0], 3, 2, &[1.0, 2.0, 3.0]).is_none()
);
assert!(lstsq_householder(&[1.0, 2.0], 1, 2, &[1.0]).is_none());
assert!(lstsq_householder(&a, 4, 2, &[1.0, 2.0]).is_none());
}
#[test]
fn layout_and_products() {
let rows = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]];
let flat = flatten(&rows);
assert_eq!(flat, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
assert_eq!(to_rows(&flat, 2, 3), rows);
assert_eq!(matvec(&flat, 2, 3, &[1.0, 1.0, 1.0]), [6.0, 15.0]);
assert_eq!(dot(&[1.0, 2.0], &[3.0, 4.0]), 11.0);
assert_eq!(
dot(&[1.0, 2.0, 3.0], &[3.0, 4.0]),
11.0,
"zip stops at the shorter"
);
}
}