Activated support for periodic box search

src/mykdtree.hpp

@@ -99,15 +99,16 @@ namespace CosmoTool {
   class RecursionMultipleInfo
   {
   public:
-    const typename KDDef<N,CType>::KDCoordinates& x;
+    typename KDDef<N,CType>::KDCoordinates x;
     BoundedQueue< KDCell<N,ValType,CType> *, typename KDDef<N,CType>::CoordType> queue;
     int traversed;
 
     RecursionMultipleInfo(const typename KDDef<N,CType>::KDCoordinates& rx,
 			  KDCell<N,ValType,CType> **cells,
 			  uint32_t numCells)
-      : x(rx), queue(cells, numCells, INFINITY),traversed(0)
+      : queue(cells, numCells, INFINITY),traversed(0)
     {
+      std::copy(rx, rx+N, x);
     }
   };
 
@@ -131,6 +132,12 @@ namespace CosmoTool {
     KDTree(Cell *cells, uint32_t Ncells);
     ~KDTree();
 
+    void setPeriodic(bool on, CoordType replicate)
+     {
+       periodic = on;
+       this->replicate = replicate;
+     }
+
     uint32_t getIntersection(const coords& x, CoordType r, 
 			     Cell **cells,
 			     uint32_t numCells)
@@ -185,6 +192,7 @@ namespace CosmoTool {
     Cell *base_cell;
 
     bool periodic;
+    CoordType replicate;
 
     Node *buildTree(Cell **cell0,
 		    uint32_t NumCells,

src/mykdtree.tcc

@@ -1,3 +1,4 @@
+#include "replicateGenerator.hpp"
 #include <cstring>
 #include <algorithm>
 #include <limits>
@@ -94,6 +95,19 @@ namespace CosmoTool {
     info.distances = 0;
 
     recursiveIntersectionCells<false>(info, root, 0);
+    if (periodic)
+      {
+        ReplicateGenerator<float, N> replicate(x);
+
+        do
+          {
+            coords x_new;
+            replicate.getPosition(info.x);
+            recursiveIntersectionCells<false>(info, root, 0);
+          }
+        while (replicate.next());
+      }
+
     return info.currentRank;
   }
 
@@ -115,6 +129,18 @@ namespace CosmoTool {
     info.distances = distances;
 
     recursiveIntersectionCells<false>(info, root, 0);
+    if (periodic)
+      {
+        ReplicateGenerator<float, N> replicate(x);
+
+        do
+          {
+            coords x_new;
+            replicate.getPosition(info.x);
+            recursiveIntersectionCells<false>(info, root, 0);
+          }
+        while (replicate.next());
+      }
     return info.currentRank;
   }
 
@@ -132,6 +158,19 @@ namespace CosmoTool {
     info.distances = 0;
 
     recursiveIntersectionCells<true>(info, root, 0);
+    if (periodic)
+      {
+        ReplicateGenerator<float, N> replicate(x);
+
+        do
+          {
+            coords x_new;
+            replicate.getPosition(info.x);
+            recursiveIntersectionCells<true>(info, root, 0);
+          }
+        while (replicate.next());
+      }
+
     return info.currentRank;
   }
 
@@ -369,15 +408,15 @@ namespace CosmoTool {
    recursiveNearest(root, 0, x, R2, best);
    if (periodic)
      {
-#if 0
-       ReplicateGenerator<N> replicate(x);
+       ReplicateGenerator<float, N> replicate(x);
 
        do
          {
-           recursiveNearest(root, 0, replicate.getPosition(), R2, best); 
+           coords x_new;
+           replicate.getPosition(x_new);
+           recursiveNearest(root, 0, x_new, R2, best); 
          }
        while (replicate.next());
-#endif
      }
    
    return best;
@@ -450,6 +489,18 @@ namespace CosmoTool {
      cells[i] = 0;
 
    recursiveMultipleNearest(info, root, 0);
+   if (periodic)
+    {
+       ReplicateGenerator<float, N> replicate(x);
+
+       do
+         {
+           coords x_new;
+           replicate.getPosition(info.x);
+           recursiveMultipleNearest(info, root, 0);
+         }
+       while (replicate.next());
+    }
 
    //   std::cout << "Traversed = " << info.traversed << std::endl;
  }
@@ -465,9 +516,19 @@ namespace CosmoTool {
      cells[i] = 0;
 
    recursiveMultipleNearest(info, root, 0);
-   memcpy(distances, info.queue.getPriorities(), sizeof(CoordType)*N2);
+   if (periodic)
+    {
+       ReplicateGenerator<float, N> replicate(x);
 
-   //   std::cout << "Traversed = " << info.traversed << std::endl;
+       do
+         {
+           coords x_new;
+           replicate.getPosition(info.x);
+           recursiveMultipleNearest(info, root, 0);
+         }
+       while (replicate.next());
+    }
+   memcpy(distances, info.queue.getPriorities(), sizeof(CoordType)*N2);
  }
 
 #ifdef __KD_TREE_SAVE_ON_DISK

src/replicateGenerator.hpp

@@ -1,6 +1,9 @@
 #ifndef __REPLICATE_GENERATOR_HPP
 #define __REPLICATE_GENERATOR_HPP
 
+#include <algorithm>
+#include "algo.hpp"
+
 namespace CosmoTool
 {
 
@@ -8,28 +11,50 @@ namespace CosmoTool
   class ReplicateGenerator
   {
   public:
-    ReplicateGenerator(const Coord x[N])
+    typedef Coord Coords[N];
+
+    ReplicateGenerator(const Coords& x)
     {
       face = 0; 
+      numFaces = spower<N,long>(3);
+      std::copy(x, x+N, x_base);
     }
 
     bool next()
     {
-      if (face == (2*N))
+      if (face == numFaces)
         return false;
 
       face++;
 
+      bool no_move = true;
+      int q_face = face;
+      for (int i = 0; i < N; i++)
+        {
+          int c_face;
+          c_face = q_face % 3;
+          q_face /= 3;
+          x_shifted[i] = x_base[i] + (c_face-1);
+          no_move = no_move && (c_face == 1);
+        }
+      if (no_move)
+        return next();
+      return true;
     }
 
-    Coord getPosition()
+    const Coord *getPosition()
     {
       return x_shifted;
     }
 
+    void getPosition(Coords& x_out)
+    {
+      std::copy(x_shifted, x_shifted+N, x_out);
+    }
+
   private:
-    Coord x_shifted[N];
-    int face;
+    Coord x_shifted[N], x_base[N];
+    long face, numFaces;
   };
 
 };