Add native QGKMeans and RaBitQKMeans clustering

.gitignore

@@ -48,3 +48,6 @@ __pycache__/
 
 # macOS garbages
 .DS_Store
+
+# Generated documentation
+/docs/site/

AGENTS.md

@@ -17,6 +17,7 @@ factors estimate L2 distance or inner product.
 | --- | --- |
 | `include/rabitqlib/quantization/` | Encoding, packing, reconstruction, and byte layouts |
 | `include/rabitqlib/fastscan/` | FastScan interfaces and high-accuracy scanning |
+| `include/rabitqlib/clustering/` | RaBitQKMeans flat assignment, QGKMeans graph assignment, and shared Lloyd training |
 | `include/rabitqlib/index/{ivf,hnsw,symqg}/` | Index construction, persistence, and search |
 | `include/rabitqlib/index/{query,estimator}.hpp` | Query state and distance estimation |
 | `include/rabitqlib/simd/`, `src/simd/` | Kernel declarations, implementations, and dispatch |

CONTRIBUTING.md

@@ -175,7 +175,6 @@ another. Each task can be a separate PR.
 | --- | --- | --- |
 | Explain Python search results | [Quick start](docs/docs/quick_start.md), [IVF binding](python_bindings/ivf_bindings.cpp), and [IVF tests](tests/python/test_ivf.py) | A short explanation covers result shapes, how IDs map to input rows, and what L2 distances represent, checked against the implementation. The example runs and the strict docs build passes. |
 | Add a self-contained IVF save/load example | [Python examples](sample/python/), [README quick start](README.md#python-quick-start), and [IVF tests](tests/python/test_ivf.py) | A deterministic script builds from synthetic data, saves to a temporary directory, reloads, and checks that search IDs and distances match. It needs no dataset download, cleans up its temporary files, passes Python checks, and is linked from the README. |
-| Explain the quick start's cluster assignment | [Quick start](docs/docs/quick_start.md) and [IVF guide](docs/docs/index/ivf.md) | The tutorial explains that its round-robin assignment is for a small runnable example, explains why real IVF workloads use clustering, and links to the existing clustering workflow. The strict docs build passes. |
 
 ## Feedback and further reading
 

DEVELOPMENT.md

@@ -310,10 +310,11 @@ covers its arithmetic kernels, not every scalar loop in construction and search.
 | IVF and float SymphonyQG batch correction | Complete estimator runs in the selected backend; non-float template paths remain generic |
 | FHT/Kac rotation | Complete rotation and scaling run in selected ISA translation units; shared AVX intrinsics preserve the FFHT butterfly order |
 | Float matrix rotation and PiPNN construction | Matrix products, row norms, and lower-triangle pairwise distances use isolated matrix backends |
+| RaBitQKMeans and QGKMeans centroid updates | Indexed float32 rows accumulate into float64 sums through the matrix backends, preserving per-cluster point order without additional workspace |
 | HNSW search and IVF centroid routing | Cached HNSW search selection; centroid routing uses the common raw-distance dispatcher |
 | Quantization orchestration, reconstruction, non-float utilities | Template/control code remains generic; no blanket native tuning or reduction-order rewrite |
 | Graph scheduling, candidate queues, I/O, allocation, random initialization | Generic control code; IVF one-bit candidate insertion stays in a small compiled function to avoid inlining-induced register spills; thread scheduling and seeds remain caller-owned |
-| Example KMeans training | Uses external FAISS, whose build and dispatch are independent of this package |
+| Example KMeans training | Explicit RaBitQKMeans flat or QGKMeans graph assignment; FAISS is an optional comparison dependency only |
 
 Portable wheels disable `RABITQ_ENABLE_NATIVE_OPTIMIZATION`. ARM64 builds exclude all
 x86 source groups and use standard AArch64 NEON intrinsics, without Apple-only APIs.

README.md

@@ -6,7 +6,7 @@
 
 <p>
   A research-backed C++17 library with Python bindings for 1-bit and multi-bit<br>
-  vector quantization, IVF, HNSW, and SymphonyQG.
+  vector quantization, IVF, HNSW, SymphonyQG, and clustering.
 </p>
 
 <p>
@@ -63,17 +63,16 @@ Build and search a small IVF index using synthetic data:
 
 ```python
 import numpy as np
-from rabitqlib import IvfIndex
+from rabitqlib import FinalAssignmentMode, IvfIndex, RaBitQKMeans
 
 rng = np.random.default_rng(42)
 data = rng.standard_normal((500, 64)).astype(np.float32)
 queries = rng.standard_normal((5, 64)).astype(np.float32)
 
-# Assign vectors to five clusters and calculate their centroids.
-cluster_ids = (np.arange(len(data)) % 5).astype(np.uint32)
-centroids = np.stack(
-    [data[cluster_ids == cluster].mean(axis=0) for cluster in range(5)]
-).astype(np.float32)
+clustering = RaBitQKMeans(
+    64, 5, num_threads=2, final_assignment=FinalAssignmentMode.Exact
+)
+clustering.train(data)
 
 index = IvfIndex(
     dim=64,
@@ -82,7 +81,7 @@ index = IvfIndex(
     nbits=4,
     metric="l2",
 )
-index.build(data, centroids, cluster_ids)
+index.build(data, clustering.centroids, clustering.assignments)
 
 ids, distances = index.search(queries, k=10, nprobe=5)
 print(ids.shape, distances.shape)  # (5, 10) (5, 10)
@@ -93,16 +92,20 @@ For all three indexes, `build` and `search` interpret `num_threads=0` as the
 detected hardware thread count. Larger requests are capped at that count;
 smaller positive requests are respected. Operations may use fewer workers when
 there are fewer work items. If hardware detection is unavailable, one thread is
-used. Omitting `num_threads` in Python still defaults to one thread.
+used. Python index methods default to one thread when `num_threads` is omitted.
 
-Python bindings are also available for `HnswIndex` and `SymqgIndex`. See the
-[Python examples](sample/python/README.md) for index construction, querying, and
-index persistence. IVF and HNSW examples run Faiss clustering in a separate process,
-then pass saved clusters to RaBitQ indexing so their OpenMP runtimes stay separate.
+IVF and HNSW examples save clusters for reuse across index configurations.
+Choose RaBitQKMeans for flat assignment or QGKMeans for graph assignment;
+FAISS is needed only for the optional clustering comparison.
 
 Index save/load paths are UTF-8 strings on Windows and native path bytes on POSIX
 in C++; Python paths are Unicode strings on all platforms.
 
+See the [Python examples](sample/python/README.md) for IVF, HNSW, and SymphonyQG.
+For clustering, use [RaBitQKMeans](docs/docs/clustering.md#rabitqkmeans) for small cluster
+counts or [QGKMeans](docs/docs/clustering.md#qgkmeans) for graph assignment. The
+[FAISS comparison](sample/python/compare_with_faiss.py) benchmarks both methods.
+
 <details>
 <summary>Build the Python bindings from source</summary>
 
@@ -145,6 +148,7 @@ depend on the dataset, index configuration, and search parameters.
 | **Accurate estimates** | An asymptotically optimal theoretical error bound supports reliable ordering and reranking. |
 | **Native CPU backends** | Runtime AVX2/AVX-512 selection on x86-64; NEON distance, packed-code, FastScan, rotation, query preparation, and HNSW search kernels on ARM64, with portable scalar fallbacks. |
 | **Ready for ANN search** | Use the quantizer directly or build complete IVF, HNSW, and [SymphonyQG](https://dl.acm.org/doi/abs/10.1145/3709730) indexes. |
+| **Native clustering** | Choose [RaBitQKMeans](docs/docs/clustering.md#rabitqkmeans) with flat RaBitQ assignment or [QGKMeans](docs/docs/clustering.md#qgkmeans) with SymphonyQG assignment; neither needs an external k-means package. |
 
 The library supports Euclidean distance and inner product. Cosine search is
 available by normalizing vectors before using inner product.

docs/docs/clustering.md

@@ -0,0 +1,189 @@
+# Clustering
+
+RaBitQ Library provides two implementations of Lloyd's k-means. Both compute
+centroid means from the original float32 data, accumulate in float64, and require
+no external clustering package. Choose the class explicitly; neither switches
+algorithms based on the cluster count.
+
+## Choose a method
+
+| Class | Assignment | Recommended use | Cluster count |
+| --- | --- | --- | --- |
+| [RaBitQKMeans](#rabitqkmeans) | Flat RaBitQ scan | Small cluster counts | `n >= k >= 1` |
+| [QGKMeans](#qgkmeans) | SymphonyQG centroid graph | Larger cluster counts | `n >= k > graph_degree` (default 32) |
+
+## Python
+
+Both classes have the same training and result interface:
+
+```python
+from rabitqlib import QGKMeans, RaBitQKMeans
+
+kmeans = RaBitQKMeans(d, 16, niter=25, num_threads=8)
+# For larger cluster counts, use instead:
+# kmeans = QGKMeans(d, 4096, niter=25, num_threads=8)
+last_iteration_obj = kmeans.train(x)
+
+centroids = kmeans.centroids      # shape (k, d)
+assignments = kmeans.assignments  # shape (n,)
+distances = kmeans.distances      # shape (n,)
+final_obj = kmeans.final_obj
+```
+
+`x` has shape `(n, d)` and is converted to a C-contiguous float32 array.
+
+## C++
+
+```cpp
+#include <rabitqlib/clustering/rabitqkmeans.hpp>
+#include <rabitqlib/clustering/qgkmeans.hpp>
+
+rabitqlib::rabitqkmeans::RaBitQKMeansParameters parameters;
+parameters.niter = 25;
+parameters.num_threads = 8;
+rabitqlib::rabitqkmeans::RaBitQKMeans kmeans(dimension, 16, parameters);
+
+// For graph assignment, use qgkmeans::QGKMeansParameters and qgkmeans::QGKMeans.
+kmeans.train(num_points, data);  // row-major float32 input
+
+const auto& centroids = kmeans.centroids;
+const auto& assignments = kmeans.assignments;
+const auto& distances = kmeans.distances;
+const double final_obj = kmeans.final_obj;
+```
+
+Complete examples: [RaBitQKMeans](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/rabitqkmeans.cpp)
+(`bin/rabitqkmeans`) and [QGKMeans](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/qgkmeans.cpp)
+(`bin/qgkmeans`). Run either without arguments for a synthetic demo, or save
+clusters for the C++ indexing examples:
+
+```bash
+./bin/qgkmeans data.fvecs 4096 centroids.fvecs labels.ivecs
+./bin/rabitqkmeans data.fvecs 16 centroids.fvecs labels.ivecs
+```
+
+File mode uses exact final assignment and the default training parameters.
+Append `ip` for spherical clustering of normalized input; the default is `l2`.
+`example.sh` runs the C++ clustering and indexing workflow without Python.
+
+## Parameters and input
+
+Set Python keyword arguments or fields of the corresponding C++ parameter type:
+
+| Parameter | Default | Meaning |
+| --- | --- | --- |
+| `niter` | `25` | Maximum clustering iterations. |
+| `early_stop_threshold` | `0.0` | Stop from iteration 2 when the absolute relative objective change is at most this value; range `[0, 1]`. Zero stops only on an unchanged objective. |
+| `num_threads` | `0` | Use detected hardware threads; positive requests are capped at that count. |
+| `spherical` | `false` | Normalize centroids and assign by inner product; input must already be normalized. |
+| `seed` | `42` | Seed for centroid selection and rotation. |
+| `min_points_per_centroid` | `39` | Warn below this training-points-per-centroid ratio when verbose. |
+| `verbose` | `false` | Print training progress and warnings. |
+| `final_assignment` | `Approximate` | Use the selected method for final labels, or `Exact` for nearest-centroid labels. |
+
+Training requires `64 <= d <= 65536`, finite coordinates, and
+`abs(x) <= sqrt(sqrt(FLT_MAX) / (64 * d))`; scale larger inputs before training.
+Cluster-count requirements are listed above. Training does not normalize or
+mutate the input; provide normalized vectors for spherical mode.
+
+Empty clusters are reseeded from distant points; duplicate data can still leave
+empty clusters in the final assignment. Thread scheduling can affect graph
+construction and floating-point reductions, so a fixed seed does not guarantee
+identical multithreaded results.
+
+## Method details
+
+### RaBitQKMeans
+
+The flat assignment design is inspired by
+[SuperKMeans](https://github.com/cwida/SuperKMeans). Training points are rotated
+and encoded once around a fixed dataset mean. Each centroid becomes a query
+that scans all packed point batches, and each point receives the centroid with
+the best estimated distance.
+
+Each iteration compares all `n * k` point-centroid pairs. The batch size of 32
+is not a cluster-count limit. This method uses one-bit point codes throughout;
+it has no graph parameters or `quantization_bits` setting.
+
+### QGKMeans
+
+QGKMeans builds a SymphonyQG graph over the centroids and uses previous
+assignments as search hints. Its additional parameters are:
+
+| Parameter | Default | Meaning |
+| --- | --- | --- |
+| `quantization_bits` | `0` | Graph centroid storage: raw, or `4`/`8` bits. |
+| `graph_degree` | `32` | A positive multiple of 32, strictly smaller than `k`. |
+| `ef_build`, `ef_search` | `240`, `16` | Positive graph construction and search window sizes. |
+| `graph_build_iterations` | `1` | PiPNN initialization followed by one refinement; values ≥2 request that many graph passes, ending in refinement. Independent of `niter`. |
+
+The seed also controls fallback neighbors; PiPNN uses fixed internal seeds.
+`QGAssigner` in `qgkmeans.hpp` supports repeated vector-to-centroid assignment
+without the clustering loop.
+
+## Results and final assignment
+
+Python result properties are `None` until a successful fit; a failed fit
+preserves the previous successful results.
+
+`distances` contains distances to the returned **assigned** centroids, computed
+from the original float32 vectors and centroids. For squared L2, `final_obj` is
+their sum (WCSS); spherical mode sums `1 - inner_product`. `obj` and
+`iteration_stats` record objectives before each centroid update; Python
+`train()` returns the last such objective.
+
+`FinalAssignmentMode.Approximate` (default) uses the selected class's assignment
+algorithm. The legacy name `FinalAssignmentMode.SymphonyQG` remains an alias for
+`Approximate`. Labels need not identify the exact nearest centroids, but reported
+distances and WCSS remain exact for those labels. Quantized QGKMeans final
+assignment retains a previous label when it is closer. Approximate training
+objectives need not decrease monotonically.
+
+To require exact nearest-centroid **final** assignments with either class:
+
+```python
+from rabitqlib import FinalAssignmentMode, RaBitQKMeans
+
+kmeans = RaBitQKMeans(d, 16, final_assignment=FinalAssignmentMode.Exact)
+kmeans.train(x)
+```
+
+In C++, set `parameters.final_assignment` to
+`rabitqlib::rabitqkmeans::FinalAssignmentMode::Exact` or
+`rabitqlib::qgkmeans::FinalAssignmentMode::Exact`; both names refer to the same
+enum. Training still uses the selected approximate assignment algorithm.
+
+Exact final assignment screens all centroids with blocked GEMM and refines
+possible winners in float64; ties select the smallest centroid index. Returned
+distances remain float32. Work is proportional to `n * k * d`; large coordinate
+offsets or near ties can require more refinement.
+
+## Compare with FAISS
+
+The [comparison script](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/python/compare_with_faiss.py)
+compares both methods with FAISS and requires `faiss-cpu`. Defaults are one
+million 256-dimensional vectors, 4,000 centroids, and up to 25 iterations.
+Load data with `--input data.npy`, or `--input data.bin --d D` for raw row-major
+float32 data. Choose a thread count within the available hardware limit.
+
+```bash
+python sample/python/compare_with_faiss.py --threads 48 --num-seeds 2 --repeats 2
+```
+
+All three methods receive the same full dataset, iteration limit, cluster count,
+thread budget, and seed values; FAISS subsampling is disabled. This command runs
+four fits per method with rotating order. Different initializers mean identical
+seeds need not produce identical centroids.
+
+Only `train()` is timed, including the library methods' default approximate
+final assignment. The script reports actual iterations, timing medians/ranges,
+exhaustive nearest-centroid objectives, and both library methods' returned-label
+objectives and label mismatches. Equal-distance ties can produce mismatches
+without changing the objective. `--spherical` normalizes the input copy and
+compares sums of `1 - inner_product`.
+
+Separate worker processes isolate OpenMP runtimes and share a temporary float32
+input file; allow about `4 * n * d` bytes of temporary disk space. Input staging,
+process startup, and exhaustive evaluation are outside training time. For Linux
+CPU placement, set `OMP_PLACES=threads OMP_PROC_BIND=close`; on NUMA systems,
+use `numactl --cpunodebind=0 --membind=0` with a thread count that fits the node.

docs/docs/index.md

@@ -46,16 +46,16 @@ Build an IVF index and search a batch of queries:
 
 ```python
 import numpy as np
-from rabitqlib import IvfIndex
+from rabitqlib import FinalAssignmentMode, IvfIndex, RaBitQKMeans
 
 rng = np.random.default_rng(42)
 data = rng.standard_normal((500, 64)).astype(np.float32)
 queries = rng.standard_normal((5, 64)).astype(np.float32)
 
-cluster_ids = (np.arange(len(data)) % 5).astype(np.uint32)
-centroids = np.stack(
-    [data[cluster_ids == cluster].mean(axis=0) for cluster in range(5)]
-).astype(np.float32)
+clustering = RaBitQKMeans(
+    64, 5, num_threads=2, final_assignment=FinalAssignmentMode.Exact
+)
+clustering.train(data)
 
 index = IvfIndex(
     dim=64,
@@ -64,7 +64,7 @@ index = IvfIndex(
     nbits=4,
     metric="l2",
 )
-index.build(data, centroids, cluster_ids)
+index.build(data, clustering.centroids, clustering.assignments)
 
 ids, distances = index.search(queries, k=10, nprobe=5)
 print(ids.shape, distances.shape)  # (5, 10) (5, 10)

docs/docs/index/hnsw.md

@@ -28,7 +28,7 @@ HierarchicalNSW::construct(size_t cluster_num,
 
 - **data**: Pointer to the raw data vectors.
 - **data_num**: The number of data vectors.
-- **centroids**: Centroids computed by K-means clustering on the raw data vectors (we recommend `cluster_num = 16`).  
+- **centroids**: Centroids computed by [RaBitQKMeans](../clustering.md#rabitqkmeans); the examples use `cluster_num = 16`.
 - **cluster_ids**: Array of length `data_num`; every entry must be in the range `[0, cluster_num)`.
 - **num_threads**: Number of threads to use (default: 0, which auto-selects).
 - **faster**: If `true`, enables the faster quantizer.