ENH perf: numeric optimisation in nan_euclidean_distances, GaussianMixture and KNeighborsRegressor.predict (scikit-learn#34277)

Author: rth

doc/whats_new/upcoming_changes/sklearn.metrics/1140.efficiency.rst

@@ -0,0 +1,4 @@
+- :func:`metrics.nan_euclidean_distances` is now several times faster on dense data,
+  with a more moderate speed-up for :class:`impute.KNNImputer` and estimators using
+  ``metric="nan_euclidean"`` such as :class:`neighbors.NearestNeighbors`.
+  By :user:`Roman Yurchak <rth>`.

doc/whats_new/upcoming_changes/sklearn.mixture/1140.efficiency.rst

@@ -0,0 +1,3 @@
+- :class:`mixture.GaussianMixture` with ``covariance_type="tied"`` is now faster,
+  with the speed-up most visible for many components or high-dimensional data.
+  By :user:`Roman Yurchak <rth>`.

doc/whats_new/upcoming_changes/sklearn.neighbors/1140.efficiency.rst

@@ -0,0 +1,3 @@
+- :meth:`neighbors.KNeighborsRegressor.predict` is now faster for regression with
+  many output targets.
+  By :user:`Roman Yurchak <rth>`.

sklearn/metrics/pairwise.py

@@ -549,9 +549,11 @@ def nan_euclidean_distances(
         # This may not be the case due to floating point rounding errors.
         np.fill_diagonal(distances, 0.0)
 
-    present_X = 1 - missing_X
-    present_Y = present_X if Y is X else ~missing_Y
-    present_count = np.dot(present_X, present_Y.T)
+    # Cast the boolean presence masks to float so the matrix product runs
+    # through BLAS instead of the much slower integer matmul.
+    present_X = (~missing_X).astype(distances.dtype)
+    present_Y = present_X if Y is X else (~missing_Y).astype(distances.dtype)
+    present_count = present_X @ present_Y.T
     distances[present_count == 0] = np.nan
     # avoid divide by zero
     np.maximum(1, present_count, out=present_count)

sklearn/mixture/_gaussian_mixture.py

@@ -527,11 +527,16 @@ def _estimate_log_gaussian_prob(X, means, precisions_chol, covariance_type, xp=N
             log_prob[:, k] = xp.sum(xp.square(y), axis=1)
 
     elif covariance_type == "tied":
-        log_prob = xp.empty((n_samples, n_components), dtype=X.dtype, device=device_)
-        for k in range(means.shape[0]):
-            mu = means[k, :]
-            y = (X @ precisions_chol) - (mu @ precisions_chol)
-            log_prob[:, k] = xp.sum(xp.square(y), axis=1)
+        # In the tied case all components share precisions_chol, so project X
+        # and the means once and expand ||Xp - mu_proj||**2 (as in the diag and
+        # spherical branches below).
+        Xp = X @ precisions_chol
+        mu_proj = means @ precisions_chol
+        log_prob = (
+            row_norms(mu_proj, squared=True)
+            - 2.0 * (Xp @ mu_proj.T)
+            + row_norms(Xp, squared=True)[:, xp.newaxis]
+        )
 
     elif covariance_type == "diag":
         precisions = precisions_chol**2

sklearn/neighbors/_regression.py

@@ -259,12 +259,15 @@ def predict(self, X):
         if weights is None:
             y_pred = np.mean(_y[neigh_ind], axis=1)
         else:
-            y_pred = np.empty((neigh_dist.shape[0], _y.shape[1]), dtype=np.float64)
+            # Weighted sum of the neighbors' targets over the neighbor axis,
+            # for every query and output column. Equivalent to
+            #     np.sum(_y[neigh_ind] * weights[:, :, None], axis=1)
+            # but the einsum avoids materializing that full product. Here
+            # _y[neigh_ind] is (n_queries, n_neighbors, n_outputs) and weights
+            # is (n_queries, n_neighbors).
             denom = np.sum(weights, axis=1)
-
-            for j in range(_y.shape[1]):
-                num = np.sum(_y[neigh_ind, j] * weights, axis=1)
-                y_pred[:, j] = num / denom
+            num = np.einsum("ijk,ij->ik", _y[neigh_ind], weights)
+            y_pred = num / denom[:, None]
 
         if self._y.ndim == 1:
             y_pred = y_pred.ravel()