from libcpp cimport bool from libcpp cimport string as cppstring import numpy as np cimport numpy as np from cpython cimport PyObject, Py_INCREF cimport cython 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 "cic.hpp" namespace "CosmoTool": ctypedef float CICType ctypedef float Coordinates[3] cdef cppclass CICParticles: float mass Coordinates coords cdef cppclass CICFilter: CICFilter(np.uint32_t resolution, double L) nogil void resetMesh() nogil void putParticles(CICParticles* particles, np.uint32_t N) nogil void getDensityField(CICType *& field, np.uint32_t& res) nogil 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 float *Mass bool noAuto cdef const int NEED_GADGET_ID 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": 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 getTypes(self): raise NotImplementedError("getTypes 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") 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 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 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 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 getTypes(self): return self.simu.types 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 def getMasses(self): return self.simu.masses 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_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 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 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, 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 cdef Simulation simu flags = 0 if loadPosition: flags |= NEED_POSITION if loadVelocity: 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: 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) @cython.boundscheck(False) def leanCic(float[:,:] particles, float L, int Resolution): cdef CICParticles p cdef CICFilter *cic cdef np.uint64_t i cdef CICType *field cdef np.uint32_t dummyRes cdef np.ndarray[np.float64_t, ndim=3] out_field cdef np.uint64_t j cic = new CICFilter(Resolution, L) cic.resetMesh() if particles.shape[1] != 3: raise ValueError("Particles must be Nx3 array") p.mass = 1 for i in xrange(particles.shape[0]): p.coords[0] = particles[i,0] p.coords[1] = particles[i,1] p.coords[2] = particles[i,2] cic.putParticles(&p, 1) field = 0 dummyRes = 0 cic.getDensityField(field, dummyRes) out_field = np.empty((dummyRes, dummyRes, dummyRes), dtype=np.float64) for j in xrange(out_field.size): out_field[j] = field[j] del cic return out_field