Logo Search packages:      
Sourcecode: python-biopython version File versions

KDTree.cpp

#include "KDTree.h"

float KDTREE_dist(float *coord1, float *coord2, int dim)
{
      // returns the SQUARE of the distance between two points
      int i;
      float sum=0, dif=0;

      for (i=0; i<dim; i++)
      {
            dif=coord1[i]-coord2[i];
            sum+=dif*dif;
      }
      return sum;
}


// DataPoint

int DataPoint::current_dim=0;
int DataPoint::dim=3;

int operator<(const DataPoint &self, const DataPoint &other)
{
      float a, b;

      a=self._coord[DataPoint::current_dim];
      b=other._coord[DataPoint::current_dim];

      return a<b; 
}

int operator==(const DataPoint &self, const DataPoint &other)
{
      float a, b;

      a=self._coord[DataPoint::current_dim];
      b=other._coord[DataPoint::current_dim];
      
      return a==b;
}

void DataPoint::set_data(long int index, float *coord)
{
      _index=index;
      _coord=coord;
}

float *DataPoint::get_coord(void)
{
      return _coord;
}

long int DataPoint::get_index(void)
{
      return _index;
}

// Node

Node::Node(float cut_value, int cut_dim, long int start, long int end)
{
      _left=NULL;
      _right=NULL;
      _cut_value=cut_value;
      _cut_dim=cut_dim;
      // start and end index in _data_point_list
      _start=start;
      _end=end;
}

Node::~Node()
{
      delete _left;
      delete _right;
}

void Node::set_left_node(Node *node)
{
      _left=node;
}

void Node::set_right_node(Node *node)
{
      _right=node;
}

Node *Node::get_left_node(void)
{
      return _left;
}

Node *Node::get_right_node(void)
{
      return _right;
}

long int Node::get_start(void)
{
      return _start;
}

long int Node::get_end(void)
{
      return _end;
}

float Node::get_cut_value(void)
{
      return _cut_value;
}

int Node::get_cut_dim(void)
{
      return _cut_dim;
}

int Node::is_leaf(void)
{
      if (_left==NULL && _right==NULL)
      {
            return 1;
      }
      else
      {
            return 0;
      }
}

int Node::is_bucket(void)
{
      if (_start==_end+1)
      {
            // Node contains a single point
            return 0;
      }
      else
      {
            // Node contains several points
            return 1;
      }
}
// Region

int Region::dim=3;

Region::Region(float *left, float *right)
{
      _left=new float[Region::dim];
      _right=new float[Region::dim];

      if (left==NULL || right==NULL)
      {
            // [-INF, INF]
            
            int i;

            for (i=0; i<Region::dim; i++)
            {
                  _left[i]=-INF;
                  _right[i]=INF;
            }
      }
      else
      {
            int i;
            for (i=0; i<Region::dim; i++)
            {
                  _left[i]=left[i];
                  _right[i]=right[i];
            }
      }
}

Region::~Region()
{
      delete [] _left;
      delete [] _right;
}

Region *Region::intersect_right(float split_coord, int current_dim)
{
      float l, r;

      r=_right[current_dim];
      l=_left[current_dim];

      if (split_coord<=l)
      {
            // split point lies to the left
            return new Region(_left, _right);
      }
      else
      {
            if (split_coord<=r)
            {
                  // split point in interval
                  // adjust left
                  int i;
                  float new_left[Region::dim];

                  for (i=0; i<Region::dim; i++)
                  {
                        new_left[i]=_left[i];
                  }
                  new_left[current_dim]=split_coord;
                  return new Region(new_left, _right);
            }
            else
            {
                  // interval lies to the left of split point

                  return NULL;
            }
      }
}

Region *Region::intersect_left(float split_coord, int current_dim)
{
      float l, r;

      r=_right[current_dim];
      l=_left[current_dim];

      if (split_coord<l)
      {
            // nothing to the left
            return NULL;
      }
      else
      {
            if (split_coord<r)
            {
                  // split point in interval
                  // adjust right
                  int i;
                  float new_right[Region::dim];

                  for (i=0; i<Region::dim; i++)
                  {
                        new_right[i]=_right[i];
                  }
                  new_right[current_dim]=split_coord;
                  return new Region(_left, new_right);
            }
            else
            {
                  return new Region(_left, _right);
            }
      }
}

int Region::encloses(float *coord)
{
      int i;
      for (i=0; i<Region::dim; i++)
      {
            if (!(coord[i]>=_left[i] && coord[i]<=_right[i]))
            {
                  return 0;
            }
      }
      return 1;
}

float *Region::get_left(void)
{
      return _left;
}

float *Region::get_right(void)
{
      return _right;
}

int Region::test_intersection(Region *query_region, float radius)
{
      int status=2;

      int i;
      for (i=0; i<Region::dim; i++)
      {
            float rs, rq, ls, lq;

            rs=_right[i];
            ls=_left[i];
            rq=query_region->get_right()[i];
            lq=query_region->get_left()[i];

            if (ls-rq>radius)
            {
                  // outside
                  return 0;
            }
            else if (lq-rs>radius)
            {
                  // outside
                  return 0;
            }
            else if (rs<=rq && ls>=lq)
            {
                  // inside (at least in dim i)
                  status=min(status, 2);
            }
            else
            {
                  // overlap (at least in dim i)    
                  status=1;
            }
      }
      return status;
}

// KDTree

KDTree::KDTree(int dim, int bucket_size)
{
      // set dimension
      this->dim=dim;

      DataPoint::dim=dim;
      Region::dim=dim;

      _center_coord=new float[dim];
      _query_region=NULL;
      _root=NULL;
      _coords=NULL;
      _count=0;
      _neighbor_count=0;
      _bucket_size=bucket_size;
}

KDTree::~KDTree()
{
      // clean up KD tree
      delete _root;
      delete _query_region;
      delete [] _center_coord;
      delete [] _coords;
}

Node *KDTree::_build_tree(long int offset_begin, long int offset_end, int depth)
{
      int localdim;

      if (depth==0)
      {
            // start with [begin, end+1[
            offset_begin=0;
            offset_end=_data_point_list.size();
            localdim=0;
      }
      else
      {
            localdim=depth%dim;
      }

      if ((offset_end-offset_begin)<=_bucket_size) 
      {
            // leaf node

            return new Node(-1, localdim, offset_begin, offset_end);
      }
      else
      {
            long int offset_split;
            long int left_offset_begin, left_offset_end;
            long int right_offset_begin, right_offset_end;
            long int d;
            float cut_value;
            DataPoint data_point;
            Node *left_node, *right_node, *new_node;

            // set sort dimension
            DataPoint::current_dim=localdim;

            // sort method sorts [first, last[
            sort(_data_point_list.begin()+offset_begin, _data_point_list.begin()+offset_end);

            // calculate index of split point
            d=offset_end-offset_begin;
            offset_split=d/2+d%2;

            data_point=_data_point_list[offset_begin+offset_split-1];
            cut_value=(data_point.get_coord())[localdim];
            
            // create new node and bind to left & right nodes
            new_node=new Node(cut_value, localdim, offset_begin, offset_end);

            // left
            left_offset_begin=offset_begin;
            left_offset_end=offset_begin+offset_split;
            left_node=_build_tree(left_offset_begin, left_offset_end, depth+1);

            // right
            right_offset_begin=left_offset_end;
            right_offset_end=offset_end;
            right_node=_build_tree(right_offset_begin, right_offset_end, depth+1);

            new_node->set_left_node(left_node);
            new_node->set_right_node(right_node);

            return new_node;
      }
}

void KDTree::_add_point(long int index, float *coord)
{
      DataPoint data_point;

      data_point.set_data(index, coord);
      // add to list of points
      _data_point_list.push_back(data_point);
}

void KDTree::_set_query_region(float *left, float *right)
{
      delete _query_region;
      _query_region=new Region(left, right);
}

void KDTree::_search(Region *region, Node *node, int depth)
{
      int current_dim;

      if(depth==0)
      {
            // start with [-INF, INF] region
            
            region=new Region();

            // start with root node
            node=_root;
      }

      current_dim=depth%dim;

      if(node->is_leaf())
      {
            long int i;

            for (i=node->get_start(); i<node->get_end(); i++)
            {
                  DataPoint data_point;

                  data_point=_data_point_list[i];

                  if (_query_region->encloses(data_point.get_coord()))
                  {
                        // point is enclosed in query region - report & stop
                        _report_point(data_point.get_index(), data_point.get_coord());
                  }
            }
      }
      else
      {
            Node *left_node, *right_node;
            Region *left_region, *right_region;

            left_node=node->get_left_node();

            // LEFT HALF PLANE

            // new region
            left_region=region->intersect_left(node->get_cut_value(), current_dim);

        // left_region is NULL if no overlap
        if(left_region)
        {
                _test_region(left_node, left_region, depth);
        }

            // RIGHT HALF PLANE

            right_node=node->get_right_node();

            // new region
            right_region=region->intersect_right(node->get_cut_value(), current_dim);

        // right_region is NULL if no overlap
        if(right_region)
        {
            // test for overlap/inside/outside & do recursion/report/stop
            _test_region(right_node, right_region, depth);
        }
      }

      delete region;
}

void KDTree::_test_region(Node *node, Region *region, int depth)
{
      int intersect_flag;

      // is node region inside, outside or overlapping 
      // with query region?
      intersect_flag=region->test_intersection(_query_region);

      if (intersect_flag==2)                    
      {
            // inside - extract points
            
            _report_subtree(node);
            // end of recursion
            // get rid of region
            delete region;
      }
      else if (intersect_flag==1)               
      {
            // overlap - recursion
            
            _search(region, node, depth+1);     
            // search does cleanup of region
      }
      else
      {
            // outside - stop
            
            // end of recursion
            // get rid of region
            delete region;    
      }
}

void KDTree::_report_subtree(Node *node)
{
      if (node->is_leaf())
      {
            // report point(s)
            long int i;

            for (i=node->get_start(); i<node->get_end(); i++)
            {
                  DataPoint data_point;
                  data_point=_data_point_list[i];
                  _report_point(data_point.get_index(), data_point.get_coord());
            }
      }
      else
      {
            // find points in subtrees via recursion 
            _report_subtree(node->get_left_node());
            _report_subtree(node->get_right_node());
      }
}

void KDTree::_report_point(long int index, float *coord)
{
      float r;

      r=KDTREE_dist(_center_coord, coord, KDTree::dim);

      if (r<=_radius_sq)
      {
            _index_list.push_back(index);
            // note use of sqrt - only calculated if necessary
            _radius_list.push_back(sqrt(r));
            _count++;
      }
}
            
void KDTree::set_data(float *coords, long int nr_points)
{
      long int i;

      DataPoint::dim=dim;
      Region::dim=dim;

      // clean up stuff from previous use
      delete _root;
      delete [] _coords;
      _index_list.clear();
      _radius_list.clear();
      _count=0;
      // keep pointer to coords to delete it
      _coords=coords;
            
      for (i=0; i<nr_points; i++)
      {
            _add_point(i, coords+i*dim);
      }

      // build KD tree
      _root=_build_tree();
}

      
void KDTree::search_center_radius(float *coord, float radius)
{
      int i;
      float left[dim], right[dim];

      DataPoint::dim=dim;
      Region::dim=dim;

      _index_list.clear();
      _radius_list.clear();
      _count=0;

      _radius=radius;
      // use of r^2 to avoid sqrt use
      _radius_sq=radius*radius;

      for (i=0; i<dim; i++)
      {
            left[i]=coord[i]-radius;
            right[i]=coord[i]+radius;
            // set center of query
            _center_coord[i]=coord[i];
      }

      // clean up! 
      delete [] coord;

      _set_query_region(left, right);

      _search();
}

long int KDTree::get_count(void)
{
      return _count;
}

void KDTree::copy_indices(long *indices)
{
      long int i;

      if (_count==0)
      {
            return;
      }

      for(i=0; i<_count; i++)
      {
            indices[i]=_index_list[i];
      }
}

void KDTree::copy_radii(float *radii)
{
      long int i;

      if (_count==0)
      {
            return;
      }

      for(i=0; i<_count; i++)
      {
            radii[i]=_radius_list[i];
      }
}

void KDTree::neighbor_copy_indices(long int *indices)
{
      long int i;

      if (_neighbor_count==0)
      {
            return;
      }

      for(i=0; i<_neighbor_count*2; i++)
      {
            indices[i]=_neighbor_index_list[i];
      }
}

void KDTree::neighbor_copy_radii(float *radii)
{
      long int i;

      if (_neighbor_count==0)
      {
            return;
      }

      for(i=0; i<_neighbor_count; i++)
      {
            radii[i]=_neighbor_radius_list[i];
      }
}

long int KDTree::neighbor_get_count(void)
{
      return _neighbor_count;
}

void KDTree::neighbor_search(float neighbor_radius)
{
      Region *region;

      DataPoint::dim=dim;
      Region::dim=dim;

      _neighbor_index_list.clear();
      _neighbor_radius_list.clear();
      // note the use of r^2 to avoid use of sqrt
      _neighbor_radius=neighbor_radius;
      _neighbor_radius_sq=neighbor_radius*neighbor_radius;
      _neighbor_count=0;

      // start with [-INF, INF]
      region=new Region();

      if (_root->is_leaf())
      {
            // this is a boundary condition
            // bucket_size>nr of points
            _search_neighbors_in_bucket(_root);
      }
      else
      {
            // "normal" situation
            _neighbor_search(_root, region, 0);
      }
      delete region;
}

void KDTree::_neighbor_search(Node *node, Region *region, int depth)
{
      Node *left, *right;
      Region *left_region, *right_region; 
      int localdim;

      localdim=depth%dim;

      left=node->get_left_node();
      right=node->get_right_node();

      // planes of left and right nodes
      left_region=region->intersect_left(node->get_cut_value(), localdim);
      right_region=region->intersect_right(node->get_cut_value(), localdim);

      if (!left->is_leaf())
      {
            // search for pairs in this half plane
            _neighbor_search(left, left_region, depth+1);
      }
      else
      {
            _search_neighbors_in_bucket(left);
      }
      
      if (!right->is_leaf())
      {
            // search for pairs in this half plane
            _neighbor_search(right, right_region, depth+1);
      }
      else
      {
            _search_neighbors_in_bucket(right);
      }

      // search for pairs between the half planes
      _neighbor_search_pairs(left, left_region, right, right_region, depth+1);

      // cleanup
      delete left_region;
      delete right_region;
}

void KDTree::_test_neighbors(DataPoint &p1, DataPoint &p2)
{
      float r;

      r=KDTREE_dist(p1.get_coord(), p2.get_coord(), dim);

      if(r<=_neighbor_radius_sq)
      {
            // we found a neighbor pair!
            _neighbor_index_list.push_back(p1.get_index());
            _neighbor_index_list.push_back(p2.get_index());
            // note sqrt
            _neighbor_radius_list.push_back(sqrt(r));
            _neighbor_count++;
      }
}

void KDTree::_search_neighbors_in_bucket(Node *node)
{
      long int i;

      for(i=node->get_start(); i<node->get_end(); i++)
      {
            DataPoint p1;
            long int j;

            p1=_data_point_list[i];

            for (j=i+1; j<node->get_end(); j++)
            {
                  DataPoint p2;

                  p2=_data_point_list[j];

                  _test_neighbors(p1, p2);
            }
      }
}

void KDTree::_search_neighbors_between_buckets(Node *node1, Node *node2)
{
      long int i;

      for(i=node1->get_start(); i<node1->get_end(); i++)
      {
            DataPoint p1;
            long int j;

            p1=_data_point_list[i];

            for (j=node2->get_start(); j<node2->get_end(); j++)
            {
                  DataPoint p2;

                  p2=_data_point_list[j];

                  _test_neighbors(p1, p2);
            }
      }
}

void KDTree::_neighbor_search_pairs(Node *down, Region *down_region, 
            Node *up, Region *up_region, int depth)
{
      int down_is_leaf, up_is_leaf;
      int localdim;

      // if regions do not overlap - STOP
      if (!down || !up || !down_region || !up_region) 
      {
            // STOP
            return;
      }
      
      if (down_region->test_intersection(up_region, _neighbor_radius)==0)
      {
            // regions cannot contain neighbors
            return;
      }

      // dim
      localdim=depth%dim;

      // are they leaves?
      up_is_leaf=up->is_leaf();
      down_is_leaf=down->is_leaf();

      if (up_is_leaf && down_is_leaf)
      {
            // two leaf nodes
            _search_neighbors_between_buckets(down, up);
      }
      else
      {
            // one or no leaf nodes

            Node *up_right, *up_left, *down_left, *down_right;
            Region *up_left_region, *up_right_region, 
                  *down_left_region, *down_right_region;  
            
            if (down_is_leaf)
            {
                  down_left=down;
                  // make a copy of down_region
                  down_left_region=new Region(down_region->get_left(), down_region->get_right());
                  down_right=NULL;
                  down_right_region=NULL;
            }
            else
            {
                  float cut_value;

                  cut_value=down->get_cut_value();

                  down_left=down->get_left_node();
                  down_right=down->get_right_node();
                  down_left_region=down_region->intersect_left(cut_value, localdim); 
                  down_right_region=down_region->intersect_right(cut_value, localdim); 
            }

            if (up_is_leaf)
            {
                  up_left=up;
                  // make a copy of up_region
                  up_left_region=new Region(up_region->get_left(), up_region->get_right());
                  up_right=NULL;
                  up_right_region=NULL;
            }
            else
            {
                  float cut_value;

                  cut_value=up->get_cut_value();

                  up_left=up->get_left_node();
                  up_right=up->get_right_node();
                  up_left_region=up_region->intersect_left(cut_value, localdim); 
                  up_right_region=up_region->intersect_right(cut_value, localdim); 
            }

            _neighbor_search_pairs(up_left, up_left_region, down_left, down_left_region, depth+1);
            _neighbor_search_pairs(up_left, up_left_region, down_right, down_right_region, depth+1);
            _neighbor_search_pairs(up_right, up_right_region, down_left, down_left_region, depth+1);
            _neighbor_search_pairs(up_right, up_right_region, down_right, down_right_region, depth+1);

            delete down_left_region;
            delete down_right_region;
            delete up_left_region;
            delete up_right_region;
      }
}

void KDTree::neighbor_simple_search(float radius)
{
      long int i;

      DataPoint::dim=dim;
      Region::dim=dim;

      _neighbor_radius=radius;
      _neighbor_radius_sq=radius*radius;

      _neighbor_count=0;
      _neighbor_index_list.clear();
      _neighbor_radius_list.clear();

      DataPoint::current_dim=0;
      sort(_data_point_list.begin(), _data_point_list.end());

      for (i=0; i<_data_point_list.size(); i++)
      {
            float x1;
            long int j;
            DataPoint p1;

            p1=_data_point_list[i];
            x1=p1.get_coord()[0];

            for (j=i+1; j<_data_point_list.size(); j++)
            {
                  DataPoint p2;
                  float x2;

                  p2=_data_point_list[j];
                  x2=p2.get_coord()[0];

                  if (fabs(x2-x1)<=radius)
                  {
                        _test_neighbors(p1, p2);
                  }
                  else
                  {
                        break;
                  }
            }
      }
}




Generated by  Doxygen 1.6.0   Back to index