New functions to save the built KDTree on disk

sample/testkd.cpp

@@ -1,10 +1,12 @@
+#define __KD_TREE_SAVE_ON_DISK
 #include <ctime>
 #include <cassert>
 #include <cstdlib>
 #include <iostream>
+#include <fstream>
 #include "mykdtree.hpp"
 
-#define NTRY 5000000
+#define NTRY 50000
 #define ND 2
 
 using namespace std;
@@ -28,6 +30,7 @@ MyCell *findNearest(MyTree::coords& xc, MyCell *cells, uint32_t Ncells)
       if (d2 < R2)
 	{
 	  near2 = &cells[i];
+	  near2->val = i;
 	  R2 = d2;
 	}
     }
@@ -36,7 +39,7 @@ MyCell *findNearest(MyTree::coords& xc, MyCell *cells, uint32_t Ncells)
 
 int main()
 {
-  uint32_t Ncells = 100000;
+  uint32_t Ncells = 3000;
   MyCell *cells = new MyCell[Ncells];
   
   for (int i = 0; i < Ncells; i++)
@@ -72,6 +75,7 @@ int main()
   
   for (int k = 0; k < NTRY; k++) {
     MyCell *near = tree.getNearestNeighbour(xc[k]);
+    near->val = 0;
   }
 
   clock_t endTimer = clock();
@@ -90,5 +94,30 @@ int main()
 
   cout << "Direct search/sec = " << myTime/NTRY << " us" << endl;
 
+  {
+    ofstream f("kd.dat");
+    tree.saveTree(f);
+  }
+
+  cout << "Trying to reload the tree" << endl;
+  {
+    ifstream f("kd.dat");
+    MyTree tree2(f, cells, Ncells);
+    cout << "Check timing..." << endl;
+    // Check timing
+    clock_t startTimer = clock();
+    
+    for (int k = 0; k < NTRY; k++) {
+      MyCell *near = tree.getNearestNeighbour(xc[k]);
+    }
+    
+    clock_t endTimer = clock();
+    clock_t delta = endTimer-startTimer;
+    double myTime = delta*1.0/CLOCKS_PER_SEC * 1000000.0;
+    
+    cout << "KDTree search/sec = " << myTime/NTRY << " us" << endl;
+  
+  }
+  
   return 0;
 }

src/mykdtree.hpp

@@ -29,6 +29,13 @@ namespace CosmoTool {
     ~NotEnoughCells() throw () {}
   };
 
+  class InvalidOnDiskKDTree : public Exception
+  {
+  public:
+    InvalidOnDiskKDTree() : Exception() {}
+    ~InvalidOnDiskKDTree() throw () {}
+  };
+
   template<int N, typename ValType, typename CType = ComputePrecision>
   struct KDTreeNode
   {
@@ -118,7 +125,12 @@ namespace CosmoTool {
     }
 #endif
 
+#ifdef __KD_TREE_SAVE_ON_DISK
+    KDTree(std::istream& i, Cell *cells, uint32_t Ncells)
+      throw (InvalidOnDiskKDTree);
 
+    void saveTree(std::ostream& o) const;
+#endif
   protected:
     Node *nodes;
     uint32_t numNodes;
@@ -126,6 +138,7 @@ namespace CosmoTool {
 
     Node *root;
     Cell **sortingHelper;
+    Cell *base_cell;
 
     Node *buildTree(Cell **cell0,
 		    uint32_t NumCells,

src/mykdtree.tcc

@@ -32,6 +32,7 @@ namespace CosmoTool {
   KDTree<N,ValType,CType>::KDTree(Cell *cells, uint32_t Ncells)
   {
     
+    base_cell = cells;
     numNodes = Ncells;
     nodes = new Node[numNodes];
 
@@ -429,5 +430,119 @@ namespace CosmoTool {
    //   std::cout << "Traversed = " << info.traversed << std::endl;
  }
 
+#ifdef __KD_TREE_SAVE_ON_DISK
+#define KDTREE_DISK_SIGNATURE "KDTREE"
+#define KDTREE_DISK_SIGNATURE_LEN 7
+  
+  template<int N, typename CType>
+  struct KDTreeOnDisk
+  {
+    long cell_id;
+    long children_node[2];
+    typename KDDef<N, CType>::KDCoordinates minBound, maxBound;
+  };  
+
+  struct KDTreeHeader
+  {
+    char id[KDTREE_DISK_SIGNATURE_LEN];
+    long nodesUsed, numCells;
+  };
+
+  template<int N, typename ValType, typename CType>
+  void KDTree<N,ValType,CType>::saveTree(std::ostream& o) const
+  {
+    KDTreeHeader h;
+
+    strncpy(h.id, KDTREE_DISK_SIGNATURE, KDTREE_DISK_SIGNATURE_LEN);
+    h.nodesUsed = lastNode;
+    h.numCells = numNodes;
+    o.write((char*)&h, sizeof(h));
+
+    for (long i = 0; i < lastNode; i++)
+      {
+	KDTreeOnDisk<N,CType> node_on_disk;
+
+	node_on_disk.cell_id = nodes[i].value - base_cell;
+	if (nodes[i].children[0] == 0)
+	  node_on_disk.children_node[0] = -1;
+	else
+	  node_on_disk.children_node[0] = nodes[i].children[0] - nodes;
+	if (nodes[i].children[1] == 0)
+	  node_on_disk.children_node[1] = -1;
+	else
+	  node_on_disk.children_node[1] = nodes[i].children[1] - nodes;
+	memcpy(node_on_disk.minBound, nodes[i].minBound, sizeof(coords));
+	memcpy(node_on_disk.maxBound, nodes[i].maxBound, sizeof(coords));
+
+	o.write((char *)&node_on_disk, sizeof(node_on_disk));
+      }
+  }
+
+  template<int N, typename ValType, typename CType>
+  KDTree<N,ValType,CType>::KDTree(std::istream& in, Cell *cells, uint32_t Ncells)
+    throw (InvalidOnDiskKDTree)
+  {
+    KDTreeHeader h;
+
+    if (!in)
+      throw InvalidOnDiskKDTree();
+
+    in.read((char *)&h, sizeof(h));
+    if (!in || strncmp(h.id, KDTREE_DISK_SIGNATURE, KDTREE_DISK_SIGNATURE_LEN) != 0)
+      throw InvalidOnDiskKDTree();
+
+    if (h.numCells != Ncells || h.nodesUsed < 0)
+      throw InvalidOnDiskKDTree();
+
+    base_cell = cells;
+    nodes = new Node[h.nodesUsed];
+    lastNode = h.nodesUsed;
+    numNodes = Ncells;
+
+    for (long i = 0; i < lastNode; i++)
+      {
+	KDTreeOnDisk<N,CType> node_on_disk;
+
+	in.read((char *)&node_on_disk, sizeof(node_on_disk));
+
+	if (!in || node_on_disk.cell_id > numNodes || node_on_disk.cell_id < 0 ||
+	    node_on_disk.children_node[0] > lastNode || node_on_disk.children_node[0] < -1 ||
+	    node_on_disk.children_node[1] > lastNode || node_on_disk.children_node[1] < -1)
+	  {
+	    delete[] nodes;
+	    throw InvalidOnDiskKDTree();
+	  }
+
+	nodes[i].value = base_cell + node_on_disk.cell_id;
+	if (node_on_disk.children_node[0] == -1)
+	  nodes[i].children[0] = 0;
+	else
+	  nodes[i].children[0] = nodes + node_on_disk.children_node[0];
+
+	if (node_on_disk.children_node[1] == -1)
+	  nodes[i].children[1] = 0;
+	else
+	  nodes[i].children[1] = nodes + node_on_disk.children_node[1];
+
+	memcpy(nodes[i].minBound, node_on_disk.minBound, sizeof(coords));
+	memcpy(nodes[i].maxBound, node_on_disk.maxBound, sizeof(coords));
+
+	int c;
+	for (c = 0; c < N; c++)
+	  if (nodes[i].value->coord[c] < nodes[i].minBound[c] ||
+	      nodes[i].value->coord[c] > nodes[i].maxBound[c])
+	    break;
+	if (c != N)
+	  {
+	    delete[] nodes;
+	    throw InvalidOnDiskKDTree();
+	  }
+      }
+    
+    sortingHelper = new Cell *[Ncells];
+    for (uint32_t i = 0; i < Ncells; i++)
+      sortingHelper[i] = &cells[i];    
+  }
+#endif
 
 };