Some fixes for parallel pairwise distances when n_jobs > 1 (scikit-learn#13877)

jeremiedbb · jnothman · commit 0c9b9847820e · 2019-05-15T09:48:17.000+10:00
diff --git a/doc/whats_new/v0.21.rst b/doc/whats_new/v0.21.rst
@@ -23,6 +23,12 @@ Changelog
   ``Y == None``.
   :issue:`13864` by :user:`Paresh Mathur <rick2047>`.
 
+- |Fix| Fixed two bugs in :class:`metrics.pairwise_distances` when
+  ``n_jobs > 1``. First it used to return a distance matrix with same dtype as
+  input, even for integer dtype. Then the diagonal was not zeros for euclidean
+  metric when ``Y`` is ``X``. :issue:`13877` by
+  :user:`Jérémie du Boisberranger <jeremiedbb>`.
+
 :mod:`sklearn.neighbors`
 ......................
 
diff --git a/sklearn/metrics/pairwise.py b/sklearn/metrics/pairwise.py
@@ -1198,17 +1198,23 @@ def _parallel_pairwise(X, Y, func, n_jobs, **kwds):
 
     if Y is None:
         Y = X
+    X, Y, dtype = _return_float_dtype(X, Y)
 
     if effective_n_jobs(n_jobs) == 1:
         return func(X, Y, **kwds)
 
     # enforce a threading backend to prevent data communication overhead
     fd = delayed(_dist_wrapper)
-    ret = np.empty((X.shape[0], Y.shape[0]), dtype=X.dtype, order='F')
+    ret = np.empty((X.shape[0], Y.shape[0]), dtype=dtype, order='F')
     Parallel(backend="threading", n_jobs=n_jobs)(
         fd(func, ret, s, X, Y[s], **kwds)
         for s in gen_even_slices(_num_samples(Y), effective_n_jobs(n_jobs)))
 
+    if (X is Y or Y is None) and func is euclidean_distances:
+        # zeroing diagonal for euclidean norm.
+        # TODO: do it also for other norms.
+        np.fill_diagonal(ret, 0)
+
     return ret
 
 
diff --git a/sklearn/metrics/tests/test_pairwise.py b/sklearn/metrics/tests/test_pairwise.py
@@ -231,31 +231,6 @@ def test_pairwise_precomputed_non_negative():
                          metric='precomputed')
 
 
-def check_pairwise_parallel(func, metric, kwds):
-    rng = np.random.RandomState(0)
-    for make_data in (np.array, csr_matrix):
-        X = make_data(rng.random_sample((5, 4)))
-        Y = make_data(rng.random_sample((3, 4)))
-
-        try:
-            S = func(X, metric=metric, n_jobs=1, **kwds)
-        except (TypeError, ValueError) as exc:
-            # Not all metrics support sparse input
-            # ValueError may be triggered by bad callable
-            if make_data is csr_matrix:
-                assert_raises(type(exc), func, X, metric=metric,
-                              n_jobs=2, **kwds)
-                continue
-            else:
-                raise
-        S2 = func(X, metric=metric, n_jobs=2, **kwds)
-        assert_array_almost_equal(S, S2)
-
-        S = func(X, Y, metric=metric, n_jobs=1, **kwds)
-        S2 = func(X, Y, metric=metric, n_jobs=2, **kwds)
-        assert_array_almost_equal(S, S2)
-
-
 _wminkowski_kwds = {'w': np.arange(1, 5).astype('double', copy=False), 'p': 1}
 
 
@@ -272,8 +247,30 @@ def callable_rbf_kernel(x, y, **kwds):
          (pairwise_distances, 'wminkowski', _wminkowski_kwds),
          (pairwise_kernels, 'polynomial', {'degree': 1}),
          (pairwise_kernels, callable_rbf_kernel, {'gamma': .1})])
-def test_pairwise_parallel(func, metric, kwds):
-    check_pairwise_parallel(func, metric, kwds)
+@pytest.mark.parametrize('array_constr', [np.array, csr_matrix])
+@pytest.mark.parametrize('dtype', [np.float64, int])
+def test_pairwise_parallel(func, metric, kwds, array_constr, dtype):
+    rng = np.random.RandomState(0)
+    X = array_constr(5 * rng.random_sample((5, 4)), dtype=dtype)
+    Y = array_constr(5 * rng.random_sample((3, 4)), dtype=dtype)
+
+    try:
+        S = func(X, metric=metric, n_jobs=1, **kwds)
+    except (TypeError, ValueError) as exc:
+        # Not all metrics support sparse input
+        # ValueError may be triggered by bad callable
+        if array_constr is csr_matrix:
+            with pytest.raises(type(exc)):
+                func(X, metric=metric, n_jobs=2, **kwds)
+            return
+        else:
+            raise
+    S2 = func(X, metric=metric, n_jobs=2, **kwds)
+    assert_allclose(S, S2)
+
+    S = func(X, Y, metric=metric, n_jobs=1, **kwds)
+    S2 = func(X, Y, metric=metric, n_jobs=2, **kwds)
+    assert_allclose(S, S2)
 
 
 def test_pairwise_callable_nonstrict_metric():
@@ -551,6 +548,16 @@ def test_pairwise_distances_chunked_diagonal(metric):
     assert_array_almost_equal(np.diag(np.vstack(chunks)), 0, decimal=10)
 
 
+@pytest.mark.parametrize(
+        'metric',
+        ('euclidean', 'l2', 'sqeuclidean'))
+def test_parallel_pairwise_distances_diagonal(metric):
+    rng = np.random.RandomState(0)
+    X = rng.normal(size=(1000, 10), scale=1e10)
+    distances = pairwise_distances(X, metric=metric, n_jobs=2)
+    assert_allclose(np.diag(distances), 0, atol=1e-10)
+
+
 @ignore_warnings
 def test_pairwise_distances_chunked():
     # Test the pairwise_distance helper function.
