flow_knn/
lib.rs

1//! # flow-knn
2//!
3//! Algorithm-agnostic k-nearest-neighbour graphs for large-n scientific data.
4//!
5//! Build once with [`compute_knn`], reuse the [`KnnGraph`] across embedders
6//! (PaCMAP, UMAP, …) without recomputing neighbours.
7
8mod config;
9mod error;
10mod exact;
11mod graph;
12mod io;
13mod select;
14
15#[cfg(feature = "gpu")]
16mod gpu_ann;
17#[cfg(feature = "ann-search")]
18mod hnsw_ann;
19#[cfg(feature = "hnsw")]
20mod hnsw_usearch;
21
22pub use config::{DistanceMetric, HnswParams, KnnMethod, Quantization};
23#[cfg(feature = "gpu")]
24pub use config::{IvfGpuParams, NnDescentGpuParams};
25pub use error::KnnError;
26pub use exact::exact_knn;
27pub use graph::{KnnGraph, NeighborList};
28pub use io::{read_knn_graph, write_knn_graph};
29#[cfg(feature = "gpu")]
30pub use gpu_ann::gpu_adapter_available;
31pub use select::{
32    PerfRecord, RecommendOpts, builtin_matrix, load_matrix, parse_matrix_jsonl, recommend_method,
33    recommend_method_with_matrix,
34};
35
36use config::KnnMethod as Method;
37
38fn method_provenance(method: &Method) -> String {
39    match method {
40        #[cfg(feature = "hnsw")]
41        Method::Hnsw(_) => "Hnsw".to_string(),
42        Method::Exact => "Exact".to_string(),
43        #[cfg(feature = "kdtree")]
44        Method::KdTree => "KdTree".to_string(),
45        #[cfg(feature = "ann-search")]
46        Method::AnnSearchHnsw(_) => "AnnSearchHnsw".to_string(),
47        #[cfg(feature = "gpu")]
48        Method::GpuExact => "GpuExact".to_string(),
49        #[cfg(feature = "gpu")]
50        Method::GpuIvf(_) => "GpuIvf".to_string(),
51        #[cfg(feature = "gpu")]
52        Method::GpuNnDescent(_) => "GpuNnDescent".to_string(),
53        Method::Annoy => "Annoy".to_string(),
54    }
55}
56
57/// Compute k nearest neighbours for all `n` points in row-major `data`.
58pub fn compute_knn(
59    data: &[f32],
60    n: usize,
61    d: usize,
62    k: usize,
63    method: &KnnMethod,
64    metric: DistanceMetric,
65) -> Result<KnnGraph, KnnError> {
66    if n < 2 {
67        return Err(KnnError::DatasetTooSmall { n });
68    }
69    if data.len() != n * d {
70        return Err(KnnError::DimensionMismatch {
71            len: data.len(),
72            d,
73        });
74    }
75    let k_capped = k.min(n - 1);
76    let neighbors = match method {
77        #[cfg(feature = "hnsw")]
78        KnnMethod::Hnsw(params) => hnsw_usearch::hnsw_knn(data, n, d, k_capped, params, metric)?,
79        KnnMethod::Exact => exact_knn(data, n, d, k_capped, metric)?,
80        #[cfg(feature = "kdtree")]
81        KnnMethod::KdTree => exact_knn(data, n, d, k_capped, metric)?,
82        #[cfg(feature = "ann-search")]
83        KnnMethod::AnnSearchHnsw(params) => {
84            hnsw_ann::ann_search_hnsw_knn(data, n, d, k_capped, params, metric)?
85        }
86        #[cfg(feature = "gpu")]
87        KnnMethod::GpuExact => gpu_ann::exact_gpu_knn(data, n, d, k_capped, metric)?,
88        #[cfg(feature = "gpu")]
89        KnnMethod::GpuIvf(params) => gpu_ann::ivf_gpu_knn(data, n, d, k_capped, params, metric)?,
90        #[cfg(feature = "gpu")]
91        KnnMethod::GpuNnDescent(params) => {
92            gpu_ann::nndescent_gpu_knn(data, n, d, k_capped, params, metric)?
93        }
94        KnnMethod::Annoy => {
95            return Err(KnnError::MethodNotImplemented {
96                method: "Annoy".to_string(),
97            });
98        }
99        #[allow(unreachable_patterns)]
100        _ => {
101            return Err(KnnError::MethodNotImplemented {
102                method: "unknown (feature disabled)".to_string(),
103            });
104        }
105    };
106
107    Ok(KnnGraph {
108        neighbors,
109        n,
110        k: k_capped,
111        metric,
112        provenance: Some(method_provenance(method)),
113    })
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    fn make_grid(n: usize) -> Vec<f32> {
121        (0..n).flat_map(|i| [i as f32, 0.0]).collect()
122    }
123
124    #[test]
125    fn exact_knn_nearest_neighbour() {
126        let data = make_grid(10);
127        let knn = exact_knn(&data, 10, 2, 2, DistanceMetric::Euclidean).unwrap();
128        let nbrs = &knn[5].indices;
129        assert!(nbrs.contains(&4u32) || nbrs.contains(&6u32));
130    }
131
132    #[test]
133    fn compute_knn_returns_graph_metadata() {
134        let data = make_grid(10);
135        let graph = compute_knn(
136            &data,
137            10,
138            2,
139            3,
140            &KnnMethod::Exact,
141            DistanceMetric::Euclidean,
142        )
143        .unwrap();
144        assert_eq!(graph.n, 10);
145        assert_eq!(graph.k, 3);
146        assert_eq!(graph.metric, DistanceMetric::Euclidean);
147        assert_eq!(graph.provenance.as_deref(), Some("Exact"));
148    }
149
150    #[test]
151    fn validate_rejects_mismatches() {
152        let data = make_grid(80);
153        let graph = compute_knn(
154            &data,
155            80,
156            2,
157            3,
158            &KnnMethod::Exact,
159            DistanceMetric::Euclidean,
160        )
161        .unwrap();
162        assert!(matches!(
163            graph.validate_for_pacmap(20, 5, DistanceMetric::Euclidean),
164            Err(KnnError::GraphSizeMismatch { .. })
165        ));
166        assert!(matches!(
167            graph.validate_for_pacmap(80, 5, DistanceMetric::Euclidean),
168            Err(KnnError::GraphInsufficientK { .. })
169        ));
170    }
171
172    #[cfg(feature = "gpu")]
173    #[test]
174    fn gpu_exact_runs_when_adapter_available() {
175        if !gpu_adapter_available() {
176            eprintln!("skip gpu_exact: no WGPU adapter");
177            return;
178        }
179        let data = make_grid(64);
180        let graph = compute_knn(
181            &data,
182            64,
183            2,
184            5,
185            &KnnMethod::GpuExact,
186            DistanceMetric::Euclidean,
187        )
188        .expect("gpu exact");
189        assert_eq!(graph.n, 64);
190        assert_eq!(graph.k, 5);
191        assert_eq!(graph.provenance.as_deref(), Some("GpuExact"));
192        assert_eq!(graph.neighbors[0].indices.len(), 5);
193    }
194}
flow_knn/lib.rs

flow_knn/
lib.rs
