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)#

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)

See also

NearestNeighbors, IVFIndex, IVFPQIndex, IVFSQIndex, IVFRaBitQIndex

User guide: Neighbors and vector search

Choosing an estimator

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)#

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)#

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.

classmethod load(path)#

Load an index written by save.

The arrays are read back byte for byte; a file whose arrays have a different dtype, or that was saved by another class, is refused.

Parameters:

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

Returns:

index – A fitted index of this class, ready to search.

Raises:

ValueError – If the file was saved by another class or does not hold this class’s parameters.

numeric_mode_used()#

The numeric mode this estimator runs in.

Returns:

str – 'identical' or 'fast', read from the compiled binary the estimator uses, not from the argument it was given.

save(path)#

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

The file holds the constructor parameters, the fitted sizes, the index arrays exactly as fit left them, and the numeric mode. Reload it with load; a loaded index searches without rebuilding.

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.

set_params(**params)#

Set constructor parameters. The estimator is rebuilt through its constructor, so every constructor check runs again and any fitted state is dropped, as a new setting invalidates it.

vendor_used()#

The GPU backend this estimator runs on.

Returns:

str – 'metal', 'cuda' or 'hip', read from the compiled binary the estimator uses.