fastlib2/fastlib/tree/kdtree.h

@@ -44,13 +44,16 @@ namespace tree {
    * @param new_from_old pointer to an unitialized arraylist; it will map
    *        original indexes to new indices
    */
-  template<typename TKdTree, typename T>
-  TKdTree *MakeKdTreeMidpoint(GenMatrix<T>& matrix, index_t leaf_size,
+  
+
+  template<typename TKdTree>
+  TKdTree *MakeKdTreeMidpointSelective(Matrix& matrix, Vector split_dimensions,
+      index_t leaf_size,
       ArrayList<index_t> *old_from_new = NULL,
       ArrayList<index_t> *new_from_old = NULL) {
     TKdTree *node = new TKdTree();
     index_t *old_from_new_ptr;
-
+   
     if (old_from_new) {
       old_from_new->Init(matrix.n_cols());
       
@@ -64,23 +67,39 @@ namespace tree {
     }
       
     node->Init(0, matrix.n_cols());
-    node->bound().Init(matrix.n_rows());
-    tree_kdtree_private::FindBoundFromMatrix(matrix,
+    node->bound().Init(split_dimensions.length());
+    tree_kdtree_private::SelectFindBoundFromMatrix(matrix, split_dimensions,
         0, matrix.n_cols(), &node->bound());
 
-    tree_kdtree_private::SplitKdTreeMidpoint(matrix, node, leaf_size,
-        old_from_new_ptr);
+    tree_kdtree_private::SelectSplitKdTreeMidpoint(matrix, split_dimensions, 
+	node, leaf_size, old_from_new_ptr);
     
     if (new_from_old) {
       new_from_old->Init(matrix.n_cols());
       for (index_t i = 0; i < matrix.n_cols(); i++) {
         (*new_from_old)[(*old_from_new)[i]] = i;
       }
-    }
-    
+    }    
     return node;
   }
 
+  template<typename TKdTree>
+    TKdTree *MakeKdTreeMidpoint(Matrix& matrix, index_t leaf_size,
+				ArrayList<index_t> *old_from_new = NULL,
+				ArrayList<index_t> *new_from_old = NULL) {  
+    Vector split_dimensions;
+    split_dimensions.Init(matrix.n_rows());
+    int i;
+    for (i = 0; i < matrix.n_rows(); i++){
+      split_dimensions[i] = i;
+    }
+    TKdTree *result;
+    result = MakeKdTreeMidpointSelective<TKdTree>(matrix, split_dimensions,
+		   leaf_size, old_from_new, new_from_old);
+    return result;
+  }
+  
+
   /**
    * Loads a KD tree from a command-line parameter,
    * creating a KD tree if necessary.
@@ -119,10 +138,10 @@ namespace tree {
    *        the matrix returned to the original data point indices
    * @return SUCCESS_PASS or SUCCESS_FAIL
    */
-  template<typename TKdTree, typename T>
-  success_t LoadKdTree(datanode *module, GenMatrix<T> *matrix, 
-		       TKdTree **tree_pp, ArrayList<index_t> *old_from_new) {
-
+  template<typename TKdTree>
+  success_t LoadKdTree(datanode *module,
+      Matrix *matrix, TKdTree **tree_pp,
+      ArrayList<index_t> *old_from_new) {
     const char *type = fx_param_str(module, "type", "text");
     const char *fname = fx_param_str(module, "", NULL);
     success_t success = SUCCESS_PASS;
@@ -135,8 +154,11 @@ namespace tree {
       success = data::Load(fname, matrix);
       fx_timer_stop(module, "load_matrix");
 
+      //if (fx_param_exists("do_pca")) {}
+
       fx_timer_start(module, "make_tree");
-      *tree_pp = MakeKdTreeMidpoint<TKdTree>(*matrix, leaflen, old_from_new);
+      *tree_pp = MakeKdTreeMidpoint<TKdTree>(
+          *matrix, leaflen, old_from_new);
       fx_timer_stop(module, "make_tree");
     }
     fx_timer_stop(module, "load");

fastlib2/fastlib/tree/kdtree_impl.h

@@ -1,25 +1,68 @@
 /* Implementation for the regular pointer-style kd-tree builder. */
 
 namespace tree_kdtree_private {
+  Vector MakeBoundVector(Vector point, Vector bound_dimensions){
+    int i;
+    Vector bound_vector;
+    bound_vector.Init(bound_dimensions.length());
+    for (i = 0; i < bound_dimensions.length(); i++){
+      bound_vector[i] = point[(int)bound_dimensions[i]];
+    }
+    return bound_vector;
+  }
 
-  template<typename TBound, typename T>
-  void FindBoundFromMatrix(const GenMatrix<T>& matrix,
-      index_t first, index_t count, TBound *bounds) {
+
+  template<typename TBound>
+    void SelectFindBoundFromMatrix(const Matrix& matrix,
+      Vector split_dimensions, index_t first, index_t count, TBound *bounds){
     index_t end = first + count;
     for (index_t i = first; i < end; i++) {
-      GenVector<T> col;
-
+      Vector col;
       matrix.MakeColumnVector(i, &col);
-      *bounds |= col;
+      if (split_dimensions.length() == matrix.n_rows()){
+	*bounds |= col;
+      } else {
+	Vector foo = MakeBoundVector(col, split_dimensions);	
+	*bounds |= foo;
+      }          
     }
   }
 
-  template<typename TBound, typename T>
-  index_t MatrixPartition(GenMatrix<T>& matrix, index_t dim, double splitvalue,
-			  index_t first, index_t count,
-			  TBound* left_bound, TBound* right_bound,
-			  index_t *old_from_new) {
+  template<typename TBound>
+  void FindBoundFromMatrix(const Matrix& matrix,
+      index_t first, index_t count, TBound *bounds){
+    Vector split_dimensions;
+    split_dimensions.Init(matrix.n_rows());
+    int i;
+    for (i = 0; i < matrix.n_rows(); i++){
+      split_dimensions[i] = i;
+    }
+    SelectFindBoundFromMatrix(matrix, split_dimensions, first, count, bounds); 
+  }
 
+   template<typename TBound>
+   index_t MatrixPartition(
+      Matrix& matrix, index_t dim, double splitvalue,
+      index_t first, index_t count,
+      TBound* left_bound, TBound* right_bound,
+      index_t *old_from_new) {
+     Vector split_dimensions;
+     split_dimensions.Init(matrix.n_rows());
+     int i;
+     for (i = 0; i < matrix.n_rows(); i++){
+       split_dimensions[i] = i;
+     }
+     index_t split_point = SelectMatrixPartition(matrix, split_dimensions, 
+      dim, splitvalue, first, count, left_bound, right_bound, old_from_new);
+     return split_point;
+   }
+
+  template<typename TBound>
+  index_t SelectMatrixPartition(
+      Matrix& matrix, Vector split_dimensions, index_t dim, double splitvalue,
+      index_t first, index_t count,
+      TBound* left_bound, TBound* right_bound,
+      index_t *old_from_new) {
     index_t left = first;
     index_t right = first + count - 1;
     
@@ -30,16 +73,24 @@ namespace tree_kdtree_private {
      */
     for (;;) {
       while (matrix.get(dim, left) < splitvalue && likely(left <= right)) {
-        GenVector<T> left_vector;
+        Vector left_vector;
         matrix.MakeColumnVector(left, &left_vector);
-        *left_bound |= left_vector;
+	if (split_dimensions.length() == matrix.n_rows()){
+	  *left_bound |= left_vector;
+	} else {
+	  *left_bound |= MakeBoundVector(left_vector, split_dimensions);
+	}
         left++;
       }
 
       while (matrix.get(dim, right) >= splitvalue && likely(left <= right)) {
-        GenVector<T> right_vector;
+        Vector right_vector;
         matrix.MakeColumnVector(right, &right_vector);
-        *right_bound |= right_vector;
+	if (split_dimensions.length() == matrix.n_rows()){
+	  *right_bound |= right_vector;
+	} else {
+	  *right_bound |= MakeBoundVector(right_vector, split_dimensions);
+	}        
         right--;
       }
 
@@ -48,16 +99,26 @@ namespace tree_kdtree_private {
         break;
       }
 
-      GenVector<T> left_vector;
-      GenVector<T> right_vector;
+      Vector left_vector;
+      Vector right_vector;
 
       matrix.MakeColumnVector(left, &left_vector);
       matrix.MakeColumnVector(right, &right_vector);
 
       left_vector.SwapValues(&right_vector);
 
-      *left_bound |= left_vector;
-      *right_bound |= right_vector;
+      if (split_dimensions.length() == matrix.n_rows()){
+	  *left_bound |= left_vector;
+      } else {
+	  *left_bound |= MakeBoundVector(left_vector, split_dimensions);
+      }  
+
+      if (split_dimensions.length() == matrix.n_rows()){
+	  *right_bound |= right_vector;
+      } else {
+	  *right_bound |= MakeBoundVector(right_vector, split_dimensions);
+      }  
+     
       
       if (old_from_new) {
         index_t t = old_from_new[left];
@@ -79,24 +140,37 @@ namespace tree_kdtree_private {
     return left;
   }
 
-  template<typename TKdTree, typename T>
-  void SplitKdTreeMidpoint(GenMatrix<T>& matrix, TKdTree *node, 
-			   index_t leaf_size, index_t *old_from_new) {
+  template<typename TKdTree>
+  void SplitKdTreeMidpoint(Matrix& matrix,
+      TKdTree *node, index_t leaf_size, index_t *old_from_new){
+    Vector split_dimensions;
+    split_dimensions.Init(matrix.n_rows());
+    int i;
+    for (i = 0; i < matrix.n_rows(); i++){
+      split_dimensions[i] = i;
+    }
+    SelectSplitKdTreeMidpoint(matrix, split_dimensions, node, 
+			      leaf_size, old_from_new);
+  }
+  
 
+  template<typename TKdTree>
+    void SelectSplitKdTreeMidpoint(Matrix& matrix, Vector& split_dimensions,
+      TKdTree *node, index_t leaf_size, index_t *old_from_new) {
     TKdTree *left = NULL;
     TKdTree *right = NULL;
     
-    //FindBoundFromMatrix(matrix, node->begin(), node->count(),
-    //    &node->bound());
+    SelectFindBoundFromMatrix(matrix, split_dimensions, node->begin(), 
+			      node->count(), &node->bound());
 
     if (node->count() > leaf_size) {
       index_t split_dim = BIG_BAD_NUMBER;
       double max_width = -1;
 
-      for (index_t d = 0; d < matrix.n_rows(); d++) {
+      for (index_t d = 0; d < split_dimensions.length(); d++) {
         double w = node->bound().get(d).width();
-
-        if (unlikely(w > max_width)) {
+	
+        if (w > max_width) {	
           max_width = w;
           split_dim = d;
         }
@@ -109,19 +183,21 @@ namespace tree_kdtree_private {
         // same.  We have to give up.
       } else {
         left = new TKdTree();
-        left->bound().Init(matrix.n_rows());
+        left->bound().Init(split_dimensions.length());
 
         right = new TKdTree();
-        right->bound().Init(matrix.n_rows());
+        right->bound().Init(split_dimensions.length());
 
-        index_t split_col = MatrixPartition(matrix, split_dim, split_val,
+        index_t split_col = SelectMatrixPartition(matrix, split_dimensions, 
+	     (int)split_dimensions[split_dim], split_val,
             node->begin(), node->count(),
             &left->bound(), &right->bound(),
             old_from_new);
         
         VERBOSE_MSG(3.0,"split (%d,[%d],%d) dim %d on %f (between %f, %f)",
             node->begin(), split_col,
-            node->begin() + node->count(), split_dim, split_val,
+            node->begin() + node->count(), (int)split_dimensions[split_dim], 
+		    split_val,
             node->bound().get(split_dim).lo,
             node->bound().get(split_dim).hi);
 
@@ -131,8 +207,10 @@ namespace tree_kdtree_private {
         // This should never happen if max_width > 0
         DEBUG_ASSERT(left->count() != 0 && right->count() != 0);
 
-        SplitKdTreeMidpoint(matrix, left, leaf_size, old_from_new);
-        SplitKdTreeMidpoint(matrix, right, leaf_size, old_from_new);
+        SelectSplitKdTreeMidpoint(matrix, split_dimensions, left, leaf_size, 
+				  old_from_new);
+        SelectSplitKdTreeMidpoint(matrix, split_dimensions, right, leaf_size, 
+				  old_from_new);
       }
     }