flow_pacmap/
pca.rs

1//! PCA initialisation via faer SVD on the d×d covariance matrix.
2//!
3//! This approach costs O(n·d²) rather than O(n²) — at n=10M and d=46,
4//! the covariance matrix is 46×46 = 2116 f32 elements, trivially small.
5
6use crate::error::PaCMAPError;
7use faer::{Mat, linalg::solvers::Svd};
8
9/// Project `n × d` row-major data onto its top 2 principal components.
10///
11/// Returns `n` (x, y) pairs where x and y are the PC1 and PC2 scores,
12/// normalised to the range used by the paper's initialisation.
13pub fn pca_init(data: &[f32], n: usize, d: usize) -> Result<Vec<[f32; 2]>, PaCMAPError> {
14    debug_assert_eq!(data.len(), n * d);
15
16    // Step 1: compute column means
17    let mut mean = vec![0.0_f32; d];
18    for row in data.chunks_exact(d) {
19        for (j, &v) in row.iter().enumerate() {
20            mean[j] += v;
21        }
22    }
23    for m in &mut mean {
24        *m /= n as f32;
25    }
26
27    // Step 2: build d×d covariance C = (X - mean)^T (X - mean) / n
28    // C is symmetric; only compute upper triangle then mirror
29    let mut cov = Mat::<f32>::zeros(d, d);
30    for row in data.chunks_exact(d) {
31        for i in 0..d {
32            let xi = row[i] - mean[i];
33            for j in i..d {
34                let xj = row[j] - mean[j];
35                cov[(i, j)] += xi * xj;
36            }
37        }
38    }
39    let inv_n = 1.0 / n as f32;
40    for i in 0..d {
41        for j in i..d {
42            cov[(i, j)] *= inv_n;
43            if i != j {
44                cov[(j, i)] = cov[(i, j)];
45            }
46        }
47    }
48
49    // Step 3: thin SVD of symmetric covariance C = U S V^T
50    // For a symmetric PSD matrix U ≈ V; top-2 columns of U are the principal components.
51    let svd = Svd::<f32>::new(cov.as_ref())
52        .map_err(|e| PaCMAPError::Pca(format!("{e:?}")))?;
53
54    let u = svd.U();
55    // Top-2 eigenvectors: columns 0 and 1 of U (sorted by descending singular value)
56    let pc1: Vec<f32> = (0..d).map(|r| *u.get(r, 0)).collect();
57    let pc2: Vec<f32> = if d >= 2 { (0..d).map(|r| *u.get(r, 1)).collect() } else { vec![0.0; d] };
58
59    // Step 4: project each row onto PC1, PC2
60    let mut embedding = Vec::with_capacity(n);
61    for row in data.chunks_exact(d) {
62        let mut s1 = 0.0_f32;
63        let mut s2 = 0.0_f32;
64        for j in 0..d {
65            let v = row[j] - mean[j];
66            s1 += v * pc1[j];
67            s2 += v * pc2[j];
68        }
69        embedding.push([s1, s2]);
70    }
71
72    Ok(embedding)
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn pca_separates_axis_aligned_clusters() {
81        // Two clusters separated along dim 0; PCA should put that on PC1
82        let mut data: Vec<f32> = Vec::new();
83        for _ in 0..50 {
84            data.extend_from_slice(&[0.0_f32, 0.0]);
85        }
86        for _ in 0..50 {
87            data.extend_from_slice(&[10.0_f32, 0.0]);
88        }
89        let emb = pca_init(&data, 100, 2).unwrap();
90        // First 50 points should have similar x; second 50 should differ
91        let mean_left_x = emb[..50].iter().map(|p| p[0]).sum::<f32>() / 50.0;
92        let mean_right_x = emb[50..].iter().map(|p| p[0]).sum::<f32>() / 50.0;
93        assert!(
94            (mean_left_x - mean_right_x).abs() > 1.0,
95            "PCA should separate the two clusters along PC1"
96        );
97    }
98}
flow_pacmap/pca.rs

flow_pacmap/
pca.rs
