Python: scipy.spatial.KDTree test

this is a short script to test scipy.spatial.KDTree that

1. synthesizes a set of database points
2. builds a kd-tree
3. synthesizes a set of query points
4. runs kNN search using the built kd-tree
5. plots the found kNN points

The result plot is like

[code lang="python"] import numpy as np import matplotlib.pyplot as plt from scipy import spatial

def synthesizeData(size,dim,scale=1.0): data = np.empty*1 for d in range(dim): data[:,d] = scale*np.random.random(size) return data

def main(): # settings scale = 100 size = 1000 dim = 2

# build kd-tree
data = synthesizeData(size, dim, scale)
tree = spatial.KDTree(data)

# search
for numQuery in range(1,5):
    print 'search ', numQuery, 'queries'
    query = synthesizeData(numQuery, dim, scale)
    knnIndex = tree.query_ball_point(query, 15.0)
    plt.plot(data[:,0], data[:,1], 'ro')
    for k in range(numQuery):
        print '    plot ', k, '-th data'
        if k == 0:
            plt.plot(data[knnIndex[k],0], data[knnIndex[k],1], 'bo')
        if k == 1:
            plt.plot(data[knnIndex[k],0], data[knnIndex[k],1], 'go')
        if k == 2:
            plt.plot(data[knnIndex[k],0], data[knnIndex[k],1], 'yo')
        if k == 3:
            plt.plot(data[knnIndex[k],0], data[knnIndex[k],1], 'ko')

if name == 'main': main() [/code]

*1:size, dim