mojolearn.IVFPQIndex#

class mojolearn.IVFPQIndex(n_lists, n_probes, pq_dim=4, pq_bits=8, n_neighbors=8, kmeans_n_iters=20, pq_kmeans_n_iters=20, random_state=0)[source]#

Bases: _AnnResident, _AnnSaved, NumericModeMixin

Approximate nearest-neighbor index with inverted lists and product quantization (IVF-PQ).

Follows cuVS’s ivf_pq. fit trains a k-means coarse quantizer that splits the rows into n_lists lists, then encodes each row with product quantization: the features are split into pq_dim subspaces and each subspace is stored as one of 2**pq_bits codebook entries. search scans the n_probes lists nearest each query and returns approximate squared Euclidean distances computed from the codes. Codebook training and code sums run in a fixed order and ties are broken by row id, so the results are reproducible. On a GPU the fitted index stays on the device between searches; each search uploads only its queries.

Parameters:

  • n_lists (int) – Number of inverted lists (k-means clusters). Must be in [1, n_samples].

  • n_probes (int) – Number of lists scanned per query. Required, with no default; a value greater than n_lists raises an error rather than being clamped.

  • pq_dim (int, default 4) – Number of PQ subspaces. Must be in [1, n_features]; the features are zero-padded up to a multiple of pq_dim.

  • pq_bits (int, default 8) – Bits per code, so each subspace has 2**pq_bits codebook entries. Must be in [1, 8].

  • n_neighbors (int, default 8) – Number of neighbors search returns per query.

  • kmeans_n_iters (int, default 20) – k-means iterations for the coarse quantizer.

  • pq_kmeans_n_iters (int, default 20) – k-means iterations for each subspace codebook.

  • random_state (int, default 0) – Seed for the k-means initializations.

  • numeric_mode (str or None, default None) – Numeric mode for this estimator; None uses the package default set with mojolearn.set_numeric_mode.

Variables:

  • centers_ (ndarray of shape (n_lists, n_features), float32) – Coarse quantizer centers.

  • list_offsets_ (ndarray of shape (n_lists + 1,), int32) – Start of each list in list_indices_.

  • list_indices_ (ndarray of shape (n_samples,), int32) – Original row ids, grouped by list.

  • codebooks_ (ndarray of shape (pq_dim, 2**pq_bits, pq_len_), float32) – One codebook per subspace.

  • codes_ (ndarray of shape (n_samples, pq_dim), int32) – PQ code of every indexed row.

  • n_features_in_ (int) – Number of features seen during fit.

  • n_rows_ (int) – Number of indexed rows.

  • pq_bits_ (n_lists_, pq_dim_,) – The fitted values of n_lists, pq_dim and pq_bits.

  • pq_len_ (int) – Features per subspace, ceil(n_features / pq_dim).

  • n_candidates_ (ndarray of shape (n_queries,), int32) – Set by search: how many indexed rows each query scanned.

Raises:

  • TypeError – If an integer parameter is not an int.

  • ValueError – If n_lists, pq_dim or pq_bits is out of range.

See also

IVFIndex

IVF without compression (exact distances within probed lists).

refine

Re-rank candidates with exact distances.

Examples

>>> import numpy as np, mojolearn
>>> X = np.random.default_rng(0).random((1000, 8), dtype=np.float32)
>>> index = mojolearn.IVFPQIndex(n_lists=16, n_probes=4, n_neighbors=5).fit(X)
>>> distances, indices = index.search(X[:10])
>>> indices.shape
(10, 5)
__init__(n_lists, n_probes, pq_dim=4, pq_bits=8, n_neighbors=8, kmeans_n_iters=20, pq_kmeans_n_iters=20, random_state=0)[source]#

Methods

fit(X[, y])

Build the index from X.

get_params([deep])

load(path)

Load an index written by save.

numeric_mode_used()

The numeric mode this estimator runs in.

save(path)

Write the fitted index to path as an .npz file.

search(queries[, filter])

Find the approximate nearest indexed rows of each query.

set_params(**params)

Set constructor parameters.

vendor_used()

The GPU backend this estimator runs on.
fit(X, y=None)[source]#

Build the index from X.

Parameters:

  • X (array-like of shape (n_samples, n_features)) – Rows to index, converted to float32.

  • y (None) – Ignored.

Returns:

self (IVFPQIndex) – The fitted index.

search(queries, filter=None)[source]#

Find the approximate nearest indexed rows of each query.

Parameters:

  • queries (array-like of shape (n_queries, n_features)) – Query rows, converted to float32.

  • filter (array-like of bool, shape (n_rows_,), optional) – One flag per indexed row; a row flagged False is never returned (cuVS’s sample filter, applied before scoring).

Returns:

  • distances (ndarray of shape (n_queries, n_neighbors), float32) – Approximate squared Euclidean distances, ascending.

  • indices (ndarray of shape (n_queries, n_neighbors), int32) – Original row ids of the neighbors. A query whose probed lists hold fewer than n_neighbors rows is padded with distance inf and id -1.

Raises:

ValueError – If the index is not fitted, the queries have a different number of features, or filter has the wrong shape or dtype.