DarkRoom/core/dr-pano/src/linalg.rs

//! The small dense linear algebra the geometry needs, and nothing more.
//!
//! Hand-rolled rather than pulled in, and the decision was made on purpose
//! (2026-09-19): the largest system this crate ever solves is a rotation
//! per frame plus one focal length — forty unknowns for a dozen frames —
//! and everything else is three-vectors. A general linear-algebra crate
//! would be the largest dependency in `dr-pano` by an order of magnitude,
//! for a Cholesky factorisation that is thirty lines.
//!
//! `f64` throughout. The geometry is solved once per merge on a few thousand
//! matches; there is no reason to give up precision for speed here, and the
//! bundle adjustment's normal equations are poorly conditioned enough near
//! convergence that `f32` would stall it.

use std::ops::{Add, Index, IndexMut, Mul, Neg, Sub};

/// A vector in three dimensions.
#[derive(Debug, Clone, Copy, PartialEq, Default)]
pub struct Vec3(pub [f64; 3]);

impl Vec3 {
    pub const fn new(x: f64, y: f64, z: f64) -> Self {
        Vec3([x, y, z])
    }

    pub fn dot(self, o: Vec3) -> f64 {
        self.0[0] * o.0[0] + self.0[1] * o.0[1] + self.0[2] * o.0[2]
    }

    pub fn cross(self, o: Vec3) -> Vec3 {
        Vec3([
            self.0[1] * o.0[2] - self.0[2] * o.0[1],
            self.0[2] * o.0[0] - self.0[0] * o.0[2],
            self.0[0] * o.0[1] - self.0[1] * o.0[0],
        ])
    }

    pub fn norm(self) -> f64 {
        self.dot(self).sqrt()
    }

    /// The unit vector along `self`, or `self` unchanged if it is zero.
    pub fn normalised(self) -> Vec3 {
        let n = self.norm();
        if n > 0.0 {
            self * (1.0 / n)
        } else {
            self
        }
    }

    pub fn x(self) -> f64 {
        self.0[0]
    }
    pub fn y(self) -> f64 {
        self.0[1]
    }
    pub fn z(self) -> f64 {
        self.0[2]
    }
}

impl Add for Vec3 {
    type Output = Vec3;
    fn add(self, o: Vec3) -> Vec3 {
        Vec3([self.0[0] + o.0[0], self.0[1] + o.0[1], self.0[2] + o.0[2]])
    }
}

impl Sub for Vec3 {
    type Output = Vec3;
    fn sub(self, o: Vec3) -> Vec3 {
        Vec3([self.0[0] - o.0[0], self.0[1] - o.0[1], self.0[2] - o.0[2]])
    }
}

impl Mul<f64> for Vec3 {
    type Output = Vec3;
    fn mul(self, s: f64) -> Vec3 {
        Vec3([self.0[0] * s, self.0[1] * s, self.0[2] * s])
    }
}

impl Neg for Vec3 {
    type Output = Vec3;
    fn neg(self) -> Vec3 {
        Vec3([-self.0[0], -self.0[1], -self.0[2]])
    }
}

/// A 3×3 matrix, row-major.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Mat3(pub [[f64; 3]; 3]);

impl Mat3 {
    pub const IDENTITY: Mat3 = Mat3([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]);

    /// The matrix whose columns are `a`, `b`, `c`.
    pub fn from_columns(a: Vec3, b: Vec3, c: Vec3) -> Mat3 {
        Mat3([
            [a.0[0], b.0[0], c.0[0]],
            [a.0[1], b.0[1], c.0[1]],
            [a.0[2], b.0[2], c.0[2]],
        ])
    }

    pub fn transpose(self) -> Mat3 {
        let m = self.0;
        Mat3([
            [m[0][0], m[1][0], m[2][0]],
            [m[0][1], m[1][1], m[2][1]],
            [m[0][2], m[1][2], m[2][2]],
        ])
    }

    pub fn column(self, i: usize) -> Vec3 {
        Vec3([self.0[0][i], self.0[1][i], self.0[2][i]])
    }

    pub fn trace(self) -> f64 {
        self.0[0][0] + self.0[1][1] + self.0[2][2]
    }

    /// The rotation about `axis` (any length) by `angle` radians — Rodrigues.
    pub fn rotation(axis: Vec3, angle: f64) -> Mat3 {
        let k = axis.normalised();
        let (s, c) = angle.sin_cos();
        let t = 1.0 - c;
        let (x, y, z) = (k.0[0], k.0[1], k.0[2]);
        Mat3([
            [t * x * x + c, t * x * y - s * z, t * x * z + s * y],
            [t * x * y + s * z, t * y * y + c, t * y * z - s * x],
            [t * x * z - s * y, t * y * z + s * x, t * z * z + c],
        ])
    }

    /// The rotation whose axis-angle vector is `w` (direction is the axis,
    /// length is the angle). The exponential map; [`Self::log`] inverts it.
    pub fn exp(w: Vec3) -> Mat3 {
        let angle = w.norm();
        if angle < 1e-12 {
            // First-order: I + [w]×, which is what the limit is and avoids
            // dividing by the angle.
            let (x, y, z) = (w.0[0], w.0[1], w.0[2]);
            return Mat3([[1.0, -z, y], [z, 1.0, -x], [-y, x, 1.0]]);
        }
        Mat3::rotation(w, angle)
    }

    /// The axis-angle vector of a rotation matrix. Inverse of [`Self::exp`].
    pub fn log(self) -> Vec3 {
        let m = self.0;
        let cos = ((self.trace() - 1.0) * 0.5).clamp(-1.0, 1.0);
        let axis = Vec3([m[2][1] - m[1][2], m[0][2] - m[2][0], m[1][0] - m[0][1]]);
        if cos > 1.0 - 1e-6 {
            // Small angle: `acos` near 1 loses everything below ~1e-8 to
            // rounding, but the antisymmetric part is `2 sin θ · axis` and
            // keeps it. First order, exact to the precision that matters.
            return axis * 0.5;
        }
        let angle = cos.acos();
        if angle > std::f64::consts::PI - 1e-6 {
            // Near π the antisymmetric part vanishes; take the axis from the
            // symmetric part instead. Rare for a panorama, but the solver may
            // pass through it on a bad start and must not return NaN.
            let d = Vec3([
                ((m[0][0] + 1.0) * 0.5).max(0.0).sqrt(),
                ((m[1][1] + 1.0) * 0.5).max(0.0).sqrt(),
                ((m[2][2] + 1.0) * 0.5).max(0.0).sqrt(),
            ]);
            return d.normalised() * angle;
        }
        axis * (angle / (2.0 * angle.sin()))
    }

    /// Re-orthonormalise a matrix that has drifted from a rotation through
    /// accumulated products. Gram–Schmidt on the columns; cheap and adequate
    /// for drift of the size floating-point products produce.
    pub fn orthonormalised(self) -> Mat3 {
        let a = self.column(0).normalised();
        let b = (self.column(1) - a * a.dot(self.column(1))).normalised();
        let c = a.cross(b);
        Mat3::from_columns(a, b, c)
    }
}

impl Mul<Vec3> for Mat3 {
    type Output = Vec3;
    fn mul(self, v: Vec3) -> Vec3 {
        let m = self.0;
        Vec3([
            m[0][0] * v.0[0] + m[0][1] * v.0[1] + m[0][2] * v.0[2],
            m[1][0] * v.0[0] + m[1][1] * v.0[1] + m[1][2] * v.0[2],
            m[2][0] * v.0[0] + m[2][1] * v.0[1] + m[2][2] * v.0[2],
        ])
    }
}

impl Mul for Mat3 {
    type Output = Mat3;
    fn mul(self, o: Mat3) -> Mat3 {
        let mut r = [[0.0; 3]; 3];
        for (i, row) in r.iter_mut().enumerate() {
            for (j, cell) in row.iter_mut().enumerate() {
                *cell = (0..3).map(|k| self.0[i][k] * o.0[k][j]).sum();
            }
        }
        Mat3(r)
    }
}

/// A dense square matrix, for the normal equations.
#[derive(Debug, Clone, PartialEq)]
pub struct DMat {
    n: usize,
    data: Vec<f64>,
}

impl DMat {
    pub fn zeros(n: usize) -> DMat {
        DMat {
            n,
            data: vec![0.0; n * n],
        }
    }

    pub fn n(&self) -> usize {
        self.n
    }

    /// Solve `self · x = b` for a symmetric positive-definite `self` by
    /// Cholesky factorisation. `None` if the matrix is not positive definite,
    /// which for the normal equations means the problem is not determined by
    /// the data — a frame with no matches, for instance — and the caller
    /// should say so rather than proceed.
    ///
    /// Destroys neither input: the factor is built in a copy. The systems
    /// here are at most a few dozen unknowns and the copy is nothing.
    pub fn solve_spd(&self, b: &[f64]) -> Option<Vec<f64>> {
        let n = self.n;
        debug_assert_eq!(b.len(), n);
        let mut l = vec![0.0; n * n];
        for j in 0..n {
            let mut d = self[(j, j)];
            for k in 0..j {
                d -= l[j * n + k] * l[j * n + k];
            }
            if d <= 0.0 || !d.is_finite() {
                return None;
            }
            let djj = d.sqrt();
            l[j * n + j] = djj;
            for i in j + 1..n {
                let mut s = self[(i, j)];
                for k in 0..j {
                    s -= l[i * n + k] * l[j * n + k];
                }
                l[i * n + j] = s / djj;
            }
        }
        // Forward: L y = b.
        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];
        }
        // Back: Lᵀ x = y.
        let mut x = vec![0.0; n];
        for i in (0..n).rev() {
            let mut s = y[i];
            for k in i + 1..n {
                s -= l[k * n + i] * x[k];
            }
            x[i] = s / l[i * n + i];
        }
        Some(x)
    }
}

impl Index<(usize, usize)> for DMat {
    type Output = f64;
    fn index(&self, (i, j): (usize, usize)) -> &f64 {
        &self.data[i * self.n + j]
    }
}

impl IndexMut<(usize, usize)> for DMat {
    fn index_mut(&mut self, (i, j): (usize, usize)) -> &mut f64 {
        &mut self.data[i * self.n + j]
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn close(a: f64, b: f64) -> bool {
        (a - b).abs() < 1e-9
    }

    #[test]
    fn exp_and_log_are_inverses() {
        for w in [
            Vec3::new(0.1, -0.2, 0.3),
            Vec3::new(1.0, 0.0, 0.0),
            Vec3::new(0.0, 0.0, 2.5),
            Vec3::new(1e-9, 0.0, 0.0),
        ] {
            let back = Mat3::exp(w).log();
            for i in 0..3 {
                assert!(close(back.0[i], w.0[i]), "{w:?} -> {back:?}");
            }
        }
    }

    #[test]
    fn a_rotation_is_orthonormal_and_preserves_length() {
        let r = Mat3::exp(Vec3::new(0.4, 0.5, -0.6));
        let rt = r.transpose() * r;
        for i in 0..3 {
            for j in 0..3 {
                assert!(close(rt.0[i][j], Mat3::IDENTITY.0[i][j]));
            }
        }
        let v = Vec3::new(1.0, 2.0, 3.0);
        assert!(close((r * v).norm(), v.norm()));
    }

    #[test]
    fn rotation_about_z_turns_x_towards_y() {
        let r = Mat3::rotation(Vec3::new(0.0, 0.0, 1.0), std::f64::consts::FRAC_PI_2);
        let v = r * Vec3::new(1.0, 0.0, 0.0);
        assert!(close(v.x(), 0.0) && close(v.y(), 1.0) && close(v.z(), 0.0));
    }

    #[test]
    fn cholesky_solves_a_small_spd_system() {
        // A = Bᵀ B for a random-ish B is SPD by construction.
        let b = [
            [2.0, 1.0, 0.0],
            [1.0, 3.0, 1.0],
            [0.0, 1.0, 4.0],
            [1.0, 1.0, 1.0],
        ];
        let mut a = DMat::zeros(3);
        for i in 0..3 {
            for j in 0..3 {
                a[(i, j)] = (0..4).map(|k| b[k][i] * b[k][j]).sum();
            }
        }
        let x_true = [1.0, -2.0, 0.5];
        let rhs: Vec<f64> = (0..3)
            .map(|i| (0..3).map(|j| a[(i, j)] * x_true[j]).sum())
            .collect();
        let x = a.solve_spd(&rhs).expect("spd");
        for i in 0..3 {
            assert!(close(x[i], x_true[i]), "{x:?}");
        }
    }

    #[test]
    fn cholesky_refuses_an_indefinite_matrix() {
        let mut a = DMat::zeros(2);
        a[(0, 0)] = 1.0;
        a[(1, 1)] = -1.0;
        assert!(a.solve_spd(&[1.0, 1.0]).is_none());
    }
}