185 lines
6.1 KiB
Python
185 lines
6.1 KiB
Python
import numexpr as ne
|
|
import multiprocessing
|
|
import pyfftw
|
|
import weakref
|
|
import numpy as np
|
|
import cosmolopy as cpy
|
|
|
|
class CubeFT(object):
|
|
def __init__(self, L, N, max_cpu=-1):
|
|
|
|
self.N = N
|
|
self.align = pyfftw.simd_alignment
|
|
self.L = L
|
|
self.max_cpu = multiprocessing.cpu_count() if max_cpu < 0 else max_cpu
|
|
self._dhat = pyfftw.n_byte_align_empty((self.N,self.N,self.N/2+1), self.align, dtype='complex64')
|
|
self._density = pyfftw.n_byte_align_empty((self.N,self.N,self.N), self.align, dtype='float32')
|
|
self._irfft = pyfftw.FFTW(self._dhat, self._density, axes=(0,1,2), direction='FFTW_BACKWARD', threads=self.max_cpu, normalize_idft=False)
|
|
self._rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False)
|
|
|
|
def rfft(self):
|
|
return ne.evaluate('c*a', local_dict={'c':self._rfft(normalise_idft=False),'a':(self.L/self.N)**3})
|
|
|
|
def irfft(self):
|
|
return ne.evaluate('c*a', local_dict={'c':self._irfft(normalise_idft=False),'a':(1/self.L)**3})
|
|
|
|
def get_dhat(self):
|
|
return self._dhat
|
|
def set_dhat(self, in_dhat):
|
|
self._dhat[:] = in_dhat
|
|
dhat = property(get_dhat, set_dhat, None)
|
|
|
|
def get_density(self):
|
|
return self._density
|
|
def set_density(self, d):
|
|
self._density[:] = d
|
|
density = property(get_density, set_density, None)
|
|
|
|
|
|
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, supersample=1, max_cpu=-1):
|
|
|
|
self.L = L
|
|
self.N = density.shape[0]
|
|
|
|
self.max_cpu = max_cpu
|
|
self.cube = CubeFT(self.L, self.N, max_cpu=max_cpu)
|
|
|
|
if not fourier:
|
|
self.cube.density = density
|
|
self.dhat = self.cube.rfft().copy()
|
|
else:
|
|
self.dhat = density.copy()
|
|
|
|
if supersample > 1:
|
|
self.upgrade_sampling(supersample)
|
|
self.ik = np.fft.fftfreq(self.N, d=L/self.N)*2*np.pi
|
|
self._kx = self.ik[:,None,None]
|
|
self._ky = self.ik[None,:,None]
|
|
self._kz = self.ik[None,None,:(self.N/2+1)]
|
|
self.cache = {}#weakref.WeakValueDictionary()
|
|
|
|
def upgrade_sampling(self, supersample):
|
|
N2 = self.N * supersample
|
|
N = self.N
|
|
dhat_new = np.zeros((N2, N2, N2/2+1), dtype=np.complex128)
|
|
|
|
hN = N/2
|
|
dhat_new[:hN, :hN, :hN+1] = self.dhat[:hN, :hN, :]
|
|
dhat_new[:hN, (N2-hN):N2, :hN+1] = self.dhat[:hN, hN:, :]
|
|
dhat_new[(N2-hN):N2, (N2-hN):N2, :hN+1] = self.dhat[hN:, hN:, :]
|
|
dhat_new[(N2-hN):N2, :hN, :hN+1] = self.dhat[hN:, :hN, :]
|
|
|
|
self.dhat = dhat_new
|
|
self.N = N2
|
|
self.cube = CubeFT(self.L, self.N, max_cpu=self.max_cpu)
|
|
|
|
def _gradient(self, phi, direction):
|
|
ne.evaluate('phi_hat * i * kv / (kx**2 + ky**2 + kz**2)', out=self.cube.dhat,
|
|
local_dict={'i':-1j, 'phi_hat':phi, 'kv':self._kdir(direction),
|
|
'kx':self._kx, 'ky':self._ky, 'kz':self._kz}
|
|
# self.cube.dhat = self._kdir(direction)*1j*phi
|
|
self.cube.dhat[0,0,0] = 0
|
|
return self.cube.irfft()
|
|
|
|
def lpt1(self, direction=0):
|
|
return self._gradient(self.dhat, 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 _do_irfft(self, array, copy=True):
|
|
if copy:
|
|
self.cube.dhat = array
|
|
return self.cube.irfft()
|
|
|
|
def _do_rfft(self, array, copy=True):
|
|
if copy:
|
|
self.cube.density = array
|
|
return self.cube.rfft()
|
|
|
|
def lpt2(self, direction=0):
|
|
# k2 = self._get_k2()
|
|
# k2[0,0,0] = 1
|
|
|
|
potgen0 = lambda i: ne.evaluate('kdir**2*d/k2',out=self.cube.dhat,local_dict={'kdir':self._kdir(i),'d':self.dhat,'k2':k2} )
|
|
potgen = lambda i,j: ne.evaluate('kdir0*kdir1*d/k2',out=self.cube.dhat,local_dict={'kdir0':self._kdir(i),'kdir1':self._kdir(j),'d':self.dhat,'k2':k2} )
|
|
|
|
if 'lpt2_potential' not in self.cache:
|
|
print("Rebuilding potential...")
|
|
div_phi2 = np.zeros((self.N,self.N,self.N), dtype=np.float64)
|
|
for j in xrange(3):
|
|
q = self._do_irfft( potgen0(j) ).copy()
|
|
for i in xrange(j+1, 3):
|
|
ne.evaluate('div + q * pot', out=div_phi2,
|
|
local_dict={'q':q,'pot':self._do_irfft( potgen0(i), copy=False ) }
|
|
)
|
|
ne.evaluate('div - pot**2',out=div_phi2,
|
|
local_dict={'div':div_phi2,'pot':self._do_irfft(potgen(i,j), copy=False) }
|
|
)
|
|
|
|
phi2_hat = self._do_rfft(div_phi2)
|
|
#self.cache['lpt2_potential'] = phi2_hat
|
|
del div_phi2
|
|
else:
|
|
phi2_hat = self.cache['lpt2_potential']
|
|
|
|
return self._gradient(phi2_hat, direction)
|