content: """
  Plot the scaling of the nearest neighbors algorithms with k, D, and N
  """
  from time import time
  
  import numpy as np
  import matplotlib.pyplot as plt
  from matplotlib import ticker
  
  from sklearn import neighbors, datasets
  
  
  def get_data(N, D, dataset='dense'):
      if dataset == 'dense':
          np.random.seed(0)
          return np.random.random((N, D))
      elif dataset == 'digits':
          X = datasets.load_digits().data
          i = np.argsort(X[0])[::-1]
          X = X[:, i]
          return X[:N, :D]
      else:
          raise ValueError("invalid dataset: %s" % dataset)
  
  
  def barplot_neighbors(Nrange=2 ** np.arange(1, 11),
                        Drange=2 ** np.arange(7),
                        krange=2 ** np.arange(10),
                        N=1000,
                        D=64,
                        k=5,
                        leaf_size=30,
                        dataset='digits'):
      algorithms = ('kd_tree', 'brute', 'ball_tree')
      fiducial_values = {'N': N,
                         'D': D,
                         'k': k}
  
      #------------------------------------------------------------
      # varying N
      N_results_build = dict([(alg, np.zeros(len(Nrange)))
                              for alg in algorithms])
      N_results_query = dict([(alg, np.zeros(len(Nrange)))
                              for alg in algorithms])
  
      for i, NN in enumerate(Nrange):
          print("N = %i (%i out of %i)" % (NN, i + 1, len(Nrange)))
          X = get_data(NN, D, dataset)
          for algorithm in algorithms:
              nbrs = neighbors.NearestNeighbors(n_neighbors=min(NN, k),
                                                algorithm=algorithm,
                                                leaf_size=leaf_size)
              t0 = time()
              nbrs.fit(X)
              t1 = time()
              nbrs.kneighbors(X)
              t2 = time()
  
              N_results_build[algorithm][i] = (t1 - t0)
              N_results_query[algorithm][i] = (t2 - t1)
  
      #------------------------------------------------------------
      # varying D
      D_results_build = dict([(alg, np.zeros(len(Drange)))
                              for alg in algorithms])
      D_results_query = dict([(alg, np.zeros(len(Drange)))
                              for alg in algorithms])
  
      for i, DD in enumerate(Drange):
          print("D = %i (%i out of %i)" % (DD, i + 1, len(Drange)))
          X = get_data(N, DD, dataset)
          for algorithm in algorithms:
              nbrs = neighbors.NearestNeighbors(n_neighbors=k,
                                                algorithm=algorithm,
                                                leaf_size=leaf_size)
              t0 = time()
              nbrs.fit(X)
              t1 = time()
              nbrs.kneighbors(X)
              t2 = time()
  
              D_results_build[algorithm][i] = (t1 - t0)
              D_results_query[algorithm][i] = (t2 - t1)
  
      #------------------------------------------------------------
      # varying k
      k_results_build = dict([(alg, np.zeros(len(krange)))
                              for alg in algorithms])
      k_results_query = dict([(alg, np.zeros(len(krange)))
                              for alg in algorithms])
  
      X = get_data(N, DD, dataset)
  
      for i, kk in enumerate(krange):
          print("k = %i (%i out of %i)" % (kk, i + 1, len(krange)))
          for algorithm in algorithms:
              nbrs = neighbors.NearestNeighbors(n_neighbors=kk,
                                                algorithm=algorithm,
                                                leaf_size=leaf_size)
              t0 = time()
              nbrs.fit(X)
              t1 = time()
              nbrs.kneighbors(X)
              t2 = time()
  
              k_results_build[algorithm][i] = (t1 - t0)
              k_results_query[algorithm][i] = (t2 - t1)
  
      plt.figure(figsize=(8, 11))
  
      for (sbplt, vals, quantity,
           build_time, query_time) in [(311, Nrange, 'N',
                                        N_results_build,
                                        N_results_query),
                                       (312, Drange, 'D',
                                        D_results_build,
                                        D_results_query),
                                       (313, krange, 'k',
                                        k_results_build,
                                        k_results_query)]:
          ax = plt.subplot(sbplt, yscale='log')
          plt.grid(True)
  
          tick_vals = []
          tick_labels = []
  
          bottom = 10 ** np.min([min(np.floor(np.log10(build_time[alg])))
                                 for alg in algorithms])
  
          for i, alg in enumerate(algorithms):
              xvals = 0.1 + i * (1 + len(vals)) + np.arange(len(vals))
              width = 0.8
  
              c_bar = plt.bar(xvals, build_time[alg] - bottom,
                              width, bottom, color='r')
              q_bar = plt.bar(xvals, query_time[alg],
                              width, build_time[alg], color='b')
  
              tick_vals += list(xvals + 0.5 * width)
              tick_labels += ['%i' % val for val in vals]
  
              plt.text((i + 0.02) / len(algorithms), 0.98, alg,
                       transform=ax.transAxes,
                       ha='left',
                       va='top',
                       bbox=dict(facecolor='w', edgecolor='w', alpha=0.5))
  
              plt.ylabel('Time (s)')
  
          ax.xaxis.set_major_locator(ticker.FixedLocator(tick_vals))
          ax.xaxis.set_major_formatter(ticker.FixedFormatter(tick_labels))
  
          for label in ax.get_xticklabels():
              label.set_rotation(-90)
              label.set_fontsize(10)
  
          title_string = 'Varying %s' % quantity
  
          descr_string = ''
  
          for s in 'NDk':
              if s == quantity:
                  pass
              else:
                  descr_string += '%s = %i, ' % (s, fiducial_values[s])
  
          descr_string = descr_string[:-2]
  
          plt.text(1.01, 0.5, title_string,
                   transform=ax.transAxes, rotation=-90,
                   ha='left', va='center', fontsize=20)
  
          plt.text(0.99, 0.5, descr_string,
                   transform=ax.transAxes, rotation=-90,
                   ha='right', va='center')
  
          plt.gcf().suptitle("%s data set" % dataset.capitalize(), fontsize=16)
  
      plt.figlegend((c_bar, q_bar), ('construction', 'N-point query'),
                    'upper right')
  
  if __name__ == '__main__':
      barplot_neighbors(dataset='digits')
      barplot_neighbors(dataset='dense')
      plt.show()
  