scikit-learn · jjerphan · Jun 23, 2021 · Jun 23, 2021 · Jun 23, 2021 · Jun 23, 2021
diff --git a/doc/whats_new/v1.1.rst b/doc/whats_new/v1.1.rst
@@ -113,6 +113,39 @@ Changelog
   :pr:`20880` by :user:`Guillaume Lemaitre <glemaitre>`
   and :user:`András Simon <simonandras>`.
 
+Miscellaneous
+.............
+
+- |Efficiency| Low-level routines for reductions on pairwise distances
+  for dense float64 datasets have been refactored. The following functions
+  and estimators now benefit from improved performances, in particular on
+  multi-cores machines:
+  - :func:`sklearn.metrics.pairwise_distances_argmin`
+  - :func:`sklearn.metrics.pairwise_distances_argmin_min`
+  - :class:`sklearn.cluster.AffinityPropagation`
+  - :class:`sklearn.cluster.Birch`
+  - :class:`sklearn.cluster.DBSCAN`
+  - :class:`sklearn.cluster.MeanShift`
+  - :class:`sklearn.cluster.OPTICS`
+  - :class:`sklearn.cluster.SpectralClustering`
+  - :func:`sklearn.feature_selection.mutual_info_regression`
+  - :class:`sklearn.neighbors.KNeighborsClassifier`
+  - :class:`sklearn.neighbors.KNeighborsRegressor`
+  - :class:`sklearn.neighbors.LocalOutlierFactor`
+  - :class:`sklearn.neighbors.NearestNeighbors`
+  - :class:`sklearn.manifold.Isomap`
+  - :class:`sklearn.manifold.LocallyLinearEmbedding`
+  - :class:`sklearn.manifold.TSNE`
+  - :func:`sklearn.manifold.trustworthiness`
+  - :class:`sklearn.semi_supervised.LabelPropagation`
+  - :class:`sklearn.semi_supervised.LabelSpreading`
+
+  For instance :class:`sklearn.neighbors.NearestNeighbors.kneighbors`
+  can be up to 20× faster than in the previous versions'.
+
+  :pr:`20254` by :user:`Julien Jerphanion <jjerphan>`.
+
+
 Code and Documentation Contributors
 -----------------------------------
 

diff --git a/sklearn/_config.py b/sklearn/_config.py
@@ -9,6 +9,9 @@
     "working_memory": int(os.environ.get("SKLEARN_WORKING_MEMORY", 1024)),
     "print_changed_only": True,
     "display": "text",
+    "pairwise_dist_chunk_size": int(
+        os.environ.get("SKLEARN_PAIRWISE_DIST_CHUNK_SIZE", 256)
+    ),
 }
 _threadlocal = threading.local()
 
@@ -40,7 +43,11 @@ def get_config():
 
 
 def set_config(
-    assume_finite=None, working_memory=None, print_changed_only=None, display=None
+    assume_finite=None,
+    working_memory=None,
+    print_changed_only=None,
+    display=None,
+    pairwise_dist_chunk_size=None,
 ):
     """Set global scikit-learn configuration
 
@@ -80,6 +87,12 @@ def set_config(
 
         .. versionadded:: 0.23
 
+    pairwise_dist_chunk_size : int, default=None
+        The number of vectors per chunk for PairwiseDistancesReduction.
+        Default is 256 (optimal for most of modern laptops' caches and architectures).
+
+        .. versionadded:: 1.1
+
     See Also
     --------
     config_context : Context manager for global scikit-learn configuration.
@@ -95,6 +108,8 @@ def set_config(
         local_config["print_changed_only"] = print_changed_only
     if display is not None:
         local_config["display"] = display
+    if pairwise_dist_chunk_size is not None:
+        local_config["pairwise_dist_chunk_size"] = pairwise_dist_chunk_size
 
 
 @contextmanager
@@ -132,6 +147,12 @@ def config_context(**new_config):
 
         .. versionadded:: 0.23
 
+    pairwise_dist_chunk_size : int, default=None
+        The number of vectors per chunk for PairwiseDistancesReduction.
+        Default is 256 (optimal for most of modern laptops' caches and architectures).
+
+        .. versionadded:: 1.1
+
     Notes
     -----
     All settings, not just those presently modified, will be returned to

diff --git a/sklearn/cluster/_affinity_propagation.py b/sklearn/cluster/_affinity_propagation.py
@@ -523,7 +523,9 @@ def predict(self, X):
 
         if self.cluster_centers_.shape[0] > 0:
             with config_context(assume_finite=True):
-                return pairwise_distances_argmin(X, self.cluster_centers_)
+                return pairwise_distances_argmin(
+                    X, self.cluster_centers_, metric="fast_euclidean"
+                )
         else:
             warnings.warn(
                 "This model does not have any cluster centers "

diff --git a/sklearn/cluster/_birch.py b/sklearn/cluster/_birch.py
@@ -676,11 +676,15 @@ def predict(self, X):
         """
         check_is_fitted(self)
         X = self._validate_data(X, accept_sparse="csr", reset=False)
-        kwargs = {"Y_norm_squared": self._subcluster_norms}
+
+        fast_euclidean_kwargs = {"Y_norm_squared": self._subcluster_norms}
 
         with config_context(assume_finite=True):
             argmin = pairwise_distances_argmin(
-                X, self.subcluster_centers_, metric_kwargs=kwargs
+                X,
+                self.subcluster_centers_,
+                metric="fast_euclidean",
+                metric_kwargs=fast_euclidean_kwargs,
             )
         return self.subcluster_labels_[argmin]
 
@@ -726,7 +730,8 @@ def _global_clustering(self, X=None):
                 "n_clusters should be an instance of ClusterMixin or an int"
             )
 
-        # To use in predict to avoid recalculation.
+        # We compute subcluster norms once here, so that we won't need to compute it
+        # again at each call of `Birch.predict`.
         self._subcluster_norms = row_norms(self.subcluster_centers_, squared=True)
 
         if clusterer is None or not_enough_centroids:

diff --git a/sklearn/cluster/_mean_shift.py b/sklearn/cluster/_mean_shift.py
@@ -512,4 +512,6 @@ def predict(self, X):
         check_is_fitted(self)
         X = self._validate_data(X, reset=False)
         with config_context(assume_finite=True):
-            return pairwise_distances_argmin(X, self.cluster_centers_)
+            return pairwise_distances_argmin(
+                X, self.cluster_centers_, metric="fast_euclidean"
+            )
diff --git a/sklearn/metrics/_dist_metrics.pxd b/sklearn/metrics/_dist_metrics.pxd
@@ -61,6 +61,22 @@ cdef class DistanceMetric:
     cdef DTYPE_t rdist(self, const DTYPE_t* x1, const DTYPE_t* x2,
                        ITYPE_t size) nogil except -1
 
+    cdef DTYPE_t csr_dist(
+        self,
+        const DTYPE_t[:] x1_data,
+        const ITYPE_t[:] x1_indices,
+        const DTYPE_t[:] x2_data,
+        const ITYPE_t[:] x2_indices,
+    ) nogil except -1
+
+    cdef DTYPE_t csr_rdist(
+        self,
+        const DTYPE_t[:] x1_data,
+        const ITYPE_t[:] x1_indices,
+        const DTYPE_t[:] x2_data,
+        const ITYPE_t[:] x2_indices,
+    ) nogil except -1
+
     cdef int pdist(self, const DTYPE_t[:, ::1] X, DTYPE_t[:, ::1] D) except -1
 
     cdef int cdist(self, const DTYPE_t[:, ::1] X, const DTYPE_t[:, ::1] Y,
@@ -69,3 +85,24 @@ cdef class DistanceMetric:
     cdef DTYPE_t _rdist_to_dist(self, DTYPE_t rdist) nogil except -1
 
     cdef DTYPE_t _dist_to_rdist(self, DTYPE_t dist) nogil except -1
+
+
+######################################################################
+# DatasetsPair base class
+cdef class DatasetsPair:
+    cdef DistanceMetric distance_metric
+
+    cdef ITYPE_t n_samples_X(self) nogil
+
+    cdef ITYPE_t n_samples_Y(self) nogil
+
+    cdef DTYPE_t dist(self, ITYPE_t i, ITYPE_t j) nogil
+
+    cdef DTYPE_t surrogate_dist(self, ITYPE_t i, ITYPE_t j) nogil
+
+
+cdef class DenseDenseDatasetsPair(DatasetsPair):
+    cdef:
+        const DTYPE_t[:, ::1] X
+        const DTYPE_t[:, ::1] Y
+        ITYPE_t d