Fix arg order and keyword the args; add tests for low_memory

umap/tests/test_umap.py

@@ -294,6 +294,25 @@ def test_nn_descent_neighbor_accuracy():
         "NN-descent did not get 99% " "accuracy on nearest neighbors",
     )
 
+def test_nn_descent_neighbor_accuracy_low_memory():
+    knn_indices, knn_dists, _ = nearest_neighbors(
+        nn_data, 10, "euclidean", {}, False, np.random, low_memory=True
+    )
+
+    tree = KDTree(nn_data)
+    true_indices = tree.query(nn_data, 10, return_distance=False)
+
+    num_correct = 0.0
+    for i in range(nn_data.shape[0]):
+        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))
+
+    percent_correct = num_correct / (nn_data.shape[0] * 10)
+    assert_greater_equal(
+        percent_correct,
+        0.89,
+        "NN-descent did not get 99% " "accuracy on nearest neighbors",
+    )
+
 
 def test_angular_nn_descent_neighbor_accuracy():
     knn_indices, knn_dists, _ = nearest_neighbors(
@@ -336,6 +355,26 @@ def test_sparse_nn_descent_neighbor_accuracy():
     )
 
 
+def test_sparse_nn_descent_neighbor_accuracy_low_memory():
+    knn_indices, knn_dists, _ = nearest_neighbors(
+        sparse_nn_data, 20, "euclidean", {}, False, np.random, low_memory=True
+    )
+
+    tree = KDTree(sparse_nn_data.todense())
+    true_indices = tree.query(sparse_nn_data.todense(), 10, return_distance=False)
+
+    num_correct = 0.0
+    for i in range(sparse_nn_data.shape[0]):
+        num_correct += np.sum(np.in1d(true_indices[i], knn_indices[i]))
+
+    percent_correct = num_correct / (sparse_nn_data.shape[0] * 10)
+    assert_greater_equal(
+        percent_correct,
+        0.90,
+        "Sparse NN-descent did not get " "99% accuracy on nearest " "neighbors",
+    )
+
+
 def test_sparse_angular_nn_descent_neighbor_accuracy():
     knn_indices, knn_dists, _ = nearest_neighbors(
         sparse_nn_data, 20, "cosine", {}, True, np.random

umap/umap_.py

@@ -325,9 +325,9 @@ def nearest_neighbors(
                     X.shape[0],
                     n_neighbors,
                     rng_state,
-                    distance_func,
-                    tuple(metric_kwds.values()),
                     max_candidates=60,
+                    dist=distance_func,
+                    dist_args=tuple(metric_kwds.values()),
                     low_memory=low_memory,
                     rp_tree_init=True,
                     leaf_array=leaf_array,
@@ -349,9 +349,9 @@ def nearest_neighbors(
                     X,
                     n_neighbors,
                     rng_state,
-                    distance_func,
-                    tuple(metric_kwds.values()),
                     max_candidates=60,
+                    dist=distance_func,
+                    dist_args=tuple(metric_kwds.values()),
                     low_memory=low_memory,
                     rp_tree_init=True,
                     leaf_array=leaf_array,