mojolearn.IVFSQIndex#

class mojolearn.IVFSQIndex(n_lists, n_probes, n_neighbors=8, kmeans_n_iters=20, random_state=0)[source]#

Bases: _AnnResident, _AnnSaved, NumericModeMixin

Approximate nearest-neighbor index with inverted lists and 8-bit scalar quantization (IVF-SQ).

Follows cuVS’s ivf_sq. fit trains the same k-means coarse quantizer as IVFPQIndex, then stores each row’s residual (row minus its list center) quantized to 8 bits per feature, using cuVS’s per-feature range with a 5% margin. search scans the decoded residuals of the n_probes lists nearest each query and returns approximate squared Euclidean distances. 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.

  • 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.

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

  • 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.

  • sq_vmin_ (ndarray of shape (n_features,), float32) – Lower end of each feature’s quantization range.

  • sq_delta_ (ndarray of shape (n_features,), float32) – Quantization step of each feature.

  • codes_ (ndarray of shape (n_samples, n_features), int32) – 8-bit code of every residual value.

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

  • n_rows_ (int) – Number of indexed rows.

  • n_lists_ (int) – The fitted value of n_lists.

  • 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 is out of range.

Examples

>>> import numpy as np, mojolearn
>>> X = np.random.default_rng(0).random((1000, 8), dtype=np.float32)
>>> index = mojolearn.IVFSQIndex(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, n_neighbors=8, 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 (IVFSQIndex) – 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.