How can I speed up a loop that queries a kd-tree?

Question

The following section of my code is taking ages to run (it's the only loop in the function, so it's the most likely culprit):

tree = KDTree(x_rest)
for i in range(len(x_lost)):
    _, idx = tree.query([x_lost[i]], k=int(np.sqrt(len(x_rest))), p=1)
    y_lost[i] = mode(y_rest[idx][0])[0][0]

Is there a way to speed this up? I have a few suggestions from Stack Overflow:

• This answer suggests using Cython. I'm not particularly familiar with it, but I'm not very against it either.
• This answer uses a multiprocessing Pool. I'm not sure how useful this will be: my current execution takes over 12h to run, and I'm hoping for at least a 5-10x speedup (though that may not be possible).

Answer 1

Here are a few notes about how you could speed this up:

1. This code loops over x_rest, and calls tree.query() with one point from x_rest at a time. However, query() supports querying multiple points at once. The loop inside query() is implemented in Cython, so I would expect it to be much faster than a loop written in Python. If you call it like this, it will return an array of matches.

2. The query() function supports a parameter called workers, which if set to a value larger than one, runs your query in parallel. Since workers is implemented using threads, it will likely be faster than a solution using multiprocessing.Pool, since it avoids pickling. See the documentation.

3. The code above doesn't define the mode() function, but I'm assuming it's scipy.stats.mode(). If that's the case, rather than calling mode() repeatedly, you can use the axis argument, which would let you take the mode of nearby points for multiple queries at once.