codemem_engine/
pca.rs

1use ndarray::{Array1, Array2};
2
3/// Power iteration with deflation to extract the top-k eigenvectors
4/// from a symmetric matrix.
5pub fn power_iteration_top_k(
6    matrix: &Array2<f64>,
7    k: usize,
8    iterations: usize,
9) -> Vec<Array1<f64>> {
10    let dim = matrix.nrows();
11    let mut deflated = matrix.clone();
12    let mut eigenvectors = Vec::with_capacity(k);
13
14    for _ in 0..k {
15        let mut v = Array1::<f64>::ones(dim);
16        for i in 0..dim {
17            if i % 2 == 0 {
18                v[i] = 1.0;
19            } else {
20                v[i] = -1.0;
21            }
22        }
23
24        for _ in 0..iterations {
25            let mv = deflated.dot(&v);
26            let norm = mv.dot(&mv).sqrt();
27            if norm < 1e-10 {
28                break;
29            }
30            v = mv / norm;
31        }
32
33        let eigenvalue = v.dot(&deflated.dot(&v));
34        let outer = outer_product(&v, &v);
35        deflated = deflated - eigenvalue * &outer;
36
37        eigenvectors.push(v);
38    }
39
40    eigenvectors
41}
42
43fn outer_product(a: &Array1<f64>, b: &Array1<f64>) -> Array2<f64> {
44    let n = a.len();
45    let m = b.len();
46    let mut result = Array2::<f64>::zeros((n, m));
47    for i in 0..n {
48        for j in 0..m {
49            result[[i, j]] = a[i] * b[j];
50        }
51    }
52    result
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58    use ndarray::Array2;
59
60    #[test]
61    fn top_k_returns_k_eigenvectors() {
62        let matrix =
63            Array2::from_shape_vec((3, 3), vec![2.0, 1.0, 0.0, 1.0, 3.0, 1.0, 0.0, 1.0, 2.0])
64                .unwrap();
65
66        let eigenvectors = power_iteration_top_k(&matrix, 2, 100);
67        assert_eq!(eigenvectors.len(), 2);
68
69        for v in &eigenvectors {
70            let norm = v.dot(v).sqrt();
71            assert!(
72                (norm - 1.0).abs() < 1e-6,
73                "Eigenvector norm {norm} should be ~1.0"
74            );
75        }
76    }
77
78    #[test]
79    fn top_k_with_identity_matrix() {
80        let matrix = Array2::eye(4);
81        let eigenvectors = power_iteration_top_k(&matrix, 3, 50);
82        assert_eq!(eigenvectors.len(), 3);
83
84        let norm0 = eigenvectors[0].dot(&eigenvectors[0]).sqrt();
85        assert!(
86            (norm0 - 1.0).abs() < 1e-6,
87            "First eigenvector should have unit norm, got {norm0}"
88        );
89    }
90
91    #[test]
92    fn top_k_with_diagonal_matrix() {
93        let mut matrix = Array2::<f64>::zeros((4, 4));
94        matrix[[0, 0]] = 4.0;
95        matrix[[1, 1]] = 3.0;
96        matrix[[2, 2]] = 2.0;
97        matrix[[3, 3]] = 1.0;
98
99        let eigenvectors = power_iteration_top_k(&matrix, 2, 100);
100        assert_eq!(eigenvectors.len(), 2);
101
102        let v0 = &eigenvectors[0];
103        assert!(
104            v0[0].abs() > 0.9,
105            "First eigenvector should point along axis 0, got {:?}",
106            v0
107        );
108    }
109
110    #[test]
111    fn top_k_requests_more_than_dimensions() {
112        let matrix = Array2::eye(2);
113        let eigenvectors = power_iteration_top_k(&matrix, 5, 50);
114        assert_eq!(eigenvectors.len(), 5);
115    }
116
117    #[test]
118    fn top_k_single_dimension() {
119        let matrix = Array2::from_shape_vec((1, 1), vec![5.0]).unwrap();
120        let eigenvectors = power_iteration_top_k(&matrix, 1, 50);
121        assert_eq!(eigenvectors.len(), 1);
122        assert!((eigenvectors[0][0].abs() - 1.0).abs() < 1e-6);
123    }
124
125    #[test]
126    fn top_k_zero_matrix() {
127        let matrix = Array2::<f64>::zeros((3, 3));
128        let eigenvectors = power_iteration_top_k(&matrix, 2, 50);
129        assert_eq!(eigenvectors.len(), 2);
130    }
131
132    #[test]
133    fn eigenvectors_are_approximately_orthogonal() {
134        let matrix =
135            Array2::from_shape_vec((3, 3), vec![5.0, 1.0, 0.5, 1.0, 3.0, 0.5, 0.5, 0.5, 1.0])
136                .unwrap();
137
138        let eigenvectors = power_iteration_top_k(&matrix, 3, 200);
139        assert_eq!(eigenvectors.len(), 3);
140
141        let dot01 = eigenvectors[0].dot(&eigenvectors[1]);
142        assert!(
143            dot01.abs() < 0.1,
144            "First two eigenvectors should be approximately orthogonal, dot product = {dot01}"
145        );
146    }
147}
codemem_engine/pca.rs

codemem_engine/
pca.rs
