This script compares the performance of three different methods for finding the nearest neighbors from a list of coordinates. The first method is a simple 'for loop'. The second method uses the KDTree object from the scipy.spatial library. The third method uses cKDTree from the scipy.spatial library. Both the KDTree and cKDTree methods are signi…

kdtree_evaluation.py

import numpy as np
import scipy.spatial as ss
from itertools import combinations
from itertools import permutations
import math
import random
from time import clock


def length(point1, point2):
    return math.sqrt((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2)


def find_node_v1(points, nodeCount, node):
    best_seg = 10000000
    for i in xrange(0, nodeCount):
	if node != i:
	    seg_len = length(points[node], points[i])
	    if seg_len < best_seg:
		best_seg = seg_len
		best_nieghbor = i
    return best_nieghbor


def get_kdtree(points):
    tree = ss.KDTree(points)
    return tree


def get_ckdtree(points):
    tree = ss.cKDTree(points)
    return tree


def query_tree(tree, points, node):
    q_result = tree.query(points[node], k=2)
    return q_result[1][1]


def main():
    nodeCount = 10000
    # create a random set of coordinates
    points = []
    for i in xrange(0, nodeCount):
	points.append((random.randint(0, nodeCount), random.randint(0, nodeCount)))

    # Run a time trial for three different methods for finding the nearest nieghbor
    # for nodes 0-1000
    
    # Method 1 using a for loop
    
    t0 = clock()
    nieghbors1 = []
    for i in xrange(0,100):
	ind = find_node_v1(points, nodeCount, i)
	nieghbors1.append(points[ind])
    t1 = clock()
    print "Method 1 (using for loops) time: ", t1-t0

    # Method 2 using KDTree

    t0 = clock()
    nieghbors2 = []
    tree = get_kdtree(points)
    for i in xrange(0,100):
	ind = query_tree(tree, points, i)
	nieghbors2.append(points[ind])
    t1 = clock()
    print "Method 2 (using kdtree) time: ", t1-t0

    # Method 3 using cKDTree

    t0 = clock()
    nieghbors3 = []
    tree = get_ckdtree(points)
    for i in xrange(0,100):
	ind = query_tree(tree, points, i)
	nieghbors3.append(points[ind])
    t1 = clock()
    print "Method 3 (using ckdtree) time: ", t1-t0

    assert nieghbors1 == nieghbors2
    assert nieghbors1 == nieghbors3


if __name__ == "__main__":
    main()