text: """
  =============================
  MNIST dataset T-SNE benchmark
  =============================
  
  """
  
  # License: BSD 3 clause
  
  import os
  import os.path as op
  from time import time
  import numpy as np
  import json
  import argparse
  from joblib import Memory
  
  from sklearn.datasets import fetch_openml
  from sklearn.manifold import TSNE
  from sklearn.neighbors import NearestNeighbors
  from sklearn.decomposition import PCA
  from sklearn.utils import check_array
  from sklearn.utils import shuffle as _shuffle
  from sklearn.utils._openmp_helpers import _openmp_effective_n_threads
  
  LOG_DIR = "mnist_tsne_output"
  if not os.path.exists(LOG_DIR):
      os.mkdir(LOG_DIR)
  
  
  memory = Memory(os.path.join(LOG_DIR, "mnist_tsne_benchmark_data"), mmap_mode="r")
  
  
  @memory.cache
  def load_data(dtype=np.float32, order="C", shuffle=True, seed=0):
      """Load the data, then cache and memmap the train/test split"""
      print("Loading dataset...")
      data = fetch_openml("mnist_784")
  
      X = check_array(data["data"], dtype=dtype, order=order)
      y = data["target"]
  
      if shuffle:
          X, y = _shuffle(X, y, random_state=seed)
  
      # Normalize features
      X /= 255
      return X, y
  
  
  def nn_accuracy(X, X_embedded, k=1):
      """Accuracy of the first nearest neighbor"""
      knn = NearestNeighbors(n_neighbors=1, n_jobs=-1)
      _, neighbors_X = knn.fit(X).kneighbors()
      _, neighbors_X_embedded = knn.fit(X_embedded).kneighbors()
      return np.mean(neighbors_X == neighbors_X_embedded)
  
  
  def tsne_fit_transform(model, data):
      transformed = model.fit_transform(data)
      return transformed, model.n_iter_
  
  
  def sanitize(filename):
      return filename.replace("/", "-").replace(" ", "_")
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser("Benchmark for t-SNE")
      parser.add_argument(
          "--order", type=str, default="C", help="Order of the input data"
      )
      parser.add_argument("--perplexity", type=float, default=30)
      parser.add_argument(
          "--bhtsne",
          action="store_true",
          help=(
              "if set and the reference bhtsne code is "
              "correctly installed, run it in the benchmark."
          ),
      )
      parser.add_argument(
          "--all",
          action="store_true",
          help=(
              "if set, run the benchmark with the whole MNIST."
              "dataset. Note that it will take up to 1 hour."
          ),
      )
      parser.add_argument(
          "--profile",
          action="store_true",
          help="if set, run the benchmark with a memory profiler.",
      )
      parser.add_argument("--verbose", type=int, default=0)
      parser.add_argument(
          "--pca-components",
          type=int,
          default=50,
          help="Number of principal components for preprocessing.",
      )
      args = parser.parse_args()
  
      print("Used number of threads: {}".format(_openmp_effective_n_threads()))
      X, y = load_data(order=args.order)
  
      if args.pca_components > 0:
          t0 = time()
          X = PCA(n_components=args.pca_components).fit_transform(X)
          print(
              "PCA preprocessing down to {} dimensions took {:0.3f}s".format(
                  args.pca_components, time() - t0
              )
          )
  
      methods = []
  
      # Put TSNE in methods
      tsne = TSNE(
          n_components=2,
          init="pca",
          perplexity=args.perplexity,
          verbose=args.verbose,
          n_iter=1000,
      )
      methods.append(("sklearn TSNE", lambda data: tsne_fit_transform(tsne, data)))
  
      if args.bhtsne:
          try:
              from bhtsne.bhtsne import run_bh_tsne
          except ImportError as e:
              raise ImportError(
                  """\
  If you want comparison with the reference implementation, build the
  binary from source (https://github.com/lvdmaaten/bhtsne) in the folder
  benchmarks/bhtsne and add an empty `__init__.py` file in the folder:
  
  $ git clone git@github.com:lvdmaaten/bhtsne.git
  $ cd bhtsne
  $ g++ sptree.cpp tsne.cpp tsne_main.cpp -o bh_tsne -O2
  $ touch __init__.py
  $ cd ..
  """
              ) from e
  
          def bhtsne(X):
              """Wrapper for the reference lvdmaaten/bhtsne implementation."""
              # PCA preprocessing is done elsewhere in the benchmark script
              n_iter = -1  # TODO find a way to report the number of iterations
              return (
                  run_bh_tsne(
                      X,
                      use_pca=False,
                      perplexity=args.perplexity,
                      verbose=args.verbose > 0,
                  ),
                  n_iter,
              )
  
          methods.append(("lvdmaaten/bhtsne", bhtsne))
  
      if args.profile:
  
          try:
              from memory_profiler import profile
          except ImportError as e:
              raise ImportError(
                  "To run the benchmark with `--profile`, you "
                  "need to install `memory_profiler`. Please "
                  "run `pip install memory_profiler`."
              ) from e
          methods = [(n, profile(m)) for n, m in methods]
  
      data_size = [100, 500, 1000, 5000, 10000]
      if args.all:
          data_size.append(70000)
  
      results = []
      basename = os.path.basename(os.path.splitext(__file__)[0])
      log_filename = os.path.join(LOG_DIR, basename + ".json")
      for n in data_size:
          X_train = X[:n]
          y_train = y[:n]
          n = X_train.shape[0]
          for name, method in methods:
              print("Fitting {} on {} samples...".format(name, n))
              t0 = time()
              np.save(
                  os.path.join(LOG_DIR, "mnist_{}_{}.npy".format("original", n)), X_train
              )
              np.save(
                  os.path.join(LOG_DIR, "mnist_{}_{}.npy".format("original_labels", n)),
                  y_train,
              )
              X_embedded, n_iter = method(X_train)
              duration = time() - t0
              precision_5 = nn_accuracy(X_train, X_embedded)
              print(
                  "Fitting {} on {} samples took {:.3f}s in {:d} iterations, "
                  "nn accuracy: {:0.3f}".format(name, n, duration, n_iter, precision_5)
              )
              results.append(dict(method=name, duration=duration, n_samples=n))
              with open(log_filename, "w", encoding="utf-8") as f:
                  json.dump(results, f)
              method_name = sanitize(name)
              np.save(
                  op.join(LOG_DIR, "mnist_{}_{}.npy".format(method_name, n)), X_embedded
              )
  