lobpcg for spectral; back to parallel

lmcinnes · Oct 6, 2018 · bdd80d3 · bdd80d3
1 parent 16e7ce6
commit bdd80d3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/umap/spectral.py b/umap/spectral.py
@@ -254,14 +254,20 @@ def spectral_layout(data, graph, dim, random_state, metric="euclidean", metric_k
     k = dim + 1
     num_lanczos_vectors = max(2 * k + 1, int(np.sqrt(graph.shape[0])))
     try:
-        eigenvalues, eigenvectors = scipy.sparse.linalg.eigsh(
+        # eigenvalues, eigenvectors = scipy.sparse.linalg.eigsh(
+        #     L,
+        #     k,
+        #     which="SM",
+        #     ncv=num_lanczos_vectors,
+        #     tol=1e-4,
+        #     v0=np.ones(L.shape[0]),
+        #     maxiter=graph.shape[0] * 5,
+        # )
+        eigenvalues, eigenvectors = scipy.sparse.linalg.lobpcg(
             L,
-            k,
-            which="SM",
-            ncv=num_lanczos_vectors,
-            tol=1e-4,
-            v0=np.ones(L.shape[0]),
-            maxiter=graph.shape[0] * 5,
+            random_state.normal(size=(L.shape[0], k)),
+            largest=False,
+            tol=1e-4
         )
         order = np.argsort(eigenvalues)[1:k]
         return eigenvectors[:, order]

diff --git a/umap/umap_.py b/umap/umap_.py
@@ -679,7 +679,7 @@ def rdist(x, y):
     return result
 
 
-@numba.njit(fastmath=True)
+@numba.njit(fastmath=True, parallel=True)
 def optimize_layout(
     head_embedding,
     tail_embedding,