Slow python code that automatically selects k for KNNG

Question

This code is used to select K for a K-nearest neighbor graph [KNNG] from a dataset with an unknown number of centroids (which is not the same as K-means clustering).

Suppose that you have a dataset of observations stored in a data matrix X[n_samples, n_features] with each row being an observation or feature vector and each column being a feature. Now suppose you want to compute the (weighted) k-Neighbors graph for points in X using sklearn.neighbors.kneighbors_graph. The first step is to pick k.

This brute force method to pick k below is very slow when the sample dataset size becomes large (ie n_neighbors_max = 10000). The FOR loop is very slow. Any tips for how to speed it up?

def autoselect_K(X, n_neighbors_max, threshold):
    # get the pairwise euclidean distance between every observation
    D = sklearn.metrics.pairwise.euclidean_distances(X, X)
    chosen_k = n_neighbors_max
    for k in range(2, n_neighbors_max):
        k_avg = []
        # loop over each row in the distance matrix
        for row in D:
            # sort the row from smallest distance to largest distance
            sorted_row = numpy.sort(row)
            # calculate the mean of the smallest k+1 distances
            k_avg.append(numpy.mean(sorted_row[0:k]))
        # find the median of the averages
        kmedian_dist = numpy.median(k_avg)
        if kmedian_dist >= threshold:
            chosen_k = k
            break
    # return the number of nearest neighbors to use
    return chosen_k

Answer 1

• Without thinking too hard about what you are actually doing (i.e., by trying to find a considerable more efficient algorithm), you could get a decent speedup by using list comprehension. Doing appends inside an explicit for loop doesn't usually lead to high performance.

• Another observation is that you are doing too much work in sorting the whole row, when you just want the first k elements in sorted order. For that, there's a more suitable function like numpy.argpartition or heapq.nsmallest.

I don't have a test program for correctness, but I wrote a simple check for performance. But even then, I have no idea about your actual use case and parameters, so take it with a grain of salt:

import numpy as np
from sklearn.metrics.pairwise import euclidean_distances
from time import time

def autoselect_K_orig(X, n_neighbors_max, threshold):
    # get the pairwise euclidean distance between every observation
    D = euclidean_distances(X, X)
    chosen_k = n_neighbors_max
    for k in range(2, n_neighbors_max):
        k_avg = []
        # loop over each row in the distance matrix
        for row in D:
            # sort the row from smallest distance to largest distance
            sorted_row = np.sort(row)
            # calculate the mean of the smallest k+1 distances
            k_avg.append(np.mean(sorted_row[0:k]))
        # find the median of the averages
        kmedian_dist = np.median(k_avg)
        if kmedian_dist >= threshold:
            chosen_k = k
            break
    # return the number of nearest neighbors to use
    return chosen_k

def autoselect_K_other(X, n_neighbors_max, threshold):
    D = euclidean_distances(X, X)
    for k in range(2, n_neighbors_max):
        kmedian_dist = np.median([np.mean(row[np.argpartition(row, k)]) for row in D])
        if kmedian_dist >= threshold:
            return k

    return n_neighbors_max

# Rudimentary performance check
X = np.random.randint(5, size=(10000, 10))

start = time()
result = autoselect_K_orig(X, 20, 4)
end = time()
total_orig = end - start
print('Original: ', total_orig)

start = time()
result = autoselect_K_other(X, 20, 4)
end = time()
total_other = end - start
print('Other: ', total_other)

print('Speedup: ', total_orig / total_other)

On my outdated hardware, a typical output is:

Original:  57.56129217147827
Other:  2.257129192352295
Speedup: 25.501992693422242

Meaning speedups of more than 25x. Not surprisingly, this speedup is strongly linked to the value of n_neighbors_max; if you make it very small, the speedup tends to go as low as 2x, but as you can see, it grows quite considerably for larger values, which is probably exactly what you want.

• I think further optimization even with your current approach is still possible. For instance, notice that everytime that you increment k, you are essentially duplicating work you have already done by again taking the first (k-1) elements, taking their mean, and so on. Instead, for step k, you should have stored the mean of step (k-1), and then you simply update this mean with a new element that now comes in, reducing the amount of computation you need to do.

• Even without such an optimization, you could still get some gains by getting rid of the explicit for loop that goes from 2 to n_neighbors_max by using a suitable generator expression that asks for more elements for as long as we've not extracted an element of at least the threshold.