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
scipy.spatial.KDTree result

import numpy as np
import matplotlib.pyplot as plt
from scipy import spatial

def synthesizeData(size,dim,scale=1.0):
    data = np.empty((size, dim))
    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()
Advertisements

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s