FIX Approximate nearest neighbors in TSNE example (scikit-learn#19809)

Author: TomDLT

examples/neighbors/approximate_nearest_neighbors.py

@@ -8,11 +8,6 @@
 replace KNeighborsTransformer and perform approximate nearest neighbors.
 These packages can be installed with `pip install annoy nmslib`.
 
-Note: Currently `TSNE(metric='precomputed')` does not modify the precomputed
-distances, and thus assumes that precomputed euclidean distances are squared.
-In future versions, a parameter in TSNE will control the optional squaring of
-precomputed distances (see #12401).
-
 Note: In KNeighborsTransformer we use the definition which includes each
 training point as its own neighbor in the count of `n_neighbors`, and for
 compatibility reasons, one extra neighbor is computed when
@@ -91,7 +86,6 @@ def fit(self, X):
         # see more metric in the manual
         # https://github.com/nmslib/nmslib/tree/master/manual
         space = {
-            'sqeuclidean': 'l2',
             'euclidean': 'l2',
             'cosine': 'cosinesimil',
             'l1': 'l1',
@@ -115,9 +109,6 @@ def transform(self, X):
         indices, distances = zip(*results)
         indices, distances = np.vstack(indices), np.vstack(distances)
 
-        if self.metric == 'sqeuclidean':
-            distances **= 2
-
         indptr = np.arange(0, n_samples_transform * n_neighbors + 1,
                            n_neighbors)
         kneighbors_graph = csr_matrix((distances.ravel(), indices.ravel(),
@@ -139,8 +130,7 @@ def __init__(self, n_neighbors=5, metric='euclidean', n_trees=10,
 
     def fit(self, X):
         self.n_samples_fit_ = X.shape[0]
-        metric = self.metric if self.metric != 'sqeuclidean' else 'euclidean'
-        self.annoy_ = annoy.AnnoyIndex(X.shape[1], metric=metric)
+        self.annoy_ = annoy.AnnoyIndex(X.shape[1], metric=self.metric)
         for i, x in enumerate(X):
             self.annoy_.add_item(i, x.tolist())
         self.annoy_.build(self.n_trees)
@@ -177,9 +167,6 @@ def _transform(self, X):
                     x.tolist(), n_neighbors, self.search_k,
                     include_distances=True)
 
-        if self.metric == 'sqeuclidean':
-            distances **= 2
-
         indptr = np.arange(0, n_samples_transform * n_neighbors + 1,
                            n_neighbors)
         kneighbors_graph = csr_matrix((distances.ravel(), indices.ravel(),
@@ -209,7 +196,7 @@ def test_transformers():
 
 def load_mnist(n_samples):
     """Load MNIST, shuffle the data, and return only n_samples."""
-    mnist = fetch_openml("mnist_784")
+    mnist = fetch_openml("mnist_784", as_frame=False)
     X, y = shuffle(mnist.data, mnist.target, random_state=2)
     return X[:n_samples] / 255, y[:n_samples]
 
@@ -222,34 +209,39 @@ def run_benchmark():
 
     n_iter = 500
     perplexity = 30
+    metric = "euclidean"
     # TSNE requires a certain number of neighbors which depends on the
     # perplexity parameter.
     # Add one since we include each sample as its own neighbor.
     n_neighbors = int(3. * perplexity + 1) + 1
 
+    tsne_params = dict(perplexity=perplexity, method="barnes_hut",
+                       random_state=42, n_iter=n_iter,
+                       square_distances=True)
+
     transformers = [
-        ('AnnoyTransformer', AnnoyTransformer(n_neighbors=n_neighbors,
-                                              metric='sqeuclidean')),
-        ('NMSlibTransformer', NMSlibTransformer(n_neighbors=n_neighbors,
-                                                metric='sqeuclidean')),
-        ('KNeighborsTransformer', KNeighborsTransformer(
-            n_neighbors=n_neighbors, mode='distance', metric='sqeuclidean')),
-        ('TSNE with AnnoyTransformer', make_pipeline(
-            AnnoyTransformer(n_neighbors=n_neighbors, metric='sqeuclidean'),
-            TSNE(metric='precomputed', perplexity=perplexity,
-                 method="barnes_hut", random_state=42, n_iter=n_iter), )),
-        ('TSNE with NMSlibTransformer', make_pipeline(
-            NMSlibTransformer(n_neighbors=n_neighbors, metric='sqeuclidean'),
-            TSNE(metric='precomputed', perplexity=perplexity,
-                 method="barnes_hut", random_state=42, n_iter=n_iter), )),
-        ('TSNE with KNeighborsTransformer', make_pipeline(
-            KNeighborsTransformer(n_neighbors=n_neighbors, mode='distance',
-                                  metric='sqeuclidean'),
-            TSNE(metric='precomputed', perplexity=perplexity,
-                 method="barnes_hut", random_state=42, n_iter=n_iter), )),
+        ('AnnoyTransformer',
+         AnnoyTransformer(n_neighbors=n_neighbors, metric=metric)),
+        ('NMSlibTransformer',
+         NMSlibTransformer(n_neighbors=n_neighbors, metric=metric)),
+        ('KNeighborsTransformer',
+         KNeighborsTransformer(n_neighbors=n_neighbors, mode='distance',
+                               metric=metric)),
+        ('TSNE with AnnoyTransformer',
+         make_pipeline(
+             AnnoyTransformer(n_neighbors=n_neighbors, metric=metric),
+             TSNE(metric='precomputed', **tsne_params))),
+        ('TSNE with NMSlibTransformer',
+         make_pipeline(
+             NMSlibTransformer(n_neighbors=n_neighbors, metric=metric),
+             TSNE(metric='precomputed', **tsne_params))),
+        ('TSNE with KNeighborsTransformer',
+         make_pipeline(
+             KNeighborsTransformer(n_neighbors=n_neighbors, mode='distance',
+                                   metric=metric),
+             TSNE(metric='precomputed', **tsne_params))),
         ('TSNE with internal NearestNeighbors',
-         TSNE(metric='sqeuclidean', perplexity=perplexity, method="barnes_hut",
-              random_state=42, n_iter=n_iter)),
+         TSNE(metric=metric, **tsne_params)),
     ]
 
     # init the plot