cosmotool/python/_cosmotool.pyx

191 lines
4.9 KiB
Cython
Raw Normal View History

2014-05-25 11:11:44 +02:00
from libcpp cimport bool
2014-05-25 10:43:06 +02:00
import numpy as np
2014-05-25 11:11:44 +02:00
cimport numpy as np
2014-05-26 14:24:53 +02:00
from cpython cimport PyObject, Py_INCREF
np.import_array()
2014-05-25 11:11:44 +02:00
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
2014-05-26 14:24:53 +02:00
np.int64_t TotalNumPart
np.int64_t NumPart
2014-05-25 11:11:44 +02:00
np.int64_t *Id
2014-05-26 14:24:53 +02:00
float *Pos[3]
float *Vel[3]
2014-05-25 11:11:44 +02:00
int *type
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 +
cdef extern from "loadRamses.hpp" namespace "CosmoTool":
SimuData *loadRamsesSimu(const char *basename, int id, int cpuid, bool dp, int flags) except +
2014-05-25 11:11:44 +02:00
cdef class Simulation:
2014-05-26 14:24:53 +02:00
cdef list positions
cdef list velocities
2014-05-30 13:21:12 +02:00
cdef object identifiers
2014-05-25 11:11:44 +02:00
cdef SimuData *data
2014-05-26 14:24:53 +02:00
property BoxSize:
def __get__(Simulation self):
return self.data.BoxSize
2014-05-30 13:21:12 +02:00
property time:
def __get__(Simulation self):
return self.data.time
2014-05-26 14:24:53 +02:00
property Hubble:
def __get__(Simulation self):
return self.data.Hubble
property Omega_M:
def __get__(Simulation self):
return self.data.Omega_M
property positions:
def __get__(Simulation self):
return self.positions
property velocities:
def __get__(Simulation self):
return self.velocities
property numParticles:
def __get__(Simulation self):
return self.data.NumPart
2014-05-25 11:11:44 +02:00
def __cinit__(Simulation self):
self.data = <SimuData *>0
def __dealloc__(Simulation self):
if self.data != <SimuData *>0:
2014-05-26 14:24:53 +02:00
print("Clearing simulation data")
2014-05-25 11:11:44 +02:00
del self.data
2014-05-26 14:24:53 +02:00
cdef class ArrayWrapper:
cdef void* data_ptr
cdef int size
2014-05-30 13:21:12 +02:00
cdef int type_array
2014-05-26 14:24:53 +02:00
2014-05-30 13:21:12 +02:00
cdef set_data(self, int size, int type_array, void* data_ptr):
2014-05-26 14:24:53 +02:00
""" 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
2014-05-30 13:21:12 +02:00
self.type_array = type_array
2014-05-26 14:24:53 +02:00
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] = <np.npy_intp> self.size
# Create a 1D array, of length 'size'
2014-05-30 13:21:12 +02:00
ndarray = np.PyArray_SimpleNewFromData(1, shape, self.type_array, self.data_ptr)
2014-05-26 14:24:53 +02:00
return ndarray
def __dealloc__(self):
""" Frees the array. This is called by Python when all the
references to the object are gone. """
pass
2014-05-30 13:21:12 +02:00
cdef object wrap_array(void *p, np.uint64_t s, int typ):
2014-05-26 14:24:53 +02:00
cdef np.ndarray ndarray
cdef ArrayWrapper wrapper
wrapper = ArrayWrapper()
2014-05-30 13:21:12 +02:00
wrapper.set_data(s, typ, p)
2014-05-26 14:24:53 +02:00
ndarray = np.array(wrapper, copy=False)
ndarray.base = <PyObject*> wrapper
Py_INCREF(wrapper)
return ndarray
2014-05-30 13:21:12 +02:00
cdef object wrap_float_array(float *p, np.uint64_t s):
return wrap_array(<void *>p, s, np.NPY_FLOAT)
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_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)]
if flags & NEED_VELOCITY:
simu.velocities = [wrap_float_array(data.Vel[i], data.NumPart) for i in xrange(3)]
2014-05-30 13:21:12 +02:00
if flags & NEED_GADGET_ID:
simu.identifiers = wrap_int64_array(data.Id, data.NumPart)
return simu
2014-05-30 13:21:12 +02:00
def loadGadget(str filename, int snapshot_id, bool loadPosition = True, bool loadVelocity = False, bool loadId = False):
2014-05-25 11:11:44 +02:00
cdef int flags
cdef SimuData *data
cdef Simulation simu
flags = 0
if loadPosition:
flags |= NEED_POSITION
if loadVelocity:
flags |= NEED_VELOCITY
2014-05-30 13:21:12 +02:00
if loadId:
flags |= NEED_GADGET_ID
2014-05-25 11:11:44 +02:00
data = loadGadgetMulti(filename, snapshot_id, flags)
2014-05-27 13:54:42 +02:00
if data == <SimuData*>0:
return None
2014-05-25 11:11:44 +02:00
return wrap_simudata(data, flags)
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
2014-05-26 14:24:53 +02:00
if loadPosition:
flags |= NEED_POSITION
2014-05-26 14:24:53 +02:00
if loadVelocity:
flags |= NEED_VELOCITY
data = loadRamsesSimu(basepath, snapshot_id, cpu_id, doublePrecision, flags)
if data == <SimuData*>0:
return None
return wrap_simudata(data, flags)