K3DTree-approx.cpp

Go to the documentation of this file.

/* 
 * Copyright (C) 2014  Daniel Haley
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include "atomprobe/algorithm/K3DTree-approx.h"

#include "atomprobe/helper/aptAssert.h"

#ifdef DEBUG
    #include <cstdlib>
    #include <iostream>
    using std::endl;
    using std::cerr;
#endif

#include <limits>
#include <cstring>

using std::vector;

namespace AtomProbe
{



//Axis compare
//==========
AxisCompare::AxisCompare() : axis(0)
{
}

void AxisCompare::setAxis(unsigned int sortAxis)
{
    axis=sortAxis;
}

//==========

//K3D node
//==========
void K3DNodeApprox::setLoc(const Point3D &locNew)
{
    loc=locNew;
}

Point3D K3DNodeApprox::getLoc() const
{
    return loc;
}

void K3DNodeApprox::deleteChildren()
{

    if(childLeft)
    {
        childLeft->deleteChildren();
        delete childLeft;
        childLeft=0;
    }
    
    if(childRight)
    {
        childRight->deleteChildren();
        delete childRight;
        childRight=0;
    }
    

}

void K3DNodeApprox::dump(std::ostream &strm, unsigned int depth) const
{
    for(unsigned int ui=0;ui<depth; ui++)
        strm << "\t";

    strm << "(" << loc.getValue(0) 
        << "," << loc.getValue(1) << "," << loc.getValue(2) 
        << ")" << std::endl;

    if(childLeft)
        childLeft->dump(strm,depth+1);
    
    if(childRight)
        childRight->dump(strm,depth+1);
}

//===========

//K3D Tree
//=============
K3DTreeApprox::K3DTreeApprox() : maxDepth(0),root(0)
{
    treeSize=0;
}


K3DTreeApprox::~K3DTreeApprox()
{
    kill();
}


void K3DTreeApprox::kill()
{
    if(root)
    {
        root->deleteChildren();
        delete root;
        root=0;
        treeSize=0;
    }
}

//Build the KD tree
void K3DTreeApprox::build(const vector<Point3D> &origPts)
{
    vector<Point3D> pts;
    pts=origPts;

    //che. to see if the pts vector is empty
    if(!pts.size())
    {
        maxDepth=0;
        return; 
    }

    if(root)
        kill();

    treeSize=pts.size();    
    maxDepth=1;
    root=buildRecurse(pts.begin(), pts.end(),0);
    
}

//Build the KD tree, shuffling original
void K3DTreeApprox::buildByRef(vector<Point3D> &pts)
{
    //che. to see if the pts vector is empty
    if(pts.empty())
    {
        maxDepth=0;
        return; 
    }

    if(root)
        kill();

    treeSize=pts.size();    
    maxDepth=1;
    root=buildRecurse(pts.begin(), pts.end(),0);
    
}

K3DNodeApprox*K3DTreeApprox::buildRecurse(vector<Point3D>::iterator pts_start, vector<Point3D>::iterator pts_end, unsigned int depth)
{

    K3DNodeApprox*node= new K3DNodeApprox;
    unsigned int curAxis=depth%3;
    unsigned int ptsSize=pts_end - pts_start - 1;//pts.size()-1
    node->setAxis(curAxis);
    
    //if we are deeper, then record that
    if(depth > maxDepth)
        maxDepth=depth;
    
    unsigned int median =(ptsSize)/2;

    //set up initial node
    AxisCompare axisCmp;
    axisCmp.setAxis(curAxis);

    //Find the median value in the current axis
    sort(pts_start,pts_end,axisCmp);
    
    //allocate node (this stores a copy of the point) and set.
    node->setLoc(*(pts_start + median));
    

    if(median)
        node->setLeft(buildRecurse(pts_start,pts_start + median,depth+1));
    else
        node->setLeft(0);   

    if(median!=ptsSize)
        node->setRight(buildRecurse(pts_start + median + 1, pts_end,depth+1));
    else
        node->setRight(0);

    return node;    

}

void K3DTreeApprox::dump( std::ostream &strm) const
{
    if(root)
        root->dump(strm,0);
}


const Point3D *K3DTreeApprox::findNearest(const Point3D &searchPt, const BoundCube &domainCube,
                                    const float deadDistSqr) const
{
    enum { NODE_FIRST_VISIT, //First visit is when you descend the tree
        NODE_SECOND_VISIT, //Second visit is when you come back from ->Left()
        NODE_THIRD_VISIT // Third visit is when you come back from ->Right()
        };


    //The stacks are for the nodes above the current Node.
    const K3DNodeApprox*nodeStack[maxDepth+1];
    float domainStack[maxDepth+1][2];
    unsigned int visitStack[maxDepth+1];

    const Point3D *bestPoint;
    const K3DNodeApprox*curNode;

    BoundCube curDomain;
    unsigned int visit;
    unsigned int stackTop;
    unsigned int curAxis;
    
    float bestDistSqr;
    float tmpEdge;

    //Set the root as the best estimate
    
    bestPoint =0; 
    bestDistSqr =std::numeric_limits<float>::max();
    curDomain=domainCube;
    visit=NODE_FIRST_VISIT;
    curAxis=0;

    stackTop=0;
/*  
    nodeStack[0]=root;  
    visitStack[0]=NODE_FIRST_VISIT;
    domainStack[0][0]=curDomain.bounds[0][0];
    domainStack[0][1]=curDomain.bounds[0][1];
*/  
    curNode=root;   

    do
    {

        switch(visit)
        {
            //Examine left branch
            case NODE_FIRST_VISIT:
            {
                if(searchPt[curAxis] < curNode->getLocVal(curAxis))
                {
                    if(curNode->left())
                    {
                        //Check bounding box when shrunk overlaps best
                        //estimate sphere
                        tmpEdge= curDomain.bounds[curAxis][1];
                        curDomain.bounds[curAxis][1] = curNode->getLocVal(curAxis);
                        if(bestPoint && !curDomain.intersects(*bestPoint,bestDistSqr))
                        {
                            curDomain.bounds[curAxis][1] = tmpEdge; 
                            visit++;
                            continue;       
                        }   
                        
                        //Preserve our current state.
                        nodeStack[stackTop]=curNode;
                        visitStack[stackTop] = NODE_SECOND_VISIT; //Oh, It will be. It will be.
                        domainStack[stackTop][1] = tmpEdge;
                        domainStack[stackTop][0]= curDomain.bounds[curAxis][0];
                        stackTop++;
                        
                        //Update the current information
                        curNode=curNode->left();
                        visit=NODE_FIRST_VISIT;
                        curAxis++;
                        curAxis%=3;
                        continue;   
                    }
                }
                else
                {
                    if(curNode->right())
                    {
                        //Check bounding box when shrunk overlaps best
                        //estimate sphere
                        tmpEdge= curDomain.bounds[curAxis][0];
                        curDomain.bounds[curAxis][0] = curNode->getLocVal(curAxis);
                        
                        if(bestPoint && !curDomain.intersects(*bestPoint,bestDistSqr))
                        {
                            curDomain.bounds[curAxis][0] =tmpEdge; 
                            visit++;
                            continue;       
                        }   
                        //Preserve our current state.
                        nodeStack[stackTop]=curNode;
                        visitStack[stackTop] = NODE_SECOND_VISIT; //Oh, It will be. It will be.
                        domainStack[stackTop][0] = tmpEdge;
                        domainStack[stackTop][1]= curDomain.bounds[curAxis][1];
                        stackTop++;

                        //Update the information
                        curNode=curNode->right();
                        visit=NODE_FIRST_VISIT;
                        curAxis++;
                        curAxis%=3;
                        continue;   
                    }

                }
                visit++;
                //Fall through
            }
            //Examine right branch
            case NODE_SECOND_VISIT:
            {
                if(searchPt[curAxis]< curNode->getLocVal(curAxis))
                {
                    if(curNode->right())
                    {
                        //Check bounding box when shrunk overlaps best
                        //estimate sphere
                        tmpEdge= curDomain.bounds[curAxis][0];
                        curDomain.bounds[curAxis][0] = curNode->getLocVal(curAxis);
                        
                        if(bestPoint && !curDomain.intersects(*bestPoint,bestDistSqr))
                        {
                            curDomain.bounds[curAxis][0] = tmpEdge; 
                            visit++;
                            continue;       
                        }
    
                        nodeStack[stackTop]=curNode;
                        visitStack[stackTop] = NODE_THIRD_VISIT; //Oh, It will be. It will be.
                        domainStack[stackTop][0] = tmpEdge;
                        domainStack[stackTop][1]= curDomain.bounds[curAxis][1];
                        stackTop++;
                        
                        //Update the information
                        curNode=curNode->right();
                        visit=NODE_FIRST_VISIT;
                        curAxis++;
                        curAxis%=3;
                        continue;   

                    }
                }
                else
                {
                    if(curNode->left())
                    {
                        //Check bounding box when shrunk overlaps best
                        //estimate sphere
                        tmpEdge= curDomain.bounds[curAxis][1];
                        curDomain.bounds[curAxis][1] = curNode->getLocVal(curAxis);
                        
                        if(bestPoint && !curDomain.intersects(*bestPoint,bestDistSqr))
                        {
                            curDomain.bounds[curAxis][1] = tmpEdge; 
                            visit++;
                            continue;       
                        }   
                        //Preserve our current state.
                        nodeStack[stackTop]=curNode;
                        visitStack[stackTop] = NODE_THIRD_VISIT; //Oh, It will be. It will be.
                        domainStack[stackTop][1] = tmpEdge;
                        domainStack[stackTop][0]= curDomain.bounds[curAxis][0];
                        stackTop++;
                        
                        //Update the information
                        curNode=curNode->left();
                        visit=NODE_FIRST_VISIT;
                        curAxis++;
                        curAxis%=3;
                        continue;   

                    }
                }
                visit++;
                //Fall through
            }
            //Go up
            case NODE_THIRD_VISIT:
            {
                float tmpDistSqr;
                tmpDistSqr = curNode->sqrDist(searchPt); 
                if(tmpDistSqr < bestDistSqr &&  tmpDistSqr > deadDistSqr)
                {
                    bestDistSqr  = tmpDistSqr;
                    bestPoint=curNode->getLocRef();
                }

                //DEBUG
                ASSERT(stackTop%3 == curAxis)
                //
                if(curAxis)
                    curAxis--;
                else
                    curAxis=2;


                
                ASSERT(stackTop < maxDepth+1);  
                if(stackTop)
                {
                    stackTop--;
                    visit=visitStack[stackTop];
                    curNode=nodeStack[stackTop];
                    curDomain.bounds[curAxis][0]=domainStack[stackTop][0];
                    curDomain.bounds[curAxis][1]=domainStack[stackTop][1];
                }   
                ASSERT((stackTop)%3 == curAxis)
                break;
            }
            default:
                ASSERT(false);


        }
        
    //Keep going until we meet the root nde for the third time (one left, one right, one finish)    
    }while(!(curNode==root &&  visit== NODE_THIRD_VISIT));

    //Check the root node, as we miss this on the way up
    float rootDist = root->sqrDist(searchPt);
    if(rootDist < bestDistSqr && rootDist > deadDistSqr)
        bestPoint = root->getLocRef();
    
    return bestPoint;   

}


void K3DTreeApprox::findKNearest(const Point3D &searchPt, const BoundCube &domainCube,
                    unsigned int num, vector<const Point3D *> &bestPts,
                float deadDistSqr) const
{
    //find the N nearest points
    bestPts.clear();
    bestPts.reserve(num);

    const Point3D *p;
    for(unsigned int ui=0; ui<num; ui++)
    {
        float sqrDist;
        p= findNearest(searchPt, domainCube,
                        deadDistSqr);

        if(!p)
            return;
        else
            bestPts.push_back(p);

        sqrDist = p->sqrDist(searchPt);
        deadDistSqr = sqrDist+std::numeric_limits<float>::epsilon();

    }
    
}
//=============

#ifdef DEBUG
bool K3DTreeApprox::test()
{
    vector<Point3D> pts;

    pts.push_back(Point3D(0,0,0));
    pts.push_back(Point3D(0,0.5,0.5));
    pts.push_back(Point3D(100,0,0));

    BoundCube bc(pts);

    K3DTreeApprox k3dTree;
    k3dTree.build(pts);

    const Point3D *p;
    p = k3dTree.findNearest(pts[0],bc,0.00001);

    TEST( p->sqrDist(pts[1]) < 0.01, "findNearest check");

    return true;
}
#endif

}