mojolearn.IVFIndex#

class mojolearn.IVFIndex(n_lists, n_probes, n_neighbors=8, kmeans_n_iters=20, metric='sqeuclidean', random_state=0)[source]#

Bases: NumericModeMixin

Approximate nearest-neighbor index with inverted lists (IVF-Flat).

Follows cuVS’s ivf_flat. fit trains a k-means coarse quantizer that splits the rows into n_lists lists and stores the rows uncompressed, grouped by list. search scans the n_probes lists nearest each query and returns exact distances to the rows in those lists, so recall depends only on n_probes. Results are reproducible. On a GPU the fitted index stays on the device between searches; each search uploads only its queries. IVFFlat is another name for this class.

Parameters:

  • n_lists (int) – Number of inverted lists (k-means clusters). Must be at least 1.

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

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

  • metric ({‘sqeuclidean’, ‘euclidean’}, default ‘sqeuclidean’) – ‘sqeuclidean’ (cuVS L2Expanded) returns squared Euclidean distances; ‘euclidean’ (cuVS L2SqrtExpanded) returns their square roots. ‘l2_expanded’ is accepted for ‘sqeuclidean’, and ‘l2_sqrt_expanded’ and ‘l2’ for ‘euclidean’. The metric is fixed when the index is built: on one index the ids are the same under both, but the k-means quantizer trains differently under each, so two indexes built with different metrics can return different neighbors when n_probes < n_lists.

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

  • center_norms_ (ndarray of shape (n_lists,), float32) – Squared norm of each center.

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

  • list_indices_ (ndarray of shape (n_rows_,), int32) – Original row ids, grouped by list and ascending within each list.

  • list_data_ (ndarray of shape (n_rows_, n_features), float32) – The indexed rows, in the order of list_indices_.

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

  • n_rows_ (int) – Number of indexed rows, including rows added by extend.

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

  • metric_code_ (int) – The metric the index was built under: 0 for ‘sqeuclidean’, 1 for ‘euclidean’.

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

  • extend_labels_ (ndarray of shape (n_new,), int32) – Set by extend: the list each new row joined.

Raises:

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

  • ValueError – If n_lists is less than 1 or metric is unknown.

See also

IVFPQIndex, IVFSQIndex, IVFRaBitQIndex

CagraIndex

Graph-based index.

Examples

>>> import numpy as np, mojolearn
>>> X = np.random.default_rng(0).random((1000, 8), dtype=np.float32)
>>> index = mojolearn.IVFIndex(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, metric='sqeuclidean', random_state=0)[source]#

Methods

extend(X)

Add the rows of X to the built index.

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 built index to path as an .npz file.

search(queries[, filter])

Find the nearest indexed rows of each query within the probed lists.

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.

Building needs a GPU install; on a CPU-only install, load a saved index instead.

Parameters:

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

  • y (None) – Ignored.

Returns:

self (IVFIndex) – The fitted index.

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

Find the nearest indexed rows of each query within the probed lists.

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, by original row id; a row flagged False is never scored, returned or counted in n_candidates_ (cuVS’s sample filter, applied before any distance is computed).

Returns:

  • distances (ndarray of shape (n_queries, n_neighbors), float32) – Distances under metric, ascending.

  • indices (ndarray of shape (n_queries, n_neighbors), int32) – Original row ids of the neighbors.

Raises:

ValueError – If the index is not fitted or loaded, the queries have a different number of features, metric no longer names the metric the index was built under, or filter has the wrong shape or dtype.

Notes

Unlike the compressed IVF indexes, a query whose probed lists hold (after filtering) fewer than n_neighbors rows raises an error rather than being padded. Raise n_probes if this happens.

extend(X)[source]#

Add the rows of X to the built index.

Follows cuVS’s ivf_flat::extend with adaptive_centers=False: each new row joins the list of its nearest center (an exact tie goes to the lower list id), and the centers do not move. The new rows take the ids n_rows_, ``n_rows_ + 1, …`` in the order given; unlike cuVS, custom ids are not accepted. Extending by a set of rows in one call or in several calls, in the same order, gives the same index. Works on a CPU-only install, since it trains nothing.

Parameters:

X (array-like of shape (n_new, n_features)) – Rows to add, converted to float32.

Returns:

self (IVFIndex) – The extended index. extend_labels_ holds the list each new row joined.

Raises:

ValueError – If the index is not fitted or loaded, X has a different number of features, or metric no longer names the metric the index was built under.

save(path)[source]#

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

The file holds the index arrays, the parameters and the numeric mode. Reload it with load; a loaded index searches without rebuilding, including on a CPU-only install.

Parameters:

path (str or path-like) – Destination file.

Returns:

path (str or path-like) – The path written.

Raises:

RuntimeError – If the index has not been fitted.

classmethod load(path)[source]#

Load an index written by save.

Every array’s dtype and shape is checked on load; the list layout is checked at the first search.

Parameters:

path (str or path-like) – File written by save.

Returns:

index (IVFIndex) – A fitted index, ready to search or extend.

Raises:

ValueError – If the file was saved by another class or its contents do not match the IVF-Flat layout.