Nearest Neighbours Using k-d Trees

This program calculates the set of nearest neighbours for each point in the input data set using the algorithm descibed in[1]. A k-d tree is built from the data set in O(n log n) time which allows the computation of the nearest neighbours for a particular query point in O(log n) time.

Reading In The Data

First the preliminaries - reading the data file. The following assumes a file format of
 "<input 1> <input 2> ... <input n>, <output 1> ... <output n>,", which allows multiple inputs and outputs:

filename = "Math/Gamma/hen100.asc";

filedata = Partition[ReadList[filename, Word, 
            RecordLists -> True,
            WordSeparators -> {" "},
            RecordSeparators -> {",", "\n"}], 2];
(* WARNING USING THIS SYNTAX IS CONVENIENT BUT THE RESULTS
 * ARE STRINGS NOT NUMBERS. TO FIX THIS WE DO THE FOLLOWING
 *)
data = Map[ToExpression, filedata];

This package usefully provides all the statistical tools needed

Needs["Statistics`DescriptiveStatistics`"]

Now the data is ready to use. It can be seen that the data consists of a list of input-output pairs. The nearest neighbours calculation is only to be performed on the input set, so the output set can be removed to simplify future processing. The following does this:

data = Table[data[[i]][[1]], {i, 1, Length[data]}];

Constructing the kd-Tree

The next step is to build the k-d tree. This is done recursively using two basic steps at each stage. First the dimension (or key) with the largest spread of values over the records in this branch is found, and the median of the record values for this dimension is found. Second, the data is partitioned into two sets about the median previously found; the tree building then continues with each of the smaller sets. The recursion halts when the number of records to be processed falls below some preset value; the records become a terminal or leaf node in the tree.

In order to do this,  the representation of the tree structure in Mathematica must be decided. At each non-terminal node, the dimension split on and the associated median value are recorded, together with the left and right branch of the recursively created subtrees. This leads to the following:

        Tree = Terminal Node |
                {Dimension Number, Partition Value, Left Tree, Right Tree}

This is complicted somewhat by the need to return a list of {record index, distance} pairs as the nearest neighbours for each node. As the tree-building mechanism destroys the ordering of the nodes, the simplest way to retain each node's index in the original list is to tag it with it. Each data point then becomes a pair {value vector, index number}.

In order to manipulate these without cluttering the code to much, a few data access functions have been defined. Given a list of index-tagged data points, "getValueFor" returns the value for dimension (d) at position (i) in the list, "getIndex" returns the index for the record at position (i) and "getPoint" returns the data vector (without index) for the record at position (i).

getValueFor[in_, i_, d_] := in[[i]][[1]][[d]];
getIndex[in_, i_] := in[[i]][[2]];
getPoint[in_, i_] := in[[i]][[1]];

The following function is used to tag each point in the data set with its index number. This is called before the tree is to be built:

addIndex[data_] := Table[{data[[i]], i}, {i,1,Length[data]}];

The following module is used to create the k-d tree. It needs to be passed the data set and the dimensionality of the data (the variable "k"). The local variable "b" controls the number of records in each terminal node  - the default is 1, but the paper by Bentley[1] mentions that: "... increasing the bucket size from one record per bucket considerably improves the performance of the search...". The paper by White and Jain[2] mentions that spliting the data on the dimension with the highest variance gives better query performance, and this should also be faster to compute. It might be interesting to compares these approaches.

buildTree[k_, inData_]  := 
        Module[
        (* "b" is the number of records per bucket *)
        {bucketsize = 1,  maxspread = 0, i, dimension,
         hi, low, thisHi, thisLow, medianval, pair, temp},
         
        If[Length[inData] <= bucketsize,
                (* This is a terminal node, so simply return the
                 * list as it was passed *)
                Return[inData],
                (* Find the dimension with the largest spread *)
                For[i = 1, i <= k, i++,
                        hi = -Infinity; low = Infinity;
                        For [j = 1, j <= Length[inData], j++,
                                temp = getValueFor[inData, j, i];
                                If[hi < temp, hi = temp, ];
                                If[low > temp, low = temp, ]
                        ];
                        If[hi - low >= maxspread,
                                maxspread = hi - low;
                                thisHi = hi; thisLow = low;
                                dimension = i, ]
                ];
                
                (* Extract values at [dimension] and take median *)
                medianval = Median[Map[Part[Part[#,1],dimension]&,
                                                        inData]];
                                                        
                (* If the median is the same as the hi or low values
                 * for the dimension, then repeated values must exist
                 * in the data set. Return the set without splitting
                 *)
                If[thisLow == medianval || thisHi == medianval,
                        Return[inData], 
                (* Else return the node in the tree with the correct
                 * subtrees *)
                        pair = split[dimension, medianval, inData];
                        Return[{dimension, medianval,
                                buildTree[k, First[pair]],
                                buildTree[k, Last[pair]]}]
                ]
        ]
];

The final conditional is necessary to deal correctly with data sets where  one or vectors are equal. If a particular vecotr is repeated enough, it is possible that the median can be
skewed to the maximum or minium of the range for the dimension. If this occurs, then the entire dataset is returned as a terminal node - no partitioning is possible as one branch will become empty, leading to an infinite loop.

The following is used during the creation of the tree to split the records around a particular value (v) for a particular dimension (d). It may be quicker to do this "in place" in a programming language such as C.

split[d_, v_, list_] := 
        Module[
        {i, hi = {}, low = {}},
        For[i = 1, i <= Length[list], i++,
                If [getValueFor[list, i, d] >= v,
                        hi = Append[hi, list[[i]]],
                        low = Append[low, list[[i]]]
                        ]
        ];
        Return[{low, hi}]
];

Finding the Nearest Neighbours

Once the tree has been built, the task is then to search through it for the nearest neighbours of each point. The module takes parameters for the query vector to look for (query), the dimensionality of the data (k) and the previously constructed k-d tree (tree). Several "global" variables are used:
        nearest:This is best set of nearest neighbours found at any one point in the 
                search. It is a pair of lists, one for the record indexes, the other for the
                distances to the query record - both lists are maintained in
                increasing distance from the query node.
        lowerBounds:    This is the set of lower bounds currently in use to define the 
                region being searched. There is one for each dimension.
        upperBounds:    These are the equivilent set for the upper bounds on the
                search region

The nearest neighbours are found by only searching those parts of the tree which contain data points "close" to the query point. The distance to the furthest nearest neighbour found so far and the size of the geometrical region defined by the tree branches traversed previously are compared to determine whether or not a branch should be examined. The comparison is carried out in the "boundsOverlapBall" procedure.

The search can terminate whenever the "ballWithinBounds" test returns True. Note that this test is not strictly required for the algorithm to work correctly - it is only used to indicate that the solution has been found.

searchTree[query_, k_, tree_] :=
        Module[
        {medianval, dimension, temp, i, distance, maxloc, maxdist},
        
        Print["upper = ", upperBounds, " lower = ", lowerBounds,
                " tree = ", tree];
        If[Not[NumberQ[tree[[1]]]], 
                (* Node is terminal - update nearest neighbours *)
                For[i = 1, i <= Length[tree], i++,
                        (* Add points in bucket if closer than furthest
                         * found so far, so long as this is not the query
                         * record *)
                        maxdist = Max[nearest[[2]]];
                        distance = norm[query - tree[[i]][[1]]];
                        If [(distance < maxdist) &&
                                (query != getPoint[tree, i]),
                                maxloc = Flatten[Position[nearest[[2]], maxdist]][[1]];
                                nearest[[1, maxloc]] = getIndex[tree, i];
                                nearest[[2, maxloc]] = distance, ]
                ];
                (* Sort nearest neighbours list *)
                nearest = Transpose[Sort[Transpose[nearest],
                                        (Part[#1, 2] < Part[#2, 2])&]];
                
                (* Check for termination and return *)
                If[ballWithinBounds[query, k],
                        Print["**done**"]; Return[ ],
                        Return[ ]],
        ];
        
        (* Node not terminal - call searchTree recursively *)   
        dimension = tree[[1]];
        medianval = tree[[2]];
        (* Recursive call on closer son *)
        If[query[[dimension]] <= medianval,
                temp = upperBounds[[dimension]];
                upperBounds[[dimension]] =  medianval;
                searchTree[query, k, tree[[3]]];
                upperBounds[[dimension]] = temp, 
        (* else *)
                temp = lowerBounds[[dimension]];
                lowerBounds[[dimension]] = medianval;
                searchTree[query, k, tree[[4]]];
                lowerBounds[[dimension]] = temp
        ];
        
        (* Recursive call on further son if necessary *)
        If[query[[dimension]] <= medianval,
                temp = lowerBounds[[dimension]];
                lowerBounds[[dimension]] = medianval;
                If[boundsOverlapBall[query, k],
                        searchTree[query, k, tree[[4]]],
                ];
                lowerBounds[[dimension]] = temp,
        (* else *)
                temp = upperBounds[[dimension]];
                upperBounds[[dimension]] = medianval;
                If[boundsOverlapBall[query, k],
                        searchTree[query, k, tree[[3]]],
                ];
                upperBounds[[dimension]] = temp
        ];
        
        (* See if we should terminate *)
        If[ballWithinBounds[query, k],
                Print["**done**"]; Return[ ],
                Return[ ]]      
];
norm[x_] :=     N[Sqrt[x.x]];

The following two procedures are used to guide the search process. "ballWithinBounds" tests to see if the coordinate distance from the query record to the closer boundary along each dimension is less than the radius of the furthest nearest neighbour found so far. The test succeeds if all the coordinate distances are greater than the radius.

ballWithinBounds[query_, k_] := 
        Module[
        {d},
        For [d = 1, d <= k, d++,
                (* "Last[nearest[[2]]]" is the furthest
                 * nearest neighbour *)
                If [((query[[d]] - lowerBounds[[d]])^2 <=
                                Last[nearest[[2]]]) ||
                        ((query[[d]] - upperBounds[[d]])^2 <=
                                Last[nearest[[2]]]),
                        Return[False],
                ]
        ];
        Return[True]
];

The "boundsOverlapBall" procedure is used on backtracking to decide if it is necessary to search down the opposite branch of the tree to the one previously searched. It determines whether the geometric boundaries of the subfile under consideration overlap a ball of radius equal to the distance of the furthest nearest neighbour found so far. The smallest distance between the query record and the bounded region is found - if this is less than the radius, then the branch of the subtree can be eliminated, as there is no overlap between the domain to be searched and the nearest neighbours.

boundsOverlapBall[query_, k_] :=
        Module[
        {total = 0, d},
        (* "Last[nearest[[2]]]" is the furthest
         * nearest neighbour *)
        For[d = 1, d <= k, d++,
                If[query[[d]] < lowerBounds[[d]],
                        total = total + (query[[d]] - lowerBounds[[d]])^2;
                                If[total > (Last[nearest[[2]]])^2,
                                Return[False],
                                (* else nothing *)
                        ],
                (* else *)
                        If[query[[d]] > upperBounds[[d]],
                                total = total + (query[[d]] - upperBounds[[d]])^2;
                                If[total > (Last[nearest[[2]]])^2,
                                        Return[False],
                                        (* else nothing *)
                                ],
                        (* else nothing *)
                        ]
                ]
        ];
        Return[True]    
];

The following module initialises the global variables to be used in the search. The number of nearest neighbours to be found is needed (numNearest), as is the query record (query), so that the dimensionality of the data cen be found. This initialisation must be repeated for each query point.

(* Initialisations for globals *)
initialise[query_, numNearest_] :=
        Module[
        {},
        nearest = {Table[-1, {numNearest}],Table[Infinity, {numNearest}]};
        upperBounds = Table[Infinity, {Length[query]}];
        lowerBounds = Table[-Infinity, {Length[query]}]
];

The following calculates the (numNearest) nearest neighbours of the input data entered (data). As an example, the 30 nearest neighbours of the points in "data" can be found by entering "NNarray = nearestNeighbours[data, 30]"

nearestNeighbours[inData_, numNearest_] :=
        Module[
        {tree, query, indexedData, currentList = {}, i},
        indexedData = addIndex[inData];
        Print["Building Search Tree..."];
        tree = buildTree[Length[inData[[1]]], indexedData];
        Print["Starting Nearest Neighbour Search..."];
        For[i = 1, i <= Length[inData], i++,
                Print["Processing ", i];
                query = indexedData[[i]][[1]];
                initialise[query, numNearest];
                searchTree[query, Length[query], tree];
                currentList = Append[currentList, nearest]
        ];
        Return[currentList]
];

Experimental Results

The algorithm above has been implemented in C, using long double arthimetic, togther with an naive O(n^2) algorithm to provide a comparison. Both were run initially on datafiles of sizes 100, 200, 400, 800, 1600 and 3200 3-dimensional points. Further runs were carried out using bucket-sizes of 4 and 8 on the same set of points. The results are presented below (in a fixed height cell):

n2 = {{8,8,8}, {27,27,27}, {101,102,103}, {387,387,386}, 
        {1490, 1490, 1490}, {5887, 5888, 5884}};
nlogn1 = {{8,8,8}, {16,16,16}, {34,34,34}, {67,67,67},
        {145,145,146}, {290,289,290}};
nlogn4 = {{6,6,6}, {12,12,12}, {27,27,27}, {55,55,55},
        {120,120,121}, {241,241,241}};
nlogn8 = {{6,6,6}, {12,12,12}, {26,26,26}, {53,53,53},
        {116,116,117}, {234,233,233}};
xs = {100, 200, 400, 800, 1600, 3200};
n2bar = Transpose[{xs, Map[Mean, n2]}];
nlogn1bar = Transpose[{xs, Map[Mean, nlogn1]}];
nlogn4bar = Transpose[{xs, Map[Mean, nlogn4]}];
nlogn8bar = Transpose[{xs, Map[Mean, nlogn8]}];
p1 = ListPlot[n2bar, PlotJoined->True,
        DisplayFunction->Identity,
        AxesLabel->{"(n)", "Time"}];
p2 = ListPlot[nlogn1bar, PlotJoined->True,
        DisplayFunction->Identity,
        AxesLabel->{"(n)", "Time"}];
p3 = ListPlot[nlogn4bar, PlotJoined->True,
        DisplayFunction->Identity,
        AxesLabel->{"(n)", "Time"}];
p4 = ListPlot[nlogn8bar, PlotJoined->True,
        DisplayFunction->Identity,
        AxesLabel->{"(n)", "Time"}];

a1 = Show[{p1, p2}, DisplayFunction->$DisplayFunction,
        PlotLabel->"kd-Tree compared to n^2",
        PlotRange->All,
        Epilog->{Text["n^2", {2500, 5000}],
                Text["kd-tree", {3000, 600}]}];

a2 = Show[{p2, p3, p4},
        PlotLabel->"kd-tree bucket size comparison",
        DisplayFunction->$DisplayFunction,
        PlotRange->All,
        Epilog->{Text["1", {3300, 280}],
                        Text["4", {3300, 250}],
                        Text["8", {3300, 230}]}];

The first graph compares the performance of the O(n^2) algorithm and the kd-tree method. It can be seen that the kd-tree algorithm does indeed have O(n log n) performance. At 3200 points, the O(n^2) algorithm took on average 1 hour 38 minutes to complete, whereas the kd-tree took just under 5 minutes.

The second graph compares different setting for the bucket size on the same set of points. The top line is bucket size = 1, the middle is bucket size = 4, the lowest line if bucket size  = 8. It can be seen that increasing the size does give a performance increase, particularly noticable in the jump from 1 to 4. The increase from 4 to 8 is less marked.

In order to determine how the algorithm scales  with increasing dimension, a series of tests were performed where the dimension ranged from 2 to 40 and the number of points ranged from 100 to 2000 (ie. 100 tests in all). All tests were perfomed on an
Alpha workstation. The raw results are in the following fixed height cell:

rawdata = {{2, 200, 0}, {2, 400, 0}, {2, 600, 0}, {2, 800, 0}, {2, 1000, 0}, {2, 1200, 0},
 {2, 1400, 0}, {2, 1600, 0}, {2, 1800,  1}, {2, 2000, 1}, {4, 200,  0},{4,  400,
 0}, {4, 600, 0}, {4, 800, 0}, {4,  1000,  1}, {4, 1200, 1}, {4,  1400,  2}, {4,
 1600,  2}, {4,  1800,  3}, {4, 2000, 3}, {6, 200, 0}, {6, 400, 0}, {6, 600, 1},
 {6, 800, 2}, {6, 1000, 3}, {6, 1200, 5}, {6, 1400, 6}, {6, 1600, 7}, {6,  1800,
 9}, {6,  2000,  10}, {8, 200, 0}, {8, 400, 0}, {8, 600,  2}, {8,  800,  5}, {8,
 1000, 7}, {8, 1200, 11}, {8, 1400, 14}, {8, 1600, 18}, {8, 1800, 21}, {8, 2000,
 25}, {10, 200, 0}, {10, 400, 1}, {10,  600, 4}, {10, 800, 7}, {10,  1000,  11},
 {10, 1200, 18}, {10, 1400, 23}, {10, 1600, 29}, {10, 1800, 36}, {10, 2000, 43},
 {12, 200, 0}, {12, 400, 1}, {12, 600, 4}, {12, 800, 9}, {12,  1000,  13},  {12,
 1200, 22}, {12, 1400, 29}, {12, 1600, 39}, {12, 1800, 44}, {12, 2000, 56}, {14,
 200, 0}, {14, 400, 1}, {14, 600, 5}, {14, 800, 9}, {14, 1000,  15}, {14,  1200,
 24}, {14, 1400, 32}, {14, 1600, 45}, {14,  1800, 51}, {14, 2000, 62}, {16, 200,
 1}, {16, 400, 2}, {16, 600, 6}, {16, 800, 11}, {16, 1000, 17}, {16, 1200,  27},
 {16,  1400, 42}, {16, 1600, 47}, {16,  1800, 62}, {16, 2000, 71}, {18, 200, 1},
 {18, 400, 2}, {18, 600, 6}, {18, 800, 12}, {18, 1000, 18}, {18, 1200, 29}, {18,
 1400,  42}, {18, 1600, 52}, {18, 1800,  62}, {18, 2000, 77}, {20, 200, 0}, {20,
 400, 2}, {20, 600, 9}, {20, 800, 13}, {20,  1000,  20}, {20,  1200,  32},  {20,
 1400,  44}, {20, 1600, 56}, {20, 1800,  71}, {20, 2000, 82}, {22, 200, 0}, {22,
 400, 3}, {22, 600, 8}, {22, 800, 15}, {22,  1000,  22}, {22,  1200,  35},  {22,
 1400,  49}, {22, 1600, 61}, {22, 1800,  74}, {22, 2000, 89}, {24, 200, 1}, {24,
 400, 3}, {24, 600, 9}, {24, 800, 17}, {24,  1000,  24}, {24,  1200,  39},  {24,
 1400,  52}, {24, 1600, 70}, {24, 1800,  79}, {24, 2000, 96}, {26, 200, 0}, {26,
 400, 3}, {26, 600, 10}, {26, 800, 17}, {26,  1000,  25}, {26,  1200,  40}, {26,
 1400, 55}, {26, 1600, 75}, {26, 1800, 87}, {26, 2000,  104}, {28, 200, 0}, {28,
 400, 4}, {28, 600, 11}, {28, 800, 18}, {28,  1000,  27}, {28,  1200,  43}, {28,
 1400, 61}, {28, 1600, 80}, {28, 1800, 91}, {28, 2000,  119}, {30, 200, 1}, {30,
 400, 4}, {30, 600, 11}, {30, 800, 19}, {30,  1000,  28}, {30,  1200,  47}, {30,
 1400, 69}, {30, 1600, 83}, {30, 1800, 100}, {30, 2000, 119}, {32, 200, 0}, {32,
 400, 4}, {32, 600, 12}, {32, 800, 22}, {32,  1000,  33}, {32,  1200,  50}, {32,
 1400, 67}, {32, 1600, 88}, {32, 1800, 116}, {32, 2000, 138}, {34, 200, 0}, {34,
 400, 5}, {34, 600, 12}, {34, 800, 22}, {34,  1000,  32}, {34,  1200,  52}, {34,
 1400, 71}, {34, 1600, 87}, {34, 1800, 114}, {34, 2000, 134}, {36, 200, 0}, {36,
 400, 5}, {36, 600, 13}, {36, 800, 23}, {36,  1000,  33}, {36,  1200,  53}, {36,
 1400, 89}, {36, 1600,  105}, {36, 1800,  159}, {36, 2000,  135},  {38, 200, 0},
 {38, 400, 6}, {38, 600, 14}, {38, 800, 24}, {38,  1000,  36}, {38,  1200,  66},
 {38, 1400, 94}, {38, 1600,  130}, {38, 1800,  162}, {38, 2000,  180}, {40, 200,
 1}, {40, 400, 10}, {40, 600, 27}, {40, 800, 49}, {40,  1000,  67},  {40,  1200,
 105}, {40, 1400, 115}, {40, 1600, 156}, {40, 1800, 168}, {40, 2000, 154}};

Needs["Graphics`Graphics3D`"];
plotdata = Partition[
                Table[rawdata[[i,3]], {i,1,Length[rawdata]}],
                10];
a3 = ListPlot3D[plotdata,
        AxesLabel->{"Dimension", "Num Vectors", Runtime},
        Ticks->{Table[{i,4*i},{i,0,10}],
                        Table[{2j,200*j},{j,1,10}],
                        Automatic},
        PlotLabel->"Surface Plot of the Runtime in Seconds"]

The above graph is a surface plot of the results. The x-axis is  the number of dimensions, the y-axis ("into" the page) is the number of popints, and the vertical axis the runtime required. It can be seen once again that the run-time is approximately linear with the number of points, but scales approximately as the square of the number of dimensions.

References

[1]     "An Algorithm for Finding Best Matches in Logarthimic Expected Time"
        J.H Friedman, J.L Bentley, R.A Finkel
        ACM Transactions on Mathematical Software, Vol 3 Number 3, 1977
        Pages 200-226
        
 [2]    "Similarity Indexing: Algorithms and Performance"
        D.A White, R Jain
        Technical Report 1996, available at "http://vision.ucsd.edu/papers"