$ show lmcinnes/pynndescent
PyNNDescent is a Python library for fast approximate nearest neighbor search, implementing the NN-descent algorithm. The core pipeline builds a k-nearest neighbor graph via random projection tree initialization followed by iterative graph refinement, with all hot loops JIT-compiled using Numba with explicit types and fastmath. The public API follows scikit-learn conventions: NNDescent constructor for training, .query() for approximate nearest neighbor queries, .neighbor_graph for extracting the training-set graph, and .update() for incremental insertion. It supports dense and sparse data, multiple distance metrics, and parallel execution via Numba's prange. The library is a dependency of UMAP and OpenTSNE, which consume the training-set neighbor graph for manifold learning and density-based clustering.
The LLM identified 3 files worth highlighting for reuse:
Contains a Numba-compatible reimplementation of scipy's rankdata, required because Numba cannot call scipy directly. The function is self-contained, tested separately, and demonstrates how to hand-roll scientific computing primitives for JIT compilation. Reusable in any Numba-accelerated Python project needing rank-based statistics without a scipy dependency.
A self-contained optimal transport solver adapted from the Python Optimal Transport (POT) library and the EGRES project, used for graph diversification in the NN-descent algorithm. It is isolated with its own license and provenance, making it extractable for any project needing a compact OT solver. Demonstrates non-obvious optimization techniques from the combinatorial optimization literature.
A helper function defined in the tutorial notebook that downloads and caches ann-benchmarks HDF5 datasets. It solves the recurring problem of loading standard benchmark datasets for nearest-neighbor evaluation, with local caching to avoid re-downloads. The pattern is directly transferable to any ANN benchmarking or evaluation pipeline.