scikit-learn/benchmarks/bench_plot_neighbors.py

"""
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(return_X_y=True)
        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 = {alg: np.zeros(len(Nrange)) for alg in algorithms}
    N_results_query = {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 = {alg: np.zeros(len(Drange)) for alg in algorithms}
    D_results_query = {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 = {alg: np.zeros(len(krange)) for alg in algorithms}
    k_results_query = {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()