ffi-kdtree 0.0.2
Sign up to get free protection for your applications and to get access to all the features.
- data/.gitignore +11 -0
- data/README.md +26 -0
- data/examples/test.rb +34 -0
- data/ext/ffi-kdtree/extconf.rb +5 -0
- data/ext/ffi-kdtree/kdtree.c +836 -0
- data/ext/ffi-kdtree/kdtree.h +129 -0
- data/ffi-kdtree.gemspec +23 -0
- data/lib/ffi-kdtree.rb +3 -0
- data/lib/ffi-kdtree/data_conversion.rb +45 -0
- data/lib/ffi-kdtree/dimensions.rb +47 -0
- data/lib/ffi-kdtree/ffi/libkdtree.rb +74 -0
- data/lib/ffi-kdtree/finders.rb +31 -0
- data/lib/ffi-kdtree/helpers.rb +15 -0
- data/lib/ffi-kdtree/tree.rb +27 -0
- data/spec/data/zips.txt +29470 -0
- data/spec/find_nearest_spec.rb +40 -0
- data/spec/spec_helper.rb +15 -0
- data/spec/tree_spec.rb +12 -0
- metadata +130 -0
data/.gitignore
ADDED
data/README.md
ADDED
@@ -0,0 +1,26 @@
|
|
1
|
+
FFI KDTree
|
2
|
+
==========
|
3
|
+
|
4
|
+
This is a simple FFI wrapped version of [libkdtree](https://github.com/sdeming/kdtree) for Ruby.
|
5
|
+
|
6
|
+
License
|
7
|
+
-------
|
8
|
+
|
9
|
+
[New BSD license][license]
|
10
|
+
|
11
|
+
Requirements
|
12
|
+
------------
|
13
|
+
|
14
|
+
* [libkdtree](https://github.com/sdeming/kdtree)
|
15
|
+
* ffi gem
|
16
|
+
|
17
|
+
Contributing
|
18
|
+
------------
|
19
|
+
|
20
|
+
1. Fork it.
|
21
|
+
2. Create a branch
|
22
|
+
3. Commit your changes
|
23
|
+
4. Push to the branch
|
24
|
+
5. Fire off a pull request
|
25
|
+
|
26
|
+
[license]: http://www.opensource.org/licenses/bsd-license.php
|
data/examples/test.rb
ADDED
@@ -0,0 +1,34 @@
|
|
1
|
+
NUM_DIMENSIONS = 10
|
2
|
+
NUM_ENTRIES = 100000
|
3
|
+
SEARCH_RADIUS = 0.8
|
4
|
+
|
5
|
+
require File.expand_path(File.join(File.dirname(__FILE__), '..', 'lib', 'ffi-kdtree'))
|
6
|
+
|
7
|
+
begin
|
8
|
+
num_dimensions = NUM_DIMENSIONS
|
9
|
+
num_entries = NUM_ENTRIES
|
10
|
+
search_radius = SEARCH_RADIUS
|
11
|
+
|
12
|
+
puts "* Initializing tree with #{num_dimensions} dimensions"
|
13
|
+
kdtree = KdTree::Tree.new(num_dimensions, :integer)
|
14
|
+
|
15
|
+
puts "* Adding #{num_entries} entries to tree with random coordinates"
|
16
|
+
1.upto(num_entries).each do |id|
|
17
|
+
kdtree.add(1.upto(num_dimensions).collect { rand * 2.0 - 1.0}, id)
|
18
|
+
end
|
19
|
+
|
20
|
+
pos = 1.upto(num_dimensions).collect { rand * 2.0 - 1.0 }
|
21
|
+
puts "* Chose a random position at #{pos.join(',')}"
|
22
|
+
|
23
|
+
puts "* Find its nearest neighbor"
|
24
|
+
where, data = kdtree.find_nearest(pos)
|
25
|
+
|
26
|
+
puts "* Found: #{data} at a distance of #{kdtree.distance(pos, where)} at #{where.join(',')}"
|
27
|
+
|
28
|
+
neighbors = kdtree.find_nearest_range(pos, search_radius)
|
29
|
+
puts "* Found #{neighbors.size} neighbors within radius #{search_radius}"
|
30
|
+
neighbors.each do |neighbor|
|
31
|
+
where, data = *neighbor
|
32
|
+
puts "* Found neighbor: #{data}, distance of #{kdtree.distance(pos, where)}"
|
33
|
+
end
|
34
|
+
end
|
@@ -0,0 +1,836 @@
|
|
1
|
+
/*
|
2
|
+
This file is part of ``kdtree'', a library for working with kd-trees.
|
3
|
+
Copyright (C) 2007-2011 John Tsiombikas <nuclear@member.fsf.org>
|
4
|
+
|
5
|
+
Redistribution and use in source and binary forms, with or without
|
6
|
+
modification, are permitted provided that the following conditions are met:
|
7
|
+
|
8
|
+
1. Redistributions of source code must retain the above copyright notice, this
|
9
|
+
list of conditions and the following disclaimer.
|
10
|
+
2. Redistributions in binary form must reproduce the above copyright notice,
|
11
|
+
this list of conditions and the following disclaimer in the documentation
|
12
|
+
and/or other materials provided with the distribution.
|
13
|
+
3. The name of the author may not be used to endorse or promote products
|
14
|
+
derived from this software without specific prior written permission.
|
15
|
+
|
16
|
+
THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR IMPLIED
|
17
|
+
WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
|
18
|
+
MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
|
19
|
+
EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
|
20
|
+
EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
|
21
|
+
OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
22
|
+
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
23
|
+
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
|
24
|
+
IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
|
25
|
+
OF SUCH DAMAGE.
|
26
|
+
*/
|
27
|
+
/* single nearest neighbor search written by Tamas Nepusz <tamas@cs.rhul.ac.uk> */
|
28
|
+
#include <stdio.h>
|
29
|
+
#include <stdlib.h>
|
30
|
+
#include <string.h>
|
31
|
+
#include <math.h>
|
32
|
+
#include "kdtree.h"
|
33
|
+
|
34
|
+
#if defined(WIN32) || defined(__WIN32__)
|
35
|
+
#include <malloc.h>
|
36
|
+
#endif
|
37
|
+
|
38
|
+
#ifdef USE_LIST_NODE_ALLOCATOR
|
39
|
+
|
40
|
+
#ifndef NO_PTHREADS
|
41
|
+
#include <pthread.h>
|
42
|
+
#else
|
43
|
+
|
44
|
+
#ifndef I_WANT_THREAD_BUGS
|
45
|
+
#error "You are compiling with the fast list node allocator, with pthreads disabled! This WILL break if used from multiple threads."
|
46
|
+
#endif /* I want thread bugs */
|
47
|
+
|
48
|
+
#endif /* pthread support */
|
49
|
+
#endif /* use list node allocator */
|
50
|
+
|
51
|
+
struct kdhyperrect {
|
52
|
+
int dim;
|
53
|
+
double *min, *max; /* minimum/maximum coords */
|
54
|
+
};
|
55
|
+
|
56
|
+
struct kdnode {
|
57
|
+
double *pos;
|
58
|
+
int dir;
|
59
|
+
void *data;
|
60
|
+
|
61
|
+
struct kdnode *left, *right; /* negative/positive side */
|
62
|
+
};
|
63
|
+
|
64
|
+
struct res_node {
|
65
|
+
struct kdnode *item;
|
66
|
+
double dist_sq;
|
67
|
+
struct res_node *next;
|
68
|
+
};
|
69
|
+
|
70
|
+
struct kdtree {
|
71
|
+
int dim;
|
72
|
+
struct kdnode *root;
|
73
|
+
struct kdhyperrect *rect;
|
74
|
+
void (*destr)(void*);
|
75
|
+
};
|
76
|
+
|
77
|
+
struct kdres {
|
78
|
+
struct kdtree *tree;
|
79
|
+
struct res_node *rlist, *riter;
|
80
|
+
int size;
|
81
|
+
};
|
82
|
+
|
83
|
+
#define SQ(x) ((x) * (x))
|
84
|
+
|
85
|
+
|
86
|
+
static void clear_rec(struct kdnode *node, void (*destr)(void*));
|
87
|
+
static int insert_rec(struct kdnode **node, const double *pos, void *data, int dir, int dim);
|
88
|
+
static int rlist_insert(struct res_node *list, struct kdnode *item, double dist_sq);
|
89
|
+
static void clear_results(struct kdres *set);
|
90
|
+
|
91
|
+
static struct kdhyperrect* hyperrect_create(int dim, const double *min, const double *max);
|
92
|
+
static void hyperrect_free(struct kdhyperrect *rect);
|
93
|
+
static struct kdhyperrect* hyperrect_duplicate(const struct kdhyperrect *rect);
|
94
|
+
static void hyperrect_extend(struct kdhyperrect *rect, const double *pos);
|
95
|
+
static double hyperrect_dist_sq(struct kdhyperrect *rect, const double *pos);
|
96
|
+
|
97
|
+
#ifdef USE_LIST_NODE_ALLOCATOR
|
98
|
+
static struct res_node *alloc_resnode(void);
|
99
|
+
static void free_resnode(struct res_node*);
|
100
|
+
#else
|
101
|
+
#define alloc_resnode() malloc(sizeof(struct res_node))
|
102
|
+
#define free_resnode(n) free(n)
|
103
|
+
#endif
|
104
|
+
|
105
|
+
|
106
|
+
|
107
|
+
struct kdtree *kd_create(int k)
|
108
|
+
{
|
109
|
+
struct kdtree *tree;
|
110
|
+
|
111
|
+
if(!(tree = malloc(sizeof *tree))) {
|
112
|
+
return 0;
|
113
|
+
}
|
114
|
+
|
115
|
+
tree->dim = k;
|
116
|
+
tree->root = 0;
|
117
|
+
tree->destr = 0;
|
118
|
+
tree->rect = 0;
|
119
|
+
|
120
|
+
return tree;
|
121
|
+
}
|
122
|
+
|
123
|
+
void kd_free(struct kdtree *tree)
|
124
|
+
{
|
125
|
+
if(tree) {
|
126
|
+
kd_clear(tree);
|
127
|
+
free(tree);
|
128
|
+
}
|
129
|
+
}
|
130
|
+
|
131
|
+
static void clear_rec(struct kdnode *node, void (*destr)(void*))
|
132
|
+
{
|
133
|
+
if(!node) return;
|
134
|
+
|
135
|
+
clear_rec(node->left, destr);
|
136
|
+
clear_rec(node->right, destr);
|
137
|
+
|
138
|
+
if(destr) {
|
139
|
+
destr(node->data);
|
140
|
+
}
|
141
|
+
free(node->pos);
|
142
|
+
free(node);
|
143
|
+
}
|
144
|
+
|
145
|
+
void kd_clear(struct kdtree *tree)
|
146
|
+
{
|
147
|
+
clear_rec(tree->root, tree->destr);
|
148
|
+
tree->root = 0;
|
149
|
+
|
150
|
+
if (tree->rect) {
|
151
|
+
hyperrect_free(tree->rect);
|
152
|
+
tree->rect = 0;
|
153
|
+
}
|
154
|
+
}
|
155
|
+
|
156
|
+
void kd_data_destructor(struct kdtree *tree, void (*destr)(void*))
|
157
|
+
{
|
158
|
+
tree->destr = destr;
|
159
|
+
}
|
160
|
+
|
161
|
+
|
162
|
+
static int insert_rec(struct kdnode **nptr, const double *pos, void *data, int dir, int dim)
|
163
|
+
{
|
164
|
+
int new_dir;
|
165
|
+
struct kdnode *node;
|
166
|
+
|
167
|
+
if(!*nptr) {
|
168
|
+
if(!(node = malloc(sizeof *node))) {
|
169
|
+
return -1;
|
170
|
+
}
|
171
|
+
if(!(node->pos = malloc(dim * sizeof *node->pos))) {
|
172
|
+
free(node);
|
173
|
+
return -1;
|
174
|
+
}
|
175
|
+
memcpy(node->pos, pos, dim * sizeof *node->pos);
|
176
|
+
node->data = data;
|
177
|
+
node->dir = dir;
|
178
|
+
node->left = node->right = 0;
|
179
|
+
*nptr = node;
|
180
|
+
return 0;
|
181
|
+
}
|
182
|
+
|
183
|
+
node = *nptr;
|
184
|
+
new_dir = (node->dir + 1) % dim;
|
185
|
+
if(pos[node->dir] < node->pos[node->dir]) {
|
186
|
+
return insert_rec(&(*nptr)->left, pos, data, new_dir, dim);
|
187
|
+
}
|
188
|
+
return insert_rec(&(*nptr)->right, pos, data, new_dir, dim);
|
189
|
+
}
|
190
|
+
|
191
|
+
int kd_insert(struct kdtree *tree, const double *pos, void *data)
|
192
|
+
{
|
193
|
+
if (insert_rec(&tree->root, pos, data, 0, tree->dim)) {
|
194
|
+
return -1;
|
195
|
+
}
|
196
|
+
|
197
|
+
if (tree->rect == 0) {
|
198
|
+
tree->rect = hyperrect_create(tree->dim, pos, pos);
|
199
|
+
} else {
|
200
|
+
hyperrect_extend(tree->rect, pos);
|
201
|
+
}
|
202
|
+
|
203
|
+
return 0;
|
204
|
+
}
|
205
|
+
|
206
|
+
int kd_insertf(struct kdtree *tree, const float *pos, void *data)
|
207
|
+
{
|
208
|
+
static double sbuf[16];
|
209
|
+
double *bptr, *buf = 0;
|
210
|
+
int res, dim = tree->dim;
|
211
|
+
|
212
|
+
if(dim > 16) {
|
213
|
+
#ifndef NO_ALLOCA
|
214
|
+
if(dim <= 256)
|
215
|
+
bptr = buf = alloca(dim * sizeof *bptr);
|
216
|
+
else
|
217
|
+
#endif
|
218
|
+
if(!(bptr = buf = malloc(dim * sizeof *bptr))) {
|
219
|
+
return -1;
|
220
|
+
}
|
221
|
+
} else {
|
222
|
+
bptr = buf = sbuf;
|
223
|
+
}
|
224
|
+
|
225
|
+
while(dim-- > 0) {
|
226
|
+
*bptr++ = *pos++;
|
227
|
+
}
|
228
|
+
|
229
|
+
res = kd_insert(tree, buf, data);
|
230
|
+
#ifndef NO_ALLOCA
|
231
|
+
if(tree->dim > 256)
|
232
|
+
#else
|
233
|
+
if(tree->dim > 16)
|
234
|
+
#endif
|
235
|
+
free(buf);
|
236
|
+
return res;
|
237
|
+
}
|
238
|
+
|
239
|
+
int kd_insert3(struct kdtree *tree, double x, double y, double z, void *data)
|
240
|
+
{
|
241
|
+
double buf[3];
|
242
|
+
buf[0] = x;
|
243
|
+
buf[1] = y;
|
244
|
+
buf[2] = z;
|
245
|
+
return kd_insert(tree, buf, data);
|
246
|
+
}
|
247
|
+
|
248
|
+
int kd_insert3f(struct kdtree *tree, float x, float y, float z, void *data)
|
249
|
+
{
|
250
|
+
double buf[3];
|
251
|
+
buf[0] = x;
|
252
|
+
buf[1] = y;
|
253
|
+
buf[2] = z;
|
254
|
+
return kd_insert(tree, buf, data);
|
255
|
+
}
|
256
|
+
|
257
|
+
static int find_nearest(struct kdnode *node, const double *pos, double range, struct res_node *list, int ordered, int dim)
|
258
|
+
{
|
259
|
+
double dist_sq, dx;
|
260
|
+
int i, ret, added_res = 0;
|
261
|
+
|
262
|
+
if(!node) return 0;
|
263
|
+
|
264
|
+
dist_sq = 0;
|
265
|
+
for(i=0; i<dim; i++) {
|
266
|
+
dist_sq += SQ(node->pos[i] - pos[i]);
|
267
|
+
}
|
268
|
+
if(dist_sq <= SQ(range)) {
|
269
|
+
if(rlist_insert(list, node, ordered ? dist_sq : -1.0) == -1) {
|
270
|
+
return -1;
|
271
|
+
}
|
272
|
+
added_res = 1;
|
273
|
+
}
|
274
|
+
|
275
|
+
dx = pos[node->dir] - node->pos[node->dir];
|
276
|
+
|
277
|
+
ret = find_nearest(dx <= 0.0 ? node->left : node->right, pos, range, list, ordered, dim);
|
278
|
+
if(ret >= 0 && fabs(dx) < range) {
|
279
|
+
added_res += ret;
|
280
|
+
ret = find_nearest(dx <= 0.0 ? node->right : node->left, pos, range, list, ordered, dim);
|
281
|
+
}
|
282
|
+
if(ret == -1) {
|
283
|
+
return -1;
|
284
|
+
}
|
285
|
+
added_res += ret;
|
286
|
+
|
287
|
+
return added_res;
|
288
|
+
}
|
289
|
+
|
290
|
+
#if 0
|
291
|
+
static int find_nearest_n(struct kdnode *node, const double *pos, double range, int num, struct rheap *heap, int dim)
|
292
|
+
{
|
293
|
+
double dist_sq, dx;
|
294
|
+
int i, ret, added_res = 0;
|
295
|
+
|
296
|
+
if(!node) return 0;
|
297
|
+
|
298
|
+
/* if the photon is close enough, add it to the result heap */
|
299
|
+
dist_sq = 0;
|
300
|
+
for(i=0; i<dim; i++) {
|
301
|
+
dist_sq += SQ(node->pos[i] - pos[i]);
|
302
|
+
}
|
303
|
+
if(dist_sq <= range_sq) {
|
304
|
+
if(heap->size >= num) {
|
305
|
+
/* get furthest element */
|
306
|
+
struct res_node *maxelem = rheap_get_max(heap);
|
307
|
+
|
308
|
+
/* and check if the new one is closer than that */
|
309
|
+
if(maxelem->dist_sq > dist_sq) {
|
310
|
+
rheap_remove_max(heap);
|
311
|
+
|
312
|
+
if(rheap_insert(heap, node, dist_sq) == -1) {
|
313
|
+
return -1;
|
314
|
+
}
|
315
|
+
added_res = 1;
|
316
|
+
|
317
|
+
range_sq = dist_sq;
|
318
|
+
}
|
319
|
+
} else {
|
320
|
+
if(rheap_insert(heap, node, dist_sq) == -1) {
|
321
|
+
return =1;
|
322
|
+
}
|
323
|
+
added_res = 1;
|
324
|
+
}
|
325
|
+
}
|
326
|
+
|
327
|
+
|
328
|
+
/* find signed distance from the splitting plane */
|
329
|
+
dx = pos[node->dir] - node->pos[node->dir];
|
330
|
+
|
331
|
+
ret = find_nearest_n(dx <= 0.0 ? node->left : node->right, pos, range, num, heap, dim);
|
332
|
+
if(ret >= 0 && fabs(dx) < range) {
|
333
|
+
added_res += ret;
|
334
|
+
ret = find_nearest_n(dx <= 0.0 ? node->right : node->left, pos, range, num, heap, dim);
|
335
|
+
}
|
336
|
+
|
337
|
+
}
|
338
|
+
#endif
|
339
|
+
|
340
|
+
static void kd_nearest_i(struct kdnode *node, const double *pos, struct kdnode **result, double *result_dist_sq, struct kdhyperrect* rect)
|
341
|
+
{
|
342
|
+
int dir = node->dir;
|
343
|
+
int i;
|
344
|
+
double dummy, dist_sq;
|
345
|
+
struct kdnode *nearer_subtree, *farther_subtree;
|
346
|
+
double *nearer_hyperrect_coord, *farther_hyperrect_coord;
|
347
|
+
|
348
|
+
/* Decide whether to go left or right in the tree */
|
349
|
+
dummy = pos[dir] - node->pos[dir];
|
350
|
+
if (dummy <= 0) {
|
351
|
+
nearer_subtree = node->left;
|
352
|
+
farther_subtree = node->right;
|
353
|
+
nearer_hyperrect_coord = rect->max + dir;
|
354
|
+
farther_hyperrect_coord = rect->min + dir;
|
355
|
+
} else {
|
356
|
+
nearer_subtree = node->right;
|
357
|
+
farther_subtree = node->left;
|
358
|
+
nearer_hyperrect_coord = rect->min + dir;
|
359
|
+
farther_hyperrect_coord = rect->max + dir;
|
360
|
+
}
|
361
|
+
|
362
|
+
if (nearer_subtree) {
|
363
|
+
/* Slice the hyperrect to get the hyperrect of the nearer subtree */
|
364
|
+
dummy = *nearer_hyperrect_coord;
|
365
|
+
*nearer_hyperrect_coord = node->pos[dir];
|
366
|
+
/* Recurse down into nearer subtree */
|
367
|
+
kd_nearest_i(nearer_subtree, pos, result, result_dist_sq, rect);
|
368
|
+
/* Undo the slice */
|
369
|
+
*nearer_hyperrect_coord = dummy;
|
370
|
+
}
|
371
|
+
|
372
|
+
/* Check the distance of the point at the current node, compare it
|
373
|
+
* with our best so far */
|
374
|
+
dist_sq = 0;
|
375
|
+
for(i=0; i < rect->dim; i++) {
|
376
|
+
dist_sq += SQ(node->pos[i] - pos[i]);
|
377
|
+
}
|
378
|
+
if (dist_sq < *result_dist_sq) {
|
379
|
+
*result = node;
|
380
|
+
*result_dist_sq = dist_sq;
|
381
|
+
}
|
382
|
+
|
383
|
+
if (farther_subtree) {
|
384
|
+
/* Get the hyperrect of the farther subtree */
|
385
|
+
dummy = *farther_hyperrect_coord;
|
386
|
+
*farther_hyperrect_coord = node->pos[dir];
|
387
|
+
/* Check if we have to recurse down by calculating the closest
|
388
|
+
* point of the hyperrect and see if it's closer than our
|
389
|
+
* minimum distance in result_dist_sq. */
|
390
|
+
if (hyperrect_dist_sq(rect, pos) < *result_dist_sq) {
|
391
|
+
/* Recurse down into farther subtree */
|
392
|
+
kd_nearest_i(farther_subtree, pos, result, result_dist_sq, rect);
|
393
|
+
}
|
394
|
+
/* Undo the slice on the hyperrect */
|
395
|
+
*farther_hyperrect_coord = dummy;
|
396
|
+
}
|
397
|
+
}
|
398
|
+
|
399
|
+
struct kdres *kd_nearest(struct kdtree *kd, const double *pos)
|
400
|
+
{
|
401
|
+
struct kdhyperrect *rect;
|
402
|
+
struct kdnode *result;
|
403
|
+
struct kdres *rset;
|
404
|
+
double dist_sq;
|
405
|
+
int i;
|
406
|
+
|
407
|
+
if (!kd) return 0;
|
408
|
+
if (!kd->rect) return 0;
|
409
|
+
|
410
|
+
/* Allocate result set */
|
411
|
+
if(!(rset = malloc(sizeof *rset))) {
|
412
|
+
return 0;
|
413
|
+
}
|
414
|
+
if(!(rset->rlist = alloc_resnode())) {
|
415
|
+
free(rset);
|
416
|
+
return 0;
|
417
|
+
}
|
418
|
+
rset->rlist->next = 0;
|
419
|
+
rset->tree = kd;
|
420
|
+
|
421
|
+
/* Duplicate the bounding hyperrectangle, we will work on the copy */
|
422
|
+
if (!(rect = hyperrect_duplicate(kd->rect))) {
|
423
|
+
kd_res_free(rset);
|
424
|
+
return 0;
|
425
|
+
}
|
426
|
+
|
427
|
+
/* Our first guesstimate is the root node */
|
428
|
+
result = kd->root;
|
429
|
+
dist_sq = 0;
|
430
|
+
for (i = 0; i < kd->dim; i++)
|
431
|
+
dist_sq += SQ(result->pos[i] - pos[i]);
|
432
|
+
|
433
|
+
/* Search for the nearest neighbour recursively */
|
434
|
+
kd_nearest_i(kd->root, pos, &result, &dist_sq, rect);
|
435
|
+
|
436
|
+
/* Free the copy of the hyperrect */
|
437
|
+
hyperrect_free(rect);
|
438
|
+
|
439
|
+
/* Store the result */
|
440
|
+
if (result) {
|
441
|
+
if (rlist_insert(rset->rlist, result, -1.0) == -1) {
|
442
|
+
kd_res_free(rset);
|
443
|
+
return 0;
|
444
|
+
}
|
445
|
+
rset->size = 1;
|
446
|
+
kd_res_rewind(rset);
|
447
|
+
return rset;
|
448
|
+
} else {
|
449
|
+
kd_res_free(rset);
|
450
|
+
return 0;
|
451
|
+
}
|
452
|
+
}
|
453
|
+
|
454
|
+
struct kdres *kd_nearestf(struct kdtree *tree, const float *pos)
|
455
|
+
{
|
456
|
+
static double sbuf[16];
|
457
|
+
double *bptr, *buf = 0;
|
458
|
+
int dim = tree->dim;
|
459
|
+
struct kdres *res;
|
460
|
+
|
461
|
+
if(dim > 16) {
|
462
|
+
#ifndef NO_ALLOCA
|
463
|
+
if(dim <= 256)
|
464
|
+
bptr = buf = alloca(dim * sizeof *bptr);
|
465
|
+
else
|
466
|
+
#endif
|
467
|
+
if(!(bptr = buf = malloc(dim * sizeof *bptr))) {
|
468
|
+
return 0;
|
469
|
+
}
|
470
|
+
} else {
|
471
|
+
bptr = buf = sbuf;
|
472
|
+
}
|
473
|
+
|
474
|
+
while(dim-- > 0) {
|
475
|
+
*bptr++ = *pos++;
|
476
|
+
}
|
477
|
+
|
478
|
+
res = kd_nearest(tree, buf);
|
479
|
+
#ifndef NO_ALLOCA
|
480
|
+
if(tree->dim > 256)
|
481
|
+
#else
|
482
|
+
if(tree->dim > 16)
|
483
|
+
#endif
|
484
|
+
free(buf);
|
485
|
+
return res;
|
486
|
+
}
|
487
|
+
|
488
|
+
struct kdres *kd_nearest3(struct kdtree *tree, double x, double y, double z)
|
489
|
+
{
|
490
|
+
double pos[3];
|
491
|
+
pos[0] = x;
|
492
|
+
pos[1] = y;
|
493
|
+
pos[2] = z;
|
494
|
+
return kd_nearest(tree, pos);
|
495
|
+
}
|
496
|
+
|
497
|
+
struct kdres *kd_nearest3f(struct kdtree *tree, float x, float y, float z)
|
498
|
+
{
|
499
|
+
double pos[3];
|
500
|
+
pos[0] = x;
|
501
|
+
pos[1] = y;
|
502
|
+
pos[2] = z;
|
503
|
+
return kd_nearest(tree, pos);
|
504
|
+
}
|
505
|
+
|
506
|
+
/* ---- nearest N search ---- */
|
507
|
+
/*
|
508
|
+
static kdres *kd_nearest_n(struct kdtree *kd, const double *pos, int num)
|
509
|
+
{
|
510
|
+
int ret;
|
511
|
+
struct kdres *rset;
|
512
|
+
|
513
|
+
if(!(rset = malloc(sizeof *rset))) {
|
514
|
+
return 0;
|
515
|
+
}
|
516
|
+
if(!(rset->rlist = alloc_resnode())) {
|
517
|
+
free(rset);
|
518
|
+
return 0;
|
519
|
+
}
|
520
|
+
rset->rlist->next = 0;
|
521
|
+
rset->tree = kd;
|
522
|
+
|
523
|
+
if((ret = find_nearest_n(kd->root, pos, range, num, rset->rlist, kd->dim)) == -1) {
|
524
|
+
kd_res_free(rset);
|
525
|
+
return 0;
|
526
|
+
}
|
527
|
+
rset->size = ret;
|
528
|
+
kd_res_rewind(rset);
|
529
|
+
return rset;
|
530
|
+
}*/
|
531
|
+
|
532
|
+
struct kdres *kd_nearest_range(struct kdtree *kd, const double *pos, double range)
|
533
|
+
{
|
534
|
+
int ret;
|
535
|
+
struct kdres *rset;
|
536
|
+
|
537
|
+
if(!(rset = malloc(sizeof *rset))) {
|
538
|
+
return 0;
|
539
|
+
}
|
540
|
+
if(!(rset->rlist = alloc_resnode())) {
|
541
|
+
free(rset);
|
542
|
+
return 0;
|
543
|
+
}
|
544
|
+
rset->rlist->next = 0;
|
545
|
+
rset->tree = kd;
|
546
|
+
|
547
|
+
if((ret = find_nearest(kd->root, pos, range, rset->rlist, 0, kd->dim)) == -1) {
|
548
|
+
kd_res_free(rset);
|
549
|
+
return 0;
|
550
|
+
}
|
551
|
+
rset->size = ret;
|
552
|
+
kd_res_rewind(rset);
|
553
|
+
return rset;
|
554
|
+
}
|
555
|
+
|
556
|
+
struct kdres *kd_nearest_rangef(struct kdtree *kd, const float *pos, float range)
|
557
|
+
{
|
558
|
+
static double sbuf[16];
|
559
|
+
double *bptr, *buf = 0;
|
560
|
+
int dim = kd->dim;
|
561
|
+
struct kdres *res;
|
562
|
+
|
563
|
+
if(dim > 16) {
|
564
|
+
#ifndef NO_ALLOCA
|
565
|
+
if(dim <= 256)
|
566
|
+
bptr = buf = alloca(dim * sizeof *bptr);
|
567
|
+
else
|
568
|
+
#endif
|
569
|
+
if(!(bptr = buf = malloc(dim * sizeof *bptr))) {
|
570
|
+
return 0;
|
571
|
+
}
|
572
|
+
} else {
|
573
|
+
bptr = buf = sbuf;
|
574
|
+
}
|
575
|
+
|
576
|
+
while(dim-- > 0) {
|
577
|
+
*bptr++ = *pos++;
|
578
|
+
}
|
579
|
+
|
580
|
+
res = kd_nearest_range(kd, buf, range);
|
581
|
+
#ifndef NO_ALLOCA
|
582
|
+
if(kd->dim > 256)
|
583
|
+
#else
|
584
|
+
if(kd->dim > 16)
|
585
|
+
#endif
|
586
|
+
free(buf);
|
587
|
+
return res;
|
588
|
+
}
|
589
|
+
|
590
|
+
struct kdres *kd_nearest_range3(struct kdtree *tree, double x, double y, double z, double range)
|
591
|
+
{
|
592
|
+
double buf[3];
|
593
|
+
buf[0] = x;
|
594
|
+
buf[1] = y;
|
595
|
+
buf[2] = z;
|
596
|
+
return kd_nearest_range(tree, buf, range);
|
597
|
+
}
|
598
|
+
|
599
|
+
struct kdres *kd_nearest_range3f(struct kdtree *tree, float x, float y, float z, float range)
|
600
|
+
{
|
601
|
+
double buf[3];
|
602
|
+
buf[0] = x;
|
603
|
+
buf[1] = y;
|
604
|
+
buf[2] = z;
|
605
|
+
return kd_nearest_range(tree, buf, range);
|
606
|
+
}
|
607
|
+
|
608
|
+
void kd_res_free(struct kdres *rset)
|
609
|
+
{
|
610
|
+
clear_results(rset);
|
611
|
+
free_resnode(rset->rlist);
|
612
|
+
free(rset);
|
613
|
+
}
|
614
|
+
|
615
|
+
int kd_res_size(struct kdres *set)
|
616
|
+
{
|
617
|
+
return (set->size);
|
618
|
+
}
|
619
|
+
|
620
|
+
void kd_res_rewind(struct kdres *rset)
|
621
|
+
{
|
622
|
+
rset->riter = rset->rlist->next;
|
623
|
+
}
|
624
|
+
|
625
|
+
int kd_res_end(struct kdres *rset)
|
626
|
+
{
|
627
|
+
return rset->riter == 0;
|
628
|
+
}
|
629
|
+
|
630
|
+
int kd_res_next(struct kdres *rset)
|
631
|
+
{
|
632
|
+
rset->riter = rset->riter->next;
|
633
|
+
return rset->riter != 0;
|
634
|
+
}
|
635
|
+
|
636
|
+
void *kd_res_item(struct kdres *rset, double *pos)
|
637
|
+
{
|
638
|
+
if(rset->riter) {
|
639
|
+
if(pos) {
|
640
|
+
memcpy(pos, rset->riter->item->pos, rset->tree->dim * sizeof *pos);
|
641
|
+
}
|
642
|
+
return rset->riter->item->data;
|
643
|
+
}
|
644
|
+
return 0;
|
645
|
+
}
|
646
|
+
|
647
|
+
void *kd_res_itemf(struct kdres *rset, float *pos)
|
648
|
+
{
|
649
|
+
if(rset->riter) {
|
650
|
+
if(pos) {
|
651
|
+
int i;
|
652
|
+
for(i=0; i<rset->tree->dim; i++) {
|
653
|
+
pos[i] = rset->riter->item->pos[i];
|
654
|
+
}
|
655
|
+
}
|
656
|
+
return rset->riter->item->data;
|
657
|
+
}
|
658
|
+
return 0;
|
659
|
+
}
|
660
|
+
|
661
|
+
void *kd_res_item3(struct kdres *rset, double *x, double *y, double *z)
|
662
|
+
{
|
663
|
+
if(rset->riter) {
|
664
|
+
if(*x) *x = rset->riter->item->pos[0];
|
665
|
+
if(*y) *y = rset->riter->item->pos[1];
|
666
|
+
if(*z) *z = rset->riter->item->pos[2];
|
667
|
+
}
|
668
|
+
return 0;
|
669
|
+
}
|
670
|
+
|
671
|
+
void *kd_res_item3f(struct kdres *rset, float *x, float *y, float *z)
|
672
|
+
{
|
673
|
+
if(rset->riter) {
|
674
|
+
if(*x) *x = rset->riter->item->pos[0];
|
675
|
+
if(*y) *y = rset->riter->item->pos[1];
|
676
|
+
if(*z) *z = rset->riter->item->pos[2];
|
677
|
+
}
|
678
|
+
return 0;
|
679
|
+
}
|
680
|
+
|
681
|
+
void *kd_res_item_data(struct kdres *set)
|
682
|
+
{
|
683
|
+
return kd_res_item(set, 0);
|
684
|
+
}
|
685
|
+
|
686
|
+
/* ---- hyperrectangle helpers ---- */
|
687
|
+
static struct kdhyperrect* hyperrect_create(int dim, const double *min, const double *max)
|
688
|
+
{
|
689
|
+
size_t size = dim * sizeof(double);
|
690
|
+
struct kdhyperrect* rect = 0;
|
691
|
+
|
692
|
+
if (!(rect = malloc(sizeof(struct kdhyperrect)))) {
|
693
|
+
return 0;
|
694
|
+
}
|
695
|
+
|
696
|
+
rect->dim = dim;
|
697
|
+
if (!(rect->min = malloc(size))) {
|
698
|
+
free(rect);
|
699
|
+
return 0;
|
700
|
+
}
|
701
|
+
if (!(rect->max = malloc(size))) {
|
702
|
+
free(rect->min);
|
703
|
+
free(rect);
|
704
|
+
return 0;
|
705
|
+
}
|
706
|
+
memcpy(rect->min, min, size);
|
707
|
+
memcpy(rect->max, max, size);
|
708
|
+
|
709
|
+
return rect;
|
710
|
+
}
|
711
|
+
|
712
|
+
static void hyperrect_free(struct kdhyperrect *rect)
|
713
|
+
{
|
714
|
+
free(rect->min);
|
715
|
+
free(rect->max);
|
716
|
+
free(rect);
|
717
|
+
}
|
718
|
+
|
719
|
+
static struct kdhyperrect* hyperrect_duplicate(const struct kdhyperrect *rect)
|
720
|
+
{
|
721
|
+
return hyperrect_create(rect->dim, rect->min, rect->max);
|
722
|
+
}
|
723
|
+
|
724
|
+
static void hyperrect_extend(struct kdhyperrect *rect, const double *pos)
|
725
|
+
{
|
726
|
+
int i;
|
727
|
+
|
728
|
+
for (i=0; i < rect->dim; i++) {
|
729
|
+
if (pos[i] < rect->min[i]) {
|
730
|
+
rect->min[i] = pos[i];
|
731
|
+
}
|
732
|
+
if (pos[i] > rect->max[i]) {
|
733
|
+
rect->max[i] = pos[i];
|
734
|
+
}
|
735
|
+
}
|
736
|
+
}
|
737
|
+
|
738
|
+
static double hyperrect_dist_sq(struct kdhyperrect *rect, const double *pos)
|
739
|
+
{
|
740
|
+
int i;
|
741
|
+
double result = 0;
|
742
|
+
|
743
|
+
for (i=0; i < rect->dim; i++) {
|
744
|
+
if (pos[i] < rect->min[i]) {
|
745
|
+
result += SQ(rect->min[i] - pos[i]);
|
746
|
+
} else if (pos[i] > rect->max[i]) {
|
747
|
+
result += SQ(rect->max[i] - pos[i]);
|
748
|
+
}
|
749
|
+
}
|
750
|
+
|
751
|
+
return result;
|
752
|
+
}
|
753
|
+
|
754
|
+
/* ---- static helpers ---- */
|
755
|
+
|
756
|
+
#ifdef USE_LIST_NODE_ALLOCATOR
|
757
|
+
/* special list node allocators. */
|
758
|
+
static struct res_node *free_nodes;
|
759
|
+
|
760
|
+
#ifndef NO_PTHREADS
|
761
|
+
static pthread_mutex_t alloc_mutex = PTHREAD_MUTEX_INITIALIZER;
|
762
|
+
#endif
|
763
|
+
|
764
|
+
static struct res_node *alloc_resnode(void)
|
765
|
+
{
|
766
|
+
struct res_node *node;
|
767
|
+
|
768
|
+
#ifndef NO_PTHREADS
|
769
|
+
pthread_mutex_lock(&alloc_mutex);
|
770
|
+
#endif
|
771
|
+
|
772
|
+
if(!free_nodes) {
|
773
|
+
node = malloc(sizeof *node);
|
774
|
+
} else {
|
775
|
+
node = free_nodes;
|
776
|
+
free_nodes = free_nodes->next;
|
777
|
+
node->next = 0;
|
778
|
+
}
|
779
|
+
|
780
|
+
#ifndef NO_PTHREADS
|
781
|
+
pthread_mutex_unlock(&alloc_mutex);
|
782
|
+
#endif
|
783
|
+
|
784
|
+
return node;
|
785
|
+
}
|
786
|
+
|
787
|
+
static void free_resnode(struct res_node *node)
|
788
|
+
{
|
789
|
+
#ifndef NO_PTHREADS
|
790
|
+
pthread_mutex_lock(&alloc_mutex);
|
791
|
+
#endif
|
792
|
+
|
793
|
+
node->next = free_nodes;
|
794
|
+
free_nodes = node;
|
795
|
+
|
796
|
+
#ifndef NO_PTHREADS
|
797
|
+
pthread_mutex_unlock(&alloc_mutex);
|
798
|
+
#endif
|
799
|
+
}
|
800
|
+
#endif /* list node allocator or not */
|
801
|
+
|
802
|
+
|
803
|
+
/* inserts the item. if dist_sq is >= 0, then do an ordered insert */
|
804
|
+
/* TODO make the ordering code use heapsort */
|
805
|
+
static int rlist_insert(struct res_node *list, struct kdnode *item, double dist_sq)
|
806
|
+
{
|
807
|
+
struct res_node *rnode;
|
808
|
+
|
809
|
+
if(!(rnode = alloc_resnode())) {
|
810
|
+
return -1;
|
811
|
+
}
|
812
|
+
rnode->item = item;
|
813
|
+
rnode->dist_sq = dist_sq;
|
814
|
+
|
815
|
+
if(dist_sq >= 0.0) {
|
816
|
+
while(list->next && list->next->dist_sq < dist_sq) {
|
817
|
+
list = list->next;
|
818
|
+
}
|
819
|
+
}
|
820
|
+
rnode->next = list->next;
|
821
|
+
list->next = rnode;
|
822
|
+
return 0;
|
823
|
+
}
|
824
|
+
|
825
|
+
static void clear_results(struct kdres *rset)
|
826
|
+
{
|
827
|
+
struct res_node *tmp, *node = rset->rlist->next;
|
828
|
+
|
829
|
+
while(node) {
|
830
|
+
tmp = node;
|
831
|
+
node = node->next;
|
832
|
+
free_resnode(tmp);
|
833
|
+
}
|
834
|
+
|
835
|
+
rset->rlist->next = 0;
|
836
|
+
}
|