108 lines
2.4 KiB
Cython
108 lines
2.4 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 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)
|
||
|
|