from libcpp cimport bool from libcpp cimport string as cppstring import numpy as np cimport numpy as np from cpython cimport PyObject, Py_INCREF np.import_array() cdef extern from "loadSimu.hpp" namespace "CosmoTool": cdef cppclass SimuData: np.float_t BoxSize np.float_t time np.float_t Hubble np.float_t Omega_M np.float_t Omega_Lambda np.int64_t TotalNumPart np.int64_t NumPart np.int64_t *Id float *Pos[3] float *Vel[3] int *type bool noAuto cdef const int NEED_GADGET_ID cdef const int NEED_POSITION cdef const int NEED_VELOCITY cdef const int NEED_TYPE cdef extern from "loadGadget.hpp" namespace "CosmoTool": SimuData *loadGadgetMulti(const char *fname, int id, int flags) except + void cxx_writeGadget "CosmoTool::writeGadget" (const char * s, SimuData *data) except + cdef extern from "loadRamses.hpp" namespace "CosmoTool": SimuData *loadRamsesSimu(const char *basename, int id, int cpuid, bool dp, int flags) except + class PySimulationBase(object): def getPositions(self): raise NotImplementedError("getPositions is not implemented") def getVelocities(self): raise NotImplementedError("getVelocities is not implemented") def getIdentifiers(self): raise NotImplementedError("getIdentifiers is not implemented") def getOmega_M(self): raise NotImplementedError("getOmega_M is not implemented") def getOmega_Lambda(self): raise NotImplementedError("getOmega_Lambda is not implemented") def getTime(self): raise NotImplementedError("getTime is not implemented") def getHubble(self): raise NotImplementedError("getHubble is not implemented") def getBoxsize(self): raise NotImplementedError("getBoxsize is not implemented") cdef class Simulation: cdef list positions cdef list velocities cdef object identifiers cdef SimuData *data property BoxSize: def __get__(Simulation self): return self.data.BoxSize property time: def __get__(Simulation self): return self.data.time property Hubble: def __get__(Simulation self): return self.data.Hubble property Omega_M: def __get__(Simulation self): return self.data.Omega_M property Omega_Lambda: def __get__(Simulation self): return self.data.Omega_Lambda property positions: def __get__(Simulation self): return self.positions property velocities: def __get__(Simulation self): return self.velocities property identifiers: def __get__(Simulation self): return self.identifiers property numParticles: def __get__(Simulation self): return self.data.NumPart def __cinit__(Simulation self): self.data = 0 def __dealloc__(Simulation self): if self.data != 0: print("Clearing simulation data") del self.data class _PySimulationAdaptor(PySimulationBase): def __init__(self,sim): self.simu = sim def getPositions(self): return self.simu.positions def getVelocities(self): return self.simu.velocities def getIdentifiers(self): return self.simu.identifiers def getTime(self): return self.simu.time def getHubble(self): return self.simu.Hubble def getOmega_M(self): return self.simu.Omega_M def getOmega_Lambda(self): return self.simu.Omega_Lambda cdef class ArrayWrapper: cdef void* data_ptr cdef int size cdef int type_array cdef set_data(self, int size, int type_array, void* data_ptr): """ Set the data of the array This cannot be done in the constructor as it must recieve C-level arguments. Parameters: ----------- size: int Length of the array. data_ptr: void* Pointer to the data """ self.data_ptr = data_ptr self.size = size self.type_array = type_array def __array__(self): """ Here we use the __array__ method, that is called when numpy tries to get an array from the object.""" cdef np.npy_intp shape[1] shape[0] = self.size # Create a 1D array, of length 'size' ndarray = np.PyArray_SimpleNewFromData(1, shape, self.type_array, self.data_ptr) return ndarray def __dealloc__(self): """ Frees the array. This is called by Python when all the references to the object are gone. """ pass cdef object wrap_array(void *p, np.uint64_t s, int typ): cdef np.ndarray ndarray cdef ArrayWrapper wrapper wrapper = ArrayWrapper() wrapper.set_data(s, typ, p) ndarray = np.array(wrapper, copy=False) ndarray.base = wrapper Py_INCREF(wrapper) return ndarray cdef object wrap_float_array(float *p, np.uint64_t s): return wrap_array(p, s, np.NPY_FLOAT32) cdef object wrap_int64_array(np.int64_t* p, np.uint64_t s): return wrap_array(p, s, np.NPY_INT64) cdef object wrap_simudata(SimuData *data, int flags): cdef Simulation simu simu = Simulation() simu.data = data if flags & NEED_POSITION: simu.positions = [wrap_float_array(data.Pos[i], data.NumPart) for i in xrange(3)] else: simu.positions = None if flags & NEED_VELOCITY: simu.velocities = [wrap_float_array(data.Vel[i], data.NumPart) for i in xrange(3)] else: simu.velocities = None if flags & NEED_GADGET_ID: simu.identifiers = wrap_int64_array(data.Id, data.NumPart) else: simu.identifiers = None return simu def loadGadget(str filename, int snapshot_id, bool loadPosition = True, bool loadVelocity = False, bool loadId = False): cdef int flags cdef SimuData *data cdef Simulation simu flags = 0 if loadPosition: flags |= NEED_POSITION if loadVelocity: flags |= NEED_VELOCITY if loadId: flags |= NEED_GADGET_ID data = loadGadgetMulti(filename, snapshot_id, flags) if data == 0: return None return _PySimulationAdaptor(wrap_simudata(data, flags)) def writeGadget(str filename, object simulation): cdef SimuData simdata cdef np.ndarray[np.float32_t, ndim=1] pos, vel cdef np.ndarray[np.int64_t, ndim=1] ids cdef np.int64_t NumPart cdef int j if not isinstance(simulation,PySimulationBase): raise TypeError("Second argument must be of type SimulationBase") NumPart = simulation.positions[0].size simdata.noAuto = True for j in xrange(3): pos = simulation.getPositions()[j] vel = simulation.getVelocities()[j] if pos.size != NumPart or vel.size != NumPart: raise ValueError("Invalid number of particles") simdata.Pos[j] = pos.data simdata.Vel[j] = vel.data ids = simulation.getIdentifiers() simdata.Id = ids.data simdata.BoxSize = simulation.getBoxsize() simdata.time = simulation.getTime() simdata.Hubble = simulation.getHubble() simdata.Omega_M = simulation.getOmega_M() simdata.Omega_Lambda = simulation.getOmega_Lambda() simdata.TotalNumPart = NumPart simdata.NumPart = NumPart cxx_writeGadget(filename, &simdata) def loadRamses(str basepath, int snapshot_id, int cpu_id, bool doublePrecision = False, bool loadPosition = True, bool loadVelocity = False): """ loadRamses(basepath, snapshot_id, cpu_id, doublePrecision=False, loadPosition=True, loadVelocity=False) Loads the indicated snapshot based on the cpu id, snapshot id and basepath. It is important to specify the correct precision in doublePrecision. """ cdef int flags cdef SimuData *data cdef Simulation simu flags = 0 if loadPosition: flags |= NEED_POSITION if loadVelocity: flags |= NEED_VELOCITY data = loadRamsesSimu(basepath, snapshot_id, cpu_id, doublePrecision, flags) if data == 0: return None return _PySimulationAdaptor(wrap_simudata(data, flags))