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
46 static int pf_kdtree_equal(pf_kdtree_t *self, int key_a[], int key_b[]);
47 
48 // Insert a node into the tree
49 static 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
53 static 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
56 static 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
65 static 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
73 pf_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
97 void 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
107 void 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.
119 void 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.
157 double 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
175 int 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
193 int 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
222 pf_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
312 pf_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 /*
345 void 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
365 void 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
414 void pf_kdtree_cluster_node(pf_kdtree_t *self, pf_kdtree_node_t *node, int depth)
415 {
416  int i;
417  int nkey[3];
418  pf_kdtree_node_t *nnode;
419 
420  for (i = 0; i < 3 * 3 * 3; i++)
421  {
422  nkey[0] = node->key[0] + (i / 9) - 1;
423  nkey[1] = node->key[1] + ((i % 9) / 3) - 1;
424  nkey[2] = node->key[2] + ((i % 9) % 3) - 1;
425 
426  nnode = pf_kdtree_find_node(self, self->root, nkey);
427  if (nnode == NULL)
428  continue;
429 
430  assert(nnode->leaf);
431 
432  // This node already has a label; skip it. The label should be
433  // consistent, however.
434  if (nnode->cluster >= 0)
435  {
436  assert(nnode->cluster == node->cluster);
437  continue;
438  }
439 
440  // Label this node and recurse
441  nnode->cluster = node->cluster;
442 
443  pf_kdtree_cluster_node(self, nnode, depth + 1);
444  }
445  return;
446 }
447 
448 
449 
450 #ifdef INCLUDE_RTKGUI
451 
452 ////////////////////////////////////////////////////////////////////////////////
453 // Draw the tree
454 void pf_kdtree_draw(pf_kdtree_t *self, rtk_fig_t *fig)
455 {
456  if (self->root != NULL)
457  pf_kdtree_draw_node(self, self->root, fig);
458  return;
459 }
460 
461 
462 ////////////////////////////////////////////////////////////////////////////////
463 // Recursively draw nodes
464 void pf_kdtree_draw_node(pf_kdtree_t *self, pf_kdtree_node_t *node, rtk_fig_t *fig)
465 {
466  double ox, oy;
467  char text[64];
468 
469  if (node->leaf)
470  {
471  ox = (node->key[0] + 0.5) * self->size[0];
472  oy = (node->key[1] + 0.5) * self->size[1];
473 
474  rtk_fig_rectangle(fig, ox, oy, 0.0, self->size[0], self->size[1], 0);
475 
476  //snprintf(text, sizeof(text), "%0.3f", node->value);
477  //rtk_fig_text(fig, ox, oy, 0.0, text);
478 
479  snprintf(text, sizeof(text), "%d", node->cluster);
480  rtk_fig_text(fig, ox, oy, 0.0, text);
481  }
482  else
483  {
484  assert(node->children[0] != NULL);
485  assert(node->children[1] != NULL);
486  pf_kdtree_draw_node(self, node->children[0], fig);
487  pf_kdtree_draw_node(self, node->children[1], fig);
488  }
489 
490  return;
491 }
492 
493 #endif