cosmotool/python/_cosmo_power.pyx

108 lines
2.4 KiB
Cython
Raw Normal View History

2014-11-20 14:04:17 +01:00
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 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)