From fc4459782978172c39f7214a1ef4409214c1c494 Mon Sep 17 00:00:00 2001 From: Guilhem Lavaux Date: Sat, 5 Jul 2014 22:05:04 +0200 Subject: [PATCH] Added support for types and masses --- python/_cosmotool.pyx | 48 +++++++++++++++++++++++++++++++++++++++++-- src/loadGadget.cpp | 36 ++++++++++++++++++++++++++++++++ src/loadSimu.hpp | 6 +++++- 3 files changed, 87 insertions(+), 3 deletions(-) diff --git a/python/_cosmotool.pyx b/python/_cosmotool.pyx index c11f743..bc032cc 100644 --- a/python/_cosmotool.pyx +++ b/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(p, s, np.NPY_INT64) +cdef object wrap_int_array(int* p, np.uint64_t s): + return wrap_array(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 == 0: diff --git a/src/loadGadget.cpp b/src/loadGadget.cpp index e121457..e166774 100644 --- a/src/loadGadget.cpp +++ b/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; diff --git a/src/loadSimu.hpp b/src/loadSimu.hpp index 7b51159..a80886c 100644 --- a/src/loadSimu.hpp +++ b/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();