From d3aa8cd34a0e1378a6a2d39f300bf2574fa4f473 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sun, 1 Nov 2026 00:14:43 +0800 Subject: [PATCH] feat(clustering): add optional k-means sampling and assignment of new vectors Add max_points_per_centroid to RaBitQKMeansParameters and QGKMeansParameters. It defaults to 0, which trains on all vectors as before; Faiss uses 256. When positive and n exceeds k * max_points_per_centroid, initialization and Lloyd iterations use a seeded random sample drawn from a stream independent of centroid initialization, and the final centroids then assign all n points with the configured final_assignment. assignments and distances keep shape (n,); obj and iteration_stats describe the sample, distances and final_obj all points. Training that does not sample is unchanged bit for bit. Sampling copies the sample, so it raises peak memory. Add assign(n, x, assignments, distances) const to both C++ classes and assign(x) -> (uint32 (m,), float32 (m,)) in Python. It uses the parameters of the last successful fit and no labels or search hints from training, so each vector's result is independent of the batch: - RaBitQKMeans stores the center its training points were encoded around and encodes new vectors around it. - QGKMeans caches a graph over the centroids, rebuilt only when the centroids or graph settings change, and searches it with per-call scratch and no query-sum buffer. assign validates x and the public centroids like training input and raises " has not been trained" before a successful train(). A sampled fit's final pass is the same assignment, so assign reproduces its labels. --- docs/docs/clustering.md | 78 ++- .../rabitqlib/clustering/detail/kmeans.hpp | 226 +++++-- include/rabitqlib/clustering/qgkmeans.hpp | 201 +++++- include/rabitqlib/clustering/rabitqkmeans.hpp | 40 +- python_bindings/__init__.py | 70 +- python_bindings/__init__.pyi | 150 ++++- python_bindings/kmeans_bindings.cpp | 30 + tests/python/test_kmeans_sampling.py | 262 ++++++++ tests/python/typing/public_api.py | 9 + .../clustering/centroid_update_test.cpp | 17 + .../clustering/kmeans_sampling_test.cpp | 634 ++++++++++++++++++ 11 files changed, 1637 insertions(+), 80 deletions(-) create mode 100644 tests/python/test_kmeans_sampling.py create mode 100644 tests/unit/rabitqlib/clustering/kmeans_sampling_test.cpp diff --git a/docs/docs/clustering.md b/docs/docs/clustering.md index 7cb16db7..8ab5fc54 100644 --- a/docs/docs/clustering.md +++ b/docs/docs/clustering.md @@ -30,6 +30,9 @@ centroids = kmeans.centroids # shape (k, d) assignments = kmeans.assignments # shape (n,) distances = kmeans.distances # shape (n,) final_obj = kmeans.final_obj + +# Assign other vectors, shape (m, d), to the trained centroids. +query_assignments, query_distances = kmeans.assign(queries) # shapes (m,), (m,) ``` `x` has shape `(n, d)` and is converted to a C-contiguous float32 array. @@ -55,6 +58,11 @@ const auto& centroids = kmeans.centroids; const auto& assignments = kmeans.assignments; const auto& distances = kmeans.distances; const double final_obj = kmeans.final_obj; + +// Caller-owned outputs with num_queries elements each. +std::vector query_assignments(num_queries); +std::vector query_distances(num_queries); +kmeans.assign(num_queries, queries, query_assignments.data(), query_distances.data()); ``` Complete examples: [RaBitQKMeans](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/cpp/rabitqkmeans.cpp) @@ -83,6 +91,7 @@ Set Python keyword arguments or fields of the corresponding C++ parameter type: | `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. | +| `max_points_per_centroid` | `0` | When positive, train on a random sample above `k * max_points_per_centroid` vectors; `0` trains on all vectors. See [sampling](#training-set-sampling). | | `verbose` | `false` | Print training progress and warnings. | | `final_assignment` | `Approximate` | Use the selected method for final labels, or `Exact` for nearest-centroid labels. | @@ -101,6 +110,35 @@ 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. +### Training-set sampling + +Sampling is disabled by default. When `max_points_per_centroid` is positive +and `n > k * max_points_per_centroid`, training draws +`k * max_points_per_centroid` distinct vectors uniformly at random, as Faiss +does; Faiss enables this by default with 256. Centroid initialization and all +Lloyd iterations use only this sample. The final centroids then assign all `n` +vectors with the selected `final_assignment` mode, as [`assign`](#assign-new-vectors) +does, so `assignments` and `distances` keep shape `(n,)` and can be passed +directly to index builders. + +`seed` selects the sample through a random stream separate from centroid +initialization. Training that does not sample, because the parameter is `0` or +`n` is at most `k * max_points_per_centroid`, returns the same results as +before this parameter existed. The selected vectors depend on the +C++ standard library's random distributions, so the same seed can choose a +different sample on another platform. + +`obj` and `iteration_stats` (and the value returned by Python `train()`) are +computed on the sampled vectors, while `distances` and `final_obj` cover all `n` +vectors. Their values are not directly comparable. Validation still checks +every input vector, including those outside the sample. + +Sampling reduces iteration work but raises peak memory. Training copies the +sample into a contiguous buffer of `k * max_points_per_centroid * d` float32 +values, 4 GiB for `k = 4096`, `max_points_per_centroid = 256` and `d = 960`, +and briefly holds `n` 8-byte row indices to select it, in addition to the +caller's input. The final assignment still processes all `n` vectors. + ## Method details ### RaBitQKMeans @@ -109,9 +147,11 @@ 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. +the best estimated distance. The model keeps this mean of the training points +(of the sample, when training samples), and [`assign`](#assign-new-vectors) +encodes new vectors around it. -Each iteration compares all `n * k` point-centroid pairs. The batch size of 32 +Each iteration compares every training point with every centroid. 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. @@ -129,7 +169,8 @@ assignments as search hints. Its additional parameters are: 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. +without the clustering loop. To assign vectors with trained QGKMeans centroids, +use [`assign`](#assign-new-vectors). ## Results and final assignment @@ -145,9 +186,10 @@ their sum (WCSS); spherical mode sums `1 - inner_product`. `obj` and `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. +distances and WCSS remain exact for those labels. Without sampling, quantized +QGKMeans final assignment retains a point's last-iteration label when it is +closer. After sampling, most points have no such label, so every point uses the +same graph search. Approximate training objectives need not decrease monotonically. To require exact nearest-centroid **final** assignments with either class: @@ -168,6 +210,30 @@ 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. +## Assign new vectors + +After a successful fit, `assign(x)` returns `(assignments, distances)` for an +`(m, d)` array: uint32 centroid IDs and float32 distances to those centroids, +each of shape `(m,)`. It uses `final_assignment` and the other parameters of the +last successful fit and checks coordinates as training does; `m` may be `0`. +Before the first successful fit, it raises `RuntimeError`. The C++ method +`assign(n, x, assignments, distances)` writes to caller-owned buffers. It throws +`std::runtime_error` until `train()` succeeds, including when `centroids` were +written directly, and `std::invalid_argument` for invalid `x` or `centroids`. + +`assign` uses no labels or search hints from training, so each vector's +result is independent of the other vectors in the call. Approximate +RaBitQKMeans encodes vectors around the stored training mean. Approximate +QGKMeans builds a graph over the centroids on first use and reuses it; +C++ callers that change `centroids` or graph parameters cause one rebuild on +the next call. Each call still validates the `k * d` centroid coordinates. + +For the training vectors, `assign` reproduces `assignments` and `distances` +with `FinalAssignmentMode.Exact`, for RaBitQKMeans, and whenever training +sampled the input, because the sampled final pass is the same assignment. +Unsampled approximate QGKMeans training also uses previous labels as graph +search hints, so its final labels can differ. + ## Compare with Faiss The [comparison script](https://github.com/VectorDB-NTU/RaBitQ-Library/blob/main/sample/python/compare_with_faiss.py) diff --git a/include/rabitqlib/clustering/detail/kmeans.hpp b/include/rabitqlib/clustering/detail/kmeans.hpp index b4603943..3810ecfc 100644 --- a/include/rabitqlib/clustering/detail/kmeans.hpp +++ b/include/rabitqlib/clustering/detail/kmeans.hpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -214,11 +215,31 @@ inline void exact_assign( failure.rethrow(); } +// Returns count sorted, distinct row indices in [0, n). The seed sequence keeps +// this stream independent of the centroid initialization stream. +[[nodiscard]] inline std::vector sample_training_points( + size_t n, size_t count, uint32_t seed +) { + assert(count <= n); + std::vector indices(n); + std::iota(indices.begin(), indices.end(), size_t{0}); + std::seed_seq sequence{seed, uint32_t{1}}; + std::mt19937 rng(sequence); + for (size_t i = 0; i < count; ++i) { + std::uniform_int_distribution pick(i, n - 1); + std::swap(indices[i], indices[pick(rng)]); + } + indices.resize(count); + std::sort(indices.begin(), indices.end()); + return indices; +} + template class LloydKMeans : public Parameters { public: using Parameters::early_stop_threshold; using Parameters::final_assignment; + using Parameters::max_points_per_centroid; using Parameters::min_points_per_centroid; using Parameters::niter; using Parameters::num_threads; @@ -242,39 +263,8 @@ class LloydKMeans : public Parameters { std::string(Assignment::kName) + " requires n >= k" ); } - if (n > static_cast(std::numeric_limits::max()) / d) { - throw std::invalid_argument( - std::string(Assignment::kName) + " input dimensions overflow ptrdiff_t" - ); - } const uint32_t num_workers = resolved_num_workers(); - // Leave headroom for products of squared distances in graph pruning, - // as well as quantized reconstruction and rotated residuals. - const double max_coordinate = std::sqrt( - std::sqrt(static_cast(std::numeric_limits::max())) / - (64.0 * static_cast(d)) - ); - bool in_range = true; - bool finite = true; -#pragma omp parallel for num_threads(num_workers) reduction(&& : finite, in_range) \ - schedule(static) - for (std::ptrdiff_t index = 0; index < static_cast(n * d); - ++index) { - const size_t i = static_cast(index); - finite = finite && std::isfinite(x[i]); - in_range = in_range && std::abs(static_cast(x[i])) <= max_coordinate; - } - if (!finite) { - throw std::invalid_argument( - std::string(Assignment::kName) + " x must contain only finite values" - ); - } - if (!in_range) { - throw std::invalid_argument( - std::string(Assignment::kName) + - " x exceeds the safe float32 coordinate range" - ); - } + validate_points(n, x, num_workers, "x"); if (verbose && min_points_per_centroid > 0 && n / k < min_points_per_centroid) { std::cerr << "WARNING " << Assignment::kName << " clustering " << n << " points to " << k << " centroids: please provide at least " @@ -291,13 +281,37 @@ class LloydKMeans : public Parameters { x = input_copy.data(); } - Assignment assigner(d, k, x, n, num_workers, *this); + // Sample when n > k * max_points_per_centroid, without forming the product. + const bool sampled = + max_points_per_centroid != 0 && max_points_per_centroid <= (n - 1) / k; + const size_t num_samples = sampled ? k * max_points_per_centroid : n; + std::vector samples; + const float* training_points = x; + if (sampled) { + if (verbose) { + std::cout << Assignment::kName << " sampling " << num_samples << " of " << n + << " training points\n"; + } + const auto indices = sample_training_points(n, num_samples, seed); + samples.resize(num_samples * d); +#pragma omp parallel for num_threads(num_workers) schedule(static) + for (std::ptrdiff_t index = 0; index < static_cast(num_samples); + ++index) { + const auto sample = static_cast(index); + std::copy_n(x + indices[sample] * d, d, samples.data() + sample * d); + } + training_points = samples.data(); + } + + std::optional assigner( + std::in_place, d, k, training_points, num_samples, num_workers, *this + ); centroids.resize(k * d); - assignments.resize(n); - distances.resize(n); + assignments.resize(num_samples); + distances.resize(num_samples); iteration_stats.clear(); - initialize_centroids(x, n); + initialize_centroids(training_points, num_samples); if (spherical) { normalize_centroids(); } @@ -316,22 +330,28 @@ class LloydKMeans : public Parameters { ++cluster) { update_worker_for_cluster[cluster] = worker; } - partition.points.reserve(n / update_workers); + partition.points.reserve(num_samples / update_workers); } double previous_obj = std::numeric_limits::infinity(); for (size_t iteration = 0; iteration < niter; ++iteration) { - assigner.assign(centroids.data(), assignments.data(), distances.data()); + assigner->assign(centroids.data(), assignments.data(), distances.data()); // Assignment has finished borrowing the current centroids. Keep them // for the shift calculation and overwrite the other buffer with means. previous_centroids.swap(centroids); const double obj = std::accumulate(distances.begin(), distances.end(), 0.0); update_centroids( - x, n, sums, cluster_sizes, update_partitions, update_worker_for_cluster + training_points, + num_samples, + sums, + cluster_sizes, + update_partitions, + update_worker_for_cluster + ); + const auto [nsplit, shift] = consolidate_centroids( + training_points, num_samples, sums, cluster_sizes, previous_centroids ); - const auto [nsplit, shift] = - consolidate_centroids(x, n, sums, cluster_sizes, previous_centroids); iteration_stats.push_back( KMeansIterationStats{iteration + 1, obj, shift, nsplit} ); @@ -351,9 +371,17 @@ class LloydKMeans : public Parameters { } } + assigner->store_model_state(assignment_state_); // Return assignments for the final centroids, not the centroids from // the beginning of the last update step. - if (final_assignment == FinalAssignmentMode::Exact) { + if (sampled) { + // Release the sample and its encoded form before assigning all points. + assigner.reset(); + std::vector().swap(samples); + assignments.resize(n); + distances.resize(n); + assign_points(n, x, num_workers, assignments.data(), distances.data()); + } else if (final_assignment == FinalAssignmentMode::Exact) { detail::exact_assign( x, centroids.data(), @@ -367,11 +395,42 @@ class LloydKMeans : public Parameters { assignments.data() ); } else { - assigner.assign(centroids.data(), assignments.data(), distances.data(), true); + assigner->assign(centroids.data(), assignments.data(), distances.data(), true); } final_obj = std::accumulate(distances.begin(), distances.end(), 0.0); } + /** + * 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. + */ + void assign( + size_t n, const float* x, PID* point_assignments, float* point_distances + ) const { + validate_parameters(); + if (centroids.size() != k * d || + !Assignment::has_model_state(assignment_state_, d)) { + throw std::runtime_error( + std::string(Assignment::kName) + " has not been trained" + ); + } + const uint32_t num_workers = resolved_num_workers(); + validate_points(k, centroids.data(), num_workers, "centroids"); + if (n == 0) { + return; + } + if (x == nullptr || point_assignments == nullptr || point_distances == nullptr) { + throw std::invalid_argument( + std::string(Assignment::kName) + + " assign input and output pointers must not be null" + ); + } + validate_points(n, x, num_workers, "x"); + assign_points(n, x, num_workers, point_assignments, point_distances); + } + size_t d; size_t k; std::vector centroids; @@ -382,6 +441,9 @@ class LloydKMeans : public Parameters { double final_obj = 0.0; private: + // Encoding center or centroid graph cache for assign(); see the Assignment type. + typename Assignment::ModelState assignment_state_; + struct UpdatePartition { size_t first_cluster; size_t last_cluster; // Exclusive; also bounds the accumulators we clear. @@ -421,6 +483,84 @@ class LloydKMeans : public Parameters { } } + // Requires a non-null x. name identifies the rows in error messages. + void validate_points( + size_t n, const float* x, uint32_t num_workers, const char* name + ) const { + if (n > static_cast(std::numeric_limits::max()) / d) { + throw std::invalid_argument( + std::string(Assignment::kName) + " input dimensions overflow ptrdiff_t" + ); + } + // Leave headroom for products of squared distances in graph pruning, + // as well as quantized reconstruction and rotated residuals. + const double max_coordinate = std::sqrt( + std::sqrt(static_cast(std::numeric_limits::max())) / + (64.0 * static_cast(d)) + ); + bool in_range = true; + bool finite = true; +#pragma omp parallel for num_threads(num_workers) reduction(&& : finite, in_range) \ + schedule(static) + for (std::ptrdiff_t index = 0; index < static_cast(n * d); + ++index) { + const size_t i = static_cast(index); + finite = finite && std::isfinite(x[i]); + in_range = in_range && std::abs(static_cast(x[i])) <= max_coordinate; + } + if (!finite) { + throw std::invalid_argument( + std::string(Assignment::kName) + " " + name + + " must contain only finite values" + ); + } + if (!in_range) { + throw std::invalid_argument( + std::string(Assignment::kName) + " " + name + + " exceeds the safe float32 coordinate range" + ); + } + } + + // Assigns validated points to the current centroids. Uses the state saved by + // train() but neither Lloyd assignments nor graph search hints, so each + // point's result is independent of the batch. Requires n > 0 and outputs + // that do not overlap x or centroids. + void assign_points( + size_t n, + const float* x, + uint32_t num_workers, + PID* point_assignments, + float* point_distances + ) const { + if (final_assignment == FinalAssignmentMode::Exact) { + detail::exact_assign( + x, + centroids.data(), + n, + d, + k, + spherical, + num_workers, + point_assignments, + point_distances + ); + } else { + Assignment::assign_with_state( + d, + k, + centroids.data(), + x, + n, + num_workers, + *this, + assignment_state_, + point_assignments, + point_distances + ); + } + } + void initialize_centroids(const float* x, size_t n) { std::vector indices(n); std::iota(indices.begin(), indices.end(), size_t{0}); diff --git a/include/rabitqlib/clustering/qgkmeans.hpp b/include/rabitqlib/clustering/qgkmeans.hpp index fd2f7d77..280c75f6 100644 --- a/include/rabitqlib/clustering/qgkmeans.hpp +++ b/include/rabitqlib/clustering/qgkmeans.hpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -32,6 +33,7 @@ using QGKMeansIterationStats = clustering::KMeansIterationStats; namespace detail { +class CentroidGraphCache; class QGAssignment; } // namespace detail @@ -58,6 +60,9 @@ struct QGKMeansParameters { uint32_t num_threads = 0; // Centroid storage in QG: 0 for raw float32, or 4/8-bit RaBitQ. uint32_t quantization_bits = 0; + // When positive, train on a random sample of k * max_points_per_centroid + // vectors if there are more, then assign all vectors; 0 trains on all vectors. + size_t max_points_per_centroid = 0; }; namespace detail { @@ -124,6 +129,7 @@ class QuerySumCache { * calls, and the previous assignment is used as the next call's search hint. */ class QGAssigner { + friend class detail::CentroidGraphCache; friend class detail::QGAssignment; public: @@ -236,9 +242,6 @@ class QGAssigner { } rebuild_centroid_graph(centroids); - const auto distance_func = graph_->metric_type() == METRIC_IP - ? dot_product_dis - : euclidean_sqr; // Stable grouping changes only the order of independent full searches. std::vector query_order; @@ -255,6 +258,58 @@ class QGAssigner { static_cast(point); } } + search_points( + centroids, + x, + n, + assignments, + distances, + query_order, + previous_assignments_.size() == n ? previous_assignments_.data() : nullptr, + query_sums, + scratch_ + ); + + previous_assignments_.assign(assignments, assignments + n); + query_sums.complete(); + } + + // Requires a graph built from centroids. Uses no hints or query-sum cache + // and allocates its own scratch, so concurrent calls only read the graph. + void search_built_graph( + const float* centroids, const float* x, size_t n, PID* assignments, float* distances + ) const { + std::vector> scratch; + scratch.reserve(num_workers_); + for (uint32_t thread = 0; thread < num_workers_; ++thread) { + scratch.emplace_back( + std::make_unique( + graph_->padded_dim(), graph_degree_, ef_search_, 1 + ) + ); + } + const detail::QuerySumCache query_sums(nullptr, false); + search_points( + centroids, x, n, assignments, distances, {}, nullptr, query_sums, scratch + ); + } + + // Searches each point independently. A non-empty query_order only changes + // scheduling; hints, if given, hold one centroid ID per point. + void search_points( + const float* centroids, + const float* x, + size_t n, + PID* assignments, + float* distances, + const std::vector& query_order, + const PID* hints, + const detail::QuerySumCache& query_sums, + const std::vector>& scratch + ) const { + const auto distance_func = graph_->metric_type() == METRIC_IP + ? dot_product_dis + : euclidean_sqr; index_detail::ParallelException failure; #pragma omp parallel for num_threads(num_workers_) schedule(dynamic, 64) for (std::ptrdiff_t index = 0; index < static_cast(n); ++index) { @@ -265,7 +320,7 @@ class QGAssigner { continue; } try { - auto& scratch = *scratch_[static_cast(thread_id)]; + auto& worker_scratch = *scratch[static_cast(thread_id)]; // Keep visited storage with the worker, as in QuantizedGraph::search. thread_local VisitedSet visited; thread_local size_t visited_size = 0; @@ -278,19 +333,18 @@ class QGAssigner { visited_size = num_clusters_; visited_capacity = capacity; } - const PID hint = - previous_assignments_.size() == n ? previous_assignments_[i] : kPidMax; + const PID hint = hints == nullptr ? kPidMax : hints[i]; symqg::detail::search_with_workspace( *graph_, x + i * dim_, 1, assignments + i, distances + i, - scratch.view(visited), + worker_scratch.view(visited), hint, query_sums.cached(i) ); - query_sums.store(i, scratch.batch_query.g_k1xsumq()); + query_sums.store(i, worker_scratch.batch_query.g_k1xsumq()); if (graph_->is_quantized()) { // Keep objectives and returned distances relative to the original // centroids; graph search scores are quantized estimates. @@ -302,9 +356,6 @@ class QGAssigner { } failure.rethrow(); - - previous_assignments_.assign(assignments, assignments + n); - query_sums.complete(); } void validate_configuration(uint32_t num_threads) const { @@ -350,6 +401,103 @@ class QGAssigner { namespace detail { +// Graph over a model's centroids for stateless assignment, rebuilt when the +// centroids or graph settings differ from those it was built from. Copies share +// the immutable graph; each object has its own lock. +class CentroidGraphCache { + struct Entry { + std::vector centroids; + QGKMeansParameters parameters; + QGAssigner assigner; + + Entry(size_t d, size_t k, const float* source, const QGKMeansParameters& cp) + : centroids(source, source + k * d) + , parameters(cp) + , assigner( + d, + k, + cp.graph_degree, + cp.ef_build, + cp.ef_search, + cp.graph_build_iterations, + cp.num_threads, + cp.seed, + cp.spherical ? METRIC_IP : METRIC_L2, + cp.quantization_bits + ) { + // The builder borrows this copy, which lives as long as the graph. + assigner.rebuild_centroid_graph(centroids.data()); + } + + [[nodiscard]] bool matches( + size_t d, size_t k, const float* source, const QGKMeansParameters& cp + ) const { + const auto& p = parameters; + return centroids.size() == k * d && + std::equal(centroids.begin(), centroids.end(), source) && + p.spherical == cp.spherical && p.seed == cp.seed && + p.graph_degree == cp.graph_degree && p.ef_build == cp.ef_build && + p.ef_search == cp.ef_search && + p.graph_build_iterations == cp.graph_build_iterations && + p.num_threads == cp.num_threads && + p.quantization_bits == cp.quantization_bits; + } + }; + + mutable std::mutex mutex_; + mutable std::shared_ptr entry_; + + [[nodiscard]] std::shared_ptr snapshot() const { + const std::lock_guard lock(mutex_); + return entry_; + } + + public: + CentroidGraphCache() = default; + + CentroidGraphCache(const CentroidGraphCache& other) : entry_(other.snapshot()) {} + + CentroidGraphCache& operator=(const CentroidGraphCache& other) { + if (this != &other) { + auto entry = other.snapshot(); + const std::lock_guard lock(mutex_); + entry_ = std::move(entry); + } + return *this; + } + + ~CentroidGraphCache() = default; + + void clear() { + const std::lock_guard lock(mutex_); + entry_.reset(); + } + + // Requires k * d validated centroid coordinates. + void assign( + size_t d, + size_t k, + const float* centroids, + const QGKMeansParameters& cp, + const float* x, + size_t n, + PID* assignments, + float* distances + ) const { + std::shared_ptr entry; + { + const std::lock_guard lock(mutex_); + if (!entry_ || !entry_->matches(d, k, centroids, cp)) { + entry_ = std::make_shared(d, k, centroids, cp); + } + entry = entry_; + } + entry->assigner.search_built_graph( + entry->centroids.data(), x, n, assignments, distances + ); + } +}; + // Retain graph-specific state for one synchronous training call. class QGAssignment { size_t dim_; @@ -434,6 +582,37 @@ class QGAssignment { } } } + + // Kept by the model after training for stateless assignment. + struct ModelState { + bool trained = false; + CentroidGraphCache graph; + }; + + void store_model_state(ModelState& state) const { + state.graph.clear(); + state.trained = true; + } + + [[nodiscard]] static bool has_model_state(const ModelState& state, size_t) { + return state.trained; + } + + // Requires validated inputs and centroids. Reuses the cached centroid graph. + static void assign_with_state( + size_t d, + size_t k, + const float* centroids, + const float* x, + size_t n, + uint32_t, + const QGKMeansParameters& cp, + const ModelState& state, + PID* assignments, + float* distances + ) { + state.graph.assign(d, k, centroids, cp, x, n, assignments, distances); + } }; } // namespace detail diff --git a/include/rabitqlib/clustering/rabitqkmeans.hpp b/include/rabitqlib/clustering/rabitqkmeans.hpp index efdb0a52..7550ed4d 100644 --- a/include/rabitqlib/clustering/rabitqkmeans.hpp +++ b/include/rabitqlib/clustering/rabitqkmeans.hpp @@ -40,12 +40,16 @@ struct RaBitQKMeansParameters { double early_stop_threshold = 0.0; FinalAssignmentMode final_assignment = FinalAssignmentMode::Approximate; uint32_t num_threads = 0; + // When positive, train on a random sample of k * max_points_per_centroid + // vectors if there are more, then assign all vectors; 0 trains on all vectors. + size_t max_points_per_centroid = 0; }; namespace detail { // Encode immutable training points once; each moving centroid supplies one LUT. // RaBitQKMeans validates inputs and keeps the borrowed x buffer alive through training. +// Points are encoded around their mean, or around a given d-element center. class RaBitQAssignment { size_t dim_; size_t num_clusters_; @@ -79,7 +83,8 @@ class RaBitQAssignment { const float* x, size_t n, uint32_t num_workers, - const RaBitQKMeansParameters& cp + const RaBitQKMeansParameters& cp, + const float* center = nullptr ) : dim_(d) , num_clusters_(k) @@ -95,7 +100,9 @@ class RaBitQAssignment { , centroid_queries_(k) , rotated_query_(padded_dim_) , lookup_table_(padded_dim_ * 4) { - { + if (center != nullptr) { + std::copy_n(center, dim_, center_.data()); + } else { std::vector> partial_sums( num_workers_, std::vector(dim_, 0.0) ); @@ -158,6 +165,35 @@ class RaBitQAssignment { failure.rethrow(); } + // Kept by the model after training: later assignments encode points around + // the training center, independently of the other points in a call. + struct ModelState { + std::vector center; + }; + + void store_model_state(ModelState& state) const { state.center = center_; } + + [[nodiscard]] static bool has_model_state(const ModelState& state, size_t d) { + return state.center.size() == d; + } + + // Requires validated inputs and centroids. + static void assign_with_state( + size_t d, + size_t k, + const float* centroids, + const float* x, + size_t n, + uint32_t num_workers, + const RaBitQKMeansParameters& cp, + const ModelState& state, + PID* assignments, + float* distances + ) { + RaBitQAssignment assigner(d, k, x, n, num_workers, cp, state.center.data()); + assigner.assign(centroids, assignments, distances); + } + void assign(const float* centroids, PID* assignments, float* distances, bool = false) { const auto distance_func = metric_type_ == METRIC_IP ? dot_product_dis : euclidean_sqr; diff --git a/python_bindings/__init__.py b/python_bindings/__init__.py index 2d454aa3..dca80bb8 100644 --- a/python_bindings/__init__.py +++ b/python_bindings/__init__.py @@ -53,8 +53,9 @@ def train(self, x): Returns ------- objective : float - Last pre-update objective. This can differ from final_obj, which uses - returned assignments and final centroids. + Last pre-update objective, over the sample when training sampled the + input. This can differ from final_obj, which uses all vectors, returned + assignments and final centroids. Notes ----- @@ -66,6 +67,43 @@ def train(self, x): self._impl = implementation return objective + def assign(self, x): + """Assign vectors to the trained centroids. + + Parameters + ---------- + x : array_like of float, shape (m, d) + Vectors in the original, unpadded coordinates, with finite + coordinates in the range accepted by train; m may be 0. Converted to + contiguous float32; a copy may be made. + + Returns + ------- + assignments : ndarray of uint32, shape (m,) + Centroid IDs. + distances : ndarray of float32, shape (m,) + Distances to the assigned centroids, as in the distances property. + + Raises + ------ + RuntimeError + If no train call has succeeded. + + Notes + ----- + Uses final_assignment and the other parameters of the last successful + train call, but no labels or search hints from it, so each vector's + result is independent of the other vectors in x. RaBitQKMeans encodes + vectors around the training center. Approximate QGKMeans builds a + centroid graph on first use and reuses it in later calls. For the + training vectors, the result reproduces assignments and distances with + Exact final assignment, with RaBitQKMeans, or when training sampled the + input; unsampled approximate QGKMeans training uses search hints. + """ + if self._impl is None: + raise RuntimeError(f"{type(self).__name__} has not been trained") + return self._impl.assign(x) + @property def centroids(self): """Final centroids, float32 array of shape (k, d), or None before training.""" @@ -78,22 +116,22 @@ def assignments(self): @property def distances(self): - """Final assignment distances, float32 array of shape (n,), or None before training.""" + """Final assignment distances for all n vectors, float32 array of shape (n,), or None before training.""" return None if self._impl is None else self._impl.distances @property def obj(self): - """Pre-update objectives, float64 array of shape (iterations,), or None before training.""" + """Pre-update objectives over the vectors used for iterations, float64 array of shape (iterations,), or None before training.""" return None if self._impl is None else self._impl.obj @property def final_obj(self): - """Objective for final assignments and centroids, or None before training.""" + """Objective for the final assignments of all n vectors, or None before training.""" return None if self._impl is None else self._impl.final_obj @property def iteration_stats(self): - """Per-iteration statistics, or None before training.""" + """Per-iteration statistics over the vectors used for iterations, or None before training.""" return None if self._impl is None else self._impl.iteration_stats @@ -119,6 +157,10 @@ class QGKMeans(_KMeans): Random seed. Default is 42. min_points_per_centroid : int, optional Warn below this many training vectors per centroid. Default is 39. + max_points_per_centroid : int, optional + When positive, train on a random sample of k * max_points_per_centroid + vectors if there are more, then assign all vectors. Default is 0, which + trains on all vectors; Faiss uses 256. early_stop_threshold : float, optional Stop when the absolute relative objective change is at most this threshold. Default is 0.0. @@ -146,6 +188,12 @@ class QGKMeans(_KMeans): Training results are None until the first successful train call. Set cp fields before training to customize subsequent runs. A failed train preserves previous results. + + When max_points_per_centroid is positive and n exceeds + k * max_points_per_centroid, initialization and iterations use a random + sample chosen by seed. obj and iteration_stats then describe the sample, + while assignments, distances and final_obj cover all n vectors, so the two + kinds of objective are not directly comparable. """ _parameters_type = QGKMeansParameters @@ -173,6 +221,10 @@ class RaBitQKMeans(_KMeans): Random seed. Default is 42. min_points_per_centroid : int, optional Warn below this many training vectors per centroid. Default is 39. + max_points_per_centroid : int, optional + When positive, train on a random sample of k * max_points_per_centroid + vectors if there are more, then assign all vectors. Default is 0, which + trains on all vectors; Faiss uses 256. early_stop_threshold : float, optional Stop when the absolute relative objective change is at most this threshold. Default is 0.0. @@ -189,6 +241,12 @@ class RaBitQKMeans(_KMeans): fields before training to customize subsequent runs. A failed train preserves previous results. + When max_points_per_centroid is positive and n exceeds + k * max_points_per_centroid, initialization and iterations use a random + sample chosen by seed. obj and iteration_stats then describe the sample, + while assignments, distances and final_obj cover all n vectors, so the two + kinds of objective are not directly comparable. + Examples -------- >>> import numpy as np diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index 42638d99..5cd2b5d3 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -868,6 +868,10 @@ class QGKMeans: Random seed. Default is 42. min_points_per_centroid : int, optional Warn below this many training vectors per centroid. Default is 39. + max_points_per_centroid : int, optional + When positive, train on a random sample of k * max_points_per_centroid + vectors if there are more, then assign all vectors. Default is 0, which + trains on all vectors; Faiss uses 256. early_stop_threshold : float, optional Stop when the absolute relative objective change is at most this threshold. Default is 0.0. @@ -895,6 +899,12 @@ class QGKMeans: Training results are None until the first successful train call. Set cp fields before training to customize subsequent runs. A failed train preserves previous results. + + When max_points_per_centroid is positive and n exceeds + k * max_points_per_centroid, initialization and iterations use a random + sample chosen by seed. obj and iteration_stats then describe the sample, + while assignments, distances and final_obj cover all n vectors, so the two + kinds of objective are not directly comparable. """ d: int @@ -911,6 +921,7 @@ class QGKMeans: spherical: bool | np.bool_ = False, seed: SupportsIndex = 42, min_points_per_centroid: SupportsIndex = 39, + max_points_per_centroid: SupportsIndex = 0, early_stop_threshold: float = 0.0, num_threads: SupportsIndex = 0, final_assignment: FinalAssignmentMode = FinalAssignmentMode.Approximate, @@ -941,6 +952,10 @@ class QGKMeans: Random seed. Default is 42. min_points_per_centroid : int, optional Warn below this many training vectors per centroid. Default is 39. + max_points_per_centroid : int, optional + When positive, train on a random sample of k * max_points_per_centroid + vectors if there are more, then assign all vectors. Default is 0, which + trains on all vectors; Faiss uses 256. early_stop_threshold : float, optional Stop when the absolute relative objective change is at most this threshold. Default is 0.0. @@ -968,6 +983,12 @@ class QGKMeans: Training results are None until the first successful train call. Set cp fields before training to customize subsequent runs. A failed train preserves previous results. + + When max_points_per_centroid is positive and n exceeds + k * max_points_per_centroid, initialization and iterations use a random + sample chosen by seed. obj and iteration_stats then describe the sample, + while assignments, distances and final_obj cover all n vectors, so the two + kinds of objective are not directly comparable. """ def train(self, x: npt.ArrayLike) -> float: @@ -982,8 +1003,9 @@ class QGKMeans: Returns ------- objective : float - Last pre-update objective. This can differ from final_obj, which uses - returned assignments and final centroids. + Last pre-update objective, over the sample when training sampled the + input. This can differ from final_obj, which uses all vectors, returned + assignments and final centroids. Notes ----- @@ -991,6 +1013,42 @@ class QGKMeans: copies. """ + def assign( + self, x: npt.ArrayLike + ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: + """Assign vectors to the trained centroids. + + Parameters + ---------- + x : array_like of float, shape (m, d) + Vectors in the original, unpadded coordinates, with finite + coordinates in the range accepted by train; m may be 0. Converted to + contiguous float32; a copy may be made. + + Returns + ------- + assignments : ndarray of uint32, shape (m,) + Centroid IDs. + distances : ndarray of float32, shape (m,) + Distances to the assigned centroids, as in the distances property. + + Raises + ------ + RuntimeError + If no train call has succeeded. + + Notes + ----- + Uses final_assignment and the other parameters of the last successful + train call, but no labels or search hints from it, so each vector's + result is independent of the other vectors in x. RaBitQKMeans encodes + vectors around the training center. Approximate QGKMeans builds a + centroid graph on first use and reuses it in later calls. For the + training vectors, the result reproduces assignments and distances with + Exact final assignment, with RaBitQKMeans, or when training sampled the + input; unsampled approximate QGKMeans training uses search hints. + """ + @property def centroids(self) -> npt.NDArray[np.float32] | None: """Final centroids, float32 array of shape (k, d), or None before training.""" @@ -1001,19 +1059,19 @@ class QGKMeans: @property def distances(self) -> npt.NDArray[np.float32] | None: - """Final assignment distances, float32 array of shape (n,), or None before training.""" + """Final assignment distances for all n vectors, float32 array of shape (n,), or None before training.""" @property def obj(self) -> npt.NDArray[np.float64] | None: - """Pre-update objectives, float64 array of shape (iterations,), or None before training.""" + """Pre-update objectives over the vectors used for iterations, float64 array of shape (iterations,), or None before training.""" @property def final_obj(self) -> float | None: - """Objective for final assignments and centroids, or None before training.""" + """Objective for the final assignments of all n vectors, or None before training.""" @property def iteration_stats(self) -> list[QGKMeansIterationStats] | None: - """Per-iteration statistics, or None before training.""" + """Per-iteration statistics over the vectors used for iterations, or None before training.""" class QGKMeansParameters: """Mutable training settings for QGKMeans. @@ -1049,6 +1107,11 @@ class QGKMeansParameters: """Warn below this many training vectors per centroid. Default is 39.""" @min_points_per_centroid.setter def min_points_per_centroid(self, value: SupportsIndex) -> None: ... + @property + def max_points_per_centroid(self) -> int: + """When positive, train on a random sample of k * max_points_per_centroid vectors if there are more, then assign all vectors. Default is 0, which trains on all vectors; Faiss uses 256.""" + @max_points_per_centroid.setter + def max_points_per_centroid(self, value: SupportsIndex) -> None: ... early_stop_threshold: float """Stop when the absolute relative objective change is at most this threshold. Default is 0.0.""" @@ -1111,6 +1174,10 @@ class RaBitQKMeans: Random seed. Default is 42. min_points_per_centroid : int, optional Warn below this many training vectors per centroid. Default is 39. + max_points_per_centroid : int, optional + When positive, train on a random sample of k * max_points_per_centroid + vectors if there are more, then assign all vectors. Default is 0, which + trains on all vectors; Faiss uses 256. early_stop_threshold : float, optional Stop when the absolute relative objective change is at most this threshold. Default is 0.0. @@ -1127,6 +1194,12 @@ class RaBitQKMeans: fields before training to customize subsequent runs. A failed train preserves previous results. + When max_points_per_centroid is positive and n exceeds + k * max_points_per_centroid, initialization and iterations use a random + sample chosen by seed. obj and iteration_stats then describe the sample, + while assignments, distances and final_obj cover all n vectors, so the two + kinds of objective are not directly comparable. + Examples -------- >>> import numpy as np @@ -1152,6 +1225,7 @@ class RaBitQKMeans: spherical: bool | np.bool_ = False, seed: SupportsIndex = 42, min_points_per_centroid: SupportsIndex = 39, + max_points_per_centroid: SupportsIndex = 0, early_stop_threshold: float = 0.0, num_threads: SupportsIndex = 0, final_assignment: FinalAssignmentMode = FinalAssignmentMode.Approximate, @@ -1176,6 +1250,10 @@ class RaBitQKMeans: Random seed. Default is 42. min_points_per_centroid : int, optional Warn below this many training vectors per centroid. Default is 39. + max_points_per_centroid : int, optional + When positive, train on a random sample of k * max_points_per_centroid + vectors if there are more, then assign all vectors. Default is 0, which + trains on all vectors; Faiss uses 256. early_stop_threshold : float, optional Stop when the absolute relative objective change is at most this threshold. Default is 0.0. @@ -1192,6 +1270,12 @@ class RaBitQKMeans: fields before training to customize subsequent runs. A failed train preserves previous results. + When max_points_per_centroid is positive and n exceeds + k * max_points_per_centroid, initialization and iterations use a random + sample chosen by seed. obj and iteration_stats then describe the sample, + while assignments, distances and final_obj cover all n vectors, so the two + kinds of objective are not directly comparable. + Examples -------- >>> import numpy as np @@ -1215,8 +1299,9 @@ class RaBitQKMeans: Returns ------- objective : float - Last pre-update objective. This can differ from final_obj, which uses - returned assignments and final centroids. + Last pre-update objective, over the sample when training sampled the + input. This can differ from final_obj, which uses all vectors, returned + assignments and final centroids. Notes ----- @@ -1224,6 +1309,42 @@ class RaBitQKMeans: copies. """ + def assign( + self, x: npt.ArrayLike + ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: + """Assign vectors to the trained centroids. + + Parameters + ---------- + x : array_like of float, shape (m, d) + Vectors in the original, unpadded coordinates, with finite + coordinates in the range accepted by train; m may be 0. Converted to + contiguous float32; a copy may be made. + + Returns + ------- + assignments : ndarray of uint32, shape (m,) + Centroid IDs. + distances : ndarray of float32, shape (m,) + Distances to the assigned centroids, as in the distances property. + + Raises + ------ + RuntimeError + If no train call has succeeded. + + Notes + ----- + Uses final_assignment and the other parameters of the last successful + train call, but no labels or search hints from it, so each vector's + result is independent of the other vectors in x. RaBitQKMeans encodes + vectors around the training center. Approximate QGKMeans builds a + centroid graph on first use and reuses it in later calls. For the + training vectors, the result reproduces assignments and distances with + Exact final assignment, with RaBitQKMeans, or when training sampled the + input; unsampled approximate QGKMeans training uses search hints. + """ + @property def centroids(self) -> npt.NDArray[np.float32] | None: """Final centroids, float32 array of shape (k, d), or None before training.""" @@ -1234,19 +1355,19 @@ class RaBitQKMeans: @property def distances(self) -> npt.NDArray[np.float32] | None: - """Final assignment distances, float32 array of shape (n,), or None before training.""" + """Final assignment distances for all n vectors, float32 array of shape (n,), or None before training.""" @property def obj(self) -> npt.NDArray[np.float64] | None: - """Pre-update objectives, float64 array of shape (iterations,), or None before training.""" + """Pre-update objectives over the vectors used for iterations, float64 array of shape (iterations,), or None before training.""" @property def final_obj(self) -> float | None: - """Objective for final assignments and centroids, or None before training.""" + """Objective for the final assignments of all n vectors, or None before training.""" @property def iteration_stats(self) -> list[QGKMeansIterationStats] | None: - """Per-iteration statistics, or None before training.""" + """Per-iteration statistics over the vectors used for iterations, or None before training.""" class RaBitQKMeansParameters: """Mutable training settings for RaBitQKMeans. @@ -1282,6 +1403,11 @@ class RaBitQKMeansParameters: """Warn below this many training vectors per centroid. Default is 39.""" @min_points_per_centroid.setter def min_points_per_centroid(self, value: SupportsIndex) -> None: ... + @property + def max_points_per_centroid(self) -> int: + """When positive, train on a random sample of k * max_points_per_centroid vectors if there are more, then assign all vectors. Default is 0, which trains on all vectors; Faiss uses 256.""" + @max_points_per_centroid.setter + def max_points_per_centroid(self, value: SupportsIndex) -> None: ... early_stop_threshold: float """Stop when the absolute relative objective change is at most this threshold. Default is 0.0.""" diff --git a/python_bindings/kmeans_bindings.cpp b/python_bindings/kmeans_bindings.cpp index bab75f8e..96077a8a 100644 --- a/python_bindings/kmeans_bindings.cpp +++ b/python_bindings/kmeans_bindings.cpp @@ -6,6 +6,7 @@ #include #include // IWYU pragma: keep (vector conversion for iteration_stats) +#include #include #include #include @@ -13,6 +14,7 @@ #include "bindings_common.hpp" #include "rabitqlib/clustering/qgkmeans.hpp" #include "rabitqlib/clustering/rabitqkmeans.hpp" +#include "rabitqlib/index/types.hpp" namespace py = pybind11; @@ -47,6 +49,28 @@ double train(Clustering& clustering, py::handle x) { return clustering.iteration_stats.back().obj; } +template +py::tuple assign(const Clustering& clustering, py::handle x) { + auto x_array = ensure_2d_array(x, "x"); + ensure_dimension(x_array, clustering.d, "x dimension does not match d"); + + const auto num_points = x_array.shape(0); + py::array_t assignments(num_points); + py::array_t distances(num_points); + PID* assignments_data = assignments.mutable_data(); + float* distances_data = distances.mutable_data(); + { + py::gil_scoped_release release; + clustering.assign( + static_cast(num_points), + x_array.data(), + assignments_data, + distances_data + ); + } + return py::make_tuple(assignments, distances); +} + template void bind_training_parameters(py::class_& binding) { binding.def(py::init<>(), "Create parameters with their documented defaults.") @@ -71,6 +95,11 @@ void bind_training_parameters(py::class_& binding) { &Parameters::min_points_per_centroid, R"doc(Warn below this many training vectors per centroid. Default is 39.)doc" ) + .def_readwrite( + "max_points_per_centroid", + &Parameters::max_points_per_centroid, + R"doc(When positive, train on a random sample of k * max_points_per_centroid vectors if there are more, then assign all vectors. Default is 0, which trains on all vectors; Faiss uses 256.)doc" + ) .def_readwrite( "early_stop_threshold", &Parameters::early_stop_threshold, @@ -93,6 +122,7 @@ void bind_clustering(py::module_& m, const char* native_name, const char* public py::class_(m, native_name) .def(py::init()) .def("train", &train, py::arg("x")) + .def("assign", &assign, py::arg("x")) .def_readonly("d", &Clustering::d) .def_readonly("k", &Clustering::k) .def_readonly("final_obj", &Clustering::final_obj) diff --git a/tests/python/test_kmeans_sampling.py b/tests/python/test_kmeans_sampling.py new file mode 100644 index 00000000..ee08fc28 --- /dev/null +++ b/tests/python/test_kmeans_sampling.py @@ -0,0 +1,262 @@ +"""Training-set sampling and assignment of new vectors for both k-means classes.""" + +import numpy as np +import pytest +from rabitqlib import ( + FinalAssignmentMode, + IvfIndex, + QGKMeans, + QGKMeansParameters, + RaBitQKMeans, + RaBitQKMeansParameters, +) + +DIM = 64 +# QGKMeans requires more centroids than its default graph degree of 32. +CLUSTERS = {QGKMeans: 33, RaBitQKMeans: 4} +MODES = [FinalAssignmentMode.Approximate, FinalAssignmentMode.Exact] + + +def make_model(model_type, **kwargs): + parameters = {"niter": 3, "num_threads": 1, "min_points_per_centroid": 1} + parameters.update(kwargs) + return model_type(DIM, CLUSTERS[model_type], **parameters) + + +def clustered(count, seed): + rng = np.random.default_rng(seed) + centers = rng.standard_normal((5, DIM)).astype(np.float32) * 4 + noise = rng.standard_normal((count, DIM), dtype=np.float32) + return centers[np.arange(count) % 5] + noise + + +def selected_distances(x, centroids, assignments): + residual = x.astype(np.float64) - centroids[assignments].astype(np.float64) + return np.einsum("ij,ij->i", residual, residual) + + +@pytest.mark.parametrize( + "parameters_type", [QGKMeansParameters, RaBitQKMeansParameters] +) +def test_max_points_per_centroid_disables_sampling_by_default(parameters_type): + parameters = parameters_type() + assert parameters.max_points_per_centroid == 0 + parameters.max_points_per_centroid = 256 + assert parameters.max_points_per_centroid == 256 + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_max_points_per_centroid_keyword_and_cp(model_type): + assert make_model(model_type).cp.max_points_per_centroid == 0 + model = make_model(model_type, max_points_per_centroid=np.int64(7)) + assert model.cp.max_points_per_centroid == 7 + model.cp.max_points_per_centroid = 0 + assert model.cp.max_points_per_centroid == 0 + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +@pytest.mark.parametrize("value", [-1, 1.5, "many"]) +def test_max_points_per_centroid_rejects_invalid_values(model_type, value): + with pytest.raises(TypeError) as error: + model_type(DIM, CLUSTERS[model_type], max_points_per_centroid=value) + assert str(error.value) == ( + f"{model_type.__name__}() could not convert parameter " + f"'max_points_per_centroid' from {type(value).__name__}" + ) + model = make_model(model_type) + with pytest.raises(TypeError): + model.cp.max_points_per_centroid = value + assert model.cp.max_points_per_centroid == 0 + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_assign_before_training_raises(model_type): + model = make_model(model_type) + with pytest.raises( + RuntimeError, match=f"^{model_type.__name__} has not been trained$" + ): + model.assign(np.zeros((1, DIM), dtype=np.float32)) + native = model._implementation_type(DIM, CLUSTERS[model_type], model.cp) + with pytest.raises( + RuntimeError, match=f"^{model_type.__name__} has not been trained$" + ): + native.assign(np.zeros((1, DIM), dtype=np.float32)) + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_assign_validates_input_like_train(model_type): + count = 2 * CLUSTERS[model_type] + valid = clustered(count, 3) + nonfinite = valid.copy() + nonfinite[-1, -1] = np.nan + out_of_range = valid.copy() + out_of_range[0, 0] = 1e8 + invalid_inputs = [ + valid[:, :-1], + valid[0], + valid[None], + nonfinite, + out_of_range, + "not numeric", + ] + model = make_model(model_type) + model.train(valid) + for invalid in invalid_inputs: + with pytest.raises((TypeError, ValueError)) as from_train: + make_model(model_type).train(invalid) + with pytest.raises(type(from_train.value)) as from_assign: + model.assign(invalid) + assert str(from_assign.value) == str(from_train.value) + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +@pytest.mark.parametrize("mode", MODES) +def test_assign_returns_documented_types_for_new_vectors(model_type, mode): + model = make_model(model_type, final_assignment=mode) + model.train(clustered(3 * CLUSTERS[model_type], 5)) + queries = clustered(45, 7) + # A float64, Fortran-ordered input exercises the converting copy. + result = model.assign(np.asfortranarray(queries, dtype=np.float64)) + assert isinstance(result, tuple) + assignments, distances = result + assert assignments.shape == distances.shape == (45,) + assert assignments.dtype == np.uint32 + assert distances.dtype == np.float32 + assert np.all(assignments < model.k) + np.testing.assert_allclose( + distances, + selected_distances(queries, model.centroids, assignments), + rtol=1e-5, + atol=1e-5, + ) + if mode == FinalAssignmentMode.Exact: + all_distances = np.stack( + [ + selected_distances(queries, model.centroids, np.full(45, cluster)) + for cluster in range(model.k) + ], + axis=1, + ) + np.testing.assert_array_equal(assignments, np.argmin(all_distances, axis=1)) + empty_assignments, empty_distances = model.assign(np.empty((0, DIM), np.float32)) + assert empty_assignments.shape == empty_distances.shape == (0,) + assert empty_assignments.dtype == np.uint32 + assert empty_distances.dtype == np.float32 + + +# assign() uses no labels or search hints from training. It reproduces train's +# final labels in Exact mode, for RaBitQKMeans (same training center) and for +# any sampled fit, whose final pass is assign() itself. Unsampled approximate +# QGKMeans training reuses graph search hints. +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +@pytest.mark.parametrize("mode", MODES) +@pytest.mark.parametrize("sampled", [False, True]) +@pytest.mark.parametrize("threads", [1, 4]) +def test_assign_on_training_data_matches_final_assignment( + model_type, mode, sampled, threads +): + clusters = CLUSTERS[model_type] + x = clustered(2 * clusters + int(sampled), 11) + model = make_model( + model_type, + final_assignment=mode, + max_points_per_centroid=2, + num_threads=threads, + ) + model.train(x) + assignments, distances = model.assign(x) + reproducible = ( + sampled or mode == FinalAssignmentMode.Exact or model_type is RaBitQKMeans + ) + if reproducible: + np.testing.assert_array_equal(assignments, model.assignments) + np.testing.assert_array_equal(distances, model.distances) + np.testing.assert_allclose( + distances, + selected_distances(x, model.centroids, assignments), + rtol=1e-5, + atol=1e-5, + ) + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +@pytest.mark.parametrize("mode", MODES) +def test_assign_is_independent_of_the_batch(model_type, mode): + model = make_model(model_type, final_assignment=mode, num_threads=4) + model.train(clustered(6 * CLUSTERS[model_type], 29)) + queries = clustered(70, 31) + assignments, distances = model.assign(queries) + for size in (1, 33): + parts = [model.assign(queries[i : i + size]) for i in range(0, 70, size)] + np.testing.assert_array_equal( + np.concatenate([p[0] for p in parts]), assignments + ) + np.testing.assert_array_equal(np.concatenate([p[1] for p in parts]), distances) + + +# C++ tests check the verbose sampling message, which is not flushed per line. +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_sampling_at_the_limit_matches_disabled_sampling(model_type): + limit = 2 * CLUSTERS[model_type] + x = clustered(limit, 13) + at_limit = make_model(model_type, max_points_per_centroid=2) + at_limit.train(x) + disabled = make_model(model_type, max_points_per_centroid=0) + disabled.train(x) + for name in ("centroids", "assignments", "distances", "obj"): + np.testing.assert_array_equal(getattr(at_limit, name), getattr(disabled, name)) + assert at_limit.final_obj == disabled.final_obj + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +@pytest.mark.parametrize("mode", MODES) +def test_sampled_fit_covers_all_points_and_is_seeded(model_type, mode): + clusters = CLUSTERS[model_type] + x = clustered(6 * clusters + 1, 17) + model = make_model(model_type, final_assignment=mode, max_points_per_centroid=2) + model.train(x) + assert model.assignments.shape == model.distances.shape == (len(x),) + assert model.final_obj == model.distances.sum(dtype=np.float64) + np.testing.assert_allclose( + model.distances, + selected_distances(x, model.centroids, model.assignments), + rtol=1e-5, + atol=1e-5, + ) + same_seed = make_model(model_type, final_assignment=mode, max_points_per_centroid=2) + same_seed.train(x) + for name in ("centroids", "assignments", "distances", "obj"): + np.testing.assert_array_equal(getattr(model, name), getattr(same_seed, name)) + other_seed = make_model( + model_type, final_assignment=mode, max_points_per_centroid=2, seed=43 + ) + other_seed.train(x) + assert not np.array_equal(other_seed.centroids, model.centroids) + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_failed_refit_keeps_previous_model_for_assign(model_type): + clusters = CLUSTERS[model_type] + x = clustered(4 * clusters, 19) + model = make_model(model_type, max_points_per_centroid=1) + model.train(x) + expected = model.assign(x) + invalid = x.copy() + invalid[-1, 0] = np.inf + with pytest.raises(ValueError, match="x must contain only finite values$"): + model.train(invalid) + for actual, previous in zip(model.assign(x), expected): + np.testing.assert_array_equal(actual, previous) + + +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_sampled_results_build_an_ivf_index(model_type): + clusters = CLUSTERS[model_type] + x = clustered(4 * clusters, 23) + model = make_model(model_type, max_points_per_centroid=1) + model.train(x) + index = IvfIndex(DIM, len(x), clusters, 4) + index.build(x, model.centroids, model.assignments) + ids, _ = index.search(x[:3], 5, nprobe=clusters) + for query, neighbors in enumerate(ids): + assert query in neighbors diff --git a/tests/python/typing/public_api.py b/tests/python/typing/public_api.py index 9696883c..4888b273 100644 --- a/tests/python/typing/public_api.py +++ b/tests/python/typing/public_api.py @@ -62,6 +62,10 @@ def check(x: npt.NDArray[np.float32], labels: npt.NDArray[np.uint32]) -> None: assert_type(model.distances, npt.NDArray[np.float32] | None) assert_type(model.obj, npt.NDArray[np.float64] | None) assert_type(model.final_obj, float | None) + assert_type( + model.assign(x), + tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]], + ) if model.iteration_stats is not None: assert_type(model.iteration_stats[0], QGKMeansIterationStats) qg.cp.final_assignment = FinalAssignmentMode.Approximate @@ -83,7 +87,9 @@ def check_parameter_reads( # the assigned NumPy scalar, although the native setter normalizes it. assert_type(parameters.niter, int) assert_type(parameters.verbose, bool) + assert_type(parameters.max_points_per_centroid, int) assert type(parameters.niter) is int + assert type(parameters.max_points_per_centroid) is int assert type(parameters.verbose) is bool @@ -133,11 +139,14 @@ def check_numpy_scalars() -> None: model.cp.niter = np.int64(1) model.cp.seed = np.uint32(42) model.cp.min_points_per_centroid = np.int64(1) + model.cp.max_points_per_centroid = np.int64(1) model.cp.num_threads = np.int64(1) model.cp.verbose = np.bool_(False) model.cp.spherical = np.bool_(False) check_parameter_reads(model.cp) assert_type(model.train(x), float) + assignments, distances = model.assign(x[:3]) + assert assignments.shape == distances.shape == (3,) qg.cp.graph_degree = np.int64(32) qg.cp.ef_build = np.int64(40) qg.cp.ef_search = np.int64(16) diff --git a/tests/unit/rabitqlib/clustering/centroid_update_test.cpp b/tests/unit/rabitqlib/clustering/centroid_update_test.cpp index 94325d28..1f9960de 100644 --- a/tests/unit/rabitqlib/clustering/centroid_update_test.cpp +++ b/tests/unit/rabitqlib/clustering/centroid_update_test.cpp @@ -47,6 +47,23 @@ class RecordedAssignment { distances[point] = 0; } } + + // These tests neither sample nor call assign() after training. + struct ModelState {}; + void store_model_state(ModelState&) const {} + static bool has_model_state(const ModelState&, size_t) { return false; } + static void assign_with_state( + size_t, + size_t, + const float*, + const float*, + size_t, + uint32_t, + const RecordedParameters&, + const ModelState&, + PID*, + float* + ) {} }; TEST(CentroidUpdateTest, ShiftIncludesNormalizedAndRefilledCentroids) { diff --git a/tests/unit/rabitqlib/clustering/kmeans_sampling_test.cpp b/tests/unit/rabitqlib/clustering/kmeans_sampling_test.cpp new file mode 100644 index 00000000..3ed8a515 --- /dev/null +++ b/tests/unit/rabitqlib/clustering/kmeans_sampling_test.cpp @@ -0,0 +1,634 @@ +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/clustering/qgkmeans.hpp" +#include "rabitqlib/clustering/rabitqkmeans.hpp" +#include "rabitqlib/index/types.hpp" + +namespace rabitqlib { +namespace { + +using clustering::FinalAssignmentMode; +using clustering::detail::sample_training_points; + +constexpr size_t kDim = 64; + +std::vector clustered_points(size_t num_points, uint32_t seed) { + constexpr size_t kGroups = 5; + std::mt19937 rng(seed); + std::normal_distribution normal(0.0F, 1.0F); + std::vector centers(kGroups * kDim); + for (auto& value : centers) { + value = 4.0F * normal(rng); + } + std::vector points(num_points * kDim); + for (size_t point = 0; point < num_points; ++point) { + for (size_t dim = 0; dim < kDim; ++dim) { + points[point * kDim + dim] = + centers[(point % kGroups) * kDim + dim] + normal(rng); + } + } + return points; +} + +rabitqkmeans::RaBitQKMeansParameters test_parameters(const rabitqkmeans::RaBitQKMeans*) { + rabitqkmeans::RaBitQKMeansParameters parameters; + parameters.niter = 4; + parameters.num_threads = 1; + parameters.min_points_per_centroid = 1; + return parameters; +} + +qgkmeans::QGKMeansParameters test_parameters(const qgkmeans::QGKMeans*) { + qgkmeans::QGKMeansParameters parameters; + parameters.niter = 4; + parameters.num_threads = 1; + parameters.min_points_per_centroid = 1; + parameters.ef_build = 64; + parameters.ef_search = 16; + return parameters; +} + +template +auto test_parameters() { + return test_parameters(static_cast(nullptr)); +} + +// QGKMeans requires more centroids than its default graph degree of 32. +template +constexpr size_t kClusters = std::is_same_v ? 33 : 7; + +template +void expect_identical_results(const Clustering& actual, const Clustering& expected) { + EXPECT_EQ(actual.centroids, expected.centroids); + EXPECT_EQ(actual.assignments, expected.assignments); + EXPECT_EQ(actual.distances, expected.distances); + EXPECT_EQ(actual.final_obj, expected.final_obj); + ASSERT_EQ(actual.iteration_stats.size(), expected.iteration_stats.size()); + for (size_t i = 0; i < actual.iteration_stats.size(); ++i) { + EXPECT_EQ(actual.iteration_stats[i].obj, expected.iteration_stats[i].obj); + EXPECT_EQ(actual.iteration_stats[i].shift, expected.iteration_stats[i].shift); + EXPECT_EQ(actual.iteration_stats[i].nsplit, expected.iteration_stats[i].nsplit); + } +} + +double reference_distance(const float* x, const float* centroid, bool spherical) { + double distance = 0; + for (size_t dim = 0; dim < kDim; ++dim) { + const double left = x[dim]; + const double right = centroid[dim]; + distance += spherical ? -left * right : (left - right) * (left - right); + } + return spherical ? 1.0 + distance : distance; +} + +// Distances belong to every input point and its returned centroid; Exact mode +// additionally returns the lowest-ID nearest centroid. +template +void expect_points_assigned( + const Clustering& clustering, + const std::vector& points, + const std::vector& assignments, + const std::vector& distances +) { + const size_t num_points = points.size() / kDim; + ASSERT_EQ(assignments.size(), num_points); + ASSERT_EQ(distances.size(), num_points); + for (size_t point = 0; point < num_points; ++point) { + const PID cluster_id = assignments[point]; + ASSERT_LT(cluster_id, clustering.k); + const float* vector = points.data() + point * kDim; + const double distance = reference_distance( + vector, clustering.centroids.data() + cluster_id * kDim, clustering.spherical + ); + EXPECT_NEAR(distances[point], distance, 1e-5 * std::max(1.0, distance)); + if (clustering.final_assignment != FinalAssignmentMode::Exact) { + continue; + } + double best = std::numeric_limits::infinity(); + PID best_id = 0; + for (size_t cluster = 0; cluster < clustering.k; ++cluster) { + const double candidate = reference_distance( + vector, clustering.centroids.data() + cluster * kDim, clustering.spherical + ); + if (candidate < best) { + best = candidate; + best_id = static_cast(cluster); + } + } + EXPECT_EQ(cluster_id, best_id) << point; + } +} + +template +void expect_all_points_assigned( + const Clustering& clustering, const std::vector& points, size_t num_points +) { + ASSERT_EQ(points.size(), num_points * kDim); + expect_points_assigned( + clustering, points, clustering.assignments, clustering.distances + ); + EXPECT_EQ( + clustering.final_obj, + std::accumulate(clustering.distances.begin(), clustering.distances.end(), 0.0) + ); +} + +template +std::pair, std::vector> assign_all( + const Clustering& clustering, const std::vector& points +) { + const size_t num_points = points.size() / kDim; + std::vector assignments(num_points); + std::vector distances(num_points); + clustering.assign(num_points, points.data(), assignments.data(), distances.data()); + return {assignments, distances}; +} + +template +void check_assign_requires_training_and_valid_input(const std::string& name) { + constexpr size_t kClusterCount = kClusters; + const auto points = clustered_points(2 * kClusterCount, 61); + Clustering clustering(kDim, kClusterCount, test_parameters()); + PID assignment = 0; + float distance = 0; + // Centroids written without train() lack the state that assignment uses. + for (const bool with_centroids : {false, true}) { + if (with_centroids) { + clustering.centroids.assign( + points.begin(), points.begin() + kClusterCount * kDim + ); + } + try { + clustering.assign(1, points.data(), &assignment, &distance); + FAIL() << "Assigned before training"; + } catch (const std::runtime_error& error) { + EXPECT_EQ(error.what(), name + " has not been trained"); + } + } + + clustering.train(2 * kClusterCount, points.data()); + EXPECT_NO_THROW(clustering.assign(0, nullptr, nullptr, nullptr)); + const auto expect_invalid = [&](const float* x, const std::string& message) { + try { + clustering.assign(1, x, &assignment, &distance); + FAIL() << "Accepted " << message; + } catch (const std::invalid_argument& error) { + EXPECT_EQ(error.what(), name + message); + } + }; + expect_invalid(nullptr, " assign input and output pointers must not be null"); + try { + clustering.assign(1, points.data(), nullptr, &distance); + FAIL() << "Accepted a null output"; + } catch (const std::invalid_argument& error) { + EXPECT_EQ( + error.what(), name + " assign input and output pointers must not be null" + ); + } + std::vector invalid(points.begin(), points.begin() + kDim); + invalid.back() = std::numeric_limits::infinity(); + expect_invalid(invalid.data(), " x must contain only finite values"); + invalid.back() = 1e8F; + expect_invalid(invalid.data(), " x exceeds the safe float32 coordinate range"); + + // Public centroids are validated as inputs, also before an empty assignment. + const auto trained = clustering.centroids; + for (const float value : {std::numeric_limits::quiet_NaN(), 1e8F}) { + clustering.centroids = trained; + clustering.centroids[kDim + 1] = value; + const std::string message = + std::isnan(value) ? " centroids must contain only finite values" + : " centroids exceeds the safe float32 coordinate range"; + for (const size_t count : {size_t{0}, size_t{1}}) { + assignment = kPidMax; + try { + clustering.assign(count, points.data(), &assignment, &distance); + FAIL() << "Accepted" << message; + } catch (const std::invalid_argument& error) { + EXPECT_EQ(error.what(), name + message); + } + EXPECT_EQ(assignment, kPidMax); + } + } + clustering.centroids = trained; + EXPECT_NO_THROW(clustering.assign(1, points.data(), &assignment, &distance)); + + // A centroid buffer that no longer matches k is not a trained model. + clustering.k = kClusterCount - 1; + EXPECT_THROW( + clustering.assign(1, points.data(), &assignment, &distance), std::runtime_error + ); +} + +// With sampling, train's final pass is assign() itself, including QGKMeans' +// cached graph. Without it, Exact assignment and RaBitQKMeans' scan around the +// stored training center also reproduce train(); graph training uses hints. +template +void check_assign_matches_final_training_assignment() { + constexpr bool kGraph = std::is_same_v; + constexpr size_t kClusterCount = kClusters; + constexpr size_t kMaxPoints = 2; + for (const uint32_t threads : {1U, 4U}) { + for (const uint32_t bits : + kGraph ? std::vector{0, 4, 8} : std::vector{0}) { + for (const size_t num_points : + {kClusterCount * kMaxPoints, 3 * kClusterCount * kMaxPoints + 1}) { + const auto points = clustered_points(num_points, 67); + for (const auto mode : + {FinalAssignmentMode::Approximate, FinalAssignmentMode::Exact}) { + SCOPED_TRACE( + ::testing::Message() << threads << "/" << bits << "/" << num_points + << "/" << static_cast(mode) + ); + auto parameters = test_parameters(); + parameters.num_threads = threads; + if constexpr (kGraph) { + parameters.quantization_bits = bits; + } + parameters.final_assignment = mode; + parameters.max_points_per_centroid = kMaxPoints; + Clustering clustering(kDim, kClusterCount, parameters); + clustering.train(num_points, points.data()); + const bool sampled = num_points > kClusterCount * kMaxPoints; + const auto [assignments, distances] = assign_all(clustering, points); + if (sampled || mode == FinalAssignmentMode::Exact || !kGraph) { + EXPECT_EQ(assignments, clustering.assignments); + EXPECT_EQ(distances, clustering.distances); + } + expect_points_assigned(clustering, points, assignments, distances); + + const auto unseen = clustered_points(45, 71); + const auto [new_assignments, new_distances] = + assign_all(clustering, unseen); + expect_points_assigned( + clustering, unseen, new_assignments, new_distances + ); + } + } + } + } +} + +// Each vector's label depends only on the model, not on the other vectors in +// a call or on earlier calls. +template +void check_assign_is_independent_of_the_batch() { + constexpr size_t kClusterCount = kClusters; + const auto points = clustered_points(6 * kClusterCount, 73); + const auto queries = clustered_points(70, 79); + for (const auto mode : {FinalAssignmentMode::Approximate, FinalAssignmentMode::Exact}) { + SCOPED_TRACE(static_cast(mode)); + auto parameters = test_parameters(); + parameters.num_threads = 4; + parameters.final_assignment = mode; + Clustering clustering(kDim, kClusterCount, parameters); + clustering.train(6 * kClusterCount, points.data()); + const auto [assignments, distances] = assign_all(clustering, queries); + for (const size_t split : {size_t{1}, size_t{33}}) { + std::vector split_assignments(queries.size() / kDim); + std::vector split_distances(split_assignments.size()); + for (size_t first = 0; first < split_assignments.size(); first += split) { + const size_t count = std::min(split, split_assignments.size() - first); + clustering.assign( + count, + queries.data() + first * kDim, + split_assignments.data() + first, + split_distances.data() + first + ); + } + EXPECT_EQ(split_assignments, assignments) << split; + EXPECT_EQ(split_distances, distances) << split; + } + } +} + +// The cached graph follows the public centroids and graph settings. A stale +// graph would report distances to the previous centroids. +void check_graph_cache_follows_centroids_and_settings() { + using qgkmeans::QGKMeans; + constexpr size_t kClusterCount = kClusters; + const auto points = clustered_points(6 * kClusterCount, 83); + const auto queries = clustered_points(60, 89); + for (const uint32_t bits : {0U, 4U}) { + SCOPED_TRACE(bits); + auto parameters = test_parameters(); + parameters.quantization_bits = bits; + QGKMeans clustering(kDim, kClusterCount, parameters); + clustering.train(6 * kClusterCount, points.data()); + QGKMeans fresh = clustering; // Copied before any graph is cached. + const auto first = assign_all(clustering, queries); + EXPECT_EQ(assign_all(clustering, queries), first); + EXPECT_EQ(assign_all(fresh, queries), first); + + fresh = clustering; + for (QGKMeans* model : {&clustering, &fresh}) { + std::swap_ranges( + model->centroids.begin(), + model->centroids.begin() + kDim, + model->centroids.begin() + kDim + ); + for (size_t cluster = 0; cluster < kClusterCount; ++cluster) { + model->centroids[cluster * kDim] += 0.5F; + } + } + const auto [moved_assignments, moved_distances] = assign_all(clustering, queries); + expect_points_assigned(clustering, queries, moved_assignments, moved_distances); + EXPECT_EQ( + assign_all(fresh, queries), std::make_pair(moved_assignments, moved_distances) + ); + + clustering.ef_search = fresh.ef_search = 33; + const auto wider = assign_all(clustering, queries); + expect_points_assigned(clustering, queries, wider.first, wider.second); + EXPECT_EQ(assign_all(fresh, queries), wider); + } +} + +template +void check_boundary_without_sampling() { + constexpr size_t kClusterCount = kClusters; + constexpr size_t kMaxPoints = 2; + const size_t num_points = kClusterCount * kMaxPoints; + const auto points = clustered_points(num_points, 11); + for (const auto mode : {FinalAssignmentMode::Approximate, FinalAssignmentMode::Exact}) { + SCOPED_TRACE(static_cast(mode)); + auto parameters = test_parameters(); + parameters.final_assignment = mode; + parameters.max_points_per_centroid = kMaxPoints; + Clustering at_limit(kDim, kClusterCount, parameters); + at_limit.train(num_points, points.data()); + parameters.max_points_per_centroid = 0; + Clustering disabled(kDim, kClusterCount, parameters); + disabled.train(num_points, points.data()); + expect_identical_results(at_limit, disabled); + expect_all_points_assigned(at_limit, points, num_points); + } +} + +// Verbose training reports sampling only when n exceeds k * max_points_per_centroid. +template +void check_verbose_sampling_threshold(const std::string& name) { + constexpr size_t kClusterCount = kClusters; + constexpr size_t kMaxPoints = 2; + constexpr size_t kLimit = kClusterCount * kMaxPoints; + const auto points = clustered_points(kLimit + 1, 13); + auto parameters = test_parameters(); + parameters.niter = 1; + parameters.verbose = true; + parameters.max_points_per_centroid = kMaxPoints; + for (const size_t num_points : {kLimit, kLimit + 1}) { + SCOPED_TRACE(num_points); + Clustering clustering(kDim, kClusterCount, parameters); + testing::internal::CaptureStdout(); + clustering.train(num_points, points.data()); + const std::string output = testing::internal::GetCapturedStdout(); + const std::string message = name + " sampling " + std::to_string(kLimit) + " of " + + std::to_string(num_points) + " training points\n"; + EXPECT_EQ(output.find(message) != std::string::npos, num_points > kLimit) << output; + EXPECT_EQ(output.find("sampling") != std::string::npos, num_points > kLimit) + << output; + } +} + +// Initialization and Lloyd iterations see only the sample, so they match +// unsampled training on the gathered sample; the final assignment covers all points. +template +void check_sampled_training() { + constexpr size_t kClusterCount = kClusters; + constexpr size_t kMaxPoints = 2; + constexpr size_t kNumSamples = kClusterCount * kMaxPoints; + for (const size_t num_points : {kNumSamples + 1, 3 * kNumSamples + 5}) { + const auto points = clustered_points(num_points, 23); + for (const bool spherical : {false, true}) { + for (const auto mode : + {FinalAssignmentMode::Approximate, FinalAssignmentMode::Exact}) { + SCOPED_TRACE( + ::testing::Message() + << num_points << "/" << spherical << "/" << static_cast(mode) + ); + auto parameters = test_parameters(); + parameters.spherical = spherical; + parameters.final_assignment = mode; + parameters.max_points_per_centroid = kMaxPoints; + Clustering sampled(kDim, kClusterCount, parameters); + sampled.train(num_points, points.data()); + + const auto indices = + sample_training_points(num_points, kNumSamples, parameters.seed); + std::vector samples; + for (const size_t index : indices) { + samples.insert( + samples.end(), + points.begin() + static_cast(index * kDim), + points.begin() + static_cast((index + 1) * kDim) + ); + } + parameters.max_points_per_centroid = 0; + Clustering reference(kDim, kClusterCount, parameters); + reference.train(kNumSamples, samples.data()); + + EXPECT_EQ(sampled.centroids, reference.centroids); + ASSERT_EQ(sampled.iteration_stats.size(), reference.iteration_stats.size()); + for (size_t i = 0; i < sampled.iteration_stats.size(); ++i) { + EXPECT_EQ( + sampled.iteration_stats[i].obj, reference.iteration_stats[i].obj + ); + EXPECT_EQ( + sampled.iteration_stats[i].nsplit, + reference.iteration_stats[i].nsplit + ); + } + expect_all_points_assigned(sampled, points, num_points); + } + } + } +} + +template +void check_reproducible_sampling() { + constexpr size_t kClusterCount = kClusters; + const size_t num_points = 6 * kClusterCount + 3; + const auto points = clustered_points(num_points, 37); + for (const auto mode : {FinalAssignmentMode::Approximate, FinalAssignmentMode::Exact}) { + SCOPED_TRACE(static_cast(mode)); + auto parameters = test_parameters(); + parameters.final_assignment = mode; + parameters.max_points_per_centroid = 2; + Clustering first(kDim, kClusterCount, parameters); + first.train(num_points, points.data()); + Clustering second(kDim, kClusterCount, parameters); + second.train(num_points, points.data()); + expect_identical_results(first, second); + + parameters.seed += 1; + Clustering reseeded(kDim, kClusterCount, parameters); + reseeded.train(num_points, points.data()); + EXPECT_NE(reseeded.centroids, first.centroids); + } +} + +// Validation covers every input point, not only those selected for training. +template +void check_rejects_nonfinite_values_outside_sample(const std::string& name) { + constexpr size_t kClusterCount = kClusters; + const size_t num_points = 3 * kClusterCount; + auto parameters = test_parameters(); + parameters.max_points_per_centroid = 1; + const auto points = clustered_points(num_points, 41); + Clustering clustering(kDim, kClusterCount, parameters); + clustering.train(num_points, points.data()); + const auto expected = clustering; + + const auto indices = sample_training_points(num_points, kClusterCount, parameters.seed); + size_t outside = 0; + while (std::binary_search(indices.begin(), indices.end(), outside)) { + ++outside; + } + auto invalid = points; + invalid[outside * kDim] = std::numeric_limits::quiet_NaN(); + try { + clustering.train(num_points, invalid.data()); + FAIL() << "Accepted nonfinite input outside the sample"; + } catch (const std::invalid_argument& error) { + EXPECT_EQ(error.what(), name + " x must contain only finite values"); + } + expect_identical_results(clustering, expected); +} + +// Retraining on a subrange of the previous centroids samples from an owned copy. +template +void check_sampled_retraining_on_own_centroids() { + constexpr size_t kClusterCount = kClusters; + constexpr size_t kInitialClusters = 2 * kClusterCount + 4; + const auto points = clustered_points(3 * kInitialClusters, 53); + auto parameters = test_parameters(); + Clustering clustering(kDim, kInitialClusters, parameters); + clustering.train(3 * kInitialClusters, points.data()); + const auto previous = clustering.centroids; + const size_t count = kInitialClusters - 1; + + clustering.max_points_per_centroid = 1; + parameters.max_points_per_centroid = 1; + Clustering reference(kDim, kClusterCount, parameters); + reference.train(count, previous.data() + kDim); + clustering.k = kClusterCount; + clustering.train(count, clustering.centroids.data() + kDim); + expect_identical_results(clustering, reference); + ASSERT_EQ(clustering.assignments.size(), count); +} + +TEST(KMeansSamplingTest, IsDisabledByDefault) { + EXPECT_EQ(rabitqkmeans::RaBitQKMeansParameters{}.max_points_per_centroid, 0U); + EXPECT_EQ(qgkmeans::QGKMeansParameters{}.max_points_per_centroid, 0U); +} + +TEST(KMeansSamplingTest, SampleIndicesAreSortedDistinctAndSeeded) { + const auto indices = sample_training_points(1000, 100, 42); + ASSERT_EQ(indices.size(), 100U); + EXPECT_TRUE(std::is_sorted(indices.begin(), indices.end())); + EXPECT_EQ(std::adjacent_find(indices.begin(), indices.end()), indices.end()); + EXPECT_LT(indices.back(), 1000U); + EXPECT_EQ(sample_training_points(1000, 100, 42), indices); + EXPECT_NE(sample_training_points(1000, 100, 43), indices); + + std::vector all(10); + std::iota(all.begin(), all.end(), size_t{0}); + EXPECT_EQ(sample_training_points(10, 10, 42), all); +} + +TEST(KMeansSamplingTest, RaBitQBoundaryWithoutSamplingMatchesDisabledSampling) { + check_boundary_without_sampling(); +} + +TEST(KMeansSamplingTest, QGBoundaryWithoutSamplingMatchesDisabledSampling) { + check_boundary_without_sampling(); +} + +TEST(KMeansSamplingTest, RaBitQSamplesOnlyAboveTheLimit) { + check_verbose_sampling_threshold("RaBitQKMeans"); +} + +TEST(KMeansSamplingTest, QGSamplesOnlyAboveTheLimit) { + check_verbose_sampling_threshold("QGKMeans"); +} + +TEST(KMeansSamplingTest, RaBitQTrainsOnSampleAndAssignsAllPoints) { + check_sampled_training(); +} + +TEST(KMeansSamplingTest, QGTrainsOnSampleAndAssignsAllPoints) { + check_sampled_training(); +} + +TEST(KMeansSamplingTest, RaBitQSamplingIsReproducibleForASeed) { + check_reproducible_sampling(); +} + +TEST(KMeansSamplingTest, QGSamplingIsReproducibleForASeed) { + check_reproducible_sampling(); +} + +TEST(KMeansSamplingTest, RaBitQRejectsNonfiniteValuesOutsideTheSample) { + check_rejects_nonfinite_values_outside_sample( + "RaBitQKMeans" + ); +} + +TEST(KMeansSamplingTest, QGRejectsNonfiniteValuesOutsideTheSample) { + check_rejects_nonfinite_values_outside_sample("QGKMeans"); +} + +TEST(KMeansAssignTest, RaBitQRequiresTrainingAndValidInput) { + check_assign_requires_training_and_valid_input( + "RaBitQKMeans" + ); +} + +TEST(KMeansAssignTest, QGRequiresTrainingAndValidInput) { + check_assign_requires_training_and_valid_input("QGKMeans"); +} + +TEST(KMeansAssignTest, RaBitQMatchesFinalTrainingAssignment) { + check_assign_matches_final_training_assignment(); +} + +TEST(KMeansAssignTest, QGMatchesFinalTrainingAssignment) { + check_assign_matches_final_training_assignment(); +} + +TEST(KMeansAssignTest, RaBitQIsIndependentOfTheBatch) { + check_assign_is_independent_of_the_batch(); +} + +TEST(KMeansAssignTest, QGIsIndependentOfTheBatch) { + check_assign_is_independent_of_the_batch(); +} + +TEST(KMeansAssignTest, QGGraphCacheFollowsCentroidsAndSettings) { + check_graph_cache_follows_centroids_and_settings(); +} + +TEST(KMeansSamplingTest, RaBitQSampledRetrainingOnOwnCentroids) { + check_sampled_retraining_on_own_centroids(); +} + +TEST(KMeansSamplingTest, QGSampledRetrainingOnOwnCentroids) { + check_sampled_retraining_on_own_centroids(); +} + +} // namespace +} // namespace rabitqlib