Fast nearest neighbours in Stata (github, do)

This is an implementation of a kd-tree in Stata's matrix programming language Mata. It facilitates the fast computation of the k nearest neighbours of a set of points according to the Euclidean distance metric. This algorithm is orders of magnitudes faster than a brute-force computation if the data are sufficiently large (i.e., thousands of data points). The function can be used e.g. to compute the average of some variable among an observation's nearest neighbours. It is only thoroughly tested for the two-dimensional case.

Syntax:

void knn(real matrix query_coords, real matrix data_coords,  real scalar k, 

         real matrix kni, real matrix knd)

where

query_coords is an M x n matrix of M n-dimensional query points,

data_coords is an N x n matrix of N n-dimensional points to query against (could be the same as query_coords),

k is the number of nearest neighbours (including self) to be computed,

kni is the M x k matrix of nearest neighbour indices for each of the M query points that will be computed, and

knd is the M x k matrix of nearest neighbour distances

Example usage:

version 15.1

mata: mata clear

mata: mata set matastrict on

run "https://raw.githubusercontent.com/robertaue/knearest/master/mata_knn.do"

mata:

    N = 10000

    k = 5

    query_coords = runiform(N,2)

    data_coords = runiform(N,2)

    knn(query_coords, data_coords, k, kni=., knd=.)

end