New octtree tool

lib/Makefile

@@ -1,5 +1,5 @@
 SHLIBS= libCosmoTool.so
-SOURCES= loadRamses.cpp yorick.cpp miniargs.cpp  fortran.cpp interpolate.cpp load_data.cpp powerSpectrum.cpp
+SOURCES= loadRamses.cpp yorick.cpp miniargs.cpp  fortran.cpp interpolate.cpp load_data.cpp powerSpectrum.cpp octTree.cpp
 LIBS= -lnetcdf_c++ -lnetcdf -lgsl -lgslcblas -lm
 
 include config.mk
@@ -8,7 +8,7 @@ VPATH=../src
 
 all: $(SHLIBS)
 
-libCosmoTool.so: loadRamses.o yorick.o miniargs.o fortran.o interpolate.o load_data.o powerSpectrum.o
+libCosmoTool.so: loadRamses.o yorick.o miniargs.o fortran.o interpolate.o load_data.o powerSpectrum.o octTree.o
 
 depend: $(SOURCES)
 	@echo "[DEPENDS] $^"

src/octTree.cpp

@@ -0,0 +1,132 @@
+#include <iostream>
+#include <cmath>
+#include <cassert>
+#include "config.hpp"
+#include "octTree.hpp"
+
+using namespace std;
+using namespace CosmoTool;
+
+//#define VERBOSE
+
+OctTree::OctTree(const FCoordinates *particles, octPtr numParticles, 
+		 uint32_t maxMeanTreeDepth, uint32_t maxAbsoluteDepth,
+		 uint32_t threshold)
+{
+  cout << "MeanTree=" << maxMeanTreeDepth << endl;
+  numCells = pow(8, maxMeanTreeDepth);
+  assert(numCells < invalidOctCell);
+  //#ifdef VERBOSE
+  cerr << "Allocating " << numCells << " octtree cells" << endl;
+  //#endif
+
+  cells = new OctCell[numCells];
+  Lbox = (float)(octCoordTypeNorm+1);
+
+  cells[0].numberLeaves = 0;
+  for (int i = 0; i < 8; i++)
+    cells[0].children[i] = emptyOctCell;
+
+  lastNode = 1;
+  this->particles = particles;
+  this->numParticles = numParticles;
+  buildTree(maxAbsoluteDepth);
+  //#ifdef VERBOSE
+  cerr << "Used " << lastNode << " cells" << endl;
+  //#endif
+}
+
+OctTree::~OctTree()
+{
+  delete cells;
+}
+
+void OctTree::buildTree(uint32_t maxAbsoluteDepth)
+{
+  for (octPtr i = 0; i < numParticles; i++)
+    {
+      OctCoords rootCenter = { octCoordCenter, octCoordCenter, octCoordCenter };
+      insertParticle(0, // root node
+		     rootCenter,
+		     octCoordCenter,
+		     i,
+		     maxAbsoluteDepth);
+    }
+}
+
+
+void OctTree::insertParticle(octPtr node, 
+			     const OctCoords& icoord,
+			     octCoordType halfNodeLength,
+			     octPtr particleId,
+			     uint32_t maxAbsoluteDepth)
+{
+  
+#ifdef VERBOSE
+  cout << "Entering " << node << " (" << icoord[0] << "," << icoord[1] << "," << icoord[2] << ")" << endl;
+#endif
+  int octPos = 0;
+  int ipos[3] = { 0,0,0};
+  octPtr newNode;
+  OctCoords newCoord;
+  
+  cells[node].numberLeaves++;
+  if (maxAbsoluteDepth == 0)
+    {
+      // All children must be invalid.
+      for (int i = 0 ; i < 8; i++)
+	cells[node].children[i] = invalidOctCell;
+      
+      return;
+    }
+  
+  for (int j = 0; j < 3; j++)
+    if ((octPtr)(particles[particleId][j]*Lbox) > icoord[j])
+      {
+	octPos |= (1 << j);
+	ipos[j] = 1;
+      }
+  
+  if (cells[node].children[octPos] == emptyOctCell)
+    {
+      // Put the particle there.
+      cells[node].children[octPos] = particleId | octParticleMarker;
+      return;
+    }
+  
+  // If it is a node, explores it.
+  if (!(cells[node].children[octPos] & octParticleMarker))
+    {
+      assert(halfNodeLength >= 2);
+      // Compute coordinates
+      for (int j = 0; j < 3; j++)
+	newCoord[j] = icoord[j]+(2*ipos[j]-1)*halfNodeLength/2;
+      insertParticle(cells[node].children[octPos], newCoord, halfNodeLength/2, 
+		     particleId, maxAbsoluteDepth-1);
+      return;
+    }
+  
+  // We have a particle there. 
+  // Make a new node and insert the old particle into this node.
+  // Insert the new particle into the node also
+  // Finally put the node in place
+  
+  newNode = lastNode++;
+  assert(lastNode != numCells);
+  
+  for (int j = 0; j < 8; j++)
+    cells[newNode].children[j] = emptyOctCell;
+  cells[newNode].numberLeaves = 0;
+  
+  // Compute coordinates
+  for (int j = 0; j < 3; j++)
+    newCoord[j] = icoord[j]+(2*ipos[j]-1)*halfNodeLength/2;
+  
+  octPtr oldPartId = cells[node].children[octPos] & octParticleMask;
+  
+  insertParticle(newNode, newCoord, halfNodeLength/2,
+		 oldPartId, maxAbsoluteDepth-1);
+  insertParticle(newNode, newCoord, halfNodeLength/2,
+		 particleId, maxAbsoluteDepth-1);  
+  cells[node].children[octPos] = newNode;
+}

src/octTree.hpp

@@ -0,0 +1,113 @@
+#ifndef __COSMOTOOL_AMR_HPP
+#define __COSMOTOOL_AMR_HPP
+
+#include "config.hpp"
+
+namespace CosmoTool
+{
+
+  typedef uint32_t octPtr;
+  typedef uint16_t octCoordType;
+  static const uint16_t octCoordTypeNorm = 0xffff;
+  static const uint16_t octCoordCenter = 0x8000;
+
+  // This is also the root cell, but this one
+  // is never referenced (really ??).
+  static const octPtr invalidOctCell = 0xffffffff;
+  static const octPtr emptyOctCell = 0;
+  static const octPtr octParticleMarker = 0x80000000;
+  static const octPtr octParticleMask = 0x7fffffff;
+
+  typedef octCoordType OctCoords[3];
+
+  struct OctCell
+  {
+    octPtr numberLeaves;
+    octPtr children[8];
+  };
+
+  class OctTree
+  {
+  public:
+    //Coordinates of particles must be in the [0:1] range
+    OctTree(const FCoordinates *particles, octPtr numParticles,
+	    uint32_t maxTreeDepth,  uint32_t maxAbsoluteDepth,
+	    uint32_t threshold = 1);
+    ~OctTree();
+
+    void buildTree(uint32_t maxAbsoluteDepth);
+    void insertParticle(octPtr node, 
+			const OctCoords& icoord,
+			octCoordType halfNodeLength,
+			octPtr particleId,
+			uint32_t maxAbsoluteDepth);
+
+
+    octPtr getNumberLeaves() const {
+      return cells[0].numberLeaves;
+    }
+
+    template<typename FunT>
+    void walkTree(FunT f)
+    {
+      OctCoords rootCenter = { octCoordCenter, octCoordCenter, octCoordCenter };
+
+      walkTreeElements(f, 0, rootCenter, octCoordCenter);
+    }
+
+  protected:
+    const FCoordinates *particles;
+    octPtr numParticles;
+    OctCell *cells;
+    float Lbox;
+    octPtr lastNode;
+    octPtr numCells;  
+
+
+    template<typename FunT>
+    void walkTreeElements(FunT f, octPtr node, 
+			  const OctCoords& icoord,
+			  octCoordType halfNodeLength)
+    {
+      OctCoords newCoord;
+      FCoordinates realCenter;
+
+      for (int j = 0; j < 3; j++)
+	realCenter[j] = icoord[j]/(2.*octCoordCenter);
+      
+      f(realCenter, cells[node].numberLeaves, halfNodeLength/(float)octCoordCenter,
+	cells[node].children[0] == invalidOctCell, // True if this is a meta-node
+	false);
+      
+      for (int i = 0; i < 8; i++)
+	{
+	  octPtr newNode = cells[node].children[i];
+	  int ipos[3] = { (i&1), (i&2)>>1, (i&4)>>2 };
+	  
+	  if (newNode == emptyOctCell || newNode == invalidOctCell)
+	    continue;
+	  	  
+	  for (int j = 0; j < 3; j++)
+	    newCoord[j] = icoord[j]+(2*ipos[j]-1)*halfNodeLength/2;
+	  
+	  if (newNode & octParticleMarker)
+	    {
+	      for (int j = 0; j < 3; j++)
+		realCenter[j] = newCoord[j]/(2.*octCoordCenter);
+
+	      f(realCenter, 
+		1, halfNodeLength/(2.*octCoordCenter),
+		false, true);
+	      continue;
+	    }
+	  walkTreeElements(f, cells[node].children[i], newCoord, halfNodeLength/2);
+	}
+      
+    }
+
+  };
+  
+};
+
+
+#endif