454 lines
11 KiB
Cython
454 lines
11 KiB
Cython
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
|
|
|
|
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 = <SimuData *>0
|
|
|
|
def __dealloc__(Simulation self):
|
|
if self.data != <SimuData *>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] = <np.npy_intp> 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 = <PyObject*> wrapper
|
|
Py_INCREF(wrapper)
|
|
|
|
return ndarray
|
|
|
|
|
|
cdef object wrap_float_array(float *p, np.uint64_t s):
|
|
return wrap_array(<void *>p, s, np.NPY_FLOAT32)
|
|
|
|
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)]
|
|
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 == <SimuData*>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] = <float *>pos.data
|
|
simdata.Vel[j] = <float *>vel.data
|
|
|
|
ids = simulation.getIdentifiers()
|
|
simdata.Id = <np.int64_t *>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 == <SimuData*>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 = <CICType*>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
|