Fawkes API  Fawkes Development Version
pf_kdtree.c
00001 
00002 /***************************************************************************
00003  *  pf_kdtree.c: KD tree functions
00004  *
00005  *  Created: Thu May 24 18:39:48 2012
00006  *  Copyright  2000  Brian Gerkey
00007  *             2000  Kasper Stoy
00008  *             2012  Tim Niemueller [www.niemueller.de]
00009  ****************************************************************************/
00010 
00011 /*  This program is free software; you can redistribute it and/or modify
00012  *  it under the terms of the GNU General Public License as published by
00013  *  the Free Software Foundation; either version 2 of the License, or
00014  *  (at your option) any later version.
00015  *
00016  *  This program is distributed in the hope that it will be useful,
00017  *  but WITHOUT ANY WARRANTY; without even the implied warranty of
00018  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
00019  *  GNU Library General Public License for more details.
00020  *
00021  *  Read the full text in the LICENSE.GPL file in the doc directory.
00022  */
00023 
00024 /*  From:
00025  *  Player - One Hell of a Robot Server (LGPL)
00026  *  Copyright (C) 2000  Brian Gerkey   &  Kasper Stoy
00027  *                      gerkey@usc.edu    kaspers@robotics.usc.edu
00028  */
00029 /**************************************************************************
00030  * Desc: KD tree functions
00031  * Author: Andrew Howard
00032  * Date: 18 Dec 2002
00033  *************************************************************************/
00034 
00035 #include <assert.h>
00036 #include <math.h>
00037 #include <stdlib.h>
00038 #include <string.h>
00039 
00040 
00041 #include "pf_vector.h"
00042 #include "pf_kdtree.h"
00043 
00044 
00045 // Compare keys to see if they are equal
00046 static int pf_kdtree_equal(pf_kdtree_t *self, int key_a[], int key_b[]);
00047 
00048 // Insert a node into the tree
00049 static pf_kdtree_node_t *pf_kdtree_insert_node(pf_kdtree_t *self, pf_kdtree_node_t *parent,
00050                                                pf_kdtree_node_t *node, int key[], double value);
00051 
00052 // Recursive node search
00053 static pf_kdtree_node_t *pf_kdtree_find_node(pf_kdtree_t *self, pf_kdtree_node_t *node, int key[]);
00054 
00055 // Recursively label nodes in this cluster
00056 static void pf_kdtree_cluster_node(pf_kdtree_t *self, pf_kdtree_node_t *node, int depth);
00057 
00058 // Recursive node printing
00059 //static void pf_kdtree_print_node(pf_kdtree_t *self, pf_kdtree_node_t *node);
00060 
00061 
00062 #ifdef INCLUDE_RTKGUI
00063 
00064 // Recursively draw nodes
00065 static void pf_kdtree_draw_node(pf_kdtree_t *self, pf_kdtree_node_t *node, rtk_fig_t *fig);
00066 
00067 #endif
00068 
00069 
00070 
00071 ////////////////////////////////////////////////////////////////////////////////
00072 // Create a tree
00073 pf_kdtree_t *pf_kdtree_alloc(int max_size)
00074 {
00075   pf_kdtree_t *self;
00076 
00077   self = calloc(1, sizeof(pf_kdtree_t));
00078 
00079   self->size[0] = 0.50;
00080   self->size[1] = 0.50;
00081   self->size[2] = (10 * M_PI / 180);
00082 
00083   self->root = NULL;
00084 
00085   self->node_count = 0;
00086   self->node_max_count = max_size;
00087   self->nodes = calloc(self->node_max_count, sizeof(pf_kdtree_node_t));
00088 
00089   self->leaf_count = 0;
00090 
00091   return self;
00092 }
00093 
00094 
00095 ////////////////////////////////////////////////////////////////////////////////
00096 // Destroy a tree
00097 void pf_kdtree_free(pf_kdtree_t *self)
00098 {
00099   free(self->nodes);
00100   free(self);
00101   return;
00102 }
00103 
00104 
00105 ////////////////////////////////////////////////////////////////////////////////
00106 // Clear all entries from the tree
00107 void pf_kdtree_clear(pf_kdtree_t *self)
00108 {
00109   self->root = NULL;
00110   self->leaf_count = 0;
00111   self->node_count = 0;
00112 
00113   return;
00114 }
00115 
00116 
00117 ////////////////////////////////////////////////////////////////////////////////
00118 // Insert a pose into the tree.
00119 void pf_kdtree_insert(pf_kdtree_t *self, pf_vector_t pose, double value)
00120 {
00121   int key[3];
00122 
00123   key[0] = floor(pose.v[0] / self->size[0]);
00124   key[1] = floor(pose.v[1] / self->size[1]);
00125   key[2] = floor(pose.v[2] / self->size[2]);
00126 
00127   self->root = pf_kdtree_insert_node(self, NULL, self->root, key, value);
00128 
00129   // Test code
00130   /*
00131   printf("find %d %d %d\n", key[0], key[1], key[2]);
00132   assert(pf_kdtree_find_node(self, self->root, key) != NULL);
00133 
00134   pf_kdtree_print_node(self, self->root);
00135 
00136   printf("\n");
00137 
00138   for (i = 0; i < self->node_count; i++)
00139   {
00140     node = self->nodes + i;
00141     if (node->leaf)
00142     {
00143       printf("find %d %d %d\n", node->key[0], node->key[1], node->key[2]);
00144       assert(pf_kdtree_find_node(self, self->root, node->key) == node);
00145     }
00146   }
00147   printf("\n\n");
00148   */
00149 
00150   return;
00151 }
00152 
00153 
00154 ////////////////////////////////////////////////////////////////////////////////
00155 // Determine the probability estimate for the given pose. TODO: this
00156 // should do a kernel density estimate rather than a simple histogram.
00157 double pf_kdtree_get_prob(pf_kdtree_t *self, pf_vector_t pose)
00158 {
00159   int key[3];
00160   pf_kdtree_node_t *node;
00161 
00162   key[0] = floor(pose.v[0] / self->size[0]);
00163   key[1] = floor(pose.v[1] / self->size[1]);
00164   key[2] = floor(pose.v[2] / self->size[2]);
00165 
00166   node = pf_kdtree_find_node(self, self->root, key);
00167   if (node == NULL)
00168     return 0.0;
00169   return node->value;
00170 }
00171 
00172 
00173 ////////////////////////////////////////////////////////////////////////////////
00174 // Determine the cluster label for the given pose
00175 int pf_kdtree_get_cluster(pf_kdtree_t *self, pf_vector_t pose)
00176 {
00177   int key[3];
00178   pf_kdtree_node_t *node;
00179 
00180   key[0] = floor(pose.v[0] / self->size[0]);
00181   key[1] = floor(pose.v[1] / self->size[1]);
00182   key[2] = floor(pose.v[2] / self->size[2]);
00183 
00184   node = pf_kdtree_find_node(self, self->root, key);
00185   if (node == NULL)
00186     return -1;
00187   return node->cluster;
00188 }
00189 
00190 
00191 ////////////////////////////////////////////////////////////////////////////////
00192 // Compare keys to see if they are equal
00193 int pf_kdtree_equal(pf_kdtree_t *self, int key_a[], int key_b[])
00194 {
00195   //double a, b;
00196 
00197   if (key_a[0] != key_b[0])
00198     return 0;
00199   if (key_a[1] != key_b[1])
00200     return 0;
00201 
00202   if (key_a[2] != key_b[2])
00203     return 0;
00204 
00205   /* TODO: make this work (pivot selection needs fixing, too)
00206   // Normalize angles
00207   a = key_a[2] * self->size[2];
00208   a = atan2(sin(a), cos(a)) / self->size[2];
00209   b = key_b[2] * self->size[2];
00210   b = atan2(sin(b), cos(b)) / self->size[2];
00211 
00212  if ((int) a != (int) b)
00213     return 0;
00214   */
00215 
00216   return 1;
00217 }
00218 
00219 
00220 ////////////////////////////////////////////////////////////////////////////////
00221 // Insert a node into the tree
00222 pf_kdtree_node_t *pf_kdtree_insert_node(pf_kdtree_t *self, pf_kdtree_node_t *parent,
00223                                         pf_kdtree_node_t *node, int key[], double value)
00224 {
00225   int i;
00226   int split, max_split;
00227 
00228   // If the node doesnt exist yet...
00229   if (node == NULL)
00230   {
00231     assert(self->node_count < self->node_max_count);
00232     node = self->nodes + self->node_count++;
00233     memset(node, 0, sizeof(pf_kdtree_node_t));
00234 
00235     node->leaf = 1;
00236 
00237     if (parent == NULL)
00238       node->depth = 0;
00239     else
00240       node->depth = parent->depth + 1;
00241 
00242     for (i = 0; i < 3; i++)
00243       node->key[i] = key[i];
00244 
00245     node->value = value;
00246     self->leaf_count += 1;
00247   }
00248 
00249   // If the node exists, and it is a leaf node...
00250   else if (node->leaf)
00251   {
00252     // If the keys are equal, increment the value
00253     if (pf_kdtree_equal(self, key, node->key))
00254     {
00255       node->value += value;
00256     }
00257 
00258     // The keys are not equal, so split this node
00259     else
00260     {
00261       // Find the dimension with the largest variance and do a mean
00262       // split
00263       max_split = 0;
00264       node->pivot_dim = -1;
00265       for (i = 0; i < 3; i++)
00266       {
00267         split = abs(key[i] - node->key[i]);
00268         if (split > max_split)
00269         {
00270           max_split = split;
00271           node->pivot_dim = i;
00272         }
00273       }
00274       assert(node->pivot_dim >= 0);
00275 
00276       node->pivot_value = (key[node->pivot_dim] + node->key[node->pivot_dim]) / 2.0;
00277 
00278       if (key[node->pivot_dim] < node->pivot_value)
00279       {
00280         node->children[0] = pf_kdtree_insert_node(self, node, NULL, key, value);
00281         node->children[1] = pf_kdtree_insert_node(self, node, NULL, node->key, node->value);
00282       }
00283       else
00284       {
00285         node->children[0] = pf_kdtree_insert_node(self, node, NULL, node->key, node->value);
00286         node->children[1] = pf_kdtree_insert_node(self, node, NULL, key, value);
00287       }
00288 
00289       node->leaf = 0;
00290       self->leaf_count -= 1;
00291     }
00292   }
00293 
00294   // If the node exists, and it has children...
00295   else
00296   {
00297     assert(node->children[0] != NULL);
00298     assert(node->children[1] != NULL);
00299 
00300     if (key[node->pivot_dim] < node->pivot_value)
00301       pf_kdtree_insert_node(self, node, node->children[0], key, value);
00302     else
00303       pf_kdtree_insert_node(self, node, node->children[1], key, value);
00304   }
00305 
00306   return node;
00307 }
00308 
00309 
00310 ////////////////////////////////////////////////////////////////////////////////
00311 // Recursive node search
00312 pf_kdtree_node_t *pf_kdtree_find_node(pf_kdtree_t *self, pf_kdtree_node_t *node, int key[])
00313 {
00314   if (node->leaf)
00315   {
00316     //printf("find  : leaf %p %d %d %d\n", node, node->key[0], node->key[1], node->key[2]);
00317 
00318     // If the keys are the same...
00319     if (pf_kdtree_equal(self, key, node->key))
00320       return node;
00321     else
00322       return NULL;
00323   }
00324   else
00325   {
00326     //printf("find  : brch %p %d %f\n", node, node->pivot_dim, node->pivot_value);
00327 
00328     assert(node->children[0] != NULL);
00329     assert(node->children[1] != NULL);
00330 
00331     // If the keys are different...
00332     if (key[node->pivot_dim] < node->pivot_value)
00333       return pf_kdtree_find_node(self, node->children[0], key);
00334     else
00335       return pf_kdtree_find_node(self, node->children[1], key);
00336   }
00337 
00338   return NULL;
00339 }
00340 
00341 
00342 ////////////////////////////////////////////////////////////////////////////////
00343 // Recursive node printing
00344 /*
00345 void pf_kdtree_print_node(pf_kdtree_t *self, pf_kdtree_node_t *node)
00346 {
00347   if (node->leaf)
00348   {
00349     printf("(%+02d %+02d %+02d)\n", node->key[0], node->key[1], node->key[2]);
00350     printf("%*s", node->depth * 11, "");
00351   }
00352   else
00353   {
00354     printf("(%+02d %+02d %+02d) ", node->key[0], node->key[1], node->key[2]);
00355     pf_kdtree_print_node(self, node->children[0]);
00356     pf_kdtree_print_node(self, node->children[1]);
00357   }
00358   return;
00359 }
00360 */
00361 
00362 
00363 ////////////////////////////////////////////////////////////////////////////////
00364 // Cluster the leaves in the tree
00365 void pf_kdtree_cluster(pf_kdtree_t *self)
00366 {
00367   int i;
00368   int queue_count, cluster_count;
00369   pf_kdtree_node_t **queue, *node;
00370 
00371   queue_count = 0;
00372   queue = calloc(self->node_count, sizeof(queue[0]));
00373 
00374   // Put all the leaves in a queue
00375   for (i = 0; i < self->node_count; i++)
00376   {
00377     node = self->nodes + i;
00378     if (node->leaf)
00379     {
00380       node->cluster = -1;
00381       assert(queue_count < self->node_count);
00382       queue[queue_count++] = node;
00383 
00384       // TESTING; remove
00385       assert(node == pf_kdtree_find_node(self, self->root, node->key));
00386     }
00387   }
00388 
00389   cluster_count = 0;
00390 
00391   // Do connected components for each node
00392   while (queue_count > 0)
00393   {
00394     node = queue[--queue_count];
00395 
00396     // If this node has already been labelled, skip it
00397     if (node->cluster >= 0)
00398       continue;
00399 
00400     // Assign a label to this cluster
00401     node->cluster = cluster_count++;
00402 
00403     // Recursively label nodes in this cluster
00404     pf_kdtree_cluster_node(self, node, 0);
00405   }
00406 
00407   free(queue);
00408   return;
00409 }
00410 
00411 
00412 ////////////////////////////////////////////////////////////////////////////////
00413 // Recursively label nodes in this cluster
00414 void pf_kdtree_cluster_node(pf_kdtree_t *self, pf_kdtree_node_t *node, int depth)
00415 {
00416   int i;
00417   int nkey[3];
00418   pf_kdtree_node_t *nnode;
00419 
00420   for (i = 0; i < 3 * 3 * 3; i++)
00421   {
00422     nkey[0] = node->key[0] + (i / 9) - 1;
00423     nkey[1] = node->key[1] + ((i % 9) / 3) - 1;
00424     nkey[2] = node->key[2] + ((i % 9) % 3) - 1;
00425 
00426     nnode = pf_kdtree_find_node(self, self->root, nkey);
00427     if (nnode == NULL)
00428       continue;
00429 
00430     assert(nnode->leaf);
00431 
00432     // This node already has a label; skip it.  The label should be
00433     // consistent, however.
00434     if (nnode->cluster >= 0)
00435     {
00436       assert(nnode->cluster == node->cluster);
00437       continue;
00438     }
00439 
00440     // Label this node and recurse
00441     nnode->cluster = node->cluster;
00442 
00443     pf_kdtree_cluster_node(self, nnode, depth + 1);
00444   }
00445   return;
00446 }
00447 
00448 
00449 
00450 #ifdef INCLUDE_RTKGUI
00451 
00452 ////////////////////////////////////////////////////////////////////////////////
00453 // Draw the tree
00454 void pf_kdtree_draw(pf_kdtree_t *self, rtk_fig_t *fig)
00455 {
00456   if (self->root != NULL)
00457     pf_kdtree_draw_node(self, self->root, fig);
00458   return;
00459 }
00460 
00461 
00462 ////////////////////////////////////////////////////////////////////////////////
00463 // Recursively draw nodes
00464 void pf_kdtree_draw_node(pf_kdtree_t *self, pf_kdtree_node_t *node, rtk_fig_t *fig)
00465 {
00466   double ox, oy;
00467   char text[64];
00468 
00469   if (node->leaf)
00470   {
00471     ox = (node->key[0] + 0.5) * self->size[0];
00472     oy = (node->key[1] + 0.5) * self->size[1];
00473 
00474     rtk_fig_rectangle(fig, ox, oy, 0.0, self->size[0], self->size[1], 0);
00475 
00476     //snprintf(text, sizeof(text), "%0.3f", node->value);
00477     //rtk_fig_text(fig, ox, oy, 0.0, text);
00478 
00479     snprintf(text, sizeof(text), "%d", node->cluster);
00480     rtk_fig_text(fig, ox, oy, 0.0, text);
00481   }
00482   else
00483   {
00484     assert(node->children[0] != NULL);
00485     assert(node->children[1] != NULL);
00486     pf_kdtree_draw_node(self, node->children[0], fig);
00487     pf_kdtree_draw_node(self, node->children[1], fig);
00488   }
00489 
00490   return;
00491 }
00492 
00493 #endif