Parallelize the "precomputed" case of knn

It is supposed to be possible to compile numpy with openmp, but it looks
like it isn't done in any of the usual packages.

It may be possible to refactor that code so there is less redundancy.

2 files changed, 11 insertions, 0 deletions

diff --git a/src/python/test/test_dtm.py b/src/python/test/test_dtm.py
index 37934fdb..bc0d3698 100755
--- a/src/python/test/test_dtm.py
+++ b/src/python/test/test_dtm.py
@@ -33,6 +33,9 @@ def test_dtm_compare_euclidean():
     dtm = DistanceToMeasure(k, metric="precomputed")
     r4 = dtm.fit_transform(d)
     assert r4 == pytest.approx(r0)
+    dtm = DistanceToMeasure(k, metric="precomputed", n_jobs=2)
+    r4b = dtm.fit_transform(d)
+    assert r4b == pytest.approx(r0)
     dtm = DistanceToMeasure(k, implementation="keops")
     r5 = dtm.fit_transform(pts)
     assert r5 == pytest.approx(r0)
diff --git a/src/python/test/test_knn.py b/src/python/test/test_knn.py
index 6aac2006..6269df54 100755
--- a/src/python/test/test_knn.py
+++ b/src/python/test/test_knn.py
@@ -52,6 +52,14 @@ def test_knn_explicit():
     r = knn.fit_transform(dist)
     assert np.array_equal(r[0], [[0, 1], [1, 0], [2, 0]])
     assert np.array_equal(r[1], [[0, 3], [0, 1], [0, 1]])
+    # Second time in parallel
+    knn = KNearestNeighbors(2, metric="precomputed", return_index=True, return_distance=False, n_jobs=2)
+    r = knn.fit_transform(dist)
+    assert np.array_equal(r, [[0, 1], [1, 0], [2, 0]])
+    knn = KNearestNeighbors(2, metric="precomputed", return_index=True, return_distance=True, n_jobs=2)
+    r = knn.fit_transform(dist)
+    assert np.array_equal(r[0], [[0, 1], [1, 0], [2, 0]])
+    assert np.array_equal(r[1], [[0, 3], [0, 1], [0, 1]])
 
 
 def test_knn_compare():