Implemented neighbour searching in KDTree (+ BoundedQueue)

src/bqueue.hpp

@@ -0,0 +1,30 @@
+#ifndef __COSMO_QUEUE_HPP
+#define __COSMO_QUEUE_HPP
+
+#include <cmath>
+
+namespace CosmoTool {
+
+  template<typename T,typename QType = int>
+  class BoundedQueue
+  {
+  public:
+    BoundedQueue(int maxSize, QType defaultMax);
+    BoundedQueue(T *pQueue, int maxSize, QType defaultMax);
+    ~BoundedQueue();
+
+    void push(T a, QType v);
+    T *getQueue() { return m_queue; }
+    QType *getPriorities() { return priority; }
+    QType getMaxPriority() { return priority[maxQueueSize-1]; }
+  private:
+    int maxQueueSize;
+    T *m_queue;
+    QType *priority;
+    bool autoFree;
+  };
+};
+
+#include "bqueue.tcc"
+
+#endif

src/bqueue.tcc

@@ -0,0 +1,67 @@
+namespace CosmoTool
+{
+  
+  template<typename T, typename QType>
+  BoundedQueue<T,QType>::BoundedQueue(int maxSize, QType defaultVal)
+  {
+    maxQueueSize = maxSize;
+    m_queue = new T[maxSize];
+    priority = new QType[maxSize];
+    autoFree = true;
+
+    for (int i = 0; i < maxSize; i++)
+      priority[i] = defaultVal;
+  }
+ 
+  template<typename T, typename QType>
+  BoundedQueue<T,QType>::BoundedQueue(T *pQueue, int maxSize, QType defaultVal)
+  {
+    maxQueueSize = maxSize;
+    m_queue = pQueue;
+    priority = new QType[maxSize];
+    autoFree = false;
+    
+    for (int i = 0; i < maxSize; i++)
+      priority[i] = defaultVal;
+
+    
+  }
+  
+  template<typename T, typename QType>
+  BoundedQueue<T,QType>::~BoundedQueue()
+  {
+    if (autoFree)
+      delete[] m_queue;
+    delete[] priority;
+  }
+
+
+  template<typename T, typename QType>
+  void BoundedQueue<T,QType>::push(T a, QType v)
+  {
+    if (v > priority[maxQueueSize-1])
+      return;
+
+    int i;
+    for (i = maxQueueSize-2; i >= 0; i--)
+      {
+	if (v > priority[i])
+	  {
+	    priority[i+1] = v;
+	    m_queue[i+1] = a;
+	    return;
+	  }
+	else
+	  {
+	    priority[i+1] = priority[i];
+	    m_queue[i+1] = m_queue[i];
+	  }
+      }
+    if (i < 0)
+      {
+	priority[0] = v;
+	m_queue[0] = a;
+      }
+  }
+   
+};

src/mykdtree.hpp

@@ -3,6 +3,7 @@
 
 #include <cmath>
 #include "config.hpp"
+#include "bqueue.hpp"
 
 namespace CosmoTool {
 
@@ -49,6 +50,22 @@ namespace CosmoTool {
     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
   {
@@ -58,7 +75,6 @@ namespace CosmoTool {
   public:
     typedef typename NGBDef<N>::CoordType CoordType;
     typedef typename NGBDef<N>::NGBCoordinates coords;
-    typedef typename CosmoQueue<NGBCell<N,ValType>*> NGBQueue;
 
     NGBTree(NGBCell<N,ValType> *cells, uint32_t Ncells);
     ~NGBTree();
@@ -106,10 +122,8 @@ namespace CosmoTool {
 			  const coords& x,
 			  CoordType& R2,
 			  NGBCell<N,ValType>*& cell);
-    void recursiveMultiNearest(NGBTreeNode<N, ValType> *node,
-			       int level,
-			       const coords& x,
-			       NGBQueue& q);
+    void recursiveMultipleNearest(RecursionMultipleInfo<N,ValType>& info, NGBTreeNode<N,ValType> *node,
+				  int level);    
   };
 
   template<int N, class T>

src/mykdtree.tcc

@@ -298,10 +298,75 @@ namespace CosmoTool {
    return best;
  }
 
+  template<int N, typename ValType>
+  void
+  NGBTree<N,ValType>::recursiveMultipleNearest(RecursionMultipleInfo<N,ValType>& info, NGBTreeNode<N,ValType> *node,
+					       int level)
+  {
+    CoordType d2 = 0;
+    int axis = level % N;
+    NGBTreeNode<N,ValType> *other, *go;
+
+    if (info.x[axis] < node->value->coord[axis])
+      {
+	// The best is potentially in 0.
+	go = node->children[0];
+	other = node->children[1];
+      }
+    else
+      {
+	// If not it is in 1.
+	go = node->children[1];
+	other = node->children[0];
+	if (go == 0)
+	  {
+	    go = other;
+	    other = 0;
+	  }
+      }
+
+    if (go != 0)
+      {
+	recursiveMultipleNearest(info, go, level+1);
+      }
+    
+    // Check if current node is not the nearest
+    CoordType thisR2 = 
+      computeDistance(node->value, info.x);
+    info.queue.push(node->value, thisR2);
+    info.traversed++;
+    if (go == 0)
+      return;
+    
+    // Now we found the best. We check whether the hypersphere
+    // intersect the hyperplane of the other branch
+
+    CoordType delta1;
+
+    delta1 = info.x[axis]-node->value->coord[axis];
+    if (delta1*delta1 < info.queue.getMaxPriority())
+      {
+	// The hypersphere intersects the hyperplane. Try the
+	// other branch
+	if (other != 0)
+	  {
+	    recursiveMultipleNearest(info, other, level+1);
+	  }
+      }    
+  }
+
  template<int N, typename ValType>
  void NGBTree<N,ValType>::getNearestNeighbours(const coords& x, uint32_t N2,
 					       NGBCell<N, ValType> **cells)
  {
+   RecursionMultipleInfo<N,ValType> info(x, cells, N2);
+   
+   for (int i = 0; i < N2; i++)
+     cells[i] = 0;
+
+   recursiveMultipleNearest(info, root, 0);
+
+   std::cout << "Traversed = " << info.traversed << std::endl;
  }
 
 };