new tutorial and scalability fixes (#71)

* new tutorial and scalability fixes

* notebook

Author: adamgayoso

docs/notebooks/large_scale.ipynb

@@ -6,4 +6,5 @@
 :maxdepth: 1
 
 notebooks/lung_example
+notebooks/large_scale
 ```

docs/tutorials.md

@@ -178,8 +178,8 @@ def kbet_per_label(
                         batches=batches_sub,
                         alpha=alpha,
                     )
-                except RuntimeError:
-                    logger.info("Not enough neighbors")
+                except ValueError:
+                    logger.info("Diffusion distance failed. Skip.")
                     score = 0  # i.e. 100% rejection
 
             else:
@@ -204,8 +204,8 @@ def kbet_per_label(
                             batches=batches_sub[idx_nonan],
                             alpha=alpha,
                         )
-                    except RuntimeError:
-                        logger.info("Not enough neighbors")
+                    except ValueError:
+                        logger.info("Diffusion distance failed. Skip.")
                         score = 0  # i.e. 100% rejection
                 else:  # if there are too many too small connected components, set kBET score to 0
                     score = 0  # i.e. 100% rejection

src/scib_metrics/_kbet.py

@@ -173,12 +173,22 @@ def prepare(self) -> None:
 
         # Compute neighbors
         for ad in tqdm(self._emb_adatas.values(), desc="Computing neighbors"):
+            # Variables from umap (https://github.com/lmcinnes/umap/blob/3f19ce19584de4cf99e3d0ae779ba13a57472cd9/umap/umap_.py#LL326-L327)
+            # which is used by scanpy under the hood
+            n_trees = min(64, 5 + int(round((ad.X.shape[0]) ** 0.5 / 20.0)))
+            n_iters = max(5, int(round(np.log2(ad.X.shape[0]))))
+            max_candidates = 60
+
             knn_search_index = NNDescent(
                 ad.X,
                 n_neighbors=max(self._neighbor_values),
                 random_state=0,
                 low_memory=True,
                 n_jobs=self._n_jobs,
+                compressed=False,
+                n_trees=n_trees,
+                n_iters=n_iters,
+                max_candidates=max_candidates,
             )
             indices, distances = knn_search_index.neighbor_graph
             for n in self._neighbor_values:

src/scib_metrics/benchmark/_core.py

@@ -116,9 +116,9 @@ def diffusion_nn(X: csr_matrix, k: int, n_comps: int = 100):
     evals, evecs = _compute_eigen(transitions, n_comps=n_comps)
     evals += 1e-8  # Avoid division by zero
     # Multiscale such that the number of steps t gets "integrated out"
-    # First eigenvalue is 1, so we start at the second one
     embedding = evecs
-    embedding[:, 1:] = (evals[1:] / (1 - evals[1:])) * embedding[:, 1:]
+    scaled_evals = np.array([e if e == 1 else e / (1 - e) for e in evals])
+    embedding *= scaled_evals
     nn_obj = pynndescent.NNDescent(embedding, n_neighbors=k + 1)
     neigh_inds, neigh_distances = nn_obj.neighbor_graph
     # We purposely ignore the first neighbor as it is the cell itself

src/scib_metrics/utils/_diffusion_nn.py

@@ -1,10 +1,13 @@
+from functools import partial
 from typing import Tuple, Union
 
 import chex
 import jax
 import jax.numpy as jnp
 import numpy as np
 
+from ._utils import get_ndarray
+
 NdArray = Union[np.ndarray, jnp.ndarray]
 
 
@@ -68,15 +71,13 @@ def _get_neighbor_probability_convergence(state):
 
 
 def _compute_simpson_index_cell(
-    knn_dists_row: jnp.ndarray, knn_row: jnp.ndarray, labels: jnp.ndarray, n_batches: int, perplexity: float, tol: float
+    knn_dists_row: jnp.ndarray, knn_labels_row: jnp.ndarray, n_batches: int, perplexity: float, tol: float
 ) -> jnp.ndarray:
     H, P = _get_neighbor_probability(knn_dists_row, perplexity, tol)
 
     def _non_zero_H_simpson():
-        knn_labels = jnp.take(labels, knn_row)
-        L = jax.nn.one_hot(knn_labels, n_batches)
-        sumP = P @ L
-        return jnp.where(knn_labels.shape[0] == P.shape[0], jnp.dot(sumP, sumP), 1)
+        sumP = jnp.bincount(knn_labels_row, weights=P, length=n_batches)
+        return jnp.where(knn_labels_row.shape[0] == P.shape[0], jnp.dot(sumP, sumP), 1)
 
     return jnp.where(H == 0, -1, _non_zero_H_simpson())
 
@@ -114,9 +115,7 @@ def compute_simpson_index(
     knn_dists = jnp.array(knn_dists)
     knn_idx = jnp.array(knn_idx)
     labels = jnp.array(labels)
-    n = knn_dists.shape[0]
-    return jax.device_get(
-        jax.vmap(
-            lambda i: _compute_simpson_index_cell(knn_dists[i, :], knn_idx[i, :], labels, n_labels, perplexity, tol)
-        )(jnp.arange(n))
-    )
+    knn_labels = labels[knn_idx]
+    simpson_fn = partial(_compute_simpson_index_cell, n_batches=n_labels, perplexity=perplexity, tol=tol)
+    out = jax.vmap(simpson_fn)(knn_dists, knn_labels)
+    return get_ndarray(out)

src/scib_metrics/utils/_lisi.py

@@ -33,16 +33,19 @@ def _silhouette_reduce(
     """
     # accumulate distances from each sample to each cluster
     D_chunk_len = D_chunk.shape[0]
-    clust_dists = jnp.zeros((D_chunk_len, len(label_freqs)), dtype=D_chunk.dtype)
 
-    def _bincount(i, _data):
-        clust_dists, D_chunk, labels, label_freqs = _data
-        clust_dists = clust_dists.at[i].set(jnp.bincount(labels, weights=D_chunk[i], length=label_freqs.shape[0]))
-        return clust_dists, D_chunk, labels, label_freqs
+    # If running into memory issues, use fori_loop instead of vmap
+    # clust_dists = jnp.zeros((D_chunk_len, len(label_freqs)), dtype=D_chunk.dtype)
+    # def _bincount(i, _data):
+    #     clust_dists, D_chunk, labels, label_freqs = _data
+    #     clust_dists = clust_dists.at[i].set(jnp.bincount(labels, weights=D_chunk[i], length=label_freqs.shape[0]))
+    #     return clust_dists, D_chunk, labels, label_freqs
 
-    clust_dists = jax.lax.fori_loop(
-        0, D_chunk_len, lambda i, _data: _bincount(i, _data), (clust_dists, D_chunk, labels, label_freqs)
-    )[0]
+    # clust_dists = jax.lax.fori_loop(
+    #     0, D_chunk_len, lambda i, _data: _bincount(i, _data), (clust_dists, D_chunk, labels, label_freqs)
+    # )[0]
+
+    clust_dists = jax.vmap(partial(jnp.bincount, length=label_freqs.shape[0]), in_axes=(None, 0))(labels, D_chunk)
 
     # intra_index selects intra-cluster distances within clust_dists
     intra_index = (jnp.arange(D_chunk_len), jax.lax.dynamic_slice(labels, (start,), (D_chunk_len,)))