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