186 lines
6.3 KiB
Python
186 lines
6.3 KiB
Python
"""
|
|
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()
|