kdtree.h

/*
    This file is part of Nori, a simple educational ray tracer
 
    Copyright (c) 2015 by Wenzel Jakob
 
    Nori is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License Version 3
    as published by the Free Software Foundation.
 
    Nori 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/>.
*/
 
#if !defined(__NORI_KDTREE_H)
#define __NORI_KDTREE_H
 
#include <nori/bbox.h>
 
NORI_NAMESPACE_BEGIN
 
template <typename _PointType, typename _DataRecord> struct GenericKDTreeNode {
    typedef _PointType                       PointType;
    typedef _DataRecord                      DataRecord;
    typedef uint32_t                         IndexType;
    typedef typename PointType::Scalar       Scalar;
 
    enum {
        ELeafFlag  =  0x10,
        EAxisMask  =  0x0F
    };
 
    PointType position;
    IndexType right;
    DataRecord data;
    uint8_t flags;
 
    GenericKDTreeNode() : position((Scalar) 0),
        right(0), data(), flags(0) { }
    GenericKDTreeNode(const PointType &position, const DataRecord &data)
        : position(position), right(0), data(data), flags(0) {}
 
    IndexType getRightIndex(IndexType self) const { return right; }
    void setRightIndex(IndexType self, IndexType value) { right = value; }
 
    IndexType getLeftIndex(IndexType self) const { return self + 1; }
    void setLeftIndex(IndexType self, IndexType value) {
        if (value != self+1)
            throw NoriException("GenericKDTreeNode::setLeftIndex(): Internal error!");
    }
 
    bool isLeaf() const { return flags & (uint8_t) ELeafFlag; }
    void setLeaf(bool value) {
        if (value)
            flags |= (uint8_t) ELeafFlag;
        else
            flags &= (uint8_t) ~ELeafFlag;
    }
 
    uint16_t getAxis() const { return flags & (uint8_t) EAxisMask; }
    void setAxis(uint8_t axis) { flags = (flags & (uint8_t) ~EAxisMask) | axis; }
 
    const PointType &getPosition() const { return position; }
    void setPosition(const PointType &value) { position = value; }
 
    DataRecord &getData() { return data; }
    const DataRecord &getData() const { return data; }
    void setData(const DataRecord &val) { data = val; }
};
 
/* Forward declaration; the implementation is at the end of this file */
template <typename DataType, typename IndexType> void permute_inplace(
        DataType *data, std::vector<IndexType> &perm);
 
template <typename _NodeType> class PointKDTree {
public:
    typedef _NodeType                        NodeType;
    typedef typename NodeType::PointType     PointType;
    typedef typename NodeType::IndexType     IndexType;
    typedef typename PointType::Scalar       Scalar;
    typedef typename PointType::VectorType   VectorType;
    typedef TBoundingBox<PointType>          BoundingBoxType;
 
    enum {
        Dimension = VectorType::RowsAtCompileTime
    };
 
    enum Heuristic {
        Balanced = 0,
 
        SlidingMidpoint
    };
 
    struct SearchResult {
        float distSquared;
        IndexType index;
 
        SearchResult() {}
 
        SearchResult(float distSquared, IndexType index)
            : distSquared(distSquared), index(index) { }
 
        std::string toString() const {
            std::ostringstream oss;
            oss << "SearchResult[distance=" << std::sqrt(distSquared)
                << ", index=" << index << "]";
            return oss.str();
        }
 
        bool operator==(const SearchResult &r) const {
            return distSquared == r.distSquared &&
                index == r.index;
        }
    };
 
public:
    PointKDTree(size_t nodes = 0, Heuristic heuristic = SlidingMidpoint)
        : m_nodes(nodes), m_heuristic(heuristic), m_depth(0) { }
 
    // =============================================================
    // =============================================================
    void clear() { m_nodes.clear(); m_bbox.reset(); }
    void resize(size_t size) { m_nodes.resize(size); }
    void reserve(size_t size) { m_nodes.reserve(size); }
    size_t size() const { return m_nodes.size(); }
    size_t capacity() const { return m_nodes.capacity(); }
    void push_back(const NodeType &node) {
        m_nodes.push_back(node);
        m_bbox.expandBy(node.getPosition());
    }
    NodeType &operator[](size_t idx) { return m_nodes[idx]; }
    const NodeType &operator[](size_t idx) const { return m_nodes[idx]; }
    // =============================================================
 
    void setBoundingBox(const BoundingBoxType &bbox) { m_bbox = bbox; }
    const BoundingBoxType &getBoundingBox() const { return m_bbox; }
    size_t getDepth() const { return m_depth; }
    void setDepth(size_t depth) { m_depth = depth; }
 
    void build(bool recomputeBoundingBox = false) {
        if (m_nodes.size() == 0) {
            std::cerr << "KDTree::build(): kd-tree is empty!" << endl;
            return;
        }
 
        cout << "Building a " << Dimension << "-dimensional kd-tree over "
             << m_nodes.size() << " data points ("
             << memString(m_nodes.size() * sizeof(NodeType)).c_str() << ") .. ";
        cout.flush();
 
        if (recomputeBoundingBox) {
            m_bbox.reset();
            for (size_t i=0; i<m_nodes.size(); ++i)
                m_bbox.expandBy(m_nodes[i].getPosition());
        }
 
        /* Instead of shuffling around the node data itself, only modify
           an indirection table initially. Once the tree construction
           is done, this table will contain a indirection that can then
           be applied to the data in one pass */
        std::vector<IndexType> indirection(m_nodes.size());
        for (size_t i=0; i<m_nodes.size(); ++i)
            indirection[i] = (IndexType) i;
 
        m_depth = 0;
        build(1, indirection.begin(), indirection.begin(), indirection.end());
        permute_inplace(&m_nodes[0], indirection);
 
        cout << "done." << endl;
    }
 
    void search(const PointType &p, float searchRadius, std::vector<IndexType> &results) const {
        if (m_nodes.size() == 0)
            return;
 
        IndexType *stack = (IndexType *) alloca((m_depth+1) * sizeof(IndexType));
        IndexType index = 0, stackPos = 1, found = 0;
        float distSquared = searchRadius*searchRadius;
        stack[0] = 0;
        results.clear();
 
        while (stackPos > 0) {
            const NodeType &node = m_nodes[index];
            IndexType nextIndex;
 
            /* Recurse on inner nodes */
            if (!node.isLeaf()) {
                float distToPlane = p[node.getAxis()]
                    - node.getPosition()[node.getAxis()];
 
                bool searchBoth = distToPlane*distToPlane <= distSquared;
 
                if (distToPlane > 0) {
                    /* The search query is located on the right side of the split.
                       Search this side first. */
                    if (hasRightChild(index)) {
                        if (searchBoth)
                            stack[stackPos++] = node.getLeftIndex(index);
                        nextIndex = node.getRightIndex(index);
                    } else if (searchBoth) {
                        nextIndex = node.getLeftIndex(index);
                    } else {
                        nextIndex = stack[--stackPos];
                    }
                } else {
                    /* The search query is located on the left side of the split.
                       Search this side first. */
                    if (searchBoth && hasRightChild(index))
                        stack[stackPos++] = node.getRightIndex(index);
 
                    nextIndex = node.getLeftIndex(index);
                }
            } else {
                nextIndex = stack[--stackPos];
            }
 
            /* Check if the current point is within the query's search radius */
            const float pointDistSquared = (node.getPosition() - p).squaredNorm();
 
            if (pointDistSquared < distSquared) {
                ++found;
                results.push_back(index);
            }
 
            index = nextIndex;
        }
    }
 
    size_t nnSearch(const PointType &p, float &_sqrSearchRadius,
            size_t k, SearchResult *results) const {
        if (m_nodes.size() == 0)
            return 0;
 
        IndexType *stack = (IndexType *) alloca((m_depth+1) * sizeof(IndexType));
        IndexType index = 0, stackPos = 1;
        float sqrSearchRadius = _sqrSearchRadius;
        size_t resultCount = 0;
        bool isHeap = false;
        stack[0] = 0;
 
        while (stackPos > 0) {
            const NodeType &node = m_nodes[index];
            IndexType nextIndex;
 
            /* Recurse on inner nodes */
            if (!node.isLeaf()) {
                float distToPlane = p[node.getAxis()] - node.getPosition()[node.getAxis()];
 
                bool searchBoth = distToPlane*distToPlane <= sqrSearchRadius;
 
                if (distToPlane > 0) {
                    /* The search query is located on the right side of the split.
                       Search this side first. */
                    if (hasRightChild(index)) {
                        if (searchBoth)
                            stack[stackPos++] = node.getLeftIndex(index);
                        nextIndex = node.getRightIndex(index);
                    } else if (searchBoth) {
                        nextIndex = node.getLeftIndex(index);
                    } else {
                        nextIndex = stack[--stackPos];
                    }
                } else {
                    /* The search query is located on the left side of the split.
                       Search this side first. */
                    if (searchBoth && hasRightChild(index))
                        stack[stackPos++] = node.getRightIndex(index);
 
                    nextIndex = node.getLeftIndex(index);
                }
            } else {
                nextIndex = stack[--stackPos];
            }
 
            /* Check if the current point is within the query's search radius */
            const float pointDistSquared = (node.getPosition() - p).squaredNorm();
 
            if (pointDistSquared < sqrSearchRadius) {
                /* Switch to a max-heap when the available search
                   result space is exhausted */
                if (resultCount < k) {
                    /* There is still room, just add the point to
                       the search result list */
                    results[resultCount++] = SearchResult(pointDistSquared, index);
                } else {
                    auto comparator = [](SearchResult &a, SearchResult &b) -> bool {
                        return a.distSquared < b.distSquared;
                    };
 
                    if (!isHeap) {
                        /* Establish the max-heap property */
                        std::make_heap(results, results + resultCount, comparator);
                        isHeap = true;
                    }
                    SearchResult *end = results + resultCount + 1;
 
                    /* Add the new point, remove the one that is farthest away */
                    results[resultCount] = SearchResult(pointDistSquared, index);
                    std::push_heap(results, end, comparator);
                    std::pop_heap(results, end, comparator);
 
                    /* Reduce the search radius accordingly */
                    sqrSearchRadius = results[0].distSquared;
                }
            }
            index = nextIndex;
        }
        _sqrSearchRadius = sqrSearchRadius;
        return resultCount;
    }
 
    size_t nnSearch(const PointType &p, size_t k,
            SearchResult *results) const {
        float searchRadiusSqr = std::numeric_limits<float>::infinity();
        return nnSearch(p, searchRadiusSqr, k, results);
    }
 
protected:
    bool hasRightChild(IndexType index) const {
        return m_nodes[index].getRightIndex(index) != 0;
    }
 
    void build(size_t depth,
              typename std::vector<IndexType>::iterator base,
              typename std::vector<IndexType>::iterator rangeStart,
              typename std::vector<IndexType>::iterator rangeEnd) {
        if (rangeEnd <= rangeStart)
            throw NoriException("Internal error!");
 
        m_depth = std::max(depth, m_depth);
 
        IndexType count = (IndexType) (rangeEnd-rangeStart);
 
        if (count == 1) {
            /* Create a leaf node */
            m_nodes[*rangeStart].setLeaf(true);
            return;
        }
 
        int axis = 0;
        typename std::vector<IndexType>::iterator split;
 
        switch (m_heuristic) {
            case Balanced: {
                    /* Build a balanced tree */
                    split = rangeStart + count/2;
                    axis = m_bbox.getLargestAxis();
                };
                break;
 
            case SlidingMidpoint: {
                    /* Sliding midpoint rule: find a split that is close to the spatial median */
                    axis = m_bbox.getLargestAxis();
 
                    Scalar midpoint = (Scalar) 0.5f
                        * (m_bbox.max[axis]+m_bbox.min[axis]);
 
                    size_t nLT = std::count_if(rangeStart, rangeEnd,
                        [&](IndexType i) {
                            return m_nodes[i].getPosition()[axis] <= midpoint;
                        }
                    );
 
                    /* Re-adjust the split to pass through a nearby point */
                    split = rangeStart + nLT;
 
                    if (split == rangeStart)
                        ++split;
                    else if (split == rangeEnd)
                        --split;
                };
                break;
        }
 
        std::nth_element(rangeStart, split, rangeEnd,
            [&](IndexType i1, IndexType i2) {
                return m_nodes[i1].getPosition()[axis] < m_nodes[i2].getPosition()[axis];
            }
        );
 
        NodeType &splitNode = m_nodes[*split];
        splitNode.setAxis(axis);
        splitNode.setLeaf(false);
 
        if (split+1 != rangeEnd)
            splitNode.setRightIndex((IndexType) (rangeStart - base),
                    (IndexType) (split + 1 - base));
        else
            splitNode.setRightIndex((IndexType) (rangeStart - base), 0);
 
        splitNode.setLeftIndex((IndexType) (rangeStart - base),
                (IndexType) (rangeStart + 1 - base));
        std::iter_swap(rangeStart, split);
 
        /* Recursively build the children */
        Scalar temp = m_bbox.max[axis],
            splitPos = splitNode.getPosition()[axis];
        m_bbox.max[axis] = splitPos;
        build(depth+1, base, rangeStart+1, split+1);
        m_bbox.max[axis] = temp;
 
        if (split+1 != rangeEnd) {
            temp = m_bbox.min[axis];
            m_bbox.min[axis] = splitPos;
            build(depth+1, base, split+1, rangeEnd);
            m_bbox.min[axis] = temp;
        }
    }
protected:
    std::vector<NodeType> m_nodes;
    BoundingBoxType m_bbox;
    Heuristic m_heuristic;
    size_t m_depth;
};
 
template <typename DataType, typename IndexType> void permute_inplace(
        DataType *data, std::vector<IndexType> &perm) {
    for (size_t i=0; i<perm.size(); i++) {
        if (perm[i] != i) {
            /* The start of a new cycle has been found. Save
               the value at this position, since it will be
               overwritten */
            IndexType j = (IndexType) i;
            DataType curval = data[i];
 
            do {
                /* Shuffle backwards */
                IndexType k = perm[j];
                data[j] = data[k];
 
                /* Also fix the permutations on the way */
                perm[j] = j;
                j = k;
 
                /* Until the end of the cycle has been found */
            } while (perm[j] != i);
 
            /* Fix the final position with the saved value */
            data[j] = curval;
            perm[j] = j;
        }
    }
}
 
NORI_NAMESPACE_END
 
#endif /* __NORI_KDTREE_H */