KdBVH.h 8.95 KB
Newer Older
xuebingbing's avatar
xuebingbing committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37
// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009 Ilya Baran <ibaran@mit.edu>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef KDBVH_H_INCLUDED
#define KDBVH_H_INCLUDED

namespace Eigen { 

namespace internal {

//internal pair class for the BVH--used instead of std::pair because of alignment
template<typename Scalar, int Dim>
struct vector_int_pair
{
EIGEN_MAKE_ALIGNED_OPERATOR_NEW_IF_VECTORIZABLE_FIXED_SIZE(Scalar, Dim)
  typedef Matrix<Scalar, Dim, 1> VectorType;

  vector_int_pair(const VectorType &v, int i) : first(v), second(i) {}

  VectorType first;
  int second;
};

//these templates help the tree initializer get the bounding boxes either from a provided
//iterator range or using bounding_box in a unified way
template<typename ObjectList, typename VolumeList, typename BoxIter>
struct get_boxes_helper {
  void operator()(const ObjectList &objects, BoxIter boxBegin, BoxIter boxEnd, VolumeList &outBoxes)
  {
    outBoxes.insert(outBoxes.end(), boxBegin, boxEnd);
    eigen_assert(outBoxes.size() == objects.size());
xuebingbing's avatar
xuebingbing committed
38
    EIGEN_ONLY_USED_FOR_DEBUG(objects);
xuebingbing's avatar
xuebingbing committed
39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223
  }
};

template<typename ObjectList, typename VolumeList>
struct get_boxes_helper<ObjectList, VolumeList, int> {
  void operator()(const ObjectList &objects, int, int, VolumeList &outBoxes)
  {
    outBoxes.reserve(objects.size());
    for(int i = 0; i < (int)objects.size(); ++i)
      outBoxes.push_back(bounding_box(objects[i]));
  }
};

} // end namespace internal


/** \class KdBVH
 *  \brief A simple bounding volume hierarchy based on AlignedBox
 *
 *  \param _Scalar The underlying scalar type of the bounding boxes
 *  \param _Dim The dimension of the space in which the hierarchy lives
 *  \param _Object The object type that lives in the hierarchy.  It must have value semantics.  Either bounding_box(_Object) must
 *                 be defined and return an AlignedBox<_Scalar, _Dim> or bounding boxes must be provided to the tree initializer.
 *
 *  This class provides a simple (as opposed to optimized) implementation of a bounding volume hierarchy analogous to a Kd-tree.
 *  Given a sequence of objects, it computes their bounding boxes, constructs a Kd-tree of their centers
 *  and builds a BVH with the structure of that Kd-tree.  When the elements of the tree are too expensive to be copied around,
 *  it is useful for _Object to be a pointer.
 */
template<typename _Scalar, int _Dim, typename _Object> class KdBVH
{
public:
  enum { Dim = _Dim };
  typedef _Object Object;
  typedef std::vector<Object, aligned_allocator<Object> > ObjectList;
  typedef _Scalar Scalar;
  typedef AlignedBox<Scalar, Dim> Volume;
  typedef std::vector<Volume, aligned_allocator<Volume> > VolumeList;
  typedef int Index;
  typedef const int *VolumeIterator; //the iterators are just pointers into the tree's vectors
  typedef const Object *ObjectIterator;

  KdBVH() {}

  /** Given an iterator range over \a Object references, constructs the BVH.  Requires that bounding_box(Object) return a Volume. */
  template<typename Iter> KdBVH(Iter begin, Iter end) { init(begin, end, 0, 0); } //int is recognized by init as not being an iterator type

  /** Given an iterator range over \a Object references and an iterator range over their bounding boxes, constructs the BVH */
  template<typename OIter, typename BIter> KdBVH(OIter begin, OIter end, BIter boxBegin, BIter boxEnd) { init(begin, end, boxBegin, boxEnd); }

  /** Given an iterator range over \a Object references, constructs the BVH, overwriting whatever is in there currently.
    * Requires that bounding_box(Object) return a Volume. */
  template<typename Iter> void init(Iter begin, Iter end) { init(begin, end, 0, 0); }

  /** Given an iterator range over \a Object references and an iterator range over their bounding boxes,
    * constructs the BVH, overwriting whatever is in there currently. */
  template<typename OIter, typename BIter> void init(OIter begin, OIter end, BIter boxBegin, BIter boxEnd)
  {
    objects.clear();
    boxes.clear();
    children.clear();

    objects.insert(objects.end(), begin, end);
    int n = static_cast<int>(objects.size());

    if(n < 2)
      return; //if we have at most one object, we don't need any internal nodes

    VolumeList objBoxes;
    VIPairList objCenters;

    //compute the bounding boxes depending on BIter type
    internal::get_boxes_helper<ObjectList, VolumeList, BIter>()(objects, boxBegin, boxEnd, objBoxes);

    objCenters.reserve(n);
    boxes.reserve(n - 1);
    children.reserve(2 * n - 2);

    for(int i = 0; i < n; ++i)
      objCenters.push_back(VIPair(objBoxes[i].center(), i));

    build(objCenters, 0, n, objBoxes, 0); //the recursive part of the algorithm

    ObjectList tmp(n);
    tmp.swap(objects);
    for(int i = 0; i < n; ++i)
      objects[i] = tmp[objCenters[i].second];
  }

  /** \returns the index of the root of the hierarchy */
  inline Index getRootIndex() const { return (int)boxes.size() - 1; }

  /** Given an \a index of a node, on exit, \a outVBegin and \a outVEnd range over the indices of the volume children of the node
    * and \a outOBegin and \a outOEnd range over the object children of the node */
  EIGEN_STRONG_INLINE void getChildren(Index index, VolumeIterator &outVBegin, VolumeIterator &outVEnd,
                                       ObjectIterator &outOBegin, ObjectIterator &outOEnd) const
  { //inlining this function should open lots of optimization opportunities to the compiler
    if(index < 0) {
      outVBegin = outVEnd;
      if(!objects.empty())
        outOBegin = &(objects[0]);
      outOEnd = outOBegin + objects.size(); //output all objects--necessary when the tree has only one object
      return;
    }

    int numBoxes = static_cast<int>(boxes.size());

    int idx = index * 2;
    if(children[idx + 1] < numBoxes) { //second index is always bigger
      outVBegin = &(children[idx]);
      outVEnd = outVBegin + 2;
      outOBegin = outOEnd;
    }
    else if(children[idx] >= numBoxes) { //if both children are objects
      outVBegin = outVEnd;
      outOBegin = &(objects[children[idx] - numBoxes]);
      outOEnd = outOBegin + 2;
    } else { //if the first child is a volume and the second is an object
      outVBegin = &(children[idx]);
      outVEnd = outVBegin + 1;
      outOBegin = &(objects[children[idx + 1] - numBoxes]);
      outOEnd = outOBegin + 1;
    }
  }

  /** \returns the bounding box of the node at \a index */
  inline const Volume &getVolume(Index index) const
  {
    return boxes[index];
  }

private:
  typedef internal::vector_int_pair<Scalar, Dim> VIPair;
  typedef std::vector<VIPair, aligned_allocator<VIPair> > VIPairList;
  typedef Matrix<Scalar, Dim, 1> VectorType;
  struct VectorComparator //compares vectors, or, more specificall, VIPairs along a particular dimension
  {
    VectorComparator(int inDim) : dim(inDim) {}
    inline bool operator()(const VIPair &v1, const VIPair &v2) const { return v1.first[dim] < v2.first[dim]; }
    int dim;
  };

  //Build the part of the tree between objects[from] and objects[to] (not including objects[to]).
  //This routine partitions the objCenters in [from, to) along the dimension dim, recursively constructs
  //the two halves, and adds their parent node.  TODO: a cache-friendlier layout
  void build(VIPairList &objCenters, int from, int to, const VolumeList &objBoxes, int dim)
  {
    eigen_assert(to - from > 1);
    if(to - from == 2) {
      boxes.push_back(objBoxes[objCenters[from].second].merged(objBoxes[objCenters[from + 1].second]));
      children.push_back(from + (int)objects.size() - 1); //there are objects.size() - 1 tree nodes
      children.push_back(from + (int)objects.size());
    }
    else if(to - from == 3) {
      int mid = from + 2;
      std::nth_element(objCenters.begin() + from, objCenters.begin() + mid,
                        objCenters.begin() + to, VectorComparator(dim)); //partition
      build(objCenters, from, mid, objBoxes, (dim + 1) % Dim);
      int idx1 = (int)boxes.size() - 1;
      boxes.push_back(boxes[idx1].merged(objBoxes[objCenters[mid].second]));
      children.push_back(idx1);
      children.push_back(mid + (int)objects.size() - 1);
    }
    else {
      int mid = from + (to - from) / 2;
      nth_element(objCenters.begin() + from, objCenters.begin() + mid,
                  objCenters.begin() + to, VectorComparator(dim)); //partition
      build(objCenters, from, mid, objBoxes, (dim + 1) % Dim);
      int idx1 = (int)boxes.size() - 1;
      build(objCenters, mid, to, objBoxes, (dim + 1) % Dim);
      int idx2 = (int)boxes.size() - 1;
      boxes.push_back(boxes[idx1].merged(boxes[idx2]));
      children.push_back(idx1);
      children.push_back(idx2);
    }
  }

  std::vector<int> children; //children of x are children[2x] and children[2x+1], indices bigger than boxes.size() index into objects.
  VolumeList boxes;
  ObjectList objects;
};

} // end namespace Eigen

#endif //KDBVH_H_INCLUDED