cosmotool/src/mykdtree.hpp

#ifndef __HV_KDTREE_HPP
#define __HV_KDTREE_HPP

#include <cmath>
#include "config.hpp"
#include "bqueue.hpp"

namespace CosmoTool {

  template<int N> 
  struct NGBDef
  {
    typedef float CoordType;
    typedef float NGBCoordinates[N];
    static const int NumCubes = 1 << (N);
  };
  
  template<int N, typename ValType>
  struct NGBCell
  {
    bool active;
    ValType val;
    typename NGBDef<N>::NGBCoordinates coord;
  };

  class NotEnoughCells: public Exception
  {
  public:
    NotEnoughCells() : Exception() {}
    ~NotEnoughCells() throw () {}
  };

  template<int N, typename ValType>
  struct NGBTreeNode
  {
    NGBCell<N,ValType> *value;
    NGBTreeNode<N,ValType> *children[2];
    typename NGBDef<N>::NGBCoordinates minBound, maxBound;
  };

  template<int N, typename ValType>
  class RecursionInfoCells
  {
  public:

    typename NGBDef<N>::NGBCoordinates x;
    typename NGBDef<N>::CoordType r, r2;
    NGBCell<N, ValType> **cells;
    uint32_t currentRank;
    uint32_t numCells;
  };

  template<int N, typename ValType>
  class RecursionMultipleInfo
  {
  public:
    const typename NGBDef<N>::NGBCoordinates& x;
    BoundedQueue< NGBCell<N,ValType> *, float> queue;
    int traversed;

    RecursionMultipleInfo(const typename NGBDef<N>::NGBCoordinates& rx,
			  NGBCell<N,ValType> **cells,
			  uint32_t numCells)
      : x(rx), queue(cells, numCells, INFINITY),traversed(0)
    {      
    }
  };

  template<int N, typename ValType>
  class NGBTree
  {
  public:
    static const int NumCubes = NGBDef<N>::NumCubes;

  public:
    typedef typename NGBDef<N>::CoordType CoordType;
    typedef typename NGBDef<N>::NGBCoordinates coords;

    NGBTree(NGBCell<N,ValType> *cells, uint32_t Ncells);
    ~NGBTree();

    uint32_t getIntersection(const coords& x, CoordType r, 
			     NGBCell<N, ValType> **cells,
			     uint32_t numCells)
      throw (NotEnoughCells);

    NGBCell<N, ValType> *getNearestNeighbour(const coords& x);
    void getNearestNeighbours(const coords& x, uint32_t N,
			      NGBCell<N, ValType> **cells);

    NGBTreeNode<N,ValType> *getRoot() { return root; }

    void optimize();

    NGBTreeNode<N,ValType> *getAllNodes() { return nodes; }
    uint32_t getNumNodes() const { return lastNode; }

    uint32_t countActives() const;

  protected:
    NGBTreeNode<N, ValType> *nodes;
    uint32_t numNodes;
    uint32_t lastNode;

    NGBTreeNode<N, ValType> *root;
    NGBCell<N, ValType> **sortingHelper;

    NGBTreeNode<N, ValType> *buildTree(NGBCell<N,ValType> **cell0,
				       uint32_t N,
				       uint32_t depth,
				       coords minBound,
				       coords maxBound);
    
    void recursiveIntersectionCells(RecursionInfoCells<N,ValType>& info,
				    NGBTreeNode<N,ValType> *node,
				    int level)
      throw (NotEnoughCells);

    CoordType computeDistance(NGBCell<N,ValType> *cell, const coords& x);
    void recursiveNearest(NGBTreeNode<N, ValType> *node,
			  int level,
			  const coords& x,
			  CoordType& R2,
			  NGBCell<N,ValType>*& cell);
    void recursiveMultipleNearest(RecursionMultipleInfo<N,ValType>& info, NGBTreeNode<N,ValType> *node,
				  int level);    
  };

  template<int N, class T>
  uint32_t gatherActiveCells(NGBCell<N,T> **cells, uint32_t numCells);

};

#include "mykdtree.tcc"

#endif