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
 
@@ -162,6 +172,14 @@ cdef class Simulation:
   property identifiers:
     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):
@@ -179,7 +197,6 @@ cdef class Simulation:
 class _PySimulationAdaptor(PySimulationBase):
 
   def __init__(self,sim):
-  
     self.simu = sim
 
   def getBoxsize(self):
@@ -187,6 +204,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: