K3DTree-exact.cpp

Go to the documentation of this file.

/* 
* K3DTree-exact.cpp - Precise KD tree implementation
* 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-exact.h"

#include <stack>
#include <queue>
#include <algorithm>
#include <cmath>
#include <utility>
#include <limits>

using namespace AtomProbe;

#include "atomprobe/helper/aptAssert.h"


using std::stack;
using std::vector;
using std::pair;
using std::queue;

const int PROGRESS_REDUCE=5000;


class AxisCompareExact
{
    private:
        unsigned int axis;
    public:
        void setAxis(unsigned int Axis){axis=Axis;};
        inline bool operator()(const std::pair<Point3D,size_t> &p1,const std::pair<Point3D,size_t> &p2) const
        {return p1.first[axis]<p2.first[axis];};
};

class NodeWalk
{
    public:
        size_t index;
        BoundCube cube;
        unsigned int depth;
        NodeWalk(unsigned int idx,const BoundCube &bc, unsigned int dpth) : 
            index(idx), cube(bc), depth(dpth) {};
};

K3DTreeExact::K3DTreeExact() : treeRoot(-1), callback(0),progress(0)
{
}

void K3DTreeExact::getBoundCube(BoundCube &b)
{
    ASSERT(treeBounds.isValid());
    b.setBounds(treeBounds);
}

const Point3D *K3DTreeExact::getPt(size_t index) const
{
    ASSERT(index < indexedPoints.size());
    return &(indexedPoints[index].first);
}

const Point3D &K3DTreeExact::getPtRef(size_t index) const
{
    ASSERT(index < indexedPoints.size());
    return (indexedPoints[index].first);
}

size_t K3DTreeExact::getOrigIndex(size_t treeIndex) const
{
    ASSERT(treeIndex <indexedPoints.size());
    return indexedPoints[treeIndex].second;
}

bool K3DTreeExact::getTag(size_t treeIndex) const
{
    ASSERT(treeIndex < nodes.size());
    return nodes[treeIndex].tagged;
}

void K3DTreeExact::tag(size_t treeIndex,bool tagVal)
{
    ASSERT(treeIndex < nodes.size());
    nodes[treeIndex].tagged=tagVal;
}

void K3DTreeExact::tag(const std::vector<size_t> &treeIndicies,bool tagVal)
{
for(unsigned int ui=0;ui<treeIndicies.size();ui++)
{
    ASSERT(treeIndicies[ui] < nodes.size());
    nodes[treeIndicies[ui]].tagged=tagVal;
}
}
size_t K3DTreeExact::size() const
{
    ASSERT(nodes.size() == indexedPoints.size());
    return indexedPoints.size();
}

void K3DTreeExact::resetPts(std::vector<Point3D> &p, bool clear)
{
    //Compute bounding box for indexedPoints
    treeBounds.setBounds(p);
    indexedPoints.resize(p.size());

    #pragma omp parallel for
    for(size_t ui=0;ui<indexedPoints.size();ui++)
    {
        indexedPoints[ui].first=p[ui];
        indexedPoints[ui].second=ui;
    }

    if(clear)
        p.clear();

    nodes.resize(indexedPoints.size());
}

void K3DTreeExact::resetPts(std::vector<IonHit> &p, bool clear)
{
    indexedPoints.resize(p.size());
    nodes.resize(p.size());

    if(p.empty())
        return;


    //Compute bounding box for indexedPoints
    treeBounds=IonHit::getIonDataLimits(p);

    #pragma omp parallel for
    for(size_t ui=0;ui<indexedPoints.size();ui++)
    {
        indexedPoints[ui].first=p[ui].getPosRef();
        indexedPoints[ui].second=ui;
    }

    if(clear)
        p.clear();

}

bool K3DTreeExact::build()
{

    using std::make_pair;

    enum
    {
        BUILT_NONE,
        BUILT_LEFT,
        BUILT_BOTH
    };

    ASSERT(callback());
    ASSERT(progress);

    //Clear any existing tags
    clearAllTags();
    maxDepth=0;

    //No indexedPoints? That was easy.
    if(indexedPoints.empty())
        return true;


    ASSERT(treeBounds.isValid());

    //Maintain a stack of nodeoffsets, and whether we have built the left hand side
    stack<pair<size_t,size_t> > limits;
    stack<char> buildStatus;
    stack<size_t> splitStack;

    //Data runs from 0 to size-1 INCLUSIVE
    limits.push(make_pair(0,indexedPoints.size()-1));
    buildStatus.push(BUILT_NONE);
    splitStack.push((size_t)-1);


    AxisCompareExact axisCmp;

    size_t numSeen=0; // for progress reporting 
    size_t splitIndex=0;




    #ifdef DEBUG
    for(size_t ui=0;ui<nodes.size();ui++)
    {
        nodes[ui].childLeft=nodes[ui].childRight=(size_t)-2;
    }
    #endif

    size_t *childPtr=0;
    do
    {

        switch(buildStatus.top())
        {
            case BUILT_NONE:
            {
                //OK, so we have not seen this data at this level before
                int curAxis=(limits.size()-1)%3;
                //First time we have seen this group? OK, we need to sort
                //along its hyper plane.
                axisCmp.setAxis(curAxis);
            
                //Sort data; note that the limits.top().second is the INCLUSIVE
                // upper end.   
                std::sort(indexedPoints.begin()+limits.top().first,
                        indexedPoints.begin() + limits.top().second+1,axisCmp);

                //Initially assume that the mid node is the median; then we slide it up 
                splitIndex=(limits.top().second+limits.top().first)/2;
                
                //Keep sliding the split towards the upper boundary until we hit a different
                //data value. This ensure that all data on the left of the sub-tree is <= 
                // to this data value for the specified sort axis
                while(splitIndex != limits.top().second
                    && indexedPoints[splitIndex].first.getValue(curAxis) ==
                        indexedPoints[splitIndex+1].first.getValue(curAxis))
                    splitIndex++;

                buildStatus.top()++; //Increment the build status to "left" case.
                        

                if(limits.size() ==1)
                {
                    //root node
                    treeRoot=splitIndex;
                }
                else
                    *childPtr=splitIndex;

                //look to see if there is any left data
                if(splitIndex >limits.top().first)
                {
                    //There is; we have to branch again
                    limits.push(make_pair(
                        limits.top().first,splitIndex-1));
                    
                    buildStatus.push(BUILT_NONE);
                    //Set the child pointer, as we don't know
                    //the correct value until the next sort.
                    childPtr=&nodes[splitIndex].childLeft;
                }
                else
                {
                    //There is not. Set the left branch to null
                    nodes[splitIndex].childLeft=(size_t)-1;
                }
                splitStack.push(splitIndex);

                break;
            }

            case BUILT_LEFT:
            {
                //Either of these cases results in use 
                //handling the right branch.
                buildStatus.top()++;
                splitIndex=splitStack.top();
                //Check to see if there is any right data
                if(splitIndex <limits.top().second)
                {
                    //There is; we have to branch again
                    limits.push(make_pair(
                        splitIndex+1,limits.top().second));
                    buildStatus.push(BUILT_NONE);

                    //Set the child pointer, as we don't know
                    //the correct value until the next sort.
                    childPtr=&nodes[splitIndex].childRight;
                }
                else
                {
                    //There is not. Set the right branch to null
                    nodes[splitIndex].childRight=(size_t)-1;
                }

                break;
            }
            case BUILT_BOTH:
            {
                ASSERT(nodes[splitStack.top()].childLeft != (size_t)-2
                    && nodes[splitStack.top()].childRight!= (size_t)-2 );
                maxDepth=std::max(maxDepth,limits.size());
                //pop limits and build status.
                limits.pop();
                buildStatus.pop();
                splitStack.pop();
                    ASSERT(limits.size() == buildStatus.size());
                    

                    numSeen++;
                    break;
                }
            }   

            if(!(numSeen%PROGRESS_REDUCE) && progress)
            {
                *progress= (unsigned int)((float)numSeen/(float)nodes.size()*100.0f);

                if(!(*callback)())
                    return false;
            }

        }while(!limits.empty());

    return true;
}

void K3DTreeExact::dump(std::ostream &strm,  size_t depth, size_t offset) const
{
    if(offset==(size_t)-1)
    {
        for(unsigned int ui=0;ui<indexedPoints.size();ui++)
        {
            strm << ui << " "<< indexedPoints[ui].first << std::endl;
        }

        strm << "----------------" << std::endl;
        offset=treeRoot;
    }

    for(size_t ui=0;ui<depth; ui++)
        strm << "\t";

    strm << offset << " : (" << indexedPoints[offset].first[0] 
        << "," << indexedPoints[offset].first[1] << "," << indexedPoints[offset].first[2]
        << ")" << std::endl;



    for(size_t ui=0;ui<depth; ui++)
        strm << "\t";
    strm << "<l>" <<std::endl;

    if(nodes[offset].childLeft!=(size_t)-1)
    {
        dump(strm,depth+1,nodes[offset].childLeft);
    }
    for(size_t ui=0;ui<depth; ui++)
        strm << "\t";
    strm << "</l>" <<std::endl;

    for(size_t ui=0;ui<depth; ui++)
        strm << "\t";
    strm << "<r>" <<std::endl;

    if(nodes[offset].childRight!=(size_t)-1)
        dump(strm,depth+1,nodes[offset].childRight);
    
    for(size_t ui=0;ui<depth; ui++)
        strm << "\t";
    strm << "</r>" <<std::endl;
}

void K3DTreeExact::ptsInSphere(const Point3D &origin, float radius,
        vector<size_t> &pts) const
{
    if(!treeBounds.intersects(origin,radius))
        return;

    //parent walking queue. This contains initial parent indices that
    // lie within the sphere.



    const float sqrRadius=radius*radius;
    //contains all completely contained points (these points and children are in sphere)
    std::queue<size_t> idxQueue;
    //queue of points whose children are partly in the sphere
    std::queue<NodeWalk> nodeQueue;
    nodeQueue.push(NodeWalk(treeRoot,treeBounds,0));

    while(!nodeQueue.empty())
    {
        size_t nodeIdx;
        BoundCube curCube;
        unsigned int depth,axis;

        nodeIdx = nodeQueue.front().index;
        curCube=nodeQueue.front().cube;
        depth = nodeQueue.front().depth;
        axis=depth %3;  
        nodeQueue.pop() ;
    

        //obtain the left and right cubes, and see if they
        // -exist
        // -intersect the spehre
        // - if intersects, are they contained entirely by the sphere 
        BoundCube leftCube,rightCube;
        if(nodes[nodeIdx].childLeft != (size_t) -1)
        {
            leftCube=curCube;
            leftCube.bounds[axis][1]=indexedPoints[nodeIdx].first[axis];
            if(leftCube.intersects(origin,sqrRadius) )
            {
                if(leftCube.containedInSphere(origin,radius))
                {
                    ASSERT(indexedPoints[nodeIdx].first.sqrDist(origin) < radius*radius);
                    idxQueue.push(nodes[nodeIdx].childLeft);
                }
                else
                    nodeQueue.push(NodeWalk(nodes[nodeIdx].childLeft,leftCube,depth+1));
            }
        }   

        if(nodes[nodeIdx].childRight != (size_t) -1)
        {
            rightCube=curCube;
            rightCube.bounds[axis][0]=indexedPoints[nodeIdx].first[axis];
            if(rightCube.intersects(origin,sqrRadius) )
            {
                if(rightCube.containedInSphere(origin,radius))
                {

                    ASSERT(indexedPoints[nodeIdx].first.sqrDist(origin) < sqrRadius);
                    idxQueue.push(nodes[nodeIdx].childRight);
                }
                else
                    nodeQueue.push(NodeWalk(nodes[nodeIdx].childRight,rightCube,depth+1));
            }
        }   

        if(indexedPoints[nodeIdx].first.sqrDist(origin) < sqrRadius)
            pts.push_back(nodeIdx);

    }   


    pts.reserve(idxQueue.size());
    //Walk the idx queue to enumerate all children that are in the sphere
    while(!idxQueue.empty())
    {
        size_t curIdx;
        curIdx=idxQueue.front();
        ASSERT(indexedPoints[curIdx].first.sqrDist(origin) < sqrRadius);
        if(nodes[curIdx].childLeft != (size_t)-1)
            idxQueue.push(nodes[curIdx].childLeft);
        
        if(nodes[curIdx].childRight !=(size_t) -1)
            idxQueue.push(nodes[curIdx].childRight);


        ASSERT(curIdx < nodes.size());
        pts.push_back(curIdx);
        idxQueue.pop();
    }
    
}

size_t K3DTreeExact::findNearestUntagged(const Point3D &searchPt,
                const BoundCube &domainCube, bool shouldTag)
{
    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()
        };
    
    ASSERT(treeRoot!=(size_t)-1);
    ASSERT(callback);
    ASSERT(progress);

    size_t nodeStack[maxDepth+1];
    float domainStack[maxDepth+1][2];
    unsigned int visitStack[maxDepth+1];

    size_t bestPoint;
    size_t curNode;

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

    if(nodes.empty())
        return -1;

    bestPoint=(size_t)-1; 
    bestDistSqr =std::numeric_limits<float>::max();
    curDomain=domainCube;
    visit=NODE_FIRST_VISIT;
    curAxis=0;
    stackTop=0;

    //Start at median of array, which is top of tree,
    //by definition
    curNode=treeRoot;

    //check root node   
    if(!nodes[curNode].tagged)
    {
        float tmpDistSqr;
        tmpDistSqr = indexedPoints[curNode].first.sqrDist(searchPt); 
        if(tmpDistSqr < bestDistSqr)
        {
            bestDistSqr  = tmpDistSqr;
            bestPoint=curNode;
        }
    }

    do
    {
        switch(visit)
        {
            //Examine left branch
            case NODE_FIRST_VISIT:
            {
                if(searchPt[curAxis] < indexedPoints[curNode].first[curAxis])
                {
                    if(nodes[curNode].childLeft!=(size_t)-1)
                    {
                        //Check bounding box when shrunk overlaps best
                        //estimate sphere
                        tmpEdge= curDomain.bounds[curAxis][1];
                        curDomain.bounds[curAxis][1] = indexedPoints[curNode].first[curAxis];
                        if(!curDomain.intersects(searchPt,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=nodes[curNode].childLeft;
                        visit=NODE_FIRST_VISIT;
                        curAxis++;
                        curAxis%=3;
                        continue;
                    }
                }   
                else
                {
                    if(nodes[curNode].childRight!=(size_t)-1)
                    {
                        //Check bounding box when shrunk overlaps best
                        //estimate sphere
                        tmpEdge= curDomain.bounds[curAxis][0];
                        curDomain.bounds[curAxis][0] = indexedPoints[curNode].first[curAxis];
                        
                        if(!curDomain.intersects(searchPt,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=nodes[curNode].childRight;
                        visit=NODE_FIRST_VISIT;
                        curAxis++;
                        curAxis%=3;
                        continue;   
                    }
                }
                visit++;
                //Fall through
            }
            //Examine right branch
            case NODE_SECOND_VISIT:
            {
                if(searchPt[curAxis]< indexedPoints[curNode].first[curAxis])
                {
                    if(nodes[curNode].childRight!=(size_t)-1)
                    {
                        //Check bounding box when shrunk overlaps best
                        //estimate sphere
                        tmpEdge= curDomain.bounds[curAxis][0];
                        curDomain.bounds[curAxis][0] = indexedPoints[curNode].first[curAxis];
                        
                        if(!curDomain.intersects(searchPt,bestDistSqr))
                        {
                            curDomain.bounds[curAxis][0] = tmpEdge; 
                            visit++;
                            continue;       
                        }
    
                        nodeStack[stackTop]=curNode;
                        visitStack[stackTop] = NODE_THIRD_VISIT; 
                        domainStack[stackTop][0] = tmpEdge;
                        domainStack[stackTop][1]= curDomain.bounds[curAxis][1];
                        stackTop++;
                        
                        //Update the information
                        curNode=nodes[curNode].childRight;
                        visit=NODE_FIRST_VISIT;
                        curAxis++;
                        curAxis%=3;
                        continue;   

                    }
                }
                else
                {
                    if(nodes[curNode].childLeft!=(size_t)-1)
                    {
                        //Check bounding box when shrunk overlaps best
                        //estimate sphere
                        tmpEdge= curDomain.bounds[curAxis][1];
                        curDomain.bounds[curAxis][1] = indexedPoints[curNode].first[curAxis];
                        
                        if(!curDomain.intersects(searchPt,bestDistSqr))
                        {
                            curDomain.bounds[curAxis][1] = tmpEdge; 
                            visit++;
                            continue;       
                        }   
                        //Preserve our current state.
                        nodeStack[stackTop]=curNode;
                        visitStack[stackTop] = NODE_THIRD_VISIT; 
                        domainStack[stackTop][1] = tmpEdge;
                        domainStack[stackTop][0]= curDomain.bounds[curAxis][0];
                        stackTop++;
                        
                        //Update the information
                        curNode=nodes[curNode].childLeft;
                        visit=NODE_FIRST_VISIT;
                        curAxis++;
                        curAxis%=3;
                        continue;   

                    }
                }
                visit++;
                //Fall through
            }
            case NODE_THIRD_VISIT:
            {
                //Decide if we should promote the current node
                //to "best" (ie nearest untagged) node.
                //To promote, it musnt be tagged, and it must
                //be closer than cur best estimate.
                if(!nodes[curNode].tagged)
                {
                    float tmpDistSqr;
                    tmpDistSqr = indexedPoints[curNode].first.sqrDist(searchPt); 
                    if(tmpDistSqr < bestDistSqr)
                    {
                        bestDistSqr  = tmpDistSqr;
                        bestPoint=curNode;
                    }
                }

                //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 node for the third time (one left, one right, one finish)   
    }while(!(curNode== treeRoot &&  visit== NODE_THIRD_VISIT));

    if(bestPoint != (size_t) -1)
        nodes[bestPoint].tagged|=shouldTag;
    return bestPoint;   

}



void K3DTreeExact::findUntaggedInRadius(const Point3D &searchPt,
                const BoundCube &domainCube, float radius, vector<size_t> &result)

{
    //TODO: Benchmark competing, threadable implementation where we track our own tags 
    // for the calculation
    result.clear();
    

    //Find all pts within sphere of "radius" from searchPt
    float sqrRad=radius*radius;
    do
    {

        //Find next untagged point
        size_t resIdx;
        resIdx =findNearestUntagged(searchPt,domainCube, radius);

        //if not a point, then we are done
        if(resIdx == (size_t)-1)
            break;

        if(searchPt.sqrDist(getPtRef(resIdx)) > sqrRad)
        {
            tag(resIdx,false);
            break;
        }

        ASSERT(resIdx < size()); 
        result.push_back(resIdx);
    }while(true);


    //The points we found were untagged intially, re-untag them
    for(size_t ui=0;ui<result.size();ui++)
    {
        tag(result[ui],false);
    }

}

size_t K3DTreeExact::tagCount() const
{
    size_t count=0;

    #pragma omp parallel for reduction(+:count)
    for(size_t ui=0;ui<nodes.size();ui++)
    {
        if(nodes[ui].tagged)
            count++;
    
    }

    return count;
}

void K3DTreeExact::clearTags(std::vector<size_t> &tagsToClear)
{
#pragma omp parallel for
    for(size_t ui=0;ui<tagsToClear.size();ui++)
        nodes[tagsToClear[ui]].tagged=false;
}

void K3DTreeExact::clearAllTags()
{
#pragma omp parallel for
    for(size_t ui=0;ui<nodes.size();ui++)
        nodes[ui].tagged=false;
}


#ifdef DEBUG
#include <iostream>
#include <atomprobe/io/dataFiles.h>
using std::cerr;
using std::endl;
bool dummyCallback()
{
    return true;
}

bool inSphereTest()
{
    const float CENTRE_POS=0.5f;

    const unsigned int NPTS=20;
    Point3D p,centre;
    centre = Point3D(CENTRE_POS,CENTRE_POS,CENTRE_POS);

    vector<Point3D> spherePts;
    vector<Point3D> allPts;

    allPts.reserve(NPTS*NPTS*NPTS);
    spherePts.reserve(allPts.size()*CENTRE_POS);

    //Build a grid of points in [0,1]
    for(size_t ui=0;ui<NPTS;ui++)
    {
        for(size_t uj=0;uj<NPTS;uj++)
        {
            for(size_t uk=0;uk<NPTS;uk++)
            {
                p= Point3D((float)ui,(float)uj,(float)uk);
                p*=1.0/(float)NPTS;
                
                if(p.sqrDist(centre) < CENTRE_POS*CENTRE_POS)
                    spherePts.push_back(p);
                allPts.push_back(p);    
            }
        }
    }


    unsigned int dummyProg=0;
    
    K3DTreeExact tree;
    tree.resetPts(allPts);
    tree.setCallback(dummyCallback);
    tree.setProgressPointer(&dummyProg);

    tree.build();
    vector<size_t> spherePtIdx;
    tree.ptsInSphere(centre,CENTRE_POS,spherePtIdx);

    TEST(spherePts.size() == spherePtIdx.size(), "Check brute force and KD methods match in size");


    //Ensure that all points lie within the sphere
    for(size_t ui=0;ui<spherePtIdx.size();ui++)
    {
        const Point3D *pt;
        pt=tree.getPt(spherePtIdx[ui]);
        ASSERT(pt->sqrDist(centre) < CENTRE_POS*CENTRE_POS); 
    }



    return true;
}

bool K3DTreeExact::runUnitTests() 
{

    TEST(inSphereTest(),"In-sphere test");

    return true;
}
#endif