Implement a modified Gram-Schmidt orthonormalization algorithm

hoomd-linear-algebra/benches/linalg.rs

@@ -11,7 +11,10 @@
 use divan::{self, Bencher, black_box, counter::ItemsCount};
 use rand::{Rng, RngExt, SeedableRng, rngs::StdRng};
 
-use hoomd_linear_algebra::{Invertible, MatMul, matrix::Matrix};
+use hoomd_linear_algebra::{
+    Invertible, MatMul,
+    matrix::{Matrix, gram_schmidt::gram_schmidt},
+};
 
 fn main() {
     divan::main();
@@ -53,6 +56,16 @@ fn det_matn<const N: usize>(bencher: Bencher) {
         .bench_local_values(|a| black_box(a.determinant()));
 }
 
+#[divan::bench(consts = SQUARE_DIMENSIONS)]
+fn gram_schmidt_matn<const N: usize>(bencher: Bencher) {
+    let mut rng = StdRng::seed_from_u64(42);
+
+    bencher
+        .counter(ItemsCount::from(1_u32))
+        .with_inputs(|| create_random_matrix::<N, N, _>(&mut rng))
+        .bench_local_values(|a| black_box(gram_schmidt(&a)));
+}
+
 /// This benchmark is included as a reference implementation for comparison with ``SquareMatrix::<3, 3>::determinant``.
 /// We want to ensure the performance of the latter is comparable to this optimal implementation.
 #[divan::bench]

hoomd-linear-algebra/src/matrix.rs

@@ -12,6 +12,9 @@ pub mod ops;
 /// ``qr`` decomposition for [`Matrix`] types.
 pub mod qr;
 
+/// ``Gram-Schmidt`` orthonormalization for [`Matrix`] types.
+pub mod gram_schmidt;
+
 pub use crate::diagonal::DiagonalMatrix;
 
 /// A lightweight representation of a diagonal matrix.
@@ -1016,10 +1019,11 @@ pub mod test_utils {
         }
     }
 
-    pub(crate) fn assert_diags_ulps_eq<const N: usize, T: Diagonal>(
-        m0: &T,
-        m1: &impl std::ops::Index<usize, Output = f64>,
-    ) {
+    pub(crate) fn assert_diags_ulps_eq<const N: usize, T0, T1>(m0: &T0, m1: &T1)
+    where
+        T0: Index<usize, Output = f64> + ?Sized,
+        T1: Index<usize, Output = f64> + ?Sized,
+    {
         for i in 0..N {
             assert_ulps_eq!(m0[i], m1[i], epsilon = EPS);
         }
@@ -1189,7 +1193,7 @@ mod tests {
         }
 
         assert_matrixes_ulps_eq::<2, 2, _, _>(&u, &faeru);
-        assert_diags_ulps_eq(&s, &faers);
+        assert_diags_ulps_eq::<2, _, _>(&s, &faers);
         // Note that faer returns V, not Vt
         assert_matrixes_ulps_eq::<2, 2, _, _>(&vt, &faerv.transpose());
     }
@@ -1223,7 +1227,7 @@ mod tests {
         let (nau, nas, navt) = (nasvd.u.unwrap(), nasvd.singular_values, nasvd.v_t.unwrap());
 
         assert_matrixes_ulps_eq::<2, 2, _, _>(&u, &nau);
-        assert_diags_ulps_eq::<2>(&s, &nas);
+        assert_diags_ulps_eq::<2, _, _>(&s, &nas);
         assert_matrixes_ulps_eq::<2, 2, _, _>(&vt, &navt);
     }
 
@@ -1256,7 +1260,7 @@ mod tests {
 
         let faers = faersvd.S();
         // Our implementation allows negative singular value
-        assert_diags_ulps_eq(
+        assert_diags_ulps_eq::<3, _, _>(
             &DiagonalMatrix {
                 elements: s.elements.map(f64::abs),
             },
@@ -1332,7 +1336,7 @@ mod tests {
         let expected_diag = DiagonalMatrix {
             elements: [1.0, 5.0, 9.0],
         };
-        assert_diags_ulps_eq(&diag, &expected_diag);
+        assert_diags_ulps_eq::<3, _, _>(&diag, &expected_diag);
 
         let from_diag = diag.to_dense();
         let expected_from_diag = Matrix {

hoomd-linear-algebra/src/matrix/gram_schmidt.rs

@@ -0,0 +1,129 @@
+// Copyright (c) 2024-2026 The Regents of the University of Michigan.
+// Part of hoomd-rs, released under the BSD 3-Clause License.
+
+use crate::matrix::Matrix;
+
+/// Tolerance for norms below which a vector is considered to be the zero vector in Gram-Schmidt.
+pub const GRAM_SCHMIDT_EPSILON: f64 = 1e-12;
+
+/// Construct an orthonormal basis from the vectors in a [`Matrix`] using the modified Gram-Schmidt procedure.
+///
+/// Implementation based on <https://www.sfu.ca/~jtmulhol/py4math/linalg/np-gramschmidt/>
+#[must_use]
+#[inline]
+pub fn gram_schmidt<const N: usize, const M: usize>(a: &Matrix<N, M>) -> Matrix<N, M> {
+    let mut a = a.clone();
+    for j in 0..M {
+        // For the vector in column k, find the perpendicular of the projection onto
+        // the previous orthogonal vectors.
+        for k in 0..j {
+            let mut j_dot_k = 0.0;
+            for i in 0..N {
+                j_dot_k += a[(i, k)] * a[(i, j)];
+            }
+            // Apply the projection
+            for i in 0..N {
+                a[(i, j)] -= a[(i, k)] * j_dot_k;
+            }
+        } // end loop over k
+
+        let mut column_j_norm_sq = 0.0;
+        for i in 0..N {
+            column_j_norm_sq += a[(i, j)] * a[(i, j)];
+        }
+        let column_j_norm = column_j_norm_sq.sqrt();
+
+        // If the initial vectors are not linearly independent, zero out the column
+        if column_j_norm > GRAM_SCHMIDT_EPSILON {
+            for i in 0..N {
+                a[(i, j)] /= column_j_norm;
+            }
+        } else {
+            for i in 0..N {
+                a[(i, j)] = 0.0;
+            }
+        }
+    } // end loop over j
+    a
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::matrix::{Matrix, test_utils::assert_matrixes_ulps_eq};
+
+    #[test]
+    fn test_gram_schmidt_3x3() {
+        // Example 1 from https://www.sfu.ca/~jtmulhol/py4math/linalg/np-gramschmidt/
+        let a = Matrix::<3, 3> {
+            rows: [[1.0, -1.0, 0.0], [1.0, 2.0, 1.0], [0.0, 1.0, 1.0]],
+        };
+        let q = gram_schmidt(&a);
+
+        let sqrt2 = 2.0f64.sqrt();
+        let sqrt11 = 11.0f64.sqrt();
+        let sqrt22 = 22.0f64.sqrt();
+
+        let expected = Matrix::<3, 3> {
+            rows: [
+                [1.0 / sqrt2, -3.0 / sqrt22, 1.0 / sqrt11],
+                [1.0 / sqrt2, 3.0 / sqrt22, -1.0 / sqrt11],
+                [0.0, 2.0 / sqrt22, 3.0 / sqrt11],
+            ],
+        };
+
+        assert_matrixes_ulps_eq::<3, 3, _, _>(&q, &expected);
+    }
+
+    #[test]
+    fn test_gram_schmidt_linearly_dependent() {
+        // Example 2 from https://www.sfu.ca/~jtmulhol/py4math/linalg/np-gramschmidt/
+        // A = [[1, 1, 2, 1], [1, 0, 1, 0], [0, 1, 1, 0], [1, 1, 2, 1]] (columns)
+        let a = Matrix::<4, 4> {
+            rows: [
+                [1.0, 1.0, 2.0, 1.0],
+                [1.0, 0.0, 1.0, 0.0],
+                [0.0, 1.0, 1.0, 0.0],
+                [1.0, 1.0, 2.0, 1.0],
+            ],
+        };
+        let q = gram_schmidt(&a);
+
+        let sqrt3 = 3.0f64.sqrt();
+        let sqrt15 = 15.0f64.sqrt();
+        let sqrt10 = 10.0f64.sqrt();
+
+        let expected = Matrix::<4, 4> {
+            rows: [
+                [1.0 / sqrt3, 1.0 / sqrt15, 0.0, 1.0 / sqrt10],
+                [1.0 / sqrt3, -2.0 / sqrt15, 0.0, -2.0 / sqrt10],
+                [0.0, 3.0 / sqrt15, 0.0, -2.0 / sqrt10],
+                [1.0 / sqrt3, 1.0 / sqrt15, 0.0, 1.0 / sqrt10],
+            ],
+        };
+
+        assert_matrixes_ulps_eq::<4, 4, _, _>(&q, &expected);
+    }
+
+    #[test]
+    fn test_gram_schmidt_subspace() {
+        // Example 3 from https://www.sfu.ca/~jtmulhol/py4math/linalg/np-gramschmidt/
+        let a = Matrix::<3, 2> {
+            rows: [[1.0, -1.0], [1.0, 2.0], [0.0, 1.0]],
+        };
+        let q = gram_schmidt(&a);
+
+        let sqrt2 = 2.0f64.sqrt();
+        let sqrt22 = 22.0f64.sqrt();
+
+        let expected = Matrix::<3, 2> {
+            rows: [
+                [1.0 / sqrt2, -3.0 / sqrt22],
+                [1.0 / sqrt2, 3.0 / sqrt22],
+                [0.0, 2.0 / sqrt22],
+            ],
+        };
+
+        assert_matrixes_ulps_eq::<3, 2, _, _>(&q, &expected);
+    }
+}