// This function creates a KD tree with the given
// points, array, and length
int KDTree::create(float *setpoints, int setnpoints, int setndim,
bool setCopy, struct KDTreeNode *setNodeMem)
{
ndim = setndim;
npoints = setnpoints;
typedef int *intptr;
// Copy the points from the original array, if necessary
copyPoints = setCopy;
if (copyPoints) {
if(points) delete[] points;
points = new float[ndim*npoints];
memcpy(points,setpoints,sizeof(float)*ndim*npoints);
}
// If we are not copying, just set the pointer
else
points = setpoints;
// Allocate some arrays;
if (workArr)
delete[]workArr;
workArr = new int[npoints];
if(!setNodeMem) {
if(m_Root) delete[] m_Root;
m_Root = new struct KDTreeNode[npoints*2+1];
nodeMemAlloc = true;
}
else {
m_Root = setNodeMem;
nodeMemAlloc = false;
}
nodeMemCnt = 0;
// Alocate array used for indexing
if(intArrMem) delete[] intArrMem;
intArrMem =
new int[(int)((float)(npoints+4)*
ceil(log((double)npoints)/log(2.0)))];
intArrMemCnt = 0;
// Create the "sortidx" array by
// sorting the range tree points on each dimension
int **sortidx = new intptr[ndim];
if(verbosity>1)
logmsg("KDTree: Sorting points\n");
float imin[3];
float imax[3];
imin[0] = imin[1] = imin[2] = 999999.f;
imax[0] = imax[1] = imax[2] = -999999.f;
for (int i = 0; i < ndim; i++) {
// Initialize the sortidx array for this
// dimension
sortidx[i] = new int[npoints];
// Initialize the "tmp" array for the sort
int *tmp = new int[npoints];
for (int j = 0; j < npoints; j++)
tmp[j] = j;
// Sort the points on dimension i, putting
// indexes in array "tmp"
heapsort(i,tmp,npoints);
// sortidx is actually the inverse of the
// index sorts
for (int j = 0; j < npoints; j++)
{
sortidx[i][tmp[j]] = j;
imin[i] = min( points[ j*3 + i ], imin[ i ] );
imax[i] = max( points[ j*3 + i ], imax[ i ] );
}
delete[] tmp;
}
if(verbosity > 1)
logmsg("KDTree: Done sorting points\n");
// Create an initial list of points that references
// all the points
int *pidx = new int[npoints];
for (int i = 0; i < npoints; i++)
pidx[i] = i;
// Build a KD Tree
AABB extents;
Vec3 vmin( imin[0] ,
imin[1],
imin[2] );
Vec3 vmax( imax[0] ,
imax[1],
imax[2] );
OutputVector( vmin, "VMin");
OutputVector( vmax, "VMax");
extents.Set( vmin,vmax );
m_Root->bounds.Set( vmin, vmax );
//add objects to this node
if( m_NodeCreationCallBack )
{
(*m_NodeCreationCallBack)( m_Root );
}
build_kdtree(sortidx, // array of sort values
0, // The current dimension
pidx, npoints, extents); // The list of points
// Delete the sort index
for (int i = 0; i < ndim; i++)
delete[]sortidx[i];
delete[] sortidx;
// Delete the initial list of points
delete[] pidx;
// Delete the sort arrays
delete[] intArrMem;
// delete the work array
if(workArr) {
delete[] workArr;
workArr = (int *)NULL;
}
if(verbosity > 1)
logmsg("KDTree: Done creating tree\n");
return 0;
} // end of create
// This function build a node of the kdtree with the
// points indexed by pidx with length "len"
// sortidx is a pre-computed array using the heapsort
// algorithm above
int KDTree::build_kdtree(int **sortidx, int dim,
int *pidx, int len, AABB ¤tBB)
{
static const Vec3 BufferOffset( KD_SPLIT_BUFFER_SIZE, KD_SPLIT_BUFFER_SIZE, KD_SPLIT_BUFFER_SIZE );
AABB b1, b2;
int ncnt = nodeMemCnt;
struct KDTreeNode *node = node_alloc();
node->bounds.Set( currentBB.m_Bounds[0] - BufferOffset, currentBB.m_Bounds[1] + BufferOffset );
node->axis = dim;
if (len <= m_KDCellSize &&
m_KDCellSize == 1 )
{
node->leftIdx = -1;
node->rightIdx = -1;
node->pntidx = pidx[0];
node->key = 0;
node->bounds.Set( currentBB.m_Bounds[0] - BufferOffset, currentBB.m_Bounds[1] + BufferOffset );
//add objects to this node
if( m_NodeCreationCallBack )
{
(*m_NodeCreationCallBack)( node );
}
return ncnt;
}
// If not, we must make a node
int pivot = -1;
int lcnt = 0, rcnt = 0;
int *larray, *rarray;
// Find the pivot (index of median point of available
// points on current dimension).
// If heapsorting the current list of points is quicker than
// iterating through all the points, just do that instead
// (heapsort if of order n*log(n)
// This test could probably use some fine tuning
if((double)len*log((double)len) < npoints) {
heapsort(dim,pidx,len);
larray = pidx;
rarray = pidx+len/2;
pivot = pidx[len/2];
lcnt = len/2;
rcnt = len/2 + (len%2==0 ? 0 : 1);
}
else
{
// Use the previously calculated sort index
// to make this a process linear in npoints
// This gets a little confusing, but it works.
// Sortidx:: sortidx[dim][idx] = val
// idx = the index to the point
// val = the order in the array
int *parray = workArr;
// Setting parray to -1 indicates we are not using
// the point
for (int i = 0; i < npoints; i++)
parray[i] = -1;
// Populate "parray" with the points that we
// are using, indexed in the order they occur
// on the current dimension
for (int i = 0; i < len; i++)
parray[sortidx[dim][pidx[i]]] = pidx[i];
int cnt = 0;
larray = int_alloc(len/2+1);
rarray = int_alloc(len/2+1);
// The middle valid value of parray is the pivot,
// the left go to a node on the left, the right
// go to a node on the right.
for (int i = 0; i < npoints; i++) {
if (parray[i] == -1)
continue;
if (cnt == len / 2) {
pivot = parray[i];
rarray[rcnt++] = parray[i];
} else if (cnt > len / 2)
rarray[rcnt++] = parray[i];
else
larray[lcnt++] = parray[i];
cnt++;
if(cnt>len)
break;
}
}
if (len <= m_KDCellSize &&
m_KDCellSize > 1 )
{
node->leftIdx = -1;
node->rightIdx = -1;
if( m_KDCellSize > 1 )
{
for( int i = 0; i < len; i++ )
{
node->fattyCellData.push_back( pidx[ i ] );
}
}
node->pntidx = pivot;
node->key = 0;
node->bounds.Set( currentBB.m_Bounds[0] - BufferOffset, currentBB.m_Bounds[1] + BufferOffset );
//add objects to this node
if( m_NodeCreationCallBack )
{
(*m_NodeCreationCallBack)( node );
}
return ncnt;
}
// Create the node
node->pntidx = -1;
node->key = points[pivot*ndim+dim] + KEY_EPSILON;
//calculate bounding boxes
if( m_NodeCreationCallBack )
{
b1 = b2 = currentBB; //b1 is for left, b2 for right
//split on the x axis
int realDim = dim%ndim;
if( realDim == 0 ){
b1.m_Bounds[1].x = node->key; //set tmax bounds of left bb to be new split point
b2.m_Bounds[0].x = node->key; //set tmin bounds of right bb to be new split point
}else if( realDim == 1 )
{
b1.m_Bounds[1].y = node->key; //set tmax bounds of left bb to be new split point
b2.m_Bounds[0].y = node->key; //set tmin bounds of right bb to be new split point
}else if( realDim == 2 )
{
b1.m_Bounds[1].z = node->key; //set tmax bounds of left bb to be new split point
b2.m_Bounds[0].z = node->key; //set tmin bounds of right bb to be new split point
}
}
// Create nodes to the left
node->leftIdx =
build_kdtree(sortidx, (dim + 1) % ndim, larray, lcnt, b1);
// Create nodes to the right
node->rightIdx =
build_kdtree(sortidx, (dim + 1) % ndim, rarray, rcnt, b2);
return ncnt;
} // end of build_kdtree
void mexFunction( int nlhs, mxArray **plhs, int nrhs, const mxArray **prhs){
Tree *tree;
double *reference, *model;
int *index, i;
unsigned int N, D, M;
double *closest_pts, *distances, *pointer_to_tree;
int SkipQueries=0;
if (nrhs <2 ){
mexErrMsgTxt("Must have at least two input arrays.");
}
#ifdef DEBUG
mexPrintf("Mex function called with %d inputs and %d explicit outputs\n",nrhs,nlhs);
#endif
reference = mxGetPr(prhs[0]);
N = mxGetM(prhs[0]);
D = mxGetN(prhs[0]);
if ((!N || !D ) && ( nrhs < 3) )
mexErrMsgTxt("You have to supply some reference points to build a k-D tree.");
#ifdef TIME
gettimeofday(&tv1,&tz);
#endif
//
//
// If the tree is not passed in as a third input, we must build it
//
//
if (nrhs < 3 ){
#ifdef DEBUG
mexPrintf("----------------------\n");
mexPrintf("Building k-D Tree ...\n");
#endif
index = (int*) malloc( sizeof(int) * N);
for (i=0; i < N; i++) index[i]=i;
if ( (tree = build_kdtree(reference,N,D,index,N,0))==NULL ){
free(index);
mexErrMsgTxt("Not enough free memory to build k-D tree\n");
} else {
tree->dims = D;
free(index);
}
#ifdef DEBUG
mexPrintf("Done Building k-D Tree\n");
mexPrintf("----------------------\n");
#endif
} else {
//
// The tree was built previously, and is now being passed in to the function.
//
//
// The tree was built previously, and is now being passed in to the function.
//
if ( (pointer_to_tree = mxGetPr(prhs[2])) == NULL )
mexErrMsgTxt("Third argument is not a valid pointer to a k-D tree\n");
if ( (tree = (Tree *) ((long) pointer_to_tree[0]))== NULL )
mexErrMsgTxt("Third argument is not a valid pointer to a k-D tree\n");
}
#ifdef TIME
gettimeofday(&tv2,&tz);
if (tv2.tv_usec - tv1.tv_usec < 0) {
tv2.tv_sec--;
tv2.tv_usec += 1000000;
}
mexPrintf("Time to Build Tree : %f\n", tv2.tv_sec -tv1.tv_sec+(tv2.tv_usec-tv1.tv_usec)/1000000.0);
#endif
#ifdef DISPLAY_TREE
mexPrintf("\nDepth first traversal of the k-D tree\n");
mexPrintf("-------------------------------------\n");
display_tree(tree->rootptr,D);
mexPrintf("-------------------------------------\n");
#endif
//
// Query section
//
//
model = mxGetPr(prhs[1]);
M = mxGetM(prhs[1]);
if (!model && !M) {
// There are no points to query
SkipQueries=1;
} else {
// Check that the model points are of the same dimension as the
// reference points in the k-d tree.
if (mxGetN(prhs[1]) != tree->dims)
mexErrMsgTxt("Reference and Model Vectors must be of the same dimension");
}
if (nlhs >=0){
plhs[0] = mxCreateDoubleMatrix(M, tree->dims,mxREAL);
closest_pts = mxGetPr(plhs[0]);
}
else{
closest_pts = (double *) malloc (sizeof(double) *M* (tree->dims));
}
if (nlhs >=2) {
plhs[1] = mxCreateDoubleMatrix(M,1,mxREAL);
distances = mxGetPr(plhs[1]);
}
else {
distances = (double *) malloc (sizeof(double)*M);
}
if (nlhs >=3) {
plhs[2] = mxCreateDoubleMatrix(1,1,mxREAL);
pointer_to_tree = mxGetPr(plhs[2]);
pointer_to_tree[0] = (long) tree;
}
if (!SkipQueries) {
#ifdef TIME
gettimeofday(&tv1,&tz);
#endif
#ifdef DEBUG
mexPrintf("--------------------\n");
mexPrintf("Running Queries...\n");
#endif
run_queries(tree->rootptr, model, M, tree->dims, closest_pts,
distances, RETURN_POINTS);
#ifdef DEBUG
mexPrintf("Done Running Queries\n");
mexPrintf("--------------------\n");
#endif
#ifdef TIME
gettimeofday(&tv2,&tz);
if (tv2.tv_usec - tv1.tv_usec < 0) {
tv2.tv_sec--;
tv2.tv_usec += 1000000;
}
mexPrintf("Time per Search : %f\n",
(tv2.tv_sec - tv1.tv_sec +
(tv2.tv_usec-tv1.tv_usec) /1000000.0 )/(double)M);
#endif
}
if (nlhs<3) {
#ifdef DEBUG
mexPrintf("-------------------------------------\n");
mexPrintf("Removing k-D Tree from system memory.\n");
#endif
free_tree(tree->rootptr);
free(tree);
#ifdef DEBUG
mexPrintf("Done.\n");
mexPrintf("-------------------------------------\n");
#endif
if (nlhs < 2){
free(distances);
}
} else{
#ifdef DEBUG
mexPrintf("--------------------------------\n");
mexPrintf("k-D Tree saved in system memory.\n");
mexPrintf("Don't forget to remove it later.\n");
mexPrintf("--------------------------------\n");
#endif
}
#ifdef DEBUG
mexPrintf("Mex function has exited normally.\n");
#endif
}
