flow_pacmap/
knn.rs

1//! K-nearest-neighbour search for PaCMAP graph construction.
2//!
3//! Returns, for each point i, its k nearest neighbours by the chosen metric,
4//! along with their distances — both used in the scaled-distance reranking step
5//! and for sigma computation (avg distance to 4th–6th neighbours).
6
7use crate::config::{DistanceMetric, KnnMethod};
8use crate::error::PaCMAPError;
9use rayon::prelude::*;
10
11/// KNN result for a single query point.
12pub struct NeighborList {
13    /// Indices of k nearest neighbours (excluding self), ascending distance order.
14    pub indices: Vec<u32>,
15    /// Distances corresponding to each index.
16    pub distances: Vec<f32>,
17}
18
19/// Compute k nearest neighbours for all n points in `data` (n×d row-major).
20///
21/// Returns a `Vec<NeighborList>` of length n.
22/// The HNSW index is built, queried in one parallel pass, then dropped —
23/// it never coexists in memory with the pair matrices.
24pub fn compute_knn(
25    data: &[f32],
26    n: usize,
27    d: usize,
28    k: usize,
29    method: &KnnMethod,
30    metric: DistanceMetric,
31) -> Result<Vec<NeighborList>, PaCMAPError> {
32    match method {
33        #[cfg(feature = "hnsw")]
34        KnnMethod::Hnsw(params) => hnsw_knn(data, n, d, k, params, metric),
35        KnnMethod::Exact => exact_knn(data, n, d, k, metric),
36        #[cfg(feature = "kdtree")]
37        KnnMethod::KdTree => kdtree_knn(data, n, d, k, metric),
38        KnnMethod::Annoy => Err(PaCMAPError::MethodNotImplemented {
39            method: "Annoy".to_string(),
40        }),
41        // When a feature is disabled, the variant does not exist; unreachable patterns
42        // are excluded by cfg. The catch-all handles any edge cases.
43        #[allow(unreachable_patterns)]
44        _ => Err(PaCMAPError::MethodNotImplemented {
45            method: "unknown (feature disabled)".to_string(),
46        }),
47    }
48}
49
50// ── Distance helpers ──────────────────────────────────────────────────────────
51
52#[inline(always)]
53fn l2_sq(a: &[f32], b: &[f32]) -> f32 {
54    a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum()
55}
56
57#[inline(always)]
58fn dist(a: &[f32], b: &[f32], metric: DistanceMetric) -> f32 {
59    match metric {
60        DistanceMetric::Euclidean => l2_sq(a, b).sqrt(),
61        DistanceMetric::EuclideanSq => l2_sq(a, b),
62        DistanceMetric::Cosine => {
63            let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
64            let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
65            let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
66            1.0 - dot / (na * nb + f32::EPSILON)
67        }
68        DistanceMetric::Manhattan => a.iter().zip(b).map(|(x, y)| (x - y).abs()).sum(),
69    }
70}
71
72// ── Exact brute-force ─────────────────────────────────────────────────────────
73
74pub fn exact_knn(
75    data: &[f32],
76    n: usize,
77    d: usize,
78    k: usize,
79    metric: DistanceMetric,
80) -> Result<Vec<NeighborList>, PaCMAPError> {
81    let k_capped = k.min(n - 1);
82    let result: Vec<NeighborList> = (0..n)
83        .into_par_iter()
84        .map(|i| {
85            let row_i = &data[i * d..(i + 1) * d];
86            // Bounded max-heap: keep the k closest seen so far.
87            // Using a Vec sorted by descending distance (max at front).
88            let mut heap: Vec<(f32, u32)> = Vec::with_capacity(k_capped + 1);
89
90            for j in 0..n {
91                if j == i {
92                    continue;
93                }
94                let row_j = &data[j * d..(j + 1) * d];
95                let d_ij = dist(row_i, row_j, metric);
96
97                if heap.len() < k_capped {
98                    heap.push((d_ij, j as u32));
99                    // Keep sorted descending by distance (max at index 0)
100                    heap.sort_unstable_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
101                } else if d_ij < heap[0].0 {
102                    heap[0] = (d_ij, j as u32);
103                    heap.sort_unstable_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
104                }
105            }
106
107            // Sort ascending by distance for output
108            heap.sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
109            NeighborList {
110                indices: heap.iter().map(|(_, idx)| *idx).collect(),
111                distances: heap.iter().map(|(d, _)| *d).collect(),
112            }
113        })
114        .collect();
115
116    Ok(result)
117}
118
119// ── HNSW via usearch ──────────────────────────────────────────────────────────
120
121#[cfg(feature = "hnsw")]
122fn hnsw_knn(
123    data: &[f32],
124    n: usize,
125    d: usize,
126    k: usize,
127    params: &crate::config::HnswParams,
128    metric: DistanceMetric,
129) -> Result<Vec<NeighborList>, PaCMAPError> {
130    use usearch::{Index, IndexOptions, MetricKind, ScalarKind};
131
132    // usearch does not support L1 natively; fall back to brute-force for Manhattan
133    if metric == DistanceMetric::Manhattan {
134        return exact_knn(data, n, d, k, metric);
135    }
136
137    let metric_kind = match metric {
138        DistanceMetric::Euclidean | DistanceMetric::EuclideanSq => MetricKind::L2sq,
139        DistanceMetric::Cosine => MetricKind::Cos,
140        DistanceMetric::Manhattan => unreachable!(),
141    };
142
143    let scalar_kind = match params.quantization {
144        crate::config::Quantization::F32 => ScalarKind::F32,
145        crate::config::Quantization::F16 => ScalarKind::F16,
146        crate::config::Quantization::I8 => ScalarKind::I8,
147    };
148
149    let options = IndexOptions {
150        dimensions: d,
151        metric: metric_kind,
152        quantization: scalar_kind,
153        connectivity: params.m,
154        expansion_add: params.ef_construction,
155        expansion_search: params.ef_search,
156        ..Default::default()
157    };
158
159    let index = Index::new(&options)
160        .map_err(|e| PaCMAPError::KnnIndex(e.to_string()))?;
161    index.reserve(n)
162        .map_err(|e| PaCMAPError::KnnIndex(e.to_string()))?;
163
164    // Parallel add — usearch Index is Send + Sync
165    (0..n).into_par_iter().try_for_each(|i| {
166        let row = &data[i * d..(i + 1) * d];
167        index.add(i as u64, row)
168            .map_err(|e| PaCMAPError::KnnIndex(e.to_string()))
169    })?;
170
171    // Parallel query all n points; exclude self-match
172    let k_fetch = k + 1; // fetch one extra to exclude self
173    let result: Vec<NeighborList> = (0..n)
174        .into_par_iter()
175        .map(|i| {
176            let row = &data[i * d..(i + 1) * d];
177            let matches = index.search(row, k_fetch).unwrap_or_else(|_| {
178                usearch::ffi::Matches {
179                    keys: vec![],
180                    distances: vec![],
181                }
182            });
183
184            let mut indices = Vec::with_capacity(k);
185            let mut distances = Vec::with_capacity(k);
186            for (&key, &dist) in matches.keys.iter().zip(matches.distances.iter()) {
187                if key == i as u64 {
188                    continue;
189                }
190                if indices.len() >= k {
191                    break;
192                }
193                indices.push(key as u32);
194                distances.push(dist);
195            }
196            NeighborList { indices, distances }
197        })
198        .collect();
199
200    // Index dropped here — HNSW memory released before pair allocation
201    drop(index);
202
203    Ok(result)
204}
205
206// ── k-d tree via kiddo ────────────────────────────────────────────────────────
207
208#[cfg(feature = "kdtree")]
209fn kdtree_knn(
210    data: &[f32],
211    n: usize,
212    d: usize,
213    k: usize,
214    _metric: DistanceMetric,
215) -> Result<Vec<NeighborList>, PaCMAPError> {
216    // kiddo's API is generic over dimension; we use a dynamic approach via ImmutableKdTree
217    // which supports arbitrary compile-time dimensions via const generics. Since d varies,
218    // we fall back to the exact brute-force for now and note that a macro-dispatch over
219    // common flow cytometry dimensions (5, 10, 20, 30, 40, 50) can be added later.
220    exact_knn(data, n, d, k, _metric)
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    fn make_grid(n: usize) -> Vec<f32> {
228        (0..n).flat_map(|i| [i as f32, 0.0]).collect()
229    }
230
231    #[test]
232    fn exact_knn_nearest_neighbour() {
233        let data = make_grid(10);
234        let knn = exact_knn(&data, 10, 2, 2, DistanceMetric::Euclidean).unwrap();
235        // Point 5 (at x=5) should have point 4 and 6 as nearest neighbours
236        let nbrs = &knn[5].indices;
237        assert!(nbrs.contains(&4u32) || nbrs.contains(&6u32));
238    }
239}
flow_pacmap/knn.rs

flow_pacmap/
knn.rs
