flow_pacmap/
lib.rs

1//! # flow-pacmap
2//!
3//! First-party implementation of PaCMAP (Pairwise Controlled Manifold
4//! Approximation Projection) from Wang et al. 2021 (JMLR 22, Algorithm 1).
5//!
6//! Designed for large-n flow cytometry data:
7//! - No ndarray version conflicts — pure `&[f32]` / `Vec<[f32;2]>` API
8//! - PCA via faer SVD on the d×d covariance matrix (O(n·d²), no large intermediates)
9//! - HNSW KNN via usearch (C++ FFI, hardware SIMD, optional f16 quantization)
10//! - All pair counts use `checked_mul`; no debug-mode overflow panics
11//! - Progress reporting via `mpsc::Sender<PaCMAPProgress>` (per phase + every 10 iters)
12//! - Cancellation via `Arc<AtomicBool>`
13
14pub mod adam;
15pub mod config;
16pub mod error;
17pub mod gradient;
18pub mod knn;
19pub mod pairs;
20pub mod pca;
21pub mod weights;
22
23pub use config::{
24    DistanceMetric, HnswParams, Init, KnnMethod, PaCMAPConfig, Quantization,
25};
26pub use error::PaCMAPError;
27
28use adam::{AdamState, adam_step};
29use gradient::compute_gradient;
30use knn::compute_knn;
31use pairs::build_pairs;
32use pca::pca_init;
33use weights::weights_at;
34
35use std::sync::{
36    Arc,
37    atomic::{AtomicBool, Ordering},
38    mpsc,
39};
40
41/// Progress event emitted during optimization.
42#[derive(Debug, Clone)]
43pub struct PaCMAPProgress {
44    /// Current optimization phase (1, 2, or 3).
45    pub phase: u8,
46    /// Current iteration number (1-indexed).
47    pub iter: usize,
48    /// Total iterations across all phases.
49    pub total_iters: usize,
50    /// Loss value at this iteration.
51    pub loss: f32,
52}
53
54/// Embed `n × d` row-major f32 data into 2 dimensions.
55///
56/// # Arguments
57/// - `data`: flat row-major slice, `len = n * d`
58/// - `n`: number of points; must be ≥ 2 and ≤ `u32::MAX`
59/// - `d`: number of dimensions per point; must be ≥ 1
60/// - `config`: algorithm configuration
61/// - `progress`: optional channel for per-iteration progress events
62/// - `cancel`: optional cancellation token; checked once per iteration
63///
64/// # Returns
65/// `n` `[f32; 2]` pairs aligned with the input rows, or a `PaCMAPError`.
66pub fn fit_transform(
67    data: &[f32],
68    n: usize,
69    d: usize,
70    config: PaCMAPConfig,
71    progress: Option<mpsc::Sender<PaCMAPProgress>>,
72    cancel: Option<Arc<AtomicBool>>,
73) -> Result<Vec<[f32; 2]>, PaCMAPError> {
74    // ── Input validation ──────────────────────────────────────────────────
75    if n < 2 {
76        return Err(PaCMAPError::DatasetTooSmall { n });
77    }
78    if data.len() != n * d {
79        return Err(PaCMAPError::DimensionMismatch { len: data.len(), d });
80    }
81    if n > u32::MAX as usize {
82        return Err(PaCMAPError::DatasetTooLarge { n });
83    }
84
85    let n_nb = config.n_neighbors.min(n - 1);
86    let n_mn = config.n_mn();
87    let n_fp = config.n_fp();
88    let total_iters = config.total_iters();
89
90    // ── KNN construction ──────────────────────────────────────────────────
91    // Request min(n_nb + 50, n-1) candidates; rerank by scaled distance inside build_pairs.
92    let k_candidates = (n_nb + 50).min(n - 1);
93    let knn = compute_knn(data, n, d, k_candidates, &config.knn_method, config.distance_metric)?;
94
95    // ── Pair sampling ─────────────────────────────────────────────────────
96    let pairs = build_pairs(&knn, data, n, d, n_nb, n_mn, n_fp, config.seed)?;
97    drop(knn); // free sigma + distance buffers before embedding allocation
98
99    // ── Initialisation ────────────────────────────────────────────────────
100    let mut embedding: Vec<[f32; 2]> = match &config.init {
101        Init::Pca => pca_init(data, n, d)?,
102        Init::Random(seed) => {
103            use rand::{RngExt, SeedableRng, rngs::SmallRng};
104            let mut rng = match seed {
105                Some(s) => SmallRng::seed_from_u64(*s),
106                None => rand::make_rng::<SmallRng>(),
107            };
108            let scale = (1e-4_f32).sqrt();
109            (0..n)
110                .map(|_| [rng.random::<f32>() * scale, rng.random::<f32>() * scale])
111                .collect()
112        }
113    };
114
115    // ── Optimization (Adam, 3-phase weight schedule) ──────────────────────
116    let mut adam = AdamState::new(n);
117    let mut global_iter = 0usize;
118
119    for (phase_idx, &phase_len) in config.phase_iters.iter().enumerate() {
120        let phase = (phase_idx + 1) as u8;
121
122        for local_iter in 0..phase_len {
123            // Cancellation check
124            if let Some(ref cancel) = cancel {
125                if cancel.load(Ordering::Relaxed) {
126                    return Err(PaCMAPError::Cancelled);
127                }
128            }
129
130            global_iter += 1;
131            let w = weights_at(global_iter, &config.phase_iters);
132
133            let (grad, loss) = compute_gradient(
134                &embedding,
135                &pairs.near,
136                &pairs.mid_near,
137                &pairs.further,
138                &w,
139                n,
140            );
141
142            adam_step(&mut embedding, &grad, &mut adam, global_iter, config.learning_rate);
143
144            // Emit progress every 10 iterations and at phase boundaries
145            if let Some(ref tx) = progress {
146                if local_iter == 0 || local_iter == phase_len - 1 || global_iter % 10 == 0 {
147                    let _ = tx.send(PaCMAPProgress {
148                        phase,
149                        iter: global_iter,
150                        total_iters,
151                        loss,
152                    });
153                }
154            }
155        }
156    }
157
158    Ok(embedding)
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164
165    /// Smoke test: two well-separated Gaussian clusters should remain separated
166    /// after embedding.
167    #[test]
168    fn two_cluster_separation() {
169        use rand::{RngExt, SeedableRng, rngs::SmallRng};
170        let mut rng = SmallRng::seed_from_u64(0);
171        let n_per_cluster = 100;
172        let n = n_per_cluster * 2;
173        let d = 10;
174
175        let mut data = Vec::with_capacity(n * d);
176        // Cluster A centred at [0, 0, ..., 0]
177        for _ in 0..n_per_cluster {
178            for _ in 0..d {
179                data.push(rng.random::<f32>() * 0.5);
180            }
181        }
182        // Cluster B centred at [5, 5, ..., 5]
183        for _ in 0..n_per_cluster {
184            for _ in 0..d {
185                data.push(5.0 + rng.random::<f32>() * 0.5);
186            }
187        }
188
189        let config = PaCMAPConfig {
190            n_neighbors: 5,
191            phase_iters: [50, 50, 100],
192            knn_method: KnnMethod::Exact,
193            init: Init::Random(Some(42)),
194            ..Default::default()
195        };
196
197        let emb = fit_transform(&data, n, d, config, None, None).unwrap();
198        assert_eq!(emb.len(), n);
199
200        // Centroids of the two clusters in the embedding should be far apart
201        let c_a: [f32; 2] = {
202            let sx: f32 = emb[..n_per_cluster].iter().map(|p| p[0]).sum();
203            let sy: f32 = emb[..n_per_cluster].iter().map(|p| p[1]).sum();
204            [sx / n_per_cluster as f32, sy / n_per_cluster as f32]
205        };
206        let c_b: [f32; 2] = {
207            let sx: f32 = emb[n_per_cluster..].iter().map(|p| p[0]).sum();
208            let sy: f32 = emb[n_per_cluster..].iter().map(|p| p[1]).sum();
209            [sx / n_per_cluster as f32, sy / n_per_cluster as f32]
210        };
211        let sep = ((c_a[0] - c_b[0]).powi(2) + (c_a[1] - c_b[1]).powi(2)).sqrt();
212        assert!(
213            sep > 0.5,
214            "cluster centroids should be separated in embedding (got {sep:.3})"
215        );
216    }
217}
flow_pacmap/lib.rs

flow_pacmap/
lib.rs
