flow-knn 0.1.1

Algorithm-agnostic k-NN graphs for flow cytometry — exact, usearch HNSW, optional ann-search-rs
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

//! HNSW via usearch.

use crate::config::{DistanceMetric, HnswParams, Quantization};
use crate::error::KnnError;
use crate::exact::exact_knn;
use crate::graph::NeighborList;
use rayon::prelude::*;
use usearch::{Index, IndexOptions, MetricKind, ScalarKind};

pub fn hnsw_knn(
    data: &[f32],
    n: usize,
    d: usize,
    k: usize,
    params: &HnswParams,
    metric: DistanceMetric,
) -> Result<Vec<NeighborList>, KnnError> {
    if metric == DistanceMetric::Manhattan {
        return exact_knn(data, n, d, k, metric);
    }

    let metric_kind = match metric {
        DistanceMetric::Euclidean | DistanceMetric::EuclideanSq => MetricKind::L2sq,
        DistanceMetric::Cosine => MetricKind::Cos,
        DistanceMetric::Manhattan => unreachable!(),
    };

    let scalar_kind = match params.quantization {
        Quantization::F32 => ScalarKind::F32,
        Quantization::F16 => ScalarKind::F16,
        Quantization::I8 => ScalarKind::I8,
    };

    let options = IndexOptions {
        dimensions: d,
        metric: metric_kind,
        quantization: scalar_kind,
        connectivity: params.m,
        expansion_add: params.ef_construction,
        expansion_search: params.ef_search,
        ..Default::default()
    };

    let index = Index::new(&options).map_err(|e| KnnError::Index(e.to_string()))?;
    index
        .reserve(n)
        .map_err(|e| KnnError::Index(e.to_string()))?;

    (0..n).into_par_iter().try_for_each(|i| {
        let row = &data[i * d..(i + 1) * d];
        index
            .add(i as u64, row)
            .map_err(|e| KnnError::Index(e.to_string()))
    })?;

    let k_fetch = k + 1;
    let result: Vec<NeighborList> = (0..n)
        .into_par_iter()
        .map(|i| {
            let row = &data[i * d..(i + 1) * d];
            let matches = index.search(row, k_fetch).unwrap_or_else(|_| usearch::ffi::Matches {
                keys: vec![],
                distances: vec![],
            });

            let mut indices = Vec::with_capacity(k);
            let mut distances = Vec::with_capacity(k);
            for (&key, &dist) in matches.keys.iter().zip(matches.distances.iter()) {
                if key == i as u64 {
                    continue;
                }
                if indices.len() >= k {
                    break;
                }
                indices.push(key as u32);
                distances.push(dist);
            }
            NeighborList { indices, distances }
        })
        .collect();

    drop(index);
    Ok(result)
}