scib_rapids.utils.compute_simpson_index

scib_rapids.utils.compute_simpson_index#

scib_rapids.utils.compute_simpson_index(knn_dists, knn_idx, row_idx, labels, n_labels, perplexity=30, tol=1e-05)[source]#

Compute the Simpson index for each cell using a fused CUDA kernel.

The entire binary search and Simpson computation runs in a single GPU kernel launch — no Python loops, no per-cell overhead.

Parameters:

  • knn_dists (ndarray | ndarray) – KNN distances of size (n_cells, n_neighbors).

  • knn_idx (ndarray | ndarray) – KNN indices of size (n_cells, n_neighbors).

  • row_idx (ndarray | ndarray) – Idx of each row (n_cells, 1).

  • labels (ndarray | ndarray) – Cell labels of size (n_cells,).

  • n_labels (int) – Number of labels.

  • perplexity (float (default: 30)) – Measure of the effective number of neighbors.

  • tol (float (default: 1e-05)) – Tolerance for binary search.

Return type:

ndarray

Returns:

simpson_index Simpson index of size (n_cells,).