cosmotool/python/_cosmotool.pyx

543 lines
15 KiB
Cython
Raw Normal View History

2014-05-25 11:11:44 +02:00
from libcpp cimport bool
2014-05-30 19:00:25 +02:00
from libcpp cimport string as cppstring
from libcpp.vector cimport vector as cppvector
from cython.parallel cimport prange
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
2014-06-12 17:16:59 +02:00
cimport cython
2014-05-26 14:24:53 +02:00
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
2014-07-05 22:05:04 +02:00
float *Mass
2014-05-25 11:11:44 +02:00
2014-05-30 19:00:25 +02:00
bool noAuto
2014-05-25 11:11:44 +02:00
cdef const int NEED_GADGET_ID
cdef const int NEED_POSITION
cdef const int NEED_VELOCITY
cdef const int NEED_TYPE
2014-07-05 22:05:04 +02:00
cdef const int NEED_MASS
2014-05-25 11:11:44 +02:00
cdef extern from "loadGadget.hpp" namespace "CosmoTool":
SimuData *loadGadgetMulti(const char *fname, int id, int flags) nogil except +
2014-05-30 19:00:25 +02:00
void cxx_writeGadget "CosmoTool::writeGadget" (const char * s, SimuData *data) except +
cdef extern from "safe_gadget.hpp":
SimuData *loadGadgetMulti_safe(cppstring.string s, int flags) nogil
SimuData **alloc_simudata(int num) nogil
void del_simudata(SimuData **d) nogil
2014-05-25 11:11:44 +02:00
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
2014-05-30 19:00:25 +02:00
class PySimulationBase(object):
2014-12-07 21:50:12 +01:00
"""
This is the base class to representation Simulation in CosmoTool/python.
"""
2014-05-30 19:00:25 +02:00
def getPositions(self):
2014-12-07 21:50:12 +01:00
"""
getPositions(self)
Returns:
A list of three arrays holding the positions of the particles.
The i-th element is the i-th coordinate of each particle.
It may be None if the positions were not requested.
"""
raise NotImplementedError("getPositions is not implemented")
2014-05-30 19:00:25 +02:00
def getVelocities(self):
2014-12-07 21:50:12 +01:00
"""
getVelocities(self)
Returns:
A list of three arrays holding the velocities of the particles.
The i-th element is the i-th coordinate of each particle.
It may be None if the velocities were not requested.
"""
raise NotImplementedError("getVelocities is not implemented")
2014-05-30 19:00:25 +02:00
def getIdentifiers(self):
2014-12-07 21:50:12 +01:00
"""
getIdentifiers(self)
Returns:
Returns an integer array that hold the unique identifiers of
each particle.
It may be None if the identifiers were not requested.
"""
raise NotImplementedError("getIdentifiers is not implemented")
2014-05-30 19:00:25 +02:00
2014-07-05 22:05:04 +02:00
def getTypes(self):
2014-12-07 21:50:12 +01:00
"""
getTypes(self)
Returns:
Returns an integer array that hold the type of
each particle.
It may be None if the types were not requested.
"""
2014-07-05 22:05:04 +02:00
raise NotImplementedError("getTypes is not implemented")
2014-05-30 19:00:25 +02:00
def getOmega_M(self):
2014-12-07 21:50:12 +01:00
"""
getOmega_M(self)
Returns:
the mean matter density in the simulation, with respect
to the critical density.
"""
raise NotImplementedError("getOmega_M is not implemented")
2014-05-30 19:00:25 +02:00
def getOmega_Lambda(self):
2014-12-07 21:50:12 +01:00
"""
getOmega_Lambda(self)
Returns:
the mean dark energy density in the simulation, with respect
to the critical density.
"""
raise NotImplementedError("getOmega_Lambda is not implemented")
2014-05-30 19:00:25 +02:00
def getTime(self):
2014-12-07 21:50:12 +01:00
"""
getTime(self)
Returns:
the time the snapshot was taken in the simulation. It can
have various units depending on the file format.
"""
raise NotImplementedError("getTime is not implemented")
2014-05-30 19:00:25 +02:00
def getHubble(self):
2014-12-07 21:50:12 +01:00
"""
getHubble(self)
Returns:
the hubble constant in unit of 100 km/s/Mpc
"""
raise NotImplementedError("getHubble is not implemented")
def getBoxsize(self):
2014-12-07 21:50:12 +01:00
"""
getBoxsize(self)
Returns:
the size of the simulation box. The length unit is not fixed,
though it is customary to have it in Mpc/h if the loader has
access to the unit normalization.
"""
raise NotImplementedError("getBoxsize is not implemented")
2014-05-30 19:00:25 +02:00
2014-07-05 22:05:04 +02:00
def getMasses(self):
2014-12-07 21:50:12 +01:00
"""
getMasses(self)
Returns:
an array with the masses of each particles, in unspecified unit that
depend on the loader.
"""
2014-07-05 22:05:04 +02:00
raise NotImplementedError("getMasses is not implemented")
2014-05-25 11:11:44 +02:00
cdef class Simulation:
2014-12-07 21:50:12 +01:00
"""
Simulation()
Class that directly manages internal loaded data obtained from a loader
"""
2014-05-25 11:11:44 +02:00
2014-12-07 21:50:12 +01:00
cdef list positions
cdef list velocities
cdef object identifiers
cdef object types
cdef object masses
2014-05-25 11:11:44 +02:00
2014-12-07 21:50:12 +01:00
cdef SimuData *data
2014-05-25 11:11:44 +02:00
2014-12-07 21:50:12 +01:00
property BoxSize:
def __get__(Simulation self):
return self.data.BoxSize
2014-05-26 14:24:53 +02:00
2014-12-07 21:50:12 +01:00
property time:
def __get__(Simulation self):
return self.data.time
2014-05-30 13:21:12 +02:00
2014-12-07 21:50:12 +01:00
property Hubble:
def __get__(Simulation self):
return self.data.Hubble
2014-05-26 14:24:53 +02:00
2014-12-07 21:50:12 +01:00
property Omega_M:
def __get__(Simulation self):
return self.data.Omega_M
2014-05-26 14:24:53 +02:00
2014-12-07 21:50:12 +01:00
property Omega_Lambda:
def __get__(Simulation self):
return self.data.Omega_Lambda
2014-12-07 21:50:12 +01:00
property positions:
def __get__(Simulation self):
return self.positions
2014-05-26 14:24:53 +02:00
2014-12-07 21:50:12 +01:00
property velocities:
def __get__(Simulation self):
return self.velocities
2014-12-07 21:50:12 +01:00
property identifiers:
def __get__(Simulation self):
return self.identifiers
2014-07-05 22:05:04 +02:00
2014-12-07 21:50:12 +01:00
property types:
def __get__(Simulation self):
return self.types
2014-07-05 22:05:04 +02:00
2014-12-07 21:50:12 +01:00
property masses:
def __get__(Simulation self):
return self.masses
2014-05-26 14:24:53 +02:00
2014-12-07 21:50:12 +01:00
property numParticles:
def __get__(Simulation self):
return self.data.NumPart
2014-05-26 14:24:53 +02:00
2014-11-03 13:56:32 +01:00
2014-12-07 21:50:12 +01:00
property totalNumParticles:
def __get__(Simulation self):
return self.data.TotalNumPart
2014-11-03 13:56:32 +01:00
2014-12-07 21:50:12 +01:00
def __cinit__(Simulation self):
self.data = <SimuData *>0
2014-05-25 11:11:44 +02:00
2014-12-07 21:50:12 +01:00
def __dealloc__(Simulation self):
if self.data != <SimuData *>0:
print("Clearing simulation data")
del self.data
2014-05-25 11:11:44 +02:00
2014-05-26 14:24:53 +02:00
2014-12-07 21:50:12 +01:00
class PySimulationAdaptor(PySimulationBase):
"""
PySimulationAdaptor(PySimulationBase_)
This class is an adaptor for an internal type to the loader. It defines
all the methods of PySimulationBase.
Attributes:
simu: a Simulation_ object
"""
2014-05-30 19:00:25 +02:00
def __init__(self,sim):
self.simu = sim
def getBoxsize(self):
return self.simu.BoxSize
2014-05-30 19:00:25 +02:00
def getPositions(self):
return self.simu.positions
2014-07-05 22:05:04 +02:00
def getTypes(self):
return self.simu.types
2014-05-30 19:00:25 +02:00
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
2014-05-30 19:00:25 +02:00
2014-07-05 22:05:04 +02:00
def getMasses(self):
return self.simu.masses
2014-05-26 14:24:53 +02:00
cdef class ArrayWrapper:
cdef void* data_ptr
2014-11-21 15:24:32 +01:00
cdef np.uint64_t size
2014-05-30 13:21:12 +02:00
cdef int type_array
2014-05-26 14:24:53 +02:00
2014-11-21 15:24:32 +01:00
cdef set_data(self, np.uint64_t size, int type_array, void* data_ptr):
2014-05-26 14:24:53 +02:00
""" Set the data of the array
2014-12-07 21:50:12 +01:00
This cannot be done in the constructor as it must recieve C-level
arguments.
2014-05-26 14:24:53 +02:00
2014-12-07 21:50:12 +01:00
Args:
size (int): Length of the array.
data_ptr (void*): Pointer to the data
"""
2014-05-26 14:24:53 +02:00
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_FLOAT32)
2014-05-30 13:21:12 +02:00
cdef object wrap_int64_array(np.int64_t* p, np.uint64_t s):
return wrap_array(<void *>p, s, np.NPY_INT64)
2014-07-05 22:05:04 +02:00
cdef object wrap_int_array(int* p, np.uint64_t s):
return wrap_array(<void *>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
2014-05-30 13:21:12 +02:00
if flags & NEED_GADGET_ID:
simu.identifiers = wrap_int64_array(data.Id, data.NumPart)
else:
simu.identifiers = None
2014-07-05 22:05:04 +02:00
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
2014-07-05 22:05:04 +02:00
def loadGadget(str filename, int snapshot_id, bool loadPosition = True, bool loadVelocity = False, bool loadId = False, bool loadType = False, bool loadMass=False):
2014-12-07 21:50:12 +01:00
"""loadGadget(filename, snapshot_id, loadPosition=True, loadVelocity=False, loadId=False, loadType=False)
This function loads Gadget-1 snapshot format.
If snapshot_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 reflect the indicated snapshot_id.
Arguments:
filename (str): input filename
snapshot_id (int): identifier of the gadget file if it is a multi-file snapshot
Keyword arguments:
loadPosition (bool): whether to load positions
loadVelocity (bool): whether to load velocities
loadId (bool): whether to load unique identifiers
loadType (bool): whether to set types to particles
loadMass (bool): whether to set the mass of particles
Returns:
an PySimulationAdaptor instance.
2014-07-05 22:05:04 +02:00
"""
2014-05-25 11:11:44 +02:00
2014-12-07 21:50:12 +01:00
cdef int flags
cdef SimuData *data
cdef Simulation simu
2014-05-25 11:11:44 +02:00
2014-12-07 21:50:12 +01:00
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
2014-05-25 11:11:44 +02:00
2014-12-07 21:50:12 +01:00
data = loadGadgetMulti(filename, snapshot_id, flags)
if data == <SimuData*>0:
return None
2014-05-25 11:11:44 +02:00
2014-12-07 21:50:12 +01:00
return PySimulationAdaptor(wrap_simudata(data, flags))
2014-05-30 19:00:25 +02:00
2014-11-21 14:21:07 +01:00
def loadParallelGadget(object filename_list, bool loadPosition = True, bool loadVelocity = False, bool loadId = False, bool loadType = False, bool loadMass=False):
2014-12-07 21:50:12 +01:00
"""loadParallelGadget(filename list, loadPosition=True, loadVelocity=False, loadId=False, loadType=False)
2014-11-21 15:24:32 +01:00
Arguments:
2014-12-07 21:50:12 +01:00
filename (list): a list or tuple of filenames to load in parallel
2014-11-21 15:24:32 +01:00
2014-12-07 21:50:12 +01:00
Keyword arguments:
loadPosition (bool): indicate to load positions
loadVelocity (bool): indicate to load velocities
loadId (bool): indicate to load id
loadType (bool): indicate to load particle types
2014-11-21 15:24:32 +01:00
2014-12-07 21:50:12 +01:00
Returns:
It loads a gadget-1 snapshot and return a cosmotool.PySimulationBase_ object.
"""
cdef int flags, i, num_files
cdef list out_arrays
cdef SimuData ** data
cdef SimuData * local_data
cdef Simulation simu
cdef cppvector[cppstring.string] filenames
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
num_files = len(filename_list)
filenames.resize(num_files)
data = alloc_simudata(num_files)
for i,l in enumerate(filename_list):
filenames[i] = l
with nogil:
for i in prange(num_files):
local_data = loadGadgetMulti_safe(filenames[i], flags)
data[i] = local_data
# data[i] = loadGadgetMulti(filenames[i].c_str(), -1, flags)
2014-11-20 14:51:39 +01:00
out_arrays = []
for i in xrange(num_files):
if data[i] == <SimuData*>0:
out_arrays.append(None)
else:
2014-12-07 21:50:12 +01:00
out_arrays.append(PySimulationAdaptor(wrap_simudata(data[i], flags)))
del_simudata(data)
return out_arrays
2014-05-30 19:00:25 +02:00
def writeGadget(str filename, object simulation):
2014-12-07 21:50:12 +01:00
"""writeGadget(filename, simulation)
This function attempts to write the content of the simulation object into
a file named `filename` using a Gadget-1 file format.
Arguments:
filename (str): output filename
simulation (PySimulationBase): a simulation object
"""
2014-05-30 19:00:25 +02:00
cdef SimuData simdata
cdef np.ndarray[np.float32_t, ndim=1] pos, vel
cdef np.ndarray[np.int64_t, ndim=1] ids
2014-05-30 19:00:25 +02:00
cdef np.int64_t NumPart
cdef int j
2014-05-30 19:00:25 +02:00
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] = <float *>pos.data
simdata.Vel[j] = <float *>vel.data
ids = simulation.getIdentifiers()
simdata.Id = <np.int64_t *>ids.data
simdata.BoxSize = simulation.getBoxsize()
2014-05-30 19:00:25 +02:00
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 otherwise the loading will fail. There is no way of auto-detecting properly the precision of the snapshot file.
Args:
basepath (str): the base directory of the snapshot
snapshot_id (int): the snapshot id
cpu_id (int): the cpu id of the file to load
Keyword args:
doublePrecision (bool): By default it is False, thus singlePrecision
loadPosition (bool): Whether to load positions
loadVelocity (bool): Whether to load velocities
Returns:
2014-12-07 21:50:12 +01:00
An object derived from PySimulationBase_.
"""
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
2014-12-07 21:50:12 +01:00
return PySimulationAdaptor(wrap_simudata(data, flags))
2014-06-02 14:28:09 +02:00