Added query_ball_point and related tests to scipy.spatial.cKDTree.

scipy.spatial.cKDTree implements only part of the more general scipy.spatial.KDTree interface in Cython.  This pull requests ports over the method query_ball_point to cKDTree and adds corresponding unit tests to verify correctness.

I have used a slightly modified version of Anne Archibald's benchmarking script to measure the resulting performance gains (see http://mail.scipy.org/pipermail/scipy-dev/2008-October/009854.html, script copied below).  The corresponding numbers show a ~ 250-350-fold speedup (this seems relatively insensitive to the query ball radius, see below):

dimension 3, 10000 points
KDTree constructed:	0.0823722
cKDTree constructed:	0.00159502
KDTree 1000 lookups:	0.574862
cKDTree 1000 lookups:	0.00470614
flat cKDTree 1000 lookups:	0.342408
KDTree 1000 ball lookups:	5.54947
cKDTree 1000 ball lookups:	0.0189772
flat cKDTree 1000 ball lookups:	0.343151
Ball lookups agree? True

dimension 8, 10000 points
KDTree constructed:	0.168782
cKDTree constructed:	0.00218606
KDTree 1000 lookups:	10.7511
cKDTree 1000 lookups:	0.091007
flat cKDTree 1000 lookups:	0.675035
KDTree 1000 ball lookups:	6.13426
cKDTree 1000 ball lookups:	0.017458
flat cKDTree 1000 ball lookups:	0.351695
Ball lookups agree? True

dimension 16, 10000 points
KDTree constructed:	0.11167
cKDTree constructed:	0.00220299
KDTree 1000 lookups:	62.2512
cKDTree 1000 lookups:	1.01098
flat cKDTree 1000 lookups:	1.44906
KDTree 1000 ball lookups:	3.89043
cKDTree 1000 ball lookups:	0.015486
flat cKDTree 1000 ball lookups:	0.352298
Ball lookups agree? True

For a probe radius of 0.5 instead of 0.2 (see script for meaning), I get the following results in 3 dimensions:

dimension 3, 10000 points
KDTree constructed:	0.0826931
cKDTree constructed:	0.00157905
KDTree 1000 lookups:	0.556555
cKDTree 1000 lookups:	0.00440192
flat cKDTree 1000 lookups:	0.338442
KDTree 1000 ball lookups:	16.3933
cKDTree 1000 ball lookups:	0.0493851
flat cKDTree 1000 ball lookups:	0.423063
Ball lookups agree? True

As far as I can tell, the issue is that KDTree.query_ball_point is written for readability, but ends up doing large numbers of list operations to build the results.

From here on, it should not be much harder to port over the remaining methods of KDTree to cKDTree, although these are not a current priority for me.  As an aside, it seemed at first glance that the public interfaces for KDTree.query() and cKDTree.query() have fallen out of sync, though I didn't check thoroughly.

All the best,
Patrick

---------------------

import numpy as np
import time

from scipy.spatial import KDTree, cKDTree



m = 3
n = 10000
r = 1000

data = np.concatenate((np.random.randn(n//2,m),
    np.random.randn(n-n//2,m)+np.ones(m)))
queries = np.concatenate((np.random.randn(r//2,m),
    np.random.randn(r-r//2,m)+np.ones(m)))

print "dimension %d, %d points" % (m,n)

t = time.time()
T1 = KDTree(data)
print "KDTree constructed:\t%g" % (time.time()-t)
t = time.time()
T2 = cKDTree(data)
print "cKDTree constructed:\t%g" % (time.time()-t)

t = time.time()
w = T1.query(queries)
print "KDTree %d lookups:\t%g" % (r, time.time()-t)
del w

t = time.time()
w = T2.query(queries)
print "cKDTree %d lookups:\t%g" % (r, time.time()-t)
del w

T3 = cKDTree(data,leafsize=n)
t = time.time()
w = T3.query(queries)
print "flat cKDTree %d lookups:\t%g" % (r, time.time()-t)
del w

t = time.time()
w1 = T1.query_ball_point(queries, 0.2)
print "KDTree %d ball lookups:\t%g" % (r, time.time()-t)

t = time.time()
w2 = T2.query_ball_point(queries, 0.2)
print "cKDTree %d ball lookups:\t%g" % (r, time.time()-t)

t = time.time()
w3 = T3.query_ball_point(queries, 0.2)
print "flat cKDTree %d ball lookups:\t%g" % (r, time.time()-t)

all_good = True
for a, b in zip(w1, w2):
    if sorted(a) != sorted(b):
        all_good = False
for a, b in zip(w1, w3):
    if sorted(a) != sorted(b):
        all_good = False

print "Ball lookups agree? %s" % str(all_good)