Added support for types and masses

python/_cosmotool.pyx

@@ -80,6 +80,7 @@ cdef extern from "loadSimu.hpp" namespace "CosmoTool":
     float *Pos[3]
     float *Vel[3]
     int *type
+    float *Mass
 
     bool noAuto
 
@@ -87,6 +88,7 @@ cdef extern from "loadSimu.hpp" namespace "CosmoTool":
   cdef const int NEED_POSITION
   cdef const int NEED_VELOCITY
   cdef const int NEED_TYPE
+  cdef const int NEED_MASS
 
 cdef extern from "loadGadget.hpp" namespace "CosmoTool":
 
@@ -108,6 +110,9 @@ class PySimulationBase(object):
   def getIdentifiers(self):
     raise NotImplementedError("getIdentifiers is not implemented")
 
+  def getTypes(self):
+    raise NotImplementedError("getTypes is not implemented")
+
   def getOmega_M(self):
     raise NotImplementedError("getOmega_M is not implemented")
   
@@ -123,11 +128,16 @@ class PySimulationBase(object):
   def getBoxsize(self):
     raise NotImplementedError("getBoxsize is not implemented")
 
+  def getMasses(self):
+    raise NotImplementedError("getMasses is not implemented")
+
 cdef class Simulation:
 
   cdef list positions
   cdef list velocities
   cdef object identifiers
+  cdef object types
+  cdef object masses
 
   cdef SimuData *data
 
@@ -163,6 +173,14 @@ cdef class Simulation:
     def __get__(Simulation self):
         return self.identifiers
 
+  property types:
+    def __get__(Simulation self):
+        return self.types
+
+  property masses:
+    def __get__(Simulation self):
+        return self.masses
+        
   property numParticles:
     def __get__(Simulation self):
         return self.data.NumPart
@@ -179,7 +197,6 @@ cdef class Simulation:
 class _PySimulationAdaptor(PySimulationBase):
 
   def __init__(self,sim):
-  
     self.simu = sim
 
   def getBoxsize(self):
@@ -188,6 +205,9 @@ class _PySimulationAdaptor(PySimulationBase):
   def getPositions(self):
     return self.simu.positions
 
+  def getTypes(self):
+    return self.simu.types
+    
   def getVelocities(self):
     return self.simu.velocities
     
@@ -206,6 +226,9 @@ class _PySimulationAdaptor(PySimulationBase):
   def getOmega_Lambda(self):
     return self.simu.Omega_Lambda
 
+  def getMasses(self):
+    return self.simu.masses
+
 cdef class ArrayWrapper:
     cdef void* data_ptr
     cdef int size
@@ -263,6 +286,9 @@ cdef object wrap_float_array(float *p, np.uint64_t s):
 cdef object wrap_int64_array(np.int64_t* p, np.uint64_t s):
     return wrap_array(<void *>p, s, np.NPY_INT64)
 
+cdef object wrap_int_array(int* p, np.uint64_t s):
+    return wrap_array(<void *>p, s, np.NPY_INT)
+
 cdef object wrap_simudata(SimuData *data, int flags):
   cdef Simulation simu
 
@@ -282,9 +308,23 @@ cdef object wrap_simudata(SimuData *data, int flags):
      simu.identifiers = wrap_int64_array(data.Id, data.NumPart)
   else:
      simu.identifiers = None
+
+  if flags & NEED_TYPE:
+     simu.types = wrap_int_array(data.type, data.NumPart)
+  else:
+     simu.types = None
+
+  if flags & NEED_MASS:
+     simu.masses = wrap_float_array(data.Mass, data.NumPart)
+  else:
+     simu.masses = None
+
   return simu
 
-def loadGadget(str filename, int snapshot_id, bool loadPosition = True, bool loadVelocity = False, bool loadId = False):
+def loadGadget(str filename, int snapshot_id, bool loadPosition = True, bool loadVelocity = False, bool loadId = False, bool loadType = False, bool loadMass=False):
+  """loadGadget(filename, cpu_id, loadPosition=True, loadVelocity=False, loadId=False, loadType=False)
+It loads a gadget-1 snapshot and return a PySimulationBase object. If cpu_id is negative then the snapshot is considered not to be part of a set of snapshots written by different cpu. Otherwise the filename is modified to reflectt this cpu_id.
+"""
 
   cdef int flags
   cdef SimuData *data
@@ -297,6 +337,10 @@ def loadGadget(str filename, int snapshot_id, bool loadPosition = True, bool loa
      flags |= NEED_VELOCITY
   if loadId:
      flags |= NEED_GADGET_ID
+  if loadType:
+     flags |= NEED_TYPE
+  if loadMass:
+     flags |= NEED_MASS
 
   data = loadGadgetMulti(filename, snapshot_id, flags)
   if data == <SimuData*>0:

src/loadGadget.cpp

@@ -261,6 +261,42 @@ SimuData *CosmoTool::loadGadgetMulti(const char *fname, int id,
       f->skip(h.npart[k]*4);
   }
 
+  if (loadflags & NEED_MASS) {
+    try
+      {
+        long l = 0;
+	f->beginCheckpoint();
+        data->Mass = new float[NumPart];
+        for (int k = 0; k < 6; k++)
+	  {
+            if (h.mass[k] == 0) {
+              for(int n = 0; n < h.npart[k]; n++)
+	        {
+                  data->Mass[l++] = f->readReal32();
+                }
+            } else {
+              for(int n = 0; n < h.npart[k]; n++)
+	        {
+                  data->Mass[l++] = h.mass[k];
+                }
+            }
+          }
+        f->endCheckpoint();
+      }
+    catch (const InvalidUnformattedAccess& e)
+      {
+	cerr << "Invalid unformatted access while reading ID" << endl;
+	delete f;
+	delete data;
+	return 0;
+      }
+  } else {
+    f->skip(2*4);
+    for (int k = 0; k < 6; k++)
+      if (h.mass[k] == 0)
+        f->skip(h.npart[k]*4);
+  }
+
   delete f;
 
   return data;

src/loadSimu.hpp

@@ -45,6 +45,7 @@ namespace CosmoTool
   static const int NEED_POSITION = 2;
   static const int NEED_VELOCITY = 4;
   static const int NEED_TYPE = 8;
+  static const int NEED_MASS = 16;
 
   struct SimuParticle
   {
@@ -78,12 +79,13 @@ namespace CosmoTool
     long *Id;
     float *Pos[3];
     float *Vel[3];
+    float *Mass;
     int *type;
 
     AttributeMap attributes;
     
   public:
-    SimuData() : Id(0),NumPart(0),type(0),noAuto(false)  { Pos[0]=Pos[1]=Pos[2]=0; Vel[0]=Vel[1]=Vel[2]=0; }
+    SimuData() : Mass(0), Id(0),NumPart(0),type(0),noAuto(false)  { Pos[0]=Pos[1]=Pos[2]=0; Vel[0]=Vel[1]=Vel[2]=0; }
     ~SimuData() 
     {
       if (!noAuto)  {
@@ -98,6 +100,8 @@ namespace CosmoTool
 	  delete[] type;
         if (Id)
 	  delete[] Id;
+        if (Mass)
+          delete[] Mass;
       }
       for (AttributeMap::iterator i = attributes.begin();
 	   i != attributes.end();