Using pairwise_distances_chunked to compute nearest neighbor search

Viewed 313

I have a long skinny data matrix (size: 250,000 x 10), which I will denote X. I also have a vector p measuring the quality of my data points. My goal is to compute the following function for each row x in my data matrix X:

r(x) = min{ ||x-y|| | p[y]>p[x], y in X }

On a smaller dataset, what I would use sklearn.metrics.pairwise_distances to precompute distances, like so:

from sklearn import metrics
n = len(X);

D_full = metrics.pairwise_distances(X);
r = np.zeros((n,1));
for i in range(n):
    r[i] = (D_full[i,p>p[i]]).min();

However, the above approach is memory-expensive, since I need to store D_full: a full n x n matrix. It seems like sklearn.metrics.pairwise_distances_chunked could be a good tool for this sort of problem since the distance matrix is only stored one chunk at a time. I was hoping to get some assistance in how to use it though, as I'm currently unfamiliar with generator objects. Suppose I call the following:

from sklearn import metrics
D = metrics.pairwise_distances_chunked(X);
D_chunk = next(D)

The above code yields D (a generator object) and D_chunk (a 536 x n array). Does D_chunkcorrespond to the first 536 rows of the matrix D_full from my earlier approach? If so, does next(D_chunk) correspond to the next 536 rows?

Thank you very much for your help.

2 Answers

This is a outline of a possible solution, but details are missing. In short, I would do the following:

Create a BallTree to query, and initialise min_quality_distance of size 250000 with say zeros.

for k=2

  1. For each vector, find the closest k neighbour (including itself).
  2. If vector with most distance within k found, has sufficient quality, update min_quality_distance for that point.
  3. For remaining, repeat with k=k+1

In each iteration, we have to query less vectors. The idea is that in each iteration you nibble a few nearest neighbors with the right condition away, and it will be easier with every step. (50% easier?) I will show how to do the first iteration, and with this is should be possible to build the loop.

You can do;

import numpy as np
size = 250000

X = np.random.random( size=(size,10))
p = np.random.random( size=size)

And create a BallTree with

from sklearn.neighbors import BallTree

tree = BallTree(X, leaf_size=10, metric='minkowski')

and query it for first iteration with (this will take about 5 minutes.)

k_nearest = 2

distances, indici = tree.query(X, k=k_nearest, return_distance=True, dualtree=False, sort_results=True)

The indici of most far away point within the nearest k is

most_far_away_indici = indici[:,-1:]

And its quality

p[most_far_away_indici]

So we can

quality_closeby = p[most_far_away_indici]

And check if it is sifficient with

indici_sufficient_quality = quality_closeby > np.expand_dims(p, axis=1)

And we have

found_closeby = np.all( indici_sufficient_quality, axis=1 )

Which is True is we have found a sufficient quality nearby.

We can update the vector with

distances_nearby = distances[:,-1:]

rx = np.zeros(size)
rx[found_closeby] = distances_nearby[found_closeby][:,0]

And we now need to take care for the remaining where we were unlucky, these are

~found_closeby

so

indici_not_found = ~found_closeby

and

distances, indici = tree.query(X[indici_not_found], k=3, return_distance=True, dualtree=False, sort_results=True)

etc..

I am sure the first few loops will take minutes, but after a few iterations the speeds will quickly go to seconds.

It is a little exercise with np.argwhere() etc to make sure the right indicis get updates.

It might not be the fastest, but it is a workable approach.

Since one cannot know the dimensions of some chunk, I suggest using np.ones_like instead of np.zeros.

Related