Skip to main content

flow_pacmap/
knn.rs

1//! Re-export and PaCMAP-facing adapters for [`flow_knn`].
2
3use crate::config::{DistanceMetric, KnnMethod};
4use crate::error::PaCMAPError;
5pub use flow_knn::{
6    KnnGraph, NeighborList, PerfRecord, RecommendOpts, builtin_matrix, load_matrix,
7    parse_matrix_jsonl, read_knn_graph, recommend_method, recommend_method_with_matrix,
8    write_knn_graph,
9};
10
11/// Compute k nearest neighbours for all n points in `data` (n×d row-major).
12pub fn compute_knn(
13    data: &[f32],
14    n: usize,
15    d: usize,
16    k: usize,
17    method: &KnnMethod,
18    metric: DistanceMetric,
19) -> Result<KnnGraph, PaCMAPError> {
20    flow_knn::compute_knn(data, n, d, k, method, metric).map_err(map_knn_error)
21}
22
23pub(crate) fn map_knn_error(err: flow_knn::KnnError) -> PaCMAPError {
24    use flow_knn::KnnError;
25    match err {
26        KnnError::DatasetTooSmall { n } => PaCMAPError::DatasetTooSmall { n },
27        KnnError::DimensionMismatch { len, d } => PaCMAPError::DimensionMismatch { len, d },
28        KnnError::MethodNotImplemented { method } => PaCMAPError::MethodNotImplemented { method },
29        KnnError::Index(msg) => PaCMAPError::KnnIndex(msg),
30        KnnError::GraphSizeMismatch {
31            graph_n,
32            neighbors_len,
33            data_n,
34        } => PaCMAPError::KnnGraphSizeMismatch {
35            graph_n,
36            neighbors_len,
37            data_n,
38        },
39        KnnError::GraphInsufficientK {
40            graph_k,
41            required_k,
42        } => PaCMAPError::KnnGraphInsufficientK {
43            graph_k,
44            required_k,
45        },
46        KnnError::GraphMetricMismatch { graph, requested } => PaCMAPError::KnnGraphMetricMismatch {
47            graph,
48            config: requested,
49        },
50        KnnError::GpuUnavailable(msg) => PaCMAPError::Gpu(msg),
51        KnnError::Io(msg) => PaCMAPError::KnnIndex(format!("knn I/O: {msg}")),
52    }
53}
54
55/// Validate a graph for PaCMAP, mapping [`flow_knn::KnnError`] into [`PaCMAPError`].
56pub fn validate_knn_for_pacmap(
57    graph: &KnnGraph,
58    data_n: usize,
59    n_neighbors: usize,
60    metric: DistanceMetric,
61) -> Result<(), PaCMAPError> {
62    graph
63        .validate_for_pacmap(data_n, n_neighbors, metric)
64        .map_err(map_knn_error)
65}