feat(kmeans): add kmeans clustering to find good entry point in the graph for the searcher

Author: Neiko2002

cpp/deglib/include/deglib/search/kmeans.h

@@ -28,6 +28,7 @@
 #include <span>
 #include <stdexcept>
 #include <type_traits>
+#include <utility>
 #include <vector>
 
 namespace deglib::search {
@@ -40,22 +41,15 @@ concept DistanceCallable = requires(F f, const float* a, const float* b) {
 };
 
 // Decodes one native graph feature vector into `dim` floats.
-inline void decode_feature_vector(
-    const deglib::distances::FloatSpace& space, const std::byte* native, float* out
-) {
+inline void decode_feature_vector(const deglib::distances::FloatSpace& space, const std::byte* native, float* out) {
     const size_t dim = space.dim();
     using deglib::distances::MetricDataType;
     switch (space.metric().get_data_type()) {
         case MetricDataType::FP32:
-            std::copy(
-                reinterpret_cast<const float*>(native),
-                reinterpret_cast<const float*>(native) + dim, out
-            );
+            std::copy(reinterpret_cast<const float*>(native), reinterpret_cast<const float*>(native) + dim, out);
             break;
         case MetricDataType::FP16:
-            deglib::distances::fp16::fp16_to_floats(
-                reinterpret_cast<const uint16_t*>(native), out, dim
-            );
+            deglib::distances::fp16::fp16_to_floats(reinterpret_cast<const uint16_t*>(native), out, dim);
             break;
         case MetricDataType::Uint8: {
             const auto* v = reinterpret_cast<const uint8_t*>(native);
@@ -89,19 +83,12 @@ inline void decode_feature_vector(
     }
 }
 
-// Lloyd k-means over caller-provided rows. Centroids are normalized means
-// (cosine-style, matching the previous Python behavior); empty or degenerate
-// clusters keep their previous centroid. Returns medoid positions into `rows`.
+// Lloyd k-means over caller-provided rows. Centroids are normalized means.
+// empty or degenerate clusters keep their previous centroid.
+// Returns medoid positions (indices) into `rows`.
 template <DistanceCallable DistFn>
-[[nodiscard]] std::vector<uint32_t> kmeans_medoids(
-    std::span<const float* const> rows,
-    uint32_t dim,
-    DistFn&& distance,
-    uint32_t n_clusters,
-    uint32_t n_iter,
-    uint32_t seed,
-    size_t threads
-) {
+[[nodiscard]] std::vector<uint32_t>
+kmeans_medoids(std::span<const float* const> rows, uint32_t dim, DistFn&& distance, uint32_t n_clusters, uint32_t n_iter, uint32_t seed, size_t threads) {
     if (rows.empty() || dim == 0) {
         throw std::invalid_argument("kmeans_medoids: empty input");
     }
@@ -128,7 +115,7 @@ template <DistanceCallable DistFn>
     std::vector<float> sums(static_cast<size_t>(n_clusters) * dim);
     std::vector<uint32_t> counts(n_clusters);
     for (uint32_t iter = 0; iter < n_iter; ++iter) {
-        // Assignment: embarrassingly parallel over rows.
+        // Assignment: parallel over rows.
         deglib::concurrent::parallel_for(0, rows.size(), threads, [&](size_t i, size_t) {
             const float* vec = rows[i];
             float best_dist = std::numeric_limits<float>::max();
@@ -203,8 +190,8 @@ template <DistanceCallable DistFn>
     uint32_t n_clusters,
     uint32_t n_iter,
     size_t sample_size,
-    uint32_t seed,
-    size_t threads
+    uint32_t seed = 7,
+    size_t threads = 1
 ) {
     const size_t n_vectors = graph.size();
     if (n_vectors == 0) {
@@ -238,20 +225,14 @@ template <DistanceCallable DistFn>
     }
 
     const bool is_l2 = space.metric().get_distance_kind() == deglib::distances::MetricDistanceKind::L2;
-    const deglib::distances::FloatSpace float_space(
-        dim, is_l2 ? deglib::distances::Metric::FP32_L2 : deglib::distances::Metric::FP32_InnerProduct
-    );
+    const deglib::distances::FloatSpace float_space(dim, is_l2 ? deglib::distances::Metric::FP32_L2 : deglib::distances::Metric::FP32_InnerProduct);
     const std::span<const float* const> row_view(rows.data(), rows.size());
     const uint32_t dim32 = static_cast<uint32_t>(dim);
     const std::vector<uint32_t> positions = float_space.compute([&](const auto& kernel) {
         // Static dispatch: the concrete SIMD kernel inlines into the loop.
         return kmeans_medoids(
             row_view, dim32,
-            [&](const float* a, const float* b) -> float {
-                return std::decay_t<decltype(kernel)>::compare(
-                    a, b, float_space.get_dist_func_param()
-                );
-            },
+            [&](const float* a, const float* b) -> float { return std::decay_t<decltype(kernel)>::compare(a, b, float_space.get_dist_func_param()); },
             n_clusters, n_iter, seed, threads
         );
     });
@@ -264,4 +245,97 @@ template <DistanceCallable DistFn>
     return medoids;
 }
 
+/**
+ * Owns the k-means entry vertices for one graph.
+ *
+ * @param graph Graph the entry vertices refer to, must outlive this selector.
+ */
+class KMeansEntrySelector {
+  public:
+    explicit KMeansEntrySelector(const deglib::graph::InternalGraph& graph) : graph_(&graph) {}
+
+    /**
+     * Select entry vertices via k-means medoids.
+     *
+     * @param n_clusters  Number of entry vertices to select.
+     * @param n_iter      Number of k-means iterations.
+     * @param sample_size Number of vertices sampled for clustering, 0 selects 3% of the graph size.
+     * @param seed        Random seed for sampling and centroid init.
+     * @param threads     Number of worker threads.
+     */
+    void optimize(uint32_t n_clusters = 128, uint32_t n_iter = 15, size_t sample_size = 0, uint32_t seed = 7, size_t threads = 1) {
+        if (sample_size == 0) {
+            const size_t n = graph_->size();
+            sample_size = n * 3 / 100;
+            if (sample_size == 0 && n > 0) sample_size = n;
+        }
+        entries_ = graph_kmeans_medoids(*graph_, n_clusters, n_iter, sample_size, seed, threads);
+    }
+
+    /**
+     * Nearest entry vertices for a query, nearest first.
+     *
+     * Falls back to vertex 0 when no usable entry exists.
+     *
+     * @param query Native query bytes sized to the graph feature size.
+     * @param count Maximum number of entries to return.
+     */
+    [[nodiscard]] std::vector<uint32_t> top_entries(std::span<const std::byte> query, size_t count = 2) const {
+        const uint32_t n = graph_->size();
+        if (n == 0 || count == 0) return {};
+        const auto dist_func = graph_->getFeatureSpace().get_dist_func();
+        const auto dist_param = graph_->getFeatureSpace().get_dist_func_param();
+        std::vector<uint32_t> valid;
+        valid.reserve(entries_.size());
+        for (auto ep : entries_) {
+            if (ep < n) valid.push_back(ep);
+        }
+        if (valid.empty()) return {0};
+        if (valid.size() <= count) return valid;
+        if (count == 2) {
+            uint32_t ep1 = valid[0], ep2 = valid[0];
+            float dist1 = std::numeric_limits<float>::max(), dist2 = std::numeric_limits<float>::max();
+            for (auto ep : valid) {
+                float d = dist_func(query.data(), graph_->getFeatureVector(ep), dist_param);
+                if (d < dist1) {
+                    dist2 = dist1;
+                    ep2 = ep1;
+                    dist1 = d;
+                    ep1 = ep;
+                } else if (d < dist2) {
+                    dist2 = d;
+                    ep2 = ep;
+                }
+            }
+            if (ep1 == ep2) return {ep1};
+            return {ep1, ep2};
+        }
+        std::vector<std::pair<float, uint32_t>> scored;
+        scored.reserve(valid.size());
+        for (auto ep : valid) {
+            scored.emplace_back(dist_func(query.data(), graph_->getFeatureVector(ep), dist_param), ep);
+        }
+        std::nth_element(scored.begin(), scored.begin() + count, scored.end(), [](const auto& a, const auto& b) { return a.first < b.first; });
+        scored.resize(count);
+        std::sort(scored.begin(), scored.end(), [](const auto& a, const auto& b) { return a.first < b.first; });
+        std::vector<uint32_t> top;
+        top.reserve(count);
+        for (const auto& s : scored) top.push_back(s.second);
+        return top;
+    }
+
+    /** Stored entry vertices as graph internal indices. */
+    [[nodiscard]] const std::vector<uint32_t>& entries() const noexcept { return entries_; }
+
+    /** True when optimize() has not produced entries yet. */
+    [[nodiscard]] bool empty() const noexcept { return entries_.empty(); }
+
+    /** Number of stored entry vertices. */
+    [[nodiscard]] size_t size() const noexcept { return entries_.size(); }
+
+  private:
+    const deglib::graph::InternalGraph* graph_ = nullptr;
+    std::vector<uint32_t> entries_;
+};
+
 }  // namespace deglib::search

cpp/deglib/include/deglib/search/searcher.h

@@ -195,7 +195,16 @@ class SearcherBase {
         bool unsorted = false
     ) const = 0;
 
-    virtual void optimize(uint32_t n_clusters = 128, uint32_t n_iter = 15, size_t sample_size = 30000, uint32_t seed = 42, size_t threads = 0) = 0;
+    /**
+     * Select entry vertices via k-means medoids.
+     *
+     * @param n_clusters  Number of entry vertices to select.
+     * @param n_iter      Number of k-means iterations.
+     * @param sample_size Number of vertices sampled for clustering, 0 selects 3% of the graph size.
+     * @param seed        Random seed for sampling and centroid init.
+     * @param threads     Number of worker threads, 0 selects a library default.
+     */
+    virtual void optimize(uint32_t n_clusters = 128, uint32_t n_iter = 15, size_t sample_size = 0, uint32_t seed = 42, size_t threads = 0) = 0;
 
     // --- Modern C++20 std::span and std::vector Convenience API ---
 
@@ -357,14 +366,23 @@ class SearcherImpl : public SearcherBase {
     const deglib::graph::InternalGraph* graph_ = nullptr;
     QuantT quantizer_;
     RefinerT refiner_;
-    std::vector<uint32_t> entry_vertex_indices_;
+    KMeansEntrySelector entry_selector_;
 
   public:
     SearcherImpl(const deglib::graph::InternalGraph& graph, QuantT quantizer, RefinerT refiner)
-        : graph_(&graph), quantizer_(std::move(quantizer)), refiner_(std::move(refiner)) {}
+        : graph_(&graph), quantizer_(std::move(quantizer)), refiner_(std::move(refiner)), entry_selector_(graph) {}
 
-    void optimize(uint32_t n_clusters = 128, uint32_t n_iter = 15, size_t sample_size = 30000, uint32_t seed = 42, size_t threads = 0) override {
-        entry_vertex_indices_ = graph_kmeans_medoids(*graph_, n_clusters, n_iter, sample_size, seed, threads);
+    /**
+     * Select entry vertices via k-means medoids.
+     *
+     * @param n_clusters  Number of entry vertices to select.
+     * @param n_iter      Number of k-means iterations.
+     * @param sample_size Number of vertices sampled for clustering, 0 selects 3% of the graph size.
+     * @param seed        Random seed for sampling and centroid init.
+     * @param threads     Number of worker threads.
+     */
+    void optimize(uint32_t n_clusters = 128, uint32_t n_iter = 15, size_t sample_size = 0, uint32_t seed = 7, size_t threads = 1) override {
+        entry_selector_.optimize(n_clusters, n_iter, sample_size, seed, threads);
     }
 
     template <typename QueryT>
@@ -381,6 +399,8 @@ class SearcherImpl : public SearcherBase {
         const uint32_t fetch_k = std::max(k, static_cast<uint32_t>(std::round(k * rerank_factor)));
         const size_t graph_feature_bytes = graph_->getFeatureSpace().get_data_size();
 
+        if (graph_->size() == 0) throw std::invalid_argument("Searcher: graph is empty");
+
         // 1. Static Query Transformation (Zero-copy bypass for NoQuantizer)
         const std::byte* query_bytes = nullptr;
         alignas(64) std::byte stack_query_bytes[512];
@@ -404,15 +424,18 @@ class SearcherImpl : public SearcherBase {
             query_bytes = q_buf;
         }
 
-        // 2. Direct Graph Search (entries come solely from optimize())
+        // 2. Entry selection via selector (falls back to vertex 0) + Direct Graph Search
         const std::span<const std::byte> query_span(query_bytes, graph_feature_bytes);
-        deglib::graph::ResultSet result = entry_vertex_indices_.empty() ? graph_->search(query_span, fetch_k, eps, nullptr, 0)
-                                                                        : graph_->search(query_span, entry_vertex_indices_, fetch_k, eps, nullptr, 0);
-        const size_t found_count = result.size();
+        const std::vector<uint32_t> entries = entry_selector_.top_entries(query_span, 2);
+        deglib::graph::ResultSet result =
+            entries.empty() ? graph_->search(query_span, fetch_k, eps, nullptr, 0) : graph_->search(query_span, entries, fetch_k, eps, nullptr, 0);
 
         // 3. Static Compile-Time Reranker Check
         if constexpr (RefinerT::enabled) {
             if (fetch_k > k) {
+                const size_t found_count = result.size();
+
+                // Stack fast path for up to 256 candidates, heap fallback beyond that without result limit.
                 uint32_t stack_cand_indices[256];
                 std::unique_ptr<uint32_t[]> heap_cand_indices;
                 uint32_t* cands = stack_cand_indices;

examples/vibe/config.yml

@@ -32,4 +32,4 @@ float:
           prune_non_rng: [false]
         query_args:
           rerank_size_factor: [1.0, 1.15, 1.35]
-          ef: [60, 80, 100, 150, 200, 250, 300, 350, 400, 500, 600, 700, 800, 845, 900, 1000]
+          search_eps: [0.0, 0.005, 0.01, 0.02, 0.04, 0.06, 0.08, 0.1, 0.12, 0.15, 0.2, 0.25, 0.3]

examples/vibe/module.py

@@ -290,7 +290,7 @@ def fit(self, X: np.ndarray, cache_dir: Path | None = None):
 
         # 6. Optimize entry vertices via K-Means cluster medoids computed in C++
         t_km = time.time()
-        self.searcher.optimize()
+        #self.searcher.optimize()
         print(f"K-Means 128 cluster medoids computed in {time.time() - t_km:.2f}s", flush=True)
 
     def set_query_arguments(self, *args, **kwargs):
@@ -331,7 +331,7 @@ def query(self, v: np.ndarray, n: int) -> np.ndarray:
             threads=1,
             return_distances=False,
             unsorted=True,
-            ef=self.ef,
+            #ef=self.ef,
         )
 
     def __str__(self) -> str:

python/src/deg_cpp/deglib_cpp.cpp

@@ -814,7 +814,7 @@ class SearcherPy {
         }
     }
 
-    void optimize(uint32_t n_clusters = 128, uint32_t n_iter = 15, size_t sample_size = 30000, uint32_t seed = 42, size_t num_threads = 0) {
+    void optimize(uint32_t n_clusters = 128, uint32_t n_iter = 15, size_t sample_size = 0, uint32_t seed = 42, size_t num_threads = 0) {
         py::gil_scoped_release release;
         searcher_->optimize(n_clusters, n_iter, sample_size, seed, num_threads);
     }
@@ -1581,7 +1581,8 @@ PYBIND11_MODULE(deglib_cpp, m) {
         .def(
             "optimize", &SearcherPy::optimize,
             py::arg("n_clusters") = 128, py::arg("n_iter") = 15,
-            py::arg("sample_size") = 30000, py::arg("seed") = 42, py::arg("num_threads") = 0
+            py::arg("sample_size") = 0, py::arg("seed") = 42, py::arg("num_threads") = 0,
+            "Select entry vertices via k-means medoids.\n\nArgs:\n    n_clusters: number of entry vertices to select.\n    n_iter: number of k-means iterations.\n    sample_size: vertices sampled for clustering, 0 selects 3% of the graph size.\n    seed: random seed for sampling and centroid init.\n    num_threads: worker threads, 0 selects a library default."
         );
 
     // graphs

python/src/deglib/search.py

@@ -127,10 +127,18 @@ def optimize(
         self,
         n_clusters: int = 128,
         n_iter: int = 15,
-        sample_size: int = 30000,
+        sample_size: int = 0,
         seed: int = 42,
         num_threads: int = 0,
     ) -> None:
+        """Select entry vertices via k-means medoids.
+
+        :param n_clusters: Number of entry vertices to select.
+        :param n_iter: Number of k-means iterations.
+        :param sample_size: Number of vertices sampled for clustering, 0 selects 3% of the graph size.
+        :param seed: Random seed for sampling and centroid init.
+        :param num_threads: Number of worker threads, 0 selects a library default.
+        """
         self.searcher_cpp.optimize(int(n_clusters), int(n_iter), int(sample_size), int(seed), int(num_threads))
 
     def search(