PBRT
/home/felix/UBC/projects/AdaptiveLightfieldSampling/pbrt_v2/src/core/kdtree.h
00001 
00002 /*
00003     pbrt source code Copyright(c) 1998-2012 Matt Pharr and Greg Humphreys.
00004 
00005     This file is part of pbrt.
00006 
00007     Redistribution and use in source and binary forms, with or without
00008     modification, are permitted provided that the following conditions are
00009     met:
00010 
00011     - Redistributions of source code must retain the above copyright
00012       notice, this list of conditions and the following disclaimer.
00013 
00014     - Redistributions in binary form must reproduce the above copyright
00015       notice, this list of conditions and the following disclaimer in the
00016       documentation and/or other materials provided with the distribution.
00017 
00018     THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
00019     IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
00020     TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
00021     PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
00022     HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
00023     SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
00024     LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
00025     DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
00026     THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
00027     (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
00028     OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
00029 
00030  */
00031 
00032 #if defined(_MSC_VER)
00033 #pragma once
00034 #endif
00035 
00036 #ifndef PBRT_CORE_KDTREE_H
00037 #define PBRT_CORE_KDTREE_H
00038 
00039 // core/kdtree.h*
00040 #include "pbrt.h"
00041 #include "geometry.h"
00042 
00043 // KdTree Declarations
00044 struct KdNode {
00045     void init(float p, uint32_t a) {
00046         splitPos = p;
00047         splitAxis = a;
00048         rightChild = (1<<29)-1;
00049         hasLeftChild = 0;
00050     }
00051     void initLeaf() {
00052         splitAxis = 3;
00053         rightChild = (1<<29)-1;
00054         hasLeftChild = 0;
00055     }
00056     // KdNode Data
00057     float splitPos;
00058     uint32_t splitAxis:2;
00059     uint32_t hasLeftChild:1, rightChild:29;
00060 };
00061 
00062 
00063 template <typename NodeData> class KdTree {
00064 public:
00065     // KdTree Public Methods
00066     KdTree(const vector<NodeData> &data);
00067     ~KdTree() {
00068         FreeAligned(nodes);
00069         FreeAligned(nodeData);
00070     }
00071     template <typename LookupProc> void Lookup(const Point &p,
00072             LookupProc &process, float &maxDistSquared) const;
00073 private:
00074     // KdTree Private Methods
00075     void recursiveBuild(uint32_t nodeNum, int start, int end,
00076         const NodeData **buildNodes);
00077     template <typename LookupProc> void privateLookup(uint32_t nodeNum,
00078         const Point &p, LookupProc &process, float &maxDistSquared) const;
00079 
00080     // KdTree Private Data
00081     KdNode *nodes;
00082     NodeData *nodeData;
00083     uint32_t nNodes, nextFreeNode;
00084 };
00085 
00086 
00087 template <typename NodeData> struct CompareNode {
00088     CompareNode(int a) { axis = a; }
00089     int axis;
00090     bool operator()(const NodeData *d1, const NodeData *d2) const {
00091         return d1->p[axis] == d2->p[axis] ? (d1 < d2) :
00092                                             d1->p[axis] < d2->p[axis];
00093     }
00094 };
00095 
00096 
00097 
00098 // KdTree Method Definitions
00099 template <typename NodeData>
00100 KdTree<NodeData>::KdTree(const vector<NodeData> &d) {
00101     nNodes = d.size();
00102     nextFreeNode = 1;
00103     nodes = AllocAligned<KdNode>(nNodes);
00104     nodeData = AllocAligned<NodeData>(nNodes);
00105     vector<const NodeData *> buildNodes(nNodes, NULL);
00106     for (uint32_t i = 0; i < nNodes; ++i)
00107         buildNodes[i] = &d[i];
00108     // Begin the KdTree building process
00109     recursiveBuild(0, 0, nNodes, &buildNodes[0]);
00110 }
00111 
00112 
00113 template <typename NodeData> void
00114 KdTree<NodeData>::recursiveBuild(uint32_t nodeNum, int start, int end,
00115         const NodeData **buildNodes) {
00116     // Create leaf node of kd-tree if we've reached the bottom
00117     if (start + 1 == end) {
00118         nodes[nodeNum].initLeaf();
00119         nodeData[nodeNum] = *buildNodes[start];
00120         return;
00121     }
00122 
00123     // Choose split direction and partition data
00124 
00125     // Compute bounds of data from _start_ to _end_
00126     BBox bound;
00127     for (int i = start; i < end; ++i)
00128         bound = Union(bound, buildNodes[i]->p);
00129     int splitAxis = bound.MaximumExtent();
00130     int splitPos = (start+end)/2;
00131     std::nth_element(&buildNodes[start], &buildNodes[splitPos],
00132                      &buildNodes[end], CompareNode<NodeData>(splitAxis));
00133 
00134     // Allocate kd-tree node and continue recursively
00135     nodes[nodeNum].init(buildNodes[splitPos]->p[splitAxis], splitAxis);
00136     nodeData[nodeNum] = *buildNodes[splitPos];
00137     if (start < splitPos) {
00138         nodes[nodeNum].hasLeftChild = 1;
00139         uint32_t childNum = nextFreeNode++;
00140         recursiveBuild(childNum, start, splitPos, buildNodes);
00141     }
00142     if (splitPos+1 < end) {
00143         nodes[nodeNum].rightChild = nextFreeNode++;
00144         recursiveBuild(nodes[nodeNum].rightChild, splitPos+1,
00145                        end, buildNodes);
00146     }
00147 }
00148 
00149 
00150 template <typename NodeData> template <typename LookupProc>
00151 void KdTree<NodeData>::Lookup(const Point &p, LookupProc &proc,
00152                               float &maxDistSquared) const {
00153     privateLookup(0, p, proc, maxDistSquared);
00154 }
00155 
00156 
00157 template <typename NodeData> template <typename LookupProc>
00158 void KdTree<NodeData>::privateLookup(uint32_t nodeNum, const Point &p,
00159         LookupProc &process, float &maxDistSquared) const {
00160     KdNode *node = &nodes[nodeNum];
00161     // Process kd-tree node's children
00162     int axis = node->splitAxis;
00163     if (axis != 3) {
00164         float dist2 = (p[axis] - node->splitPos) * (p[axis] - node->splitPos);
00165         if (p[axis] <= node->splitPos) {
00166             if (node->hasLeftChild)
00167                 privateLookup(nodeNum+1, p, process, maxDistSquared);
00168             if (dist2 < maxDistSquared && node->rightChild < nNodes)
00169                 privateLookup(node->rightChild, p, process, maxDistSquared);
00170         }
00171         else {
00172             if (node->rightChild < nNodes)
00173                 privateLookup(node->rightChild, p, process, maxDistSquared);
00174             if (dist2 < maxDistSquared && node->hasLeftChild)
00175                 privateLookup(nodeNum+1, p, process, maxDistSquared);
00176         }
00177     }
00178 
00179     // Hand kd-tree node to processing function
00180     float dist2 = DistanceSquared(nodeData[nodeNum].p, p);
00181     if (dist2 < maxDistSquared)
00182         process(p, nodeData[nodeNum], dist2, maxDistSquared);
00183 }
00184 
00185 
00186 
00187 #endif // PBRT_CORE_KDTREE_H