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 "cosmopower.hpp" namespace "CosmoTool": cdef enum CosmoFunction "CosmoTool::CosmoPower::CosmoFunction": POWER_EFSTATHIOU "CosmoTool::CosmoPower::POWER_EFSTATHIOU", HU_WIGGLES "CosmoTool::CosmoPower::HU_WIGGLES", HU_BARYON "CosmoTool::CosmoPower::HU_BARYON", OLD_POWERSPECTRUM, POWER_BARDEEN "CosmoTool::CosmoPower::POWER_BARDEEN", POWER_SUGIYAMA "CosmoTool::CosmoPower::POWER_SUGIYAMA", POWER_BDM, POWER_TEST cdef cppclass CosmoPower: double n double K0 double V_LG_CMB double CMB_VECTOR[3] double h double SIGMA8 double OMEGA_B double OMEGA_C double omega_B double omega_C double Theta_27 double OMEGA_0 double Omega double beta double OmegaEff double Gamma0 double normPower CosmoPower() void setFunction(CosmoFunction) void updateCosmology() void updatePhysicalCosmology() void normalize() void setNormalization(double) double power(double) 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 getBoxsize(self): return self.simu.BoxSize 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)) cdef class CosmologyPower: cdef CosmoPower power def __init__(self,**cosmo): self.power = CosmoPower() self.power.OMEGA_B = cosmo['omega_B_0'] self.power.OMEGA_C = cosmo['omega_M_0']-cosmo['omega_B_0'] self.power.h = cosmo['h'] if 'ns' in cosmo: self.power.n = cosmo['ns'] assert self.power.OMEGA_C > 0 self.power.updateCosmology() def normalize(self,s8): self.power.SIGMA8 = s8 self.power.normalize() def setFunction(self,funcname): cdef CosmoFunction f f = POWER_EFSTATHIOU if funcname=='EFSTATHIOU': f = POWER_EFSTATHIOU elif funcname=='HU_WIGGLES': f = HU_WIGGLES elif funcname=='HU_BARYON': f = HU_BARYON elif funcname=='BARDEEN': f = POWER_BARDEEN elif funcname=='SUGIYAMA': f = POWER_SUGIYAMA self.power.setFunction(f) cdef double _compute(self, double k): k *= self.power.h return self.power.power(k) def compute(self, k): cdef np.ndarray out cdef double kval cdef tuple i if isinstance(k, np.ndarray): out = np.empty(k.shape, dtype=np.float64) for i,kval in np.ndenumerate(k): out[i] = self._compute(kval) return out else: return self._compute(k)