103 lines
2.9 KiB
Python
103 lines
2.9 KiB
Python
import weakref
|
|
import numpy as np
|
|
import cosmolopy as cpy
|
|
|
|
class CosmoGrowth(object):
|
|
|
|
def __init__(self, **cosmo):
|
|
self.cosmo = cosmo
|
|
|
|
def D(self, a):
|
|
return cpy.perturbation.fgrowth(1/a-1, self.cosmo['omega_M_0'], unnormed=True)
|
|
|
|
def compute_E(self, a):
|
|
om = self.cosmo['omega_M_0']
|
|
ol = self.cosmo['omega_lambda_0']
|
|
ok = self.cosmo['omega_k_0']
|
|
|
|
E = np.sqrt(om/a**3 + ol + ok/a**2)
|
|
|
|
H2 = -3*om/a**4 - 2*ok/a**3
|
|
|
|
Eprime = 0.5*H2/E
|
|
|
|
return E,Eprime
|
|
|
|
def Ddot(self, a):
|
|
E,Eprime = self.compute_E(a)
|
|
D = self.D(a)
|
|
Ddot_D = Eprime/E + 2.5 * self.cosmo['omega_M_0']/(a**3*E**2*D)
|
|
Ddot_D *= a
|
|
return Ddot_D
|
|
|
|
def compute_velmul(self, a):
|
|
E,_ = self.compute_E(a)
|
|
velmul = self.Ddot(a)
|
|
velmul *= 100 * a * E
|
|
return velmul
|
|
|
|
|
|
|
|
|
|
|
|
class LagrangianPerturbation(object):
|
|
|
|
def __init__(self,density,L, fourier=False):
|
|
|
|
self.L = L
|
|
self.N = density.shape[0]
|
|
self.dhat = np.fft.rfftn(density)*(L/self.N)**3 if not fourier else density
|
|
self.ik = np.fft.fftfreq(self.N, d=L/self.N)*2*np.pi
|
|
self.cache = weakref.WeakValueDictionary()
|
|
|
|
def _gradient(self, phi, direction):
|
|
return np.fft.irfftn(self._kdir(direction)*1j*phi)*(self.N/self.L)**3
|
|
|
|
def lpt1(self, direction=0):
|
|
k2 = self._get_k2()
|
|
k2[0,0,0] = 1
|
|
|
|
return self._gradient(-self.dhat/k2, direction)
|
|
|
|
def new_shape(self,direction, q=3, half=False):
|
|
N0 = (self.N/2+1) if half else self.N
|
|
return ((1,)*direction) + (N0,) + ((1,)*(q-1-direction))
|
|
|
|
def _kdir(self, direction, q=3):
|
|
if direction != q-1:
|
|
return self.ik.reshape(self.new_shape(direction, q=q))
|
|
else:
|
|
return self.ik[:self.N/2+1].reshape(self.new_shape(direction, q=q, half=True))
|
|
|
|
def _get_k2(self, q=3):
|
|
if 'k2' in self.cache:
|
|
return self.cache['k2']
|
|
|
|
k2 = self._kdir(0, q=q)**2
|
|
for d in xrange(1,q):
|
|
k2 = k2 + self._kdir(d, q=q)**2
|
|
|
|
self.cache['k2'] = k2
|
|
return k2
|
|
|
|
def lpt2(self, direction=0):
|
|
k2 = self._get_k2()
|
|
k2[0,0,0] = 1
|
|
|
|
if 'lpt2_potential' not in self.cache:
|
|
div_phi2 = np.zeros((N,N,N), dtype=np.float64)
|
|
for j in xrange(3):
|
|
q = np.fft.irfftn( build_dir(ik, j)**2*self.dhat / k2 )
|
|
for i in xrange(j+1, 3):
|
|
div_phi2 += q * np.fft.irfftn( build_dir(ik, i)**2*self.dhat / k2 )
|
|
div_phi2 -= (np.fft.irfftn( build_dir(ik, j)*build_dir(ik, i)*self.dhat / k2 ))**2
|
|
|
|
div_phi2 *= (self.N/self.L)**3
|
|
phi2_hat = np.fft.rfftn(div_phi2) * ((L/N)**3) / k2
|
|
self.cache['lpt2_potential'] = phi2_hat
|
|
del div_phi2
|
|
else:
|
|
phi2_hat = self.cache['lpt2_potential']
|
|
|
|
return self._gradient(phi2_hat, direction)
|