New functions to save the built KDTree on disk

This commit is contained in:
Your Name 2011-02-24 18:04:42 -05:00
parent 707eaefbe1
commit f04e7dd41b
3 changed files with 176 additions and 19 deletions

View File

@ -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;
@ -15,28 +17,29 @@ typedef MyTree::Cell MyCell;
MyCell *findNearest(MyTree::coords& xc, MyCell *cells, uint32_t Ncells)
{
MyCell *near2 = 0;
double R2 = INFINITY;
for (int i = 0; i < Ncells; i++)
{
double d2 = 0;
for (int j = 0; j < ND; j++)
{
double delta = xc[j]-cells[i].coord[j];
d2 += delta*delta;
}
if (d2 < R2)
{
near2 = &cells[i];
R2 = d2;
}
}
return near2;
MyCell *near2 = 0;
double R2 = INFINITY;
for (int i = 0; i < Ncells; i++)
{
double d2 = 0;
for (int j = 0; j < ND; j++)
{
double delta = xc[j]-cells[i].coord[j];
d2 += delta*delta;
}
if (d2 < R2)
{
near2 = &cells[i];
near2->val = i;
R2 = d2;
}
}
return near2;
}
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();
@ -89,6 +93,31 @@ int main()
myTime = delta*1.0/CLOCKS_PER_SEC * 1000000.0;
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;
}

View File

@ -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,

View File

@ -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
};