feat(index): train centroids inside IVF and HNSW indexes

docs/development.md

@@ -349,7 +349,10 @@ mutating state: non-finite values and magnitudes above the clustering bound rais
 `std::invalid_argument`. The core `quantize_*` encoders and `Lut` rely on that (the
 LUT additionally rejects non-finite rotated queries) because quantization converts
 floats to integers, which is undefined for NaN. Keep the check at the boundary rather
-than inside per-vector kernels.
+than inside per-vector kernels. IVF and HNSW `train` are the exception: they pass
+training data only to k-means, which applies the same rule and bound in parallel
+before using it, so they skip the serial pass; `set_centroids` still checks the
+trained centroids.
 
 **Python conversion**
 

docs/docs/index/hnsw.md

@@ -106,6 +106,41 @@ void HierarchicalNSW::construct(
 - **num_threads**: Number of threads to use (default: every available thread; `0` also auto-selects).
 - **fast_quantization**: Use calibrated scaling instead of per-vector scale search; see [encoding paths](../rabitq/quantizer.md#implementations).
 
+### Training centroids in C++
+
+Instead of passing centroids to `construct`, train or set them on an empty
+index, then add the points:
+
+```cpp
+#include <rabitqlib/index/hnsw/hnsw.hpp>
+#include <rabitqlib/index/training.hpp>
+
+rabitqlib::hnsw::HierarchicalNSW index(max_elements, dim, 4, 16, 200);
+index.train(16, num_points, data.data(),
+            rabitqlib::default_training_parameters(rabitqlib::METRIC_L2, 100));
+// Or supply centroids: index.set_centroids(num_clusters, centroids.data());
+index.add(data.data(), num_points);
+```
+
+- `train(num_clusters, num_points, data, cp)` trains centroids with
+  [RaBitQKMeans](../clustering.md#rabitqkmeans); `num_points` must be at least
+  `num_clusters`. Points are assigned by `add`, not by training. See the
+  [IVF-RaBitQ defaults](ivf.md#training-centroids-in-c) for
+  `default_training_parameters`.
+- The first `add` into an empty graph links points with their exact vectors, as
+  `construct` does; with one thread it produces the same index as `construct`
+  with the same centroids and cluster IDs. If it fails, the graph is left empty
+  and the centroids are kept. Later adds link through the stored codes; see
+  [how added points are linked](#how-added-points-are-linked).
+- `train` and `set_centroids` replace the centroids only while the index stores
+  no points, removed ones included; otherwise they throw `std::logic_error`.
+  A failed call leaves the index unchanged. `construct` needs an untrained index,
+  so after either call it throws `std::logic_error`; use `add` instead.
+- `is_trained()` reports whether centroids are set. Before that, `add` throws
+  `std::logic_error`; searching an empty index, trained or not, returns no results.
+- `save` throws `std::logic_error` for an index without stored points, trained
+  or not. Earlier releases saved an untrained empty index; such files still load.
+
 ## Querying
 
 ```cpp
@@ -131,7 +166,8 @@ Python follows the shared [concurrency contract](usage.md#threading-and-file-pat
 ## Updating an index
 
 A constructed or loaded index supports `add()`, `resize()`,
-and `remove()` without the original data. Adding points and resizing preserve
+and `remove()` without the original data; C++ also adds to a
+[trained](#training-centroids-in-c) index. Adding points and resizing preserve
 the index file format; files with removals have
 [compatibility restrictions](#how-removal-is-stored).
 

docs/docs/index/ivf.md

@@ -199,6 +199,47 @@ For example:
 ivf.construct(data.data(), centroids.data(), cluster_ids.data(), true);
 ```
 
+### Training centroids in C++
+
+Instead of passing centroids to `construct`, create an index without a point
+count, train or set its centroids, then add the points:
+
+```c++
+#include <rabitqlib/index/ivf/ivf.hpp>
+#include <rabitqlib/index/training.hpp>
+
+rabitqlib::ivf::IVF index(dim, num_clusters, 4, rabitqlib::METRIC_L2, /*seed=*/42);
+index.train(num_points, data.data(),
+            rabitqlib::default_training_parameters(rabitqlib::METRIC_L2, 42));
+// Or supply centroids: index.set_centroids(centroids.data());
+index.add(data.data(), num_points);
+```
+
+- The constructor takes `dim`, `num_clusters`, `total_bits`, the metric and the
+  rotation seed, followed by the optional rotator and initializer types. The
+  index starts with no point storage. `construct` rebuilds an index from
+  `max_elements()` rows, so it fails on this index until `add` has stored points.
+- An index created with a point count is built by `construct`: `train` and
+  `set_centroids` throw `std::logic_error` for it.
+- `train` uses [RaBitQKMeans](../clustering.md#rabitqkmeans) for up to
+  `kMaxRaBitQKMeansClusters` (2,000) clusters and
+  [QGKMeans](../clustering.md#qgkmeans) above that, copying the fields the two
+  parameter types share and keeping QGKMeans' default graph settings. The
+  threshold is provisional and has not been measured. Only centroids are
+  trained; `add` assigns the points. `cp.num_threads` also limits building the
+  centroid routing.
+- `default_training_parameters(metric, seed)` runs 25 iterations on at most
+  256 sampled points per centroid and selects spherical k-means for inner
+  product. The k-means classes keep their own defaults, which train on every point.
+- `train` and `set_centroids` replace the centroids only while the index stores
+  no points, removed ones included; otherwise they throw `std::logic_error`.
+  A failed call leaves the index unchanged.
+- `is_trained()` reports whether centroids are set. Before that, `add` and
+  `search` throw `std::logic_error`. A trained index without points rejects
+  every `k` in `search` with `std::invalid_argument`, because `k` may not exceed
+  the stored point count; `save` throws `std::logic_error` until the index stores
+  at least one point.
+
 ## Querying
 
 To search a saved index, load it first:
@@ -271,7 +312,8 @@ Concurrent Python searches may use independent `nprobe`, `high_accuracy`, and
 ## Updating an index
 
 C++ `IVF` and Python `IvfIndex` support `add()` and `remove()` on a constructed
-or loaded index without the original dataset. Updates preserve the file format.
+or loaded index without the original dataset; C++ also adds to a
+[trained](#training-centroids-in-c) index. Updates preserve the file format.
 
 ```c++
 void IVF::add(

include/rabitqlib/clustering/detail/kmeans.hpp

@@ -234,9 +234,26 @@ inline void exact_assign(
     return indices;
 }
 
+// Copy the fields LloydKMeans reads, which every k-means parameter type shares,
+// leaving the method-specific ones of `to` unchanged. Keep this list and the
+// using-declarations below the same.
+template <typename From, typename To>
+void copy_shared_parameters(const From& from, To& to) {
+    to.early_stop_threshold = from.early_stop_threshold;
+    to.final_assignment = from.final_assignment;
+    to.max_points_per_centroid = from.max_points_per_centroid;
+    to.min_points_per_centroid = from.min_points_per_centroid;
+    to.niter = from.niter;
+    to.num_threads = from.num_threads;
+    to.seed = from.seed;
+    to.spherical = from.spherical;
+    to.verbose = from.verbose;
+}
+
 template <typename Parameters, typename Assignment>
 class LloydKMeans : public Parameters {
    public:
+    // Keep in sync with copy_shared_parameters.
     using Parameters::early_stop_threshold;
     using Parameters::final_assignment;
     using Parameters::max_points_per_centroid;
@@ -251,7 +268,19 @@ class LloydKMeans : public Parameters {
         validate_parameters();
     }
 
-    void train(size_t n, const float* x) {
+    void train(size_t n, const float* x) { fit(n, x, true); }
+
+    /**
+     * Trains the same centroids as train() but skips its final assignment of all n
+     * vectors, which can cost more than every sampled iteration together. Leaves
+     * assignments and distances empty and final_obj zero; centroids,
+     * iteration_stats and assign() match a train() call.
+     */
+    void train_centroids(size_t n, const float* x) { fit(n, x, false); }
+
+   private:
+    // assign_all selects the final assignment that train() returns.
+    void fit(size_t n, const float* x, bool assign_all) {
         validate_parameters();
         if (x == nullptr) {
             throw std::invalid_argument(
@@ -372,6 +401,12 @@ class LloydKMeans : public Parameters {
         }
 
         assigner->store_model_state(assignment_state_);
+        if (!assign_all) {
+            std::vector<PID>().swap(assignments);
+            std::vector<float>().swap(distances);
+            final_obj = 0.0;
+            return;
+        }
         // Return assignments for the final centroids, not the centroids from
         // the beginning of the last update step.
         if (sampled) {
@@ -400,11 +435,12 @@ class LloydKMeans : public Parameters {
         final_obj = std::accumulate(distances.begin(), distances.end(), 0.0);
     }
 
+   public:
     /**
      * Assigns n row-major vectors to the current centroids with final_assignment.
      * Each vector's result is independent of the others in the call and of
-     * previous labels. Requires a successful train(); the outputs hold n elements
-     * each and must not overlap x or centroids.
+     * previous labels. Requires a successful train() or train_centroids(); the
+     * outputs hold n elements each and must not overlap x or centroids.
      */
     void assign(
         size_t n, const float* x, PID* point_assignments, float* point_distances

include/rabitqlib/index/hnsw/hnsw.hpp

@@ -30,6 +30,9 @@ class VisitedSetPool;
 namespace detail {
 struct DistanceEstimate;
 }
+namespace rabitqkmeans {
+struct RaBitQKMeansParameters;
+}
 }  // namespace rabitqlib
 
 namespace rabitqlib::hnsw {
@@ -73,6 +76,9 @@ class HierarchicalNSW {
         return num_stored_points_.load() - num_removed_;
     }
     [[nodiscard]] size_t num_stored_points() const { return num_stored_points_; }
+    // Centroids are set, by train, set_centroids or construct, or the index was loaded
+    // with clusters.
+    [[nodiscard]] bool is_trained() const { return num_clusters_ != 0; }
     // Cluster of a stored record by internal ID (0 through num_stored_points() - 1;
     // these coincide with labels only for points added in order). Removed points
     // keep their cluster.
@@ -88,9 +94,39 @@ class HierarchicalNSW {
         return stored_cluster_id(internal_id);
     }
 
+    // Requires at least one stored point; a trained index must add points first.
     void save(const char* filename) const;
     void load(const char* filename);
 
+    /**
+     * @brief Replace the centroids of an initialized index without stored points.
+     *
+     * The index is unchanged if this throws. Later adds route and quantize against
+     * these centroids; the first add into the empty graph links like construct.
+     *
+     * @param num_clusters Number of centroids, from 1 through 2^31
+     * @param centroids Centroid vectors (num_clusters * dim) in the original coordinates
+     */
+    void set_centroids(size_t num_clusters, const float* centroids);
+
+    /**
+     * @brief Train centroids with RaBitQKMeans and set them as set_centroids does.
+     *
+     * Only centroids are trained: points are not assigned until add. The index is
+     * unchanged if this throws.
+     *
+     * @param num_clusters Number of centroids, at most num_points
+     * @param num_points Number of training points
+     * @param data Training points (num_points * dim)
+     * @param cp Clustering parameters, for example default_training_parameters()
+     */
+    void train(
+        size_t num_clusters,
+        size_t num_points,
+        const float* data,
+        const rabitqkmeans::RaBitQKMeansParameters& cp
+    );
+
     // Centroids and data use the original dimension; cluster_ids has num_points entries.
     // Input buffers must remain valid until this synchronous call returns or throws;
     // the index retains no pointers to them afterward.
@@ -107,10 +143,12 @@ class HierarchicalNSW {
     );
 
     /**
-     * @brief Insert points into a constructed or loaded index.
+     * @brief Insert points into a trained, constructed or loaded index.
      *
-     * Linked with the stored quantized codes, not the exact vectors construct uses,
-     * which are not kept. Centroids and rotation are never retrained. Point `i` gets
+     * Into an empty graph, points are linked with their exact vectors as construct
+     * links them, and a failure leaves the graph empty. Otherwise they are linked with
+     * the stored quantized codes, since the exact vectors are not kept. Centroids and
+     * rotation are never retrained. Point `i` gets
      * the label `num_stored_points() + i` whatever `num_threads` is. Throws rather than
      * growing past `max_elements()`. Not safe to call while another thread searches.
      * Preparation failures leave the graph and point count unchanged. A failure
@@ -168,7 +206,8 @@ class HierarchicalNSW {
 
     // Queries use the original dimension. Each result is (distance, external label),
     // where distance estimates squared L2 or 1 - dot according to the index metric.
-    // ef_search is raised to k and clamped to the stored point count.
+    // ef_search is raised to k and clamped to the stored point count. An empty index,
+    // trained or not, returns no results.
     std::vector<std::vector<std::pair<float, PID>>> search(
         const float* queries,
         size_t num_queries,
@@ -437,6 +476,16 @@ class HierarchicalNSW {
 
     void add_point(PID label, PID cluster_id, const quant::RabitqConfig& config);
 
+    // Link validated points into an empty graph that has centroids, scoring with their
+    // exact vectors. Labels are 0 through num_points - 1. A failure empties the graph.
+    void link_exact(
+        size_t num_points,
+        const float* data,
+        const PID* cluster_ids,
+        size_t num_threads,
+        bool fast_quantization
+    );
+
     // Insertion steps shared by raw and quantized scoring. distance(id) scores an
     // existing node against the node being inserted; prefetch(id) may warm its data.
     template <class Distance>