picoflann.h

Go to the documentation of this file.

#ifndef PicoFlann_H
#define PicoFlann_H
#include <vector>
#include <cassert>
#include <iostream>
#include <cstdint>
#include <limits>
#include <algorithm>
#include <cstring>
namespace picoflann
{




struct L2
{
  template <typename ElementType, typename ElementType2, typename Adapter>
  double compute_distance(const ElementType &elema, const ElementType2 &elemb,
                          const Adapter &adapter, int ndims, double worstDist) const
  {
    // compute dist
    double sqd = 0;
    for (int i = 0; i < ndims; i++)
    {
      double d = adapter(elema, i) - adapter(elemb, i);
      sqd += d * d;
      if (sqd > worstDist)
        return sqd;
    }
    return sqd;
  }
};

template <int DIMS, typename Adapter, typename DistanceType = L2>
class KdTreeIndex
{
public:
  template <typename Container>
  inline void build(const Container &container)
  {
    _index.clear();
    _index.reserve(container.size() * 2);
    _index.dims = DIMS;
    _index.nValues = container.size();
    // Create root and assign all items
    all_indices.resize(container.size());
    for (size_t i = 0; i < container.size(); i++)
      all_indices[i] = i;
    if (container.size() == 0)
      return;
    computeBoundingBox<Container>(_index.rootBBox, 0, all_indices.size(), container);
    _index.push_back(Node());
    divideTree<Container>(_index, 0, 0, all_indices.size(), _index.rootBBox, container);
  }


  inline void clear()
  {
    _index.clear();
    all_indices.clear();
  }

  // saves to a stream. Note that the container is not saved!
  inline void toStream(std::ostream &str) const;
  // reads from an stream. Note that the container is not readed!
  inline void fromStream(std::istream &str);

  template <typename Type, typename Container>
  inline std::vector<std::pair<uint32_t, double> > searchKnn(const Container &container,
                                                             const Type &val, int nn,
                                                             bool sorted = true)
  {
    std::vector<std::pair<uint32_t, double> > res;
    generalSearch<Type, Container>(res, container, val, -1, sorted, nn);
    return res;
  }


  template <typename Type, typename Container>
  inline std::vector<std::pair<uint32_t, double> > radiusSearch(const Container &container,
                                                                const Type &val, double dist,
                                                                bool sorted = true,
                                                                int maxNN = -1) const
  {
    std::vector<std::pair<uint32_t, double> > res;
    generalSearch<Type, Container>(res, container, val, dist, sorted, maxNN);
    return res;
  }


  template <typename Type, typename Container>
  inline void radiusSearch(std::vector<std::pair<uint32_t, double> > &res,
                           const Container &container, const Type &val, double dist,
                           bool sorted = true, int maxNN = -1)
  {
    generalSearch<Type, Container>(res, container, val, dist, sorted, maxNN);
  }



private:
  struct Node
  {
    inline bool isLeaf() const
    {
      return _ileft == -1 && _iright == -1;
    }
    inline void setNodesInfo(uint32_t l, uint32_t r)
    {
      _ileft = l;
      _iright = r;
    }
    double div_val;
    uint16_t col_index;  // column index of the feature vector
    std::vector<int> idx;
    float divhigh, divlow;
    int64_t _ileft = -1, _iright = -1;  // children
    void toStream(std::ostream &str) const;
    void fromStream(std::istream &str);
  };


  typedef std::vector<std::pair<double, double> > BoundingBox;

  struct Index : public std::vector<Node>
  {
    BoundingBox rootBBox;
    int dims = 0;
    int nValues = 0;  // number of elements of the set when call to build
    inline void toStream(std::ostream &str) const;
    inline void fromStream(std::istream &str);
  };
  Index _index;
  DistanceType _distance;
  Adapter adapter;
  // next are only used during build
  std::vector<uint32_t> all_indices;
  int _maxLeafSize = 10;



  // temporal used during creation of the tree
  template <typename Container>
  void divideTree(Index &index, uint64_t nodeIdx, int startIndex, int endIndex,
                  BoundingBox &bbox, const Container &container)
  {
    // std::cout<<"CREATE="<<startIndex<<"-"<<endIndex<<"|";toStream(std::cout,bbox);
    Node &currNode = index[nodeIdx];
    int count = endIndex - startIndex;
    assert(startIndex < endIndex);

    if (count <= _maxLeafSize)
    {
      currNode.idx.resize(count);
      for (int i = 0; i < count; i++)
        currNode.idx[i] = all_indices[startIndex + i];
      computeBoundingBox<Container>(bbox, startIndex, endIndex, container);
      //  std::cout<<std::endl;
      return;
    }


    currNode.setNodesInfo(index.size(), index.size() + 1);
    index.push_back(Node());
    int leftNode = index.size() - 1;
    index.push_back(Node());
    int rightNode = index.size() - 1;


    if (0)
    {
      BoundingBox _bbox;
      computeBoundingBox<Container>(_bbox, startIndex, endIndex, container);
      //        //get the dimension with highest distnaces
      double max_spread = -1;
      currNode.col_index = 0;
      for (int i = 0; i < DIMS; i++)
      {
        double spread = _bbox[i].second - _bbox[i].first;  // maxV[i]-minV[i];
        if (spread > max_spread)
        {
          max_spread = spread;
          currNode.col_index = i;
        }
      }
      // select the split val
      double split_val = (bbox[currNode.col_index].first + bbox[currNode.col_index].second) / 2;
      if (split_val < _bbox[currNode.col_index].first)
        currNode.div_val = _bbox[currNode.col_index].first;
      else if (split_val > _bbox[currNode.col_index].second)
        currNode.div_val = _bbox[currNode.col_index].second;
      else
        currNode.div_val = split_val;
    }
    else
    {
      double var[DIMS], mean[DIMS];
      // compute the variance of the features to  select the highest one
      mean_var_calculate<Container>(startIndex, endIndex, var, mean, container);
      currNode.col_index = 0;
      // select element with highest variance
      for (int i = 1; i < DIMS; i++)
        if (var[i] > var[currNode.col_index])
          currNode.col_index = i;

      // now sort all indices according to the selected value

      currNode.div_val = mean[currNode.col_index];
    }



    // compute the variance of the features to  select the highest one
    // now sort all indices according to the selected value

    // std::cout<<" CUT FEAT="<<currNode.col_index<< " VAL="<<currNode.div_val<<std::endl;
    int lim1, lim2;
    planeSplit<Container>(&all_indices[startIndex], count, currNode.col_index,
                          currNode.div_val, lim1, lim2, container);

    int split_index;

    if (lim1 > count / 2)
      split_index = lim1;
    else if (lim2 < count / 2)
      split_index = lim2;
    else
      split_index = count / 2;

    //        /* If either list is empty, it means that all remaining features
    //              * are identical. Split in the middle to maintain a balanced tree.
    //              */
    if ((lim1 == count) || (lim2 == 0))
      split_index = count / 2;
    // create partitions with at least minLeafSize elements
    if (_maxLeafSize != 1)
      if (split_index < _maxLeafSize || count - split_index < _maxLeafSize)
      {
        std::sort(all_indices.begin() + startIndex, all_indices.begin() + endIndex,
                  [&](const uint32_t &a, const uint32_t &b)
                  {
                    return adapter(container.at(a), currNode.col_index) <
                           adapter(container.at(b), currNode.col_index);
                  });
        split_index = count / 2;
        currNode.div_val =
            adapter(container.at(all_indices[startIndex + split_index]), currNode.col_index);
      }



    //  currNode.div_val=_features.ptr<float>(all_indices[split_index])[currNode.col_index];

    BoundingBox left_bbox(bbox);
    left_bbox[currNode.col_index].second = currNode.div_val;
    divideTree<Container>(index, leftNode, startIndex, startIndex + split_index,
                          left_bbox, container);
    left_bbox[currNode.col_index].second = currNode.div_val;
    assert(left_bbox[currNode.col_index].second <= currNode.div_val);
    BoundingBox right_bbox(bbox);
    right_bbox[currNode.col_index].first = currNode.div_val;
    divideTree<Container>(index, rightNode, startIndex + split_index, endIndex,
                          right_bbox, container);

    currNode.divlow = left_bbox[currNode.col_index].second;
    currNode.divhigh = right_bbox[currNode.col_index].first;
    assert(currNode.divlow <= currNode.divhigh);

    for (int i = 0; i < DIMS; ++i)
    {
      bbox[i].first = std::min(left_bbox[i].first, right_bbox[i].first);
      bbox[i].second = std::max(left_bbox[i].second, right_bbox[i].second);
    }
  }



  template <typename Container>
  void computeBoundingBox(BoundingBox &bbox, int start, int end, const Container &container)
  {
    bbox.resize(DIMS);
    for (int i = 0; i < DIMS; ++i)
      bbox[i].second = bbox[i].first = adapter(container.at(all_indices[start]), i);

    for (int k = start + 1; k < end; ++k)
    {
      for (int i = 0; i < DIMS; ++i)
      {
        float v = adapter(container.at(all_indices[k]), i);
        if (v < bbox[i].first)
          bbox[i].first = v;
        if (v > bbox[i].second)
          bbox[i].second = v;
      }
    }
  }

  template <typename Container>
  void mean_var_calculate(int startindex, int endIndex, double var[], double mean[],
                          const Container &container)
  {
    const int MAX_ELEM_MEAN = 100;
    // recompute centers
    // compute new center
    memset(mean, 0, sizeof(double) * DIMS);
    double sum2[DIMS];
    memset(sum2, 0, sizeof(double) * DIMS);
    // finish when at least MAX_ELEM_MEAN elements computed
    int cnt = 0;
    // std::min(MAX_ELEM_MEAN,endIndex-startindex );
    int increment = 1;
    if (endIndex - startindex >= 2 * MAX_ELEM_MEAN)
      increment = (endIndex - startindex) / MAX_ELEM_MEAN;
    for (int i = startindex; i < endIndex; i += increment)
    {
      for (int c = 0; c < DIMS; c++)
      {
        auto val = adapter(container.at(all_indices[i]), c);
        mean[c] += val;
        sum2[c] += val * val;
      }
      cnt++;
    }

    double invcnt = 1. / double(cnt);
    for (int c = 0; c < DIMS; c++)
    {
      mean[c] *= invcnt;
      var[c] = sum2[c] * invcnt - mean[c] * mean[c];
    }
  }


  template <typename Container>
  void planeSplit(uint32_t *ind, int count, int cutfeat, float cutval, int &lim1,
                  int &lim2, const Container &container)
  {
    /* Move vector indices for left subtree to front of list. */
    int left = 0;
    int right = count - 1;
    for (;;)
    {
      while (left <= right && adapter(container.at(ind[left]), cutfeat) < cutval)
        ++left;
      while (left <= right && adapter(container.at(ind[right]), cutfeat) >= cutval)
        --right;
      if (left > right)
        break;
      std::swap(ind[left], ind[right]);
      ++left;
      --right;
    }
    lim1 = left;
    right = count - 1;
    for (;;)
    {
      while (left <= right && adapter(container.at(ind[left]), cutfeat) <= cutval)
        ++left;
      while (left <= right && adapter(container.at(ind[right]), cutfeat) > cutval)
        --right;
      if (left > right)
        break;
      std::swap(ind[left], ind[right]);
      ++left;
      --right;
    }
    lim2 = left;
  }


  template <typename Type>
  inline double computeInitialDistances(const Type &elem, double dists[],
                                        const BoundingBox &bbox) const
  {
    float distsq = 0.0;

    for (int i = 0; i < DIMS; ++i)
    {
      double elem_i = adapter(elem, i);
      if (elem_i < bbox[i].first)
      {
        auto d = elem_i - bbox[i].first;
        dists[i] = d * d;  // distance_.accum_dist(vec[i], root_bbox_[i].first, i);
        distsq += dists[i];
      }
      if (elem_i > bbox[i].second)
      {
        auto d = elem_i - bbox[i].second;
        dists[i] = d * d;  // distance_.accum_dist(vec[i], root_bbox_[i].second, i);
        distsq += dists[i];
      }
    }
    return distsq;
  }
  // THe function that does the search in all exact methods
  template <typename Type, typename Container>
  inline void generalSearch(std::vector<std::pair<uint32_t, double> > &res,
                            const Container &container, const Type &val, double dist,
                            bool sorted = true,
                            uint32_t maxNn = std::numeric_limits<int>::max()) const
  {
    double dists[DIMS];
    memset(dists, 0, sizeof(double) * DIMS);
    res.clear();
    ResultSet hres(res, maxNn, dist > 0 ? dist * dist : -1.f);
    float distsq = computeInitialDistances<Type>(val, dists, _index.rootBBox);
    searchExactLevel<Type, Container>(_index, 0, val, hres, distsq, dists, 1, container);
    if (sorted && res.size() > 1)
      std::sort(res.begin(), res.end(),
                [](const std::pair<uint32_t, double> &a, const std::pair<uint32_t, double> &b)
                { return a.second < b.second; });
  }

  // heap having at the top the maximum element

  class ResultSet
  {
  public:
    std::vector<std::pair<uint32_t, double> > &array;
    int maxSize;
    double maxValue = std::numeric_limits<double>::max();
    bool radius_search = false;

  public:
    ResultSet(std::vector<std::pair<uint32_t, double> > &data_ref,
              uint32_t MaxSize = std::numeric_limits<uint32_t>::max(), double MaxV = -1)
      : array(data_ref)
    {
      maxSize = MaxSize;
      // set value for radius search
      if (MaxV > 0)
      {
        maxValue = MaxV;
        radius_search = true;
      }
    }


    inline void push(const std::pair<uint32_t, double> &val)
    {
      if (radius_search && val.second < maxValue)
      {
        array.push_back(val);
      }
      else
      {
        if (array.size() >= size_t(maxSize))
        {
          // check if the maxium must be replaced by this
          if (val.second < array[0].second)
          {
            swap(array.front(), array.back());
            array.pop_back();
            if (array.size() > 1)
              up(0);
          }
          else
            return;
        }
        array.push_back(val);
        if (array.size() > 1)
          down(array.size() - 1);
      }
      //            array_size++;
    }

    inline double worstDist() const
    {
      if (radius_search)
        return maxValue;  // radius search
      else if (array.size() < size_t(maxSize))
        return std::numeric_limits<double>::max();
      return array[0].second;
    }
    inline double top() const
    {
      assert(!array.empty());
      return array[0].second;
    }

  private:
    inline void down(size_t index)
    {
      if (index == 0)
        return;
      size_t parentIndex = (index - 1) / 2;
      if (array[parentIndex].second < array[index].second)
      {
        swap(array[index], array[parentIndex]);
        down(parentIndex);
      }
    }
    inline void up(size_t index)
    {
      size_t leftIndex = 2 * index + 1;   // vl_heap_left_child (index) ;
      size_t rightIndex = 2 * index + 2;  // vl_heap_right_child (index) ;

      /* no childer: stop */
      if (leftIndex >= array.size())
        return;

      /* only left childer: easy */
      if (rightIndex >= array.size())
      {
        if (array[index].second < array[leftIndex].second)
          swap(array[index], array[leftIndex]);
        return;
      }

      /* both childern */
      {
        if (array[rightIndex].second < array[leftIndex].second)
        {
          /* swap with left */
          if (array[index].second < array[leftIndex].second)
          {
            swap(array[index], array[leftIndex]);
            up(leftIndex);
          }
        }
        else
        {
          /* swap with right */
          if (array[index].second < array[rightIndex].second)
          {
            swap(array[index], array[rightIndex]);
            up(rightIndex);
          }
        }
      }
    }
  };

  template <typename Type, typename Container>
  inline void searchExactLevel(const Index &index, int64_t nodeIdx, const Type &elem,
                               ResultSet &res, double mindistsq, double dists[],
                               double epsError, const Container &container) const
  {
    const Node &currNode = index[nodeIdx];
    if (currNode.isLeaf())
    {
      double worstDist = res.worstDist();
      for (size_t i = 0; i < currNode.idx.size(); i++)
      {
        double sqd = _distance.compute_distance(elem, container.at(currNode.idx[i]),
                                                adapter, DIMS, worstDist);
        if (sqd < worstDist)
        {
          res.push({ currNode.idx[i], sqd });
          worstDist = res.worstDist();
        }
      }
    }
    else
    {
      double val = adapter(elem, currNode.col_index);
      double diff1 = val - currNode.divlow;
      double diff2 = val - currNode.divhigh;

      uint32_t bestChild;
      uint32_t otherChild;
      double cut_dist;
      if ((diff1 + diff2) < 0)
      {
        bestChild = currNode._ileft;
        otherChild = currNode._iright;
        cut_dist = diff2 * diff2;
      }
      else
      {
        bestChild = currNode._iright;
        otherChild = currNode._ileft;
        cut_dist = diff1 * diff1;
      }
      /* Call recursively to search next level down. */
      searchExactLevel<Type, Container>(index, bestChild, elem, res, mindistsq, dists,
                                        epsError, container);

      float dst = dists[currNode.col_index];
      mindistsq = mindistsq + cut_dist - dst;
      dists[currNode.col_index] = cut_dist;
      if (mindistsq * epsError <= res.worstDist())
        searchExactLevel<Type, Container>(index, otherChild, elem, res, mindistsq, dists,
                                          epsError, container);
      dists[currNode.col_index] = dst;
    }
  }
};
template <int DIMS, typename AAdapter, typename DistanceType>
void KdTreeIndex<DIMS, AAdapter, DistanceType>::Node::toStream(std::ostream &str) const
{
  str.write((char *)&div_val, sizeof(div_val));
  str.write((char *)&col_index, sizeof(col_index));
  str.write((char *)&divhigh, sizeof(divhigh));
  str.write((char *)&divlow, sizeof(divlow));
  str.write((char *)&_ileft, sizeof(_ileft));
  str.write((char *)&_iright, sizeof(_iright));
  uint64_t s = idx.size();
  str.write((char *)&s, sizeof(s));
  str.write((char *)&idx[0], sizeof(idx[0]) * idx.size());
}

template <int DIMS, typename AAdapter, typename DistanceType>
void KdTreeIndex<DIMS, AAdapter, DistanceType>::Node::fromStream(std::istream &str)
{
  str.read((char *)&div_val, sizeof(div_val));
  str.read((char *)&col_index, sizeof(col_index));
  str.read((char *)&divhigh, sizeof(divhigh));
  str.read((char *)&divlow, sizeof(divlow));
  str.read((char *)&_ileft, sizeof(_ileft));
  str.read((char *)&_iright, sizeof(_iright));
  uint64_t s;
  str.read((char *)&s, sizeof(s));
  idx.resize(s);
  str.read((char *)&idx[0], sizeof(idx[0]) * idx.size());
}

template <int DIMS, typename AAdapter, typename DistanceType>
void KdTreeIndex<DIMS, AAdapter, DistanceType>::Index::toStream(std::ostream &str) const
{
  str.write((char *)&dims, sizeof(dims));
  str.write((char *)&rootBBox[0], sizeof(rootBBox[0]) * dims);
  str.write((char *)&nValues, sizeof(nValues));

  uint64_t s = std::vector<Node>::size();
  str.write((char *)&s, sizeof(s));
  for (size_t i = 0; i < std::vector<Node>::size(); i++)
    std::vector<Node>::at(i).toStream(str);
}

template <int DIMS, typename AAdapter, typename DistanceType>
void KdTreeIndex<DIMS, AAdapter, DistanceType>::Index::fromStream(std::istream &str)
{
  str.read((char *)&dims, sizeof(dims));
  rootBBox.resize(dims);
  str.read((char *)&rootBBox[0], sizeof(rootBBox[0]) * dims);
  str.read((char *)&nValues, sizeof(nValues));


  uint64_t s;
  ;
  str.read((char *)&s, sizeof(s));
  std::vector<Node>::resize(s);
  for (size_t i = 0; i < std::vector<Node>::size(); i++)
    std::vector<Node>::at(i).fromStream(str);
  if (dims != DIMS && this->size() != 0 && nValues != 0)
    throw std::runtime_error(
        "Number of dimensions of the index in the stream is different from the number of dimensions of this");
}

template <int DIMS, typename AAdapter, typename DistanceType>
void KdTreeIndex<DIMS, AAdapter, DistanceType>::toStream(std::ostream &str) const
{
  _index.toStream(str);
}

template <int DIMS, typename AAdapter, typename DistanceType>
void KdTreeIndex<DIMS, AAdapter, DistanceType>::fromStream(std::istream &str)
{
  _index.fromStream(str);
}
}  // namespace picoflann

#endif