mojolearn.CagraIndex#

class mojolearn.CagraIndex(graph_degree=32, intermediate_graph_degree=64, n_neighbors=8, itopk_size=64, search_width=1, max_iterations=0, n_seeds=0, random_state=None)[source]#

Bases: _AnnResident, _AnnSaved, NumericModeMixin

Graph-based approximate nearest-neighbor index (CAGRA).

Follows cuVS’s cagra. fit builds an exact k-nearest-neighbor graph of degree intermediate_graph_degree, prunes it to graph_degree with cuVS’s rank-based detour pruning, and adds reverse edges. search walks the graph greedily from a set of start nodes, keeping the best itopk_size candidates, and returns squared Euclidean distances. Ties are broken by row id and the start nodes are the same for every query, so results are reproducible. The index keeps a copy of the dataset, as cuVS’s does. On a GPU the fitted index stays on the device between searches; each search uploads only its queries.

Parameters:

  • graph_degree (int, default 32) – Out-degree of the final graph. Capped at intermediate_graph_degree.

  • intermediate_graph_degree (int, default 64) – Degree of the k-NN graph built before pruning. Capped at n_samples - 1.

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

  • itopk_size (int, default 64) – Number of candidates kept during the search. Larger values give better recall and slower searches.

  • search_width (int, default 1) – Number of candidates whose neighbors are expanded per iteration.

  • max_iterations (int, default 0) – Search iterations; 0 means itopk_size.

  • n_seeds (int, default 0) – Number of start nodes; 0 means cuVS’s default, itopk_size + search_width * graph_degree (capped at the number of rows).

  • random_state (int or None, default None) – Seed for the start nodes, in the role of cuVS’s rand_xor_mask. None takes evenly spaced nodes; a non-negative int takes seeded pseudo-random nodes. Either way every query uses the same start nodes.

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

Variables:

  • dataset_ (ndarray of shape (n_samples, n_features), float32) – The indexed rows.

  • graph_ (ndarray of shape (n_samples, graph_degree_), int32) – Neighbor ids of each row in the search graph.

  • graph_degree_ (int) – The graph degree used, after capping.

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

  • n_rows_ (int) – Number of indexed rows.

Raises:

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

  • ValueError – If random_state is negative.

See also

IVFIndex, IVFPQIndex

Examples

>>> import numpy as np, mojolearn
>>> X = np.random.default_rng(0).random((1000, 8), dtype=np.float32)
>>> index = mojolearn.CagraIndex(n_neighbors=5).fit(X)
>>> distances, indices = index.search(X[:10])
>>> indices.shape
(10, 5)
__init__(graph_degree=32, intermediate_graph_degree=64, n_neighbors=8, itopk_size=64, search_width=1, max_iterations=0, n_seeds=0, random_state=None)[source]#

Methods

fit(X[, y])

Build the search graph 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 search graph from X.

Parameters:

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

  • y (None) – Ignored.

Returns:

self (CagraIndex) – 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. As in cuVS’s single-CTA search, the filter is applied after the graph walk, over the itopk_size candidates: removed rows still guide the walk, and the answer is the first n_neighbors kept candidates in (distance, id) order. Distances are the same values the unfiltered search computes.

Returns:

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

  • indices (ndarray of shape (n_queries, n_neighbors), int32) – Row ids of the neighbors. With a filter, a query with fewer kept candidates than n_neighbors 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.