Repository: YosefLab/Hotspot

Issue:
import knn graph?
Would it be useful to have a function to import and use the neighbor graph computed in scanpy (fast for big datasests) 
the function `hs.create_knn_graph` seems a bit slow on very big datasets, and if the nn graph is already being computed for other analyses in scanpy makes sense to reuse this?

Genuinely very useful tool :)

Discussion:

Patch:
diff --git a/hotspot/hotspot.py b/hotspot/hotspot.py
index a858f66..459f24b 100644
--- a/hotspot/hotspot.py
+++ b/hotspot/hotspot.py
@@ -126,7 +126,7 @@ class Hotspot:
         self.module_scores = None
 
     def create_knn_graph(
-            self, weighted_graph=False, n_neighbors=30, neighborhood_factor=3):
+            self, weighted_graph=False, n_neighbors=30, neighborhood_factor=3, approx_neighbors=True):
         """Create's the KNN graph and graph weights
 
         The resulting matrices containing the neighbors and weights are
@@ -143,12 +143,16 @@ class Hotspot:
             relative to the distances within the neighborhood.  The weight for
             a cell with a distance d will decay as exp(-d/D) where D is the distance
             to the `n_neighbors`/`neighborhood_factor`-th neighbor.
+        approx_neighbors: bool
+            Use approximate nearest neighbors or exact scikit-learn neighbors. Only
+            when hotspot initialized with `latent`.
         """
 
         if self.latent is not None:
             neighbors, weights = neighbors_and_weights(
                 self.latent, n_neighbors=n_neighbors,
-                neighborhood_factor=neighborhood_factor)
+                neighborhood_factor=neighborhood_factor,
+                approx_neighbors=approx_neighbors)
         elif self.tree is not None:
             if weighted_graph:
                 raise ValueError("When using `tree` as the metric space, `weighted_graph=True` is not supported")
diff --git a/hotspot/knn.py b/hotspot/knn.py
index d26bdf2..2831dc1 100644
--- a/hotspot/knn.py
+++ b/hotspot/knn.py
@@ -4,9 +4,10 @@ import pandas as pd
 from math import ceil
 from numba import jit
 from tqdm import tqdm
+from pynndescent import NNDescent
 
 
-def neighbors_and_weights(data, n_neighbors=30, neighborhood_factor=3):
+def neighbors_and_weights(data, n_neighbors=30, neighborhood_factor=3, approx_neighbors=True):
     """
     Computes nearest neighbors and associated weights for data
     Uses euclidean distance between rows of `data`
@@ -23,9 +24,14 @@ def neighbors_and_weights(data, n_neighbors=30, neighborhood_factor=3):
     """
 
     coords = data.values
-    nbrs = NearestNeighbors(n_neighbors=n_neighbors,
-                            algorithm="ball_tree").fit(coords)
-    dist, ind = nbrs.kneighbors()
+
+    if approx_neighbors:
+        index = NNDescent(coords, n_neighbors=n_neighbors)
+        ind, dist = index.neighbor_graph
+    else:
+        nbrs = NearestNeighbors(n_neighbors=n_neighbors,
+                        algorithm="ball_tree").fit(coords)
+        dist, ind = nbrs.kneighbors()
 
     weights = compute_weights(
         dist, neighborhood_factor=neighborhood_factor)
diff --git a/setup.py b/setup.py
index 0aee543..fde915f 100644
--- a/setup.py
+++ b/setup.py
@@ -22,6 +22,7 @@ setup(
         'tqdm>=4.32.2',
         'statsmodels>=0.9.0',
         'scikit-learn>=0.21.2',
+        'pynndescent>=0.5',
     ],
     extras_require=dict(
         test=['pytest>=5.0.0'],


Test patch:
diff --git a/tests/test_bernoulli.py b/tests/test_bernoulli.py
index d908b2a..82530f0 100644
--- a/tests/test_bernoulli.py
+++ b/tests/test_bernoulli.py
@@ -25,7 +25,7 @@ def test_local_autocorrelation_centered():
     umi_counts = pd.Series(umi_counts)
 
     neighbors, weights = neighbors_and_weights(
-        latent, n_neighbors=30, neighborhood_factor=3
+        latent, n_neighbors=30, neighborhood_factor=3, approx_neighbors=False
     )
     neighbors = neighbors.values
     weights = weights.values
@@ -145,7 +145,7 @@ def test_local_correlation_centered():
             gene_p, Gmean, EG
         )
         assert math.isclose(
-            EstdG, Gstd, rel_tol=5e-2
+            EstdG, Gstd, rel_tol=1e-1
         ), "stdG is off for gene_p={}, Actual={:.2f}, Expected={:.2f}".format(
             gene_p, Gstd, EstdG
         )
diff --git a/tests/test_local_stats.py b/tests/test_local_stats.py
index 1efea09..21f9b00 100644
--- a/tests/test_local_stats.py
+++ b/tests/test_local_stats.py
@@ -22,24 +22,33 @@ def test_moments_fast():
     latent = sim_data.sim_latent(N_CELLS, N_DIM)
     latent = pd.DataFrame(latent)
 
-    neighbors, weights = neighbors_and_weights(
-        latent, n_neighbors=30, neighborhood_factor=3)
+    def _compute(neighbors, weights, rel_tol=1e-12):
+        neighbors = neighbors.values
+        weights = weights.values
+
+        weights = make_weights_non_redundant(neighbors, weights)
 
-    neighbors = neighbors.values
-    weights = weights.values
+        # Just generate random muX, muY, x2, y2
+        muX = np.random.rand(N_CELLS) * 2 + 2
+        x2 = (np.random.rand(N_CELLS) * 2 + 1)**2 + muX**2
 
-    weights = make_weights_non_redundant(neighbors, weights)
+        EG, EG2 = compute_moments_weights_slow(
+            muX, x2, neighbors, weights
+        )
+        EG_fast, EG2_fast = compute_moments_weights(
+            muX, x2,  neighbors, weights
+        )
 
-    # Just generate random muX, muY, x2, y2
-    muX = np.random.rand(N_CELLS) * 2 + 2
-    x2 = (np.random.rand(N_CELLS) * 2 + 1)**2 + muX**2
+        assert math.isclose(EG, EG_fast, rel_tol=rel_tol)
+        assert math.isclose(EG2, EG2_fast, rel_tol=rel_tol)
+
+    neighbors, weights = neighbors_and_weights(
+        latent, n_neighbors=30, neighborhood_factor=3, approx_neighbors=False)
 
-    EG, EG2 = compute_moments_weights_slow(
-        muX, x2, neighbors, weights
-    )
-    EG_fast, EG2_fast = compute_moments_weights(
-        muX, x2,  neighbors, weights
-    )
+    _compute(neighbors, weights)
+
+
+    neighbors, weights = neighbors_and_weights(
+        latent, n_neighbors=30, neighborhood_factor=3, approx_neighbors=True)
 
-    assert math.isclose(EG, EG_fast, rel_tol=1e-12)
-    assert math.isclose(EG2, EG2_fast, rel_tol=1e-12)
+    _compute(neighbors, weights, rel_tol=0.005)
diff --git a/tests/test_local_stats_pairs.py b/tests/test_local_stats_pairs.py
index f688831..67213de 100644
--- a/tests/test_local_stats_pairs.py
+++ b/tests/test_local_stats_pairs.py
@@ -27,30 +27,37 @@ def test_pairs():
     latent = sim_data.sim_latent(N_CELLS, N_DIM)
     latent = pd.DataFrame(latent)
 
-    neighbors, weights = neighbors_and_weights(
-        latent, n_neighbors=30, neighborhood_factor=3)
+    def _compute(neighbors, weights, rel_tol=1e-12):
+        neighbors = neighbors.values
+        weights = weights.values
 
-    neighbors = neighbors.values
-    weights = weights.values
+        weights = make_weights_non_redundant(neighbors, weights)
 
-    weights = make_weights_non_redundant(neighbors, weights)
+        # Just generate random muX, muY, x2, y2
+        muX = np.random.rand(N_CELLS) * 2 + 2
+        muY = np.random.rand(N_CELLS) * .5 + 1
+        x2 = (np.random.rand(N_CELLS) * 2 + 1)**2 + muX**2
+        y2 = (np.random.rand(N_CELLS) * 3 + .5)**2 + muY**2
 
-    # Just generate random muX, muY, x2, y2
-    muX = np.random.rand(N_CELLS) * 2 + 2
-    muY = np.random.rand(N_CELLS) * .5 + 1
-    x2 = (np.random.rand(N_CELLS) * 2 + 1)**2 + muX**2
-    y2 = (np.random.rand(N_CELLS) * 3 + .5)**2 + muY**2
+        EG, EG2 = compute_moments_weights_pairs_slow(
+            muX, x2, muY, y2, neighbors, weights
+        )
+        EG_fast, EG2_fast = compute_moments_weights_pairs(
+            muX, x2, muY, y2, neighbors, weights
+        )
 
-    EG, EG2 = compute_moments_weights_pairs_slow(
-        muX, x2, muY, y2, neighbors, weights
-    )
-    EG_fast, EG2_fast = compute_moments_weights_pairs(
-        muX, x2, muY, y2, neighbors, weights
-    )
+        assert math.isclose(EG, EG_fast, rel_tol=rel_tol)
+        assert math.isclose(EG2, EG2_fast, rel_tol=rel_tol)
+
+    neighbors, weights = neighbors_and_weights(
+        latent, n_neighbors=30, neighborhood_factor=3, approx_neighbors=False)
+
+    _compute(neighbors, weights)
 
-    assert math.isclose(EG, EG_fast, rel_tol=1e-12)
-    assert math.isclose(EG2, EG2_fast, rel_tol=1e-12)
+    neighbors, weights = neighbors_and_weights(
+        latent, n_neighbors=30, neighborhood_factor=3, approx_neighbors=True)
 
+    _compute(neighbors, weights, rel_tol=0.005)
 
 def test_pairs_fast():
     """
@@ -65,29 +72,39 @@ def test_pairs_fast():
     latent = sim_data.sim_latent(N_CELLS, N_DIM)
     latent = pd.DataFrame(latent)
 
-    neighbors, weights = neighbors_and_weights(
-        latent, n_neighbors=30, neighborhood_factor=3)
+    def _compute(neighbors, weights, rel_tol=1e-12):
+        neighbors = neighbors.values
+        weights = weights.values
 
-    neighbors = neighbors.values
-    weights = weights.values
+        weights = make_weights_non_redundant(neighbors, weights)
 
-    weights = make_weights_non_redundant(neighbors, weights)
+        # Just generate random muX, muY, x2, y2
+        muX = np.random.rand(N_CELLS) * 2 + 2
+        muY = np.random.rand(N_CELLS) * .5 + 1
+        x2 = (np.random.rand(N_CELLS) * 2 + 1)**2 + muX**2
+        y2 = (np.random.rand(N_CELLS) * 3 + .5)**2 + muY**2
 
-    # Just generate random muX, muY, x2, y2
-    muX = np.random.rand(N_CELLS) * 2 + 2
-    muY = np.random.rand(N_CELLS) * .5 + 1
-    x2 = (np.random.rand(N_CELLS) * 2 + 1)**2 + muX**2
-    y2 = (np.random.rand(N_CELLS) * 3 + .5)**2 + muY**2
+        EG, EG2 = compute_moments_weights_pairs(
+            muX, x2, muY, y2, neighbors, weights
+        )
+        EG_fast, EG2_fast = compute_moments_weights_pairs_fast(
+            muX, x2, muY, y2, neighbors, weights
+        )
 
-    EG, EG2 = compute_moments_weights_pairs(
-        muX, x2, muY, y2, neighbors, weights
-    )
-    EG_fast, EG2_fast = compute_moments_weights_pairs_fast(
-        muX, x2, muY, y2, neighbors, weights
-    )
+        assert math.isclose(EG, EG_fast, rel_tol=rel_tol)
+        assert math.isclose(EG2, EG2_fast, rel_tol=rel_tol)
+
+
+    neighbors, weights = neighbors_and_weights(
+        latent, n_neighbors=30, neighborhood_factor=3, approx_neighbors=False)
+
+    _compute(neighbors, weights)
+
+    neighbors, weights = neighbors_and_weights(
+        latent, n_neighbors=30, neighborhood_factor=3, approx_neighbors=False)
+
+    _compute(neighbors, weights, rel_tol=0.005)
 
-    assert math.isclose(EG, EG_fast, rel_tol=1e-12)
-    assert math.isclose(EG2, EG2_fast, rel_tol=1e-12)
 
 
 def test_pairs_std():