UT/UT_KDTree.h

Go to the documentation of this file.

/*
 * PROPRIETARY INFORMATION.  This software is proprietary to
 * Side Effects Software Inc., and is not to be reproduced,
 * transmitted, or disclosed in any way without written permission.
 *
 * Produced by:
 *      Side Effects Software Inc
 *      477 Richmond Street West
 *      Toronto, Ontario
 *      Canada   M5V 3E7
 *      416-504-9876
 *
 * NAME:        UT_KDTree.h ( UT Library, C++)
 *
 * COMMENTS:    KD Tree Implementation
 *
 * The template class representing the nodes must have the following methods
 * defined:
 *      const UT_Vector3/UT_Vector4     getP(int idx) const;
 *      bool                            isValid(int idx) const;
 * It's ok for the area function to return 0.
 * The valid method is used when searching for valid nodes.
 * The splits are the dimension used which the node is split in.  Thus, the
 *    split direction can be encoded using a minimal 2 bits.
 *
 * TODO: We may want to convert this to a template class
 */

#ifndef __UT_KDTree__
#define __UT_KDTree__

#include "UT_API.h"
#include <SYS/SYS_Math.h>

#include "UT_IntArray.h"

class UT_Vector3;
class UT_BoundingBox;
class UT_FloatArray;
class ut_KDPQueue;
class ut_KDSplit;

// The maximum dimension must be less than 255 (unless you change the splits
// variable to a ushort)
#define UT_KD_MAXDIM    254

/// Lookup point information to be passed to the query functions
struct UT_KDQueryPt {
    UT_KDQueryPt(const float *p)
    { P = p; myData = 0; }
    UT_KDQueryPt(const float *p, const void *data)
    { P = p; myData = data; }

    const float         *P;
    const void          *myData;
};

class UT_API UT_KDTree
{
public:
    /// KD Tree balancing algorithms.  See setBalancer.
    typedef enum {
        UT_KD_MEDIAN,
        UT_KD_SAH
    } ut_KDBalancer;

public:
    UT_KDTree(int dim=3, size_t size=0)
    {
        myDim = dim;
        myList = 0;
        mySplits = 0;
        myMaxLeafNodes = 25;
        myRebalanceCount = 128;
        myBalancer = UT_KD_MEDIAN;
        myHasRadius = false;
        setEntries(size);
    }
    virtual ~UT_KDTree();

    /// setEntries instructs the KD Tree about the number of points
    /// that it should read back from its callbacks.
    /// Note that each call to setEntries has O(size) cost!  Thus,
    ///   for (i = 0; i < n; i++) tree.setEntries(i);
    /// is O(n^2).  There isn't a good reason for this (as one could instead
    /// just mark the tree as dirty until the next query, much like
    /// already happens with rebalancing) but it is the way things are so be
    // /careful.
    void                 setEntries(size_t size);
    size_t               getEntries() const     { return myFullEntries; }

    /// Sometimes, a tree might have new points added since it's construction.
    /// Rather than re-balancing the tree (which can be relatively expensive),
    /// we have allow N points to be added to the tree without re-balancing.
    /// To add entries to the tree, simply call "growEntries" with the number
    /// of points being added.  If the size of the unsorted pool exceeds the
    /// re-balance count, then the tree will be re-balanced.
    ///
    /// The rebalance count is automatically scaled as the tree is constructed,
    /// so it shouldn't be required to set the count (unless you know what
    /// you're doing).
    ///
    /// To force a re-build of the tree, simply call setEntries(), or flag the
    /// tree as dirty.
    void                 growEntries(size_t amount);
    size_t               getRebalanceCount() const
                         {
                             return myRebalanceCount;
                         }
    void                 setRebalanceCount(size_t count);

    /// Sets the KD-tree balancing algorithm.  The UT_KD_SAH (surface area
    /// heuristic) algorithm builds more efficient trees but they are
    /// slower to construct.  The default is UT_KD_MEDIAN.
    void                 setBalancer(ut_KDBalancer balance)
                         {
                             myBalancer = balance;
                         }
    /// Sets the maximum number of nodes to be stored in a leaf.  Smaller
    /// values will produce deeper and more memory-hungry trees.
    void                 setMaxLeafNodes(int max_leaf_nodes)
                         {
                             UT_ASSERT_P(max_leaf_nodes >= 3);
                             myMaxLeafNodes = SYSmax(3, max_leaf_nodes);
                         }

    /// Marks the tree as dirty.  Note, however, that this has
    /// O(size) time cost.
    void                 flagDirty()
                         {
                             // Re-build the lists
                             setEntries(myFullEntries);
                         }

    /// The compare function should compare the distance between two points
    /// (given by idx0 and idx1) in the dimension specified.  The dimension will
    /// be between 0 and 3.  For example,
    ///         comparePosition(...) {
    ///            fpreal delta = myP[idx0](dim) - myP[idx1](dim);
    ///            return delta < 0 ? -1 : (delta > 0 ? 1 : 0);
    ///         }
    virtual int                  comparePosition(int idx0,
                                                 int idx1, int dim) const = 0;
    /// Return the position associated with the given point
    virtual const fpreal        *getP(int idx) const = 0;

    /// If each point in the KD Tree has a radius associated with it, then this
    /// method should return true.  Adding a per-point radius can affect
    /// performance adversely.
    virtual bool                 pointsHaveRadius() const { return false; }
    /// Return the radius associated with the point in question
    virtual fpreal               getRadius(int /*idx*/) const { return 0; }

    /// Returns whether the given index should be considered in searches
    virtual bool         isValid(int /*idx*/) const
                         { return true; }
    virtual bool         isValid(int idx, const UT_KDQueryPt & /*pt*/) const
                         { return isValid(idx); }

    /// Find the closest node to the position specified.  This method
    /// returns -1 if no photon is found.
    /// NOTE THAT max_distance IS SQUARED DISTANCE.
    int                  findClosest(const UT_KDQueryPt &pt,
                                     fpreal max_distance);
    /// Returns the single closest point, using search parameters
    /// in a prebuilt queue.
    int                  findClosestQueue(const UT_KDQueryPt &pt,
                                     ut_KDPQueue &queue);

    /// Find the closest N (max_nodes) nodes to the position given.  Only points
    /// within the sphere defined by max_distance will be considered.
    /// NOTE THAT max_distance IS SQUARED DISTANCE.
    int                  findClosest(UT_IntArray &list,
                                     const UT_KDQueryPt &pt,
                                     fpreal max_distance,
                                     int    max_nodes,
                                     bool   sorted = true);

    /// Requires you to have pre-built your queue.
    int                  findClosestQueue(UT_IntArray &list,
                                     ut_KDPQueue &q,
                                     const  UT_KDQueryPt &pt);

    /// Find the closest N (max_nodes) nodes to the position given.  Only points
    /// within the sphere defined by max_distance will be considered.
    /// NOTE THAT max_distance IS SQUARED DISTANCE.
    int                  findClosest(UT_IntArray &list,
                                     UT_FloatArray &dist,
                                     const  UT_KDQueryPt &pt,
                                     fpreal max_distance,
                                     int    max_nodes,
                                     bool   sorted = true);

    int                  findNClosest(UT_IntArray &list,
                                     const        UT_KDQueryPt &pt,
                                     int          max_nodes)
                         {
                             return findClosest(list, pt, 1e37, max_nodes);
                         }
    int                  findAllClosest(UT_IntArray &list,
                                     const  UT_KDQueryPt &pt,
                                     fpreal max_distance)
                         {
                             return findClosest(list, pt, max_distance,
                                                myFullEntries, false);
                         }

    /// Tube queries (i.e. find the nearest N points to the line specified by
    /// the origin and direction)
    int                  findTubeClosest(const UT_Vector3 &orig,
                                         const UT_Vector3 &dir);
    int                  findTubeClosest(UT_IntArray &list,
                                         const UT_Vector3 &orig,
                                         const UT_Vector3 &dir,
                                         fpreal tube_radius,    // Radius
                                         int    max_nodes,
                                         bool   sorted=true);
    int                  findTubeClosest(UT_IntArray &list,
                                         UT_FloatArray &dlist,
                                         const UT_Vector3 &orig,
                                         const UT_Vector3 &dir,
                                         fpreal tube_radius,
                                         int max_nodes,
                                         bool   sorted=true);

    int                  findTubeNClosest(UT_IntArray &list,
                                         const UT_Vector3 &orig,
                                         const UT_Vector3 &dir,
                                         int max_nodes)
                         {
                             return findTubeClosest(list, orig, dir,
                                                    1e37, max_nodes);
                         }
    int                  findTubeAllClosest(UT_IntArray &list,
                                         const UT_Vector3 &orig,
                                         const UT_Vector3 &dir,
                                         fpreal radius)
                         {
                             return findTubeClosest(list, orig, dir, radius,
                                                    myFullEntries, false);
                         }

    const fpreal        *getBoxMin();
    const fpreal        *getBoxMax();

    /// This allows you to create a search queue so you can maintain
    /// it over time, avoiding the cost of reallocing for each search.
    /// NOTE THAT max_distance IS SQUARED DISTANCE.
    static ut_KDPQueue  *createQueue(fpreal max_dist,
                                            int max_nodes,
                                            bool sorted);
    static void          destroyQueue(ut_KDPQueue *q);

protected:
    int                  findClosest(ut_KDPQueue &list, const UT_KDQueryPt &pt);
    int                  findTubeClosest(ut_KDPQueue &list,
                                         const UT_Vector3 &orig,
                                         const UT_Vector3 &dir);

    void                 recurseFind(ut_KDPQueue &list, const UT_KDQueryPt &pt,
                                     ut_KDSplit *split, int lo, int hi) const;
    void                 recurseTubeFind(ut_KDPQueue &q,
                                        ut_KDSplit *split, int lo, int hi,
                                        const UT_Vector3 &orig,
                                        const UT_Vector3 &dir,
                                        UT_BoundingBox &box) const;
    bool                 isBalanced() const     { return myBalanced; }

    void                 computeBox(int start_index=0);
    bool                 isBoxClose(const fpreal *P, fpreal qd, fpreal r) const;

    /// Creates a KD tree from an unsorted set of points.
    /// Given a list of points, we:
    /// 1) Find the dimension in which there is most change.
    /// 2) Use the balancing algorithm to choose a splitting plane and
    ///    partition myList into left and right portions.
    /// 3) Store the splitting information in mySplits
    /// 4) Recurse on the two halves.
    /// This is repeated until myMaxLeafNodes is reached where we just
    /// leave it as a linear list.
    void                 balance();
    /// Balances a subset, returning the maximum radius of that subset.
    ut_KDSplit          *balanceSet(fpreal &radius,
                                    int *list, int entries, int axis);

    /// Stores the bounding box of all of the points in the KD tree, including
    /// their individual radii, in each of the possible dimensions.
    /// Note that during balanceSet these values are temporarily altered
    /// before recursion, so don't be too shocked.
    fpreal               myBMin[UT_KD_MAXDIM];
    fpreal               myBMax[UT_KD_MAXDIM];

    /// Nodes in a KD tree have two indices.  One is the index that
    /// they were added in.  This is the index used in getP, for example.
    /// The other is there location in the binary heap.  myList takes
    /// a binary heap location and returns the index useable for
    /// getP().
    int                 *myList;

    /// Split information stored as a tree structure.  This stores the
    /// midpoints, splitting offset, axis, and max radius.
    ut_KDSplit          *mySplits;

    /// This marks the number of entries that have been added to our tree
    size_t               myEntries;

    /// This marks the total number of entries in this data structure.
    /// The difference between this and myEntries is the unsorted
    /// chaff at the end.
    size_t               myFullEntries;
    size_t               myRebalanceCount;
    int                  myDim;
    int                  myMaxLeafNodes;
    bool                 myBalanced;
    bool                 myBoxComputed;

    /// A faster way to evaluate pointsHaveRadius() at runtime.
    bool                 myHasRadius;
    ut_KDBalancer        myBalancer;
};

#endif

---