kdtree_cartesian.cc 13.6 KB
Newer Older
Uwe Schulzweida's avatar
Uwe Schulzweida committed
1
2
3
4
/*!\file kdtree_cartesian.c
 * \brief Routines specific to the n-dimensional cartesian kd-tree
 *
 */
5
6
7
#if defined(HAVE_CONFIG_H)
#include "config.h"
#endif
Uwe Schulzweida's avatar
Uwe Schulzweida committed
8
9
10
11
12
13
14
15
16
17

#include "kdtree.h"
#include "pqueue.h"

/* ********************************************************************

   general utility functions 

   ********************************************************************* */

18
static constexpr
19
kdata_t square(const kdata_t x)
20
21
22
23
{
  return x*x;
}

24
25
static constexpr
kdata_t kd_dist_sq(const kdata_t *restrict a, const kdata_t *restrict b)
26
{
27
  return square((a[0]-b[0]))+square((a[1]-b[1]))+square((a[2]-b[2]));
28
29
}

30

31
kdata_t kd_min(kdata_t x, kdata_t y)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
{
    return x < y ? x : y;
}


/* ******************************************************************
   
   Functions for building and destroying trees 

   ******************************************************************** */


/*! \brief build kd-tree structure
 *
 * \param points an array of kd_points (struct with position vector
 * and data container).
 *
 * \param nPoints the length of the points array.
 *
 * \param constr a pointer to a void *constructor() function to include
 * the data container in the tree; optional, can be NULL
 *
 * \param destr a pointer to a void destructor() function to free()
 * the data containers in the tree; optional, can be NULL, but should
 * be given if the constr argument is non-NULL.
 *
 * \param min a vector with the minimum positions of the corners of the
 * hyperrectangle containing the data.
 *
 * \param max a vector with the maximum positions of the corners of
 * the hyperrectangle containing the data.
 *
 * \param dim the dimensionality of the data.
 *
 * \param max_threads the maximal number of threads spawned for
 * construction of the tree. The threads will be unbalanced if this is
 * not a power of 2.
 *
 * \return root node of the tree
 */
struct kdNode *
kd_buildTree(struct kd_point *points, unsigned long nPoints,
74
             kdata_t *min, kdata_t *max, int dim, int max_threads)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
75
76
77
78
79
{
    struct kd_thread_data *my_data;
    struct kdNode *tree;

    my_data = kd_buildArg(points, nPoints, min, max, 0, max_threads, dim);
Oliver Heidmann's avatar
Oliver Heidmann committed
80
    tree = (kdNode *)kd_doBuildTree(my_data);
Uwe Schulzweida's avatar
Uwe Schulzweida committed
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
    free(my_data);
    return tree;
}


/* *******************************************************************

   Functions for range searches 
  
   *******************************************************************
*/

/* Returns 1 if node is a point in the hyperrectangle defined by
   minimum and maximum vectors min and max. */
int
96
kd_isPointInRect(struct kdNode *node, kdata_t *min, kdata_t *max, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
{
    int i;

    if (node == NULL)
        return 0;

    for(i = 0; i < dim; i++) {
        if (node->location[i] < min[i] || node->location[i] > max[i])
            return 0;
    }
    return 1;
}

/* Returns 1 if the hyperrectangle of node is fully contained within
   the HR described by the minimum and maximum vectors min and
   max. Returns 0 otherwise. */
int
114
kd_isRectInRect(struct kdNode *node, kdata_t *min, kdata_t *max, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
{
    int i;

    if (node == NULL)
        return 0;

    for(i = 0; i < dim; i++) {
        if (node->min[i] < min[i] || node->max[i] > max[i])
            return 0;
    }
    return 1;
}

/* Returns 1 if the hyperrectangle of node overlaps the HR described
   the minimum and maximum vectors min and max. Returns 0
   otherwise. */
int
132
kd_rectOverlapsRect(struct kdNode *node, kdata_t *min, kdata_t *max, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
{
    int i;

    if (node == NULL)
        return 0;

    for(i = 0; i < dim; i++) {
        if ((node->min[i] > max[i] || node->max[i] < min[i]))
            return 0;
    }
    return 1;
}

/*!
 * \brief Perform orthogonal range search (get all points in a
 * hyperrectangle).
 *
 * \param node the root node of tree to be searched.
 *
 * \param min a vector with the minimum positions of the corners of the
 * hyperrectangle containing the data.
 *
 * \param max a vector with the maximum positions of the corners of
 * the hyperrectangle containing the data.
 *
 * \param dim the dimension of the data.
 *
 * \return  Pointer to a priority queue, NULL in case of problems.
*/
struct pqueue *
163
kd_ortRangeSearch(struct kdNode *node, kdata_t *min, kdata_t *max, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
{
    struct pqueue *res;
    uint32_t i;

    if ((res = pqinit(NULL, 1)) == NULL)
        return NULL;
    if (!kd_doOrtRangeSearch(node, min, max, dim, res)) {
        for(i = 0; i < res->size; i++) {
            free(res->d[i]);
        }
        free(res->d);
        free(res);
        return NULL;
    }
    return res;
}

/* This is the orthogonal range search. Returns 1 if okay, 0 in case
   of problems. */
int
184
kd_doOrtRangeSearch(struct kdNode *node, kdata_t *min, kdata_t *max,
Uwe Schulzweida's avatar
Uwe Schulzweida committed
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
                    int dim, struct pqueue *res)
{

    if (node == NULL)
        return 1;

    if (kd_isleaf(node) && kd_isPointInRect(node, min, max, dim)) {
        return kd_insertResTree(node, res);
    } else {
        if (kd_isRectInRect(node->left, min, max, dim)) {
            if (!kd_insertResTree(node->left, res))
                return 0;
        } else {
            if (kd_rectOverlapsRect(node->left, min, max, dim))
                if (!kd_doOrtRangeSearch(node->left, min, max, dim, res))
                    return 0;
        }
        if (kd_isRectInRect(node->right, min, max, dim)) {
            if (!kd_insertResTree(node->right, res))
                return 0;
        } else {
            if (kd_rectOverlapsRect(node->right, min, max, dim))
                if (!kd_doOrtRangeSearch(node->right, min, max, dim, res))
                    return 0;
        }
    }
    return 1;
}

/*! 
 * \brief Find the nearest neighbor of a point.
 *
 * \param node the root node of the tree to be searched.
 *
 * \param p a vector to the point whose nearest neighbor is sought.
 *
 * \param max_dist_sq the square of the maximum distance to the
 * nearest neighbor.
 *
 * \param dim the dimension of the data.
 *
 * \return A pointer to node containing the nearest neighbor.
 * max_dist_sq is set to the square of the distance to the nearest
 * neigbor.
 */
struct kdNode *
231
kd_nearest(struct kdNode *node, kdata_t *p, kdata_t *max_dist_sq, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
232
{
233
    if ( !node ) return NULL;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
234

235
236
237
    struct kdNode *nearer, *further;
    if ( p[node->split] < node->location[node->split] ) {
        nearer  = node->left;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
238
239
        further = node->right;
    } else {
240
        nearer  = node->right;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
241
242
        further = node->left;
    }
243
244
245
246
    
    struct kdNode *tmp = kd_nearest(nearer, p, max_dist_sq, dim);
    struct kdNode *nearest = NULL;
    if ( tmp )
Uwe Schulzweida's avatar
Uwe Schulzweida committed
247
248
249
250
        nearest = tmp;
    else
        nearest = node;

251
    kdata_t dist_sq = kd_dist_sq(nearest->location, p);
Uwe Schulzweida's avatar
Uwe Schulzweida committed
252
253
254
255
256
257
    if (*max_dist_sq > dist_sq)
	*max_dist_sq = dist_sq;

    if (!further)
        return nearest;

258
259
260
    kdata_t dx = kd_min(KDATA_ABS(p[node->split] - further->min[node->split]),
                        KDATA_ABS(p[node->split] - further->max[node->split]));
    if ( *max_dist_sq > kd_sqr(dx) ) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
261
262
263
264
        /*
         * some part of the further hyper-rectangle is in the search
         * radius, search the further node 
         */
265
        struct kdNode *tmp_nearest;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
266
        tmp = kd_nearest(further, p, max_dist_sq, dim);
267
        if ( tmp )
Uwe Schulzweida's avatar
Uwe Schulzweida committed
268
269
270
271
            tmp_nearest = tmp;
        else
            tmp_nearest = further;

272
        kdata_t tmp_dist_sq = kd_dist_sq(tmp_nearest->location, p);
273
        if ( tmp_dist_sq < dist_sq ) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
274
275
276
277
            nearest = tmp_nearest;
            dist_sq = tmp_dist_sq;
            *max_dist_sq = kd_min(dist_sq, *max_dist_sq);
        }
278
        // Uwe Schulzweida
279
        else if ( tmp_dist_sq <= dist_sq && tmp_nearest->index < nearest->index ) {
280
281
282
283
            nearest = tmp_nearest;
            dist_sq = tmp_dist_sq;
            *max_dist_sq = kd_min(dist_sq, *max_dist_sq);
        }
Uwe Schulzweida's avatar
Uwe Schulzweida committed
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
    }
    return nearest;
}

/*! 
 * \brief Return the q nearest-neighbors to a point.
 *
 * \param node the root node of the tree to be searched.
 *
 * \param p a vector to the point whose nearest neighbors are sought.
 *
 * \param max_dist_sq the square of the maximum distance to the
 * nearest neighbors.
 *
 * \param q the maximum number of points to be retured.
 *
 * \param dim the dimension of the data.
 *
 * \return A pointer to a priority queue of the points found, or NULL
 * in case of problems.
 */
struct pqueue *
306
307
kd_qnearest(struct kdNode *node, kdata_t *p,
            kdata_t *max_dist_sq, unsigned int q, int dim)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
308
{
309
310
311
312
313
    struct pqueue *res = pqinit(NULL, q + 2);
    if ( res == NULL) return NULL;
    
    if ( !kd_doQnearest(node, p, max_dist_sq, q + 1, dim, res) ) {
        for ( uint32_t i = 0; i < res->size; ++i ) free(res->d[i]);
Uwe Schulzweida's avatar
Uwe Schulzweida committed
314
315
316
317
        free(res->d);
        free(res);
        return NULL;
    }
318
    
Uwe Schulzweida's avatar
Uwe Schulzweida committed
319
320
321
322
323
324
325
326
327
328
329
330
    return res;
}

/* *
 * This is the q nearest-neighbor search.
 *
 * This is a modification of the range search in which the maximum
 * search radius is decreased to the maximum of the queue as soon as
 * the queue is filled.
 *
 * return 1 if okay, zero in case of problems
 */
331
332
// Uwe Schulzweida: extract kd_check_dist() from kd_doQnearest()
static int
333
kd_check_dist(struct kdNode *node, kdata_t *p,
334
              kdata_t *max_dist_sq, unsigned int q, struct pqueue *res)
Uwe Schulzweida's avatar
Uwe Schulzweida committed
335
{
336
    kdata_t dist_sq = kd_dist_sq(node->location, p);
337
338
339
    if ( dist_sq < *max_dist_sq && kd_isleaf(node) ) {
        struct resItem *point = (struct resItem *) kd_malloc(sizeof(struct resItem), "kd_doQnearest: ");
        if ( point == NULL) return 0;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
340
341
342
343
        point->node = node;
        point->dist_sq = dist_sq;
        pqinsert(res, point);
    }
344

345
346
    if ( res->size > q ) {
        struct resItem *item;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
347
348
        pqremove_max(res, &item);
        free(item);
349
        if ( res->size > 1 ) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
350
351
352
353
354
355
356
357
358
359
360
361
362
            /*
             * Only inspect the queue if there are items left 
             */
            pqpeek_max(res, &item);
            *max_dist_sq = item->dist_sq;
        } else {
            /*
             * Nothing was found within the max search radius 
             */
            *max_dist_sq = 0;
        }
    }

363
364
365
366
    return 1;
}

int
367
368
kd_doQnearest(struct kdNode *node, kdata_t *p,
              kdata_t *max_dist_sq, unsigned int q, int dim, struct pqueue *res)
369
{
370
    if ( !node ) return 1;
371

372
    if ( !kd_check_dist(node, p, max_dist_sq, q, res) ) return 0;
373

374
375
376
    struct kdNode *nearer, *further;
    if ( p[node->split] < node->location[node->split] ) {
        nearer  = node->left;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
377
378
        further = node->right;
    } else {
379
        nearer  = node->right;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
380
381
        further = node->left;
    }
382
    if ( !kd_doQnearest(nearer, p, max_dist_sq, q, dim, res) ) return 0;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
383

384
    if ( !further ) return 1;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
385

386
387
388
389
    kdata_t dx = kd_min(KDATA_ABS(p[node->split] - further->min[node->split]),
                        KDATA_ABS(p[node->split] - further->max[node->split]));
    
    if ( *max_dist_sq > kd_sqr(dx) ) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
390
391
392
393
        /*
         * some part of the further hyper-rectangle is in the search
         * radius, search the further node 
         */
394
        if (!kd_doQnearest(further, p, max_dist_sq, q, dim, res)) return 0;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
395

396
        if (!kd_check_dist(node, p, max_dist_sq, q, res)) return 0;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
    }
    return 1;
}


/*!
 * \brief Perform a range search around a point.
 *
 * \param node the root node of the tree to be searched.
 *
 * \param p the location of the point around which the search is carried out .
 *
 * \param max_dist_sq the square of the radius of the hypersphere.
 *
 * \param dim the dimension of the data.  \param ordered determines
 * whether the result list should be ordered in increasing distance
 * (KD_ORDERED) or unordered (KD_UNORDERED).
 *
 * \return A pointer to a priority queue containing the points found,
 * NULL in case of problems.
 */
struct pqueue *
419
kd_range(struct kdNode *node, kdata_t *p, kdata_t *max_dist_sq,
Uwe Schulzweida's avatar
Uwe Schulzweida committed
420
421
         int dim, int ordered)
{
422
423
424
425
426
    struct pqueue *res = pqinit(NULL, 1);
    if ( res == NULL ) return NULL;
    
    if ( !kd_doRange(node, p, max_dist_sq, dim, res, ordered) ) {
        for( uint32_t i = 0; i < res->size; ++i ) {
Uwe Schulzweida's avatar
Uwe Schulzweida committed
427
428
429
430
431
432
433
434
435
436
437
438
            free(res->d[i]);
        }
        free(res->d);
        free(res);
        return NULL;
    }
    return res;
}


/* This is the range search. Returns 1 if okay, 0 in case of problems */
int
439
kd_doRange(struct kdNode *node, kdata_t *p, kdata_t *max_dist_sq,
Uwe Schulzweida's avatar
Uwe Schulzweida committed
440
441
442
443
444
           int dim, struct pqueue *res, int ordered)
{

    struct kdNode *nearer, *further;
    struct resItem *point;
445
    kdata_t dist_sq, dx;
Uwe Schulzweida's avatar
Uwe Schulzweida committed
446
447
448
449

    if (!node)
        return 1;

450
    dist_sq = kd_dist_sq(node->location, p);
Uwe Schulzweida's avatar
Uwe Schulzweida committed
451
    if (dist_sq < *max_dist_sq && kd_isleaf(node)) {
Oliver Heidmann's avatar
Oliver Heidmann committed
452
        if ((point = (struct resItem *)kd_malloc(sizeof(struct resItem), "kd_doRange:"))
Uwe Schulzweida's avatar
Uwe Schulzweida committed
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
            == NULL)
            return 0;
        point->node = node;
        point->dist_sq = ordered ? dist_sq : -1;
        pqinsert(res, point);
    }

    if (p[node->split] < node->location[node->split]) {
        nearer = node->left;
        further = node->right;
    } else {
        nearer = node->right;
        further = node->left;
    }
    if (!kd_doRange(nearer, p, max_dist_sq, dim, res, ordered))
        return 0;

    if (!further)
        return 1;

473
474
    dx = kd_min(KDATA_ABS(p[node->split] - further->min[node->split]),
                KDATA_ABS(p[node->split] - further->max[node->split]));
Uwe Schulzweida's avatar
Uwe Schulzweida committed
475
476
477
478
479
480
481
    if (*max_dist_sq > kd_sqr(dx)) {
        /*
         * some part of the further hyper-rectangle is in the search
         * radius, search the further node 
         */
        if (!kd_doRange(further, p, max_dist_sq, dim, res, ordered))
            return 0;
482
        dist_sq = kd_dist_sq(node->location, p);
Uwe Schulzweida's avatar
Uwe Schulzweida committed
483
484

        if (dist_sq < *max_dist_sq && kd_isleaf(node)) {
Oliver Heidmann's avatar
Oliver Heidmann committed
485
            if ((point = (struct resItem *)kd_malloc(sizeof(struct resItem), "kd_doRange: "))
Uwe Schulzweida's avatar
Uwe Schulzweida committed
486
487
488
489
490
491
492
493
494
495
496
497
                == NULL)
                return 0;
            point->node = node;
            point->dist_sq = ordered ? dist_sq : -1;
            pqinsert(res, point);
        }

    }
    return 1;
}

/* End range searching functions */