Fawkes API Fawkes Development Version
pf_kdtree.c
1
2/***************************************************************************
3 * pf_kdtree.c: KD tree functions
4 *
5 * Created: Thu May 24 18:39:48 2012
6 * Copyright 2000 Brian Gerkey
7 * 2000 Kasper Stoy
8 * 2012 Tim Niemueller [www.niemueller.de]
9 ****************************************************************************/
10
11/* This program is free software; you can redistribute it and/or modify
12 * it under the terms of the GNU General Public License as published by
13 * the Free Software Foundation; either version 2 of the License, or
14 * (at your option) any later version.
15 *
16 * This program is distributed in the hope that it will be useful,
17 * but WITHOUT ANY WARRANTY; without even the implied warranty of
18 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
19 * GNU Library General Public License for more details.
20 *
21 * Read the full text in the LICENSE.GPL file in the doc directory.
22 */
23
24/* From:
25 * Player - One Hell of a Robot Server (LGPL)
26 * Copyright (C) 2000 Brian Gerkey & Kasper Stoy
27 * gerkey@usc.edu kaspers@robotics.usc.edu
28 */
29/**************************************************************************
30 * Desc: KD tree functions
31 * Author: Andrew Howard
32 * Date: 18 Dec 2002
33 *************************************************************************/
34
35#include <assert.h>
36#include <math.h>
37#include <stdlib.h>
38#include <string.h>
39
40
41#include "pf_vector.h"
42#include "pf_kdtree.h"
43
44
45// Compare keys to see if they are equal
46static int pf_kdtree_equal(pf_kdtree_t *self, int key_a[], int key_b[]);
47
48// Insert a node into the tree
49static pf_kdtree_node_t *pf_kdtree_insert_node(pf_kdtree_t *self, pf_kdtree_node_t *parent,
50 pf_kdtree_node_t *node, int key[], double value);
51
52// Recursive node search
53static pf_kdtree_node_t *pf_kdtree_find_node(pf_kdtree_t *self, pf_kdtree_node_t *node, int key[]);
54
55// Recursively label nodes in this cluster
56static void pf_kdtree_cluster_node(pf_kdtree_t *self, pf_kdtree_node_t *node, int depth);
57
58// Recursive node printing
59//static void pf_kdtree_print_node(pf_kdtree_t *self, pf_kdtree_node_t *node);
60
61
62#ifdef INCLUDE_RTKGUI
63
64// Recursively draw nodes
65static void pf_kdtree_draw_node(pf_kdtree_t *self, pf_kdtree_node_t *node, rtk_fig_t *fig);
66
67#endif
68
69
70
71////////////////////////////////////////////////////////////////////////////////
72// Create a tree
73pf_kdtree_t *pf_kdtree_alloc(int max_size)
74{
75 pf_kdtree_t *self;
76
77 self = calloc(1, sizeof(pf_kdtree_t));
78
79 self->size[0] = 0.50;
80 self->size[1] = 0.50;
81 self->size[2] = (10 * M_PI / 180);
82
83 self->root = NULL;
84
85 self->node_count = 0;
86 self->node_max_count = max_size;
87 self->nodes = calloc(self->node_max_count, sizeof(pf_kdtree_node_t));
88
89 self->leaf_count = 0;
90
91 return self;
92}
93
94
95////////////////////////////////////////////////////////////////////////////////
96// Destroy a tree
97void pf_kdtree_free(pf_kdtree_t *self)
98{
99 free(self->nodes);
100 free(self);
101 return;
102}
103
104
105////////////////////////////////////////////////////////////////////////////////
106// Clear all entries from the tree
107void pf_kdtree_clear(pf_kdtree_t *self)
108{
109 self->root = NULL;
110 self->leaf_count = 0;
111 self->node_count = 0;
112
113 return;
114}
115
116
117////////////////////////////////////////////////////////////////////////////////
118// Insert a pose into the tree.
119void pf_kdtree_insert(pf_kdtree_t *self, pf_vector_t pose, double value)
120{
121 int key[3];
122
123 key[0] = floor(pose.v[0] / self->size[0]);
124 key[1] = floor(pose.v[1] / self->size[1]);
125 key[2] = floor(pose.v[2] / self->size[2]);
126
127 self->root = pf_kdtree_insert_node(self, NULL, self->root, key, value);
128
129 // Test code
130 /*
131 printf("find %d %d %d\n", key[0], key[1], key[2]);
132 assert(pf_kdtree_find_node(self, self->root, key) != NULL);
133
134 pf_kdtree_print_node(self, self->root);
135
136 printf("\n");
137
138 for (i = 0; i < self->node_count; i++)
139 {
140 node = self->nodes + i;
141 if (node->leaf)
142 {
143 printf("find %d %d %d\n", node->key[0], node->key[1], node->key[2]);
144 assert(pf_kdtree_find_node(self, self->root, node->key) == node);
145 }
146 }
147 printf("\n\n");
148 */
149
150 return;
151}
152
153
154////////////////////////////////////////////////////////////////////////////////
155// Determine the probability estimate for the given pose. TODO: this
156// should do a kernel density estimate rather than a simple histogram.
157double pf_kdtree_get_prob(pf_kdtree_t *self, pf_vector_t pose)
158{
159 int key[3];
160 pf_kdtree_node_t *node;
161
162 key[0] = floor(pose.v[0] / self->size[0]);
163 key[1] = floor(pose.v[1] / self->size[1]);
164 key[2] = floor(pose.v[2] / self->size[2]);
165
166 node = pf_kdtree_find_node(self, self->root, key);
167 if (node == NULL)
168 return 0.0;
169 return node->value;
170}
171
172
173////////////////////////////////////////////////////////////////////////////////
174// Determine the cluster label for the given pose
175int pf_kdtree_get_cluster(pf_kdtree_t *self, pf_vector_t pose)
176{
177 int key[3];
178 pf_kdtree_node_t *node;
179
180 key[0] = floor(pose.v[0] / self->size[0]);
181 key[1] = floor(pose.v[1] / self->size[1]);
182 key[2] = floor(pose.v[2] / self->size[2]);
183
184 node = pf_kdtree_find_node(self, self->root, key);
185 if (node == NULL)
186 return -1;
187 return node->cluster;
188}
189
190
191////////////////////////////////////////////////////////////////////////////////
192// Compare keys to see if they are equal
193int pf_kdtree_equal(pf_kdtree_t *self, int key_a[], int key_b[])
194{
195 //double a, b;
196
197 if (key_a[0] != key_b[0])
198 return 0;
199 if (key_a[1] != key_b[1])
200 return 0;
201
202 if (key_a[2] != key_b[2])
203 return 0;
204
205 /* TODO: make this work (pivot selection needs fixing, too)
206 // Normalize angles
207 a = key_a[2] * self->size[2];
208 a = atan2(sin(a), cos(a)) / self->size[2];
209 b = key_b[2] * self->size[2];
210 b = atan2(sin(b), cos(b)) / self->size[2];
211
212 if ((int) a != (int) b)
213 return 0;
214 */
215
216 return 1;
217}
218
219
220////////////////////////////////////////////////////////////////////////////////
221// Insert a node into the tree
222pf_kdtree_node_t *pf_kdtree_insert_node(pf_kdtree_t *self, pf_kdtree_node_t *parent,
223 pf_kdtree_node_t *node, int key[], double value)
224{
225 int i;
226 int split, max_split;
227
228 // If the node doesnt exist yet...
229 if (node == NULL)
230 {
231 assert(self->node_count < self->node_max_count);
232 node = self->nodes + self->node_count++;
233 memset(node, 0, sizeof(pf_kdtree_node_t));
234
235 node->leaf = 1;
236
237 if (parent == NULL)
238 node->depth = 0;
239 else
240 node->depth = parent->depth + 1;
241
242 for (i = 0; i < 3; i++)
243 node->key[i] = key[i];
244
245 node->value = value;
246 self->leaf_count += 1;
247 }
248
249 // If the node exists, and it is a leaf node...
250 else if (node->leaf)
251 {
252 // If the keys are equal, increment the value
253 if (pf_kdtree_equal(self, key, node->key))
254 {
255 node->value += value;
256 }
257
258 // The keys are not equal, so split this node
259 else
260 {
261 // Find the dimension with the largest variance and do a mean
262 // split
263 max_split = 0;
264 node->pivot_dim = -1;
265 for (i = 0; i < 3; i++)
266 {
267 split = abs(key[i] - node->key[i]);
268 if (split > max_split)
269 {
270 max_split = split;
271 node->pivot_dim = i;
272 }
273 }
274 assert(node->pivot_dim >= 0);
275
276 node->pivot_value = (key[node->pivot_dim] + node->key[node->pivot_dim]) / 2.0;
277
278 if (key[node->pivot_dim] < node->pivot_value)
279 {
280 node->children[0] = pf_kdtree_insert_node(self, node, NULL, key, value);
281 node->children[1] = pf_kdtree_insert_node(self, node, NULL, node->key, node->value);
282 }
283 else
284 {
285 node->children[0] = pf_kdtree_insert_node(self, node, NULL, node->key, node->value);
286 node->children[1] = pf_kdtree_insert_node(self, node, NULL, key, value);
287 }
288
289 node->leaf = 0;
290 self->leaf_count -= 1;
291 }
292 }
293
294 // If the node exists, and it has children...
295 else
296 {
297 assert(node->children[0] != NULL);
298 assert(node->children[1] != NULL);
299
300 if (key[node->pivot_dim] < node->pivot_value)
301 pf_kdtree_insert_node(self, node, node->children[0], key, value);
302 else
303 pf_kdtree_insert_node(self, node, node->children[1], key, value);
304 }
305
306 return node;
307}
308
309
310////////////////////////////////////////////////////////////////////////////////
311// Recursive node search
312pf_kdtree_node_t *pf_kdtree_find_node(pf_kdtree_t *self, pf_kdtree_node_t *node, int key[])
313{
314 if (node->leaf)
315 {
316 //printf("find : leaf %p %d %d %d\n", node, node->key[0], node->key[1], node->key[2]);
317
318 // If the keys are the same...
319 if (pf_kdtree_equal(self, key, node->key))
320 return node;
321 else
322 return NULL;
323 }
324 else
325 {
326 //printf("find : brch %p %d %f\n", node, node->pivot_dim, node->pivot_value);
327
328 assert(node->children[0] != NULL);
329 assert(node->children[1] != NULL);
330
331 // If the keys are different...
332 if (key[node->pivot_dim] < node->pivot_value)
333 return pf_kdtree_find_node(self, node->children[0], key);
334 else
335 return pf_kdtree_find_node(self, node->children[1], key);
336 }
337
338 return NULL;
339}
340
341
342////////////////////////////////////////////////////////////////////////////////
343// Recursive node printing
344/*
345void pf_kdtree_print_node(pf_kdtree_t *self, pf_kdtree_node_t *node)
346{
347 if (node->leaf)
348 {
349 printf("(%+02d %+02d %+02d)\n", node->key[0], node->key[1], node->key[2]);
350 printf("%*s", node->depth * 11, "");
351 }
352 else
353 {
354 printf("(%+02d %+02d %+02d) ", node->key[0], node->key[1], node->key[2]);
355 pf_kdtree_print_node(self, node->children[0]);
356 pf_kdtree_print_node(self, node->children[1]);
357 }
358 return;
359}
360*/
361
362
363////////////////////////////////////////////////////////////////////////////////
364// Cluster the leaves in the tree
365void pf_kdtree_cluster(pf_kdtree_t *self)
366{
367 int i;
368 int queue_count, cluster_count;
369 pf_kdtree_node_t **queue, *node;
370
371 queue_count = 0;
372 queue = calloc(self->node_count, sizeof(queue[0]));
373
374 // Put all the leaves in a queue
375 for (i = 0; i < self->node_count; i++)
376 {
377 node = self->nodes + i;
378 if (node->leaf)
379 {
380 node->cluster = -1;
381 assert(queue_count < self->node_count);
382 queue[queue_count++] = node;
383
384 // TESTING; remove
385 assert(node == pf_kdtree_find_node(self, self->root, node->key));
386 }
387 }
388
389 cluster_count = 0;
390
391 // Do connected components for each node
392 while (queue_count > 0)
393 {
394 node = queue[--queue_count];
395
396 // If this node has already been labelled, skip it
397 if (node->cluster >= 0)
398 continue;
399
400 // Assign a label to this cluster
401 node->cluster = cluster_count++;
402
403 // Recursively label nodes in this cluster
404 pf_kdtree_cluster_node(self, node, 0);
405 }
406
407 free(queue);
408 return;
409}
410
411
412////////////////////////////////////////////////////////////////////////////////
413// Recursively label nodes in this cluster
414void pf_kdtree_cluster_node(pf_kdtree_t *self, pf_kdtree_node_t *node, int depth)
415{
416 int i;
417 int nkey[3];
418
419 for (i = 0; i < 3 * 3 * 3; i++)
420 {
421 nkey[0] = node->key[0] + (i / 9) - 1;
422 nkey[1] = node->key[1] + ((i % 9) / 3) - 1;
423 nkey[2] = node->key[2] + ((i % 9) % 3) - 1;
424
425 pf_kdtree_node_t *nnode = pf_kdtree_find_node(self, self->root, nkey);
426 if (nnode == NULL)
427 continue;
428
429 assert(nnode->leaf);
430
431 // This node already has a label; skip it. The label should be
432 // consistent, however.
433 if (nnode->cluster >= 0)
434 {
435 assert(nnode->cluster == node->cluster);
436 continue;
437 }
438
439 // Label this node and recurse
440 nnode->cluster = node->cluster;
441
442 pf_kdtree_cluster_node(self, nnode, depth + 1);
443 }
444 return;
445}
446
447
448
449#ifdef INCLUDE_RTKGUI
450
451////////////////////////////////////////////////////////////////////////////////
452// Draw the tree
453void pf_kdtree_draw(pf_kdtree_t *self, rtk_fig_t *fig)
454{
455 if (self->root != NULL)
456 pf_kdtree_draw_node(self, self->root, fig);
457 return;
458}
459
460
461////////////////////////////////////////////////////////////////////////////////
462// Recursively draw nodes
463void pf_kdtree_draw_node(pf_kdtree_t *self, pf_kdtree_node_t *node, rtk_fig_t *fig)
464{
465 if (node->leaf)
466 {
467 double ox = (node->key[0] + 0.5) * self->size[0];
468 double oy = (node->key[1] + 0.5) * self->size[1];
469 char text[64];
470
471 rtk_fig_rectangle(fig, ox, oy, 0.0, self->size[0], self->size[1], 0);
472
473 //snprintf(text, sizeof(text), "%0.3f", node->value);
474 //rtk_fig_text(fig, ox, oy, 0.0, text);
475
476 snprintf(text, sizeof(text), "%d", node->cluster);
477 rtk_fig_text(fig, ox, oy, 0.0, text);
478 }
479 else
480 {
481 assert(node->children[0] != NULL);
482 assert(node->children[1] != NULL);
483 pf_kdtree_draw_node(self, node->children[0], fig);
484 pf_kdtree_draw_node(self, node->children[1], fig);
485 }
486
487 return;
488}
489
490#endif