guilhem_lavaux/cosmotool
0
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

KSZ work

Browse source
This commit is contained in:
Guilhem Lavaux 2015-01-22 19:16:46 +01:00
parent 7a27dae025
commit f12c9a0c1a
5 changed files with 128 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

36
python_sample/icgen/cosmogrowth.py
View file

@ -4,6 +4,7 @@ import pyfftw
import weakref
import numpy as np
import cosmolopy as cpy
import cosmotool as ct
class CubeFT(object):
def __init__(self, L, N, max_cpu=-1):
@ -98,6 +99,7 @@ class LagrangianPerturbation(object):
self._kz = self.ik[None,None,:(self.N/2+1)]
self.cache = {}#weakref.WeakValueDictionary()
@ct.timeit_quiet
def upgrade_sampling(self, supersample):
N2 = self.N * supersample
N = self.N
@ -113,14 +115,26 @@ class LagrangianPerturbation(object):
self.N = N2
self.cube = CubeFT(self.L, self.N, max_cpu=self.max_cpu)
@ct.timeit_quiet
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),
if direction == 'all':
dirs = [0,1,2]
copy = True
else:
dirs = [direction]
copy = False
ret=[]
for dir in dirs:
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(dir),
'kx':self._kx, 'ky':self._ky, 'kz':self._kz},casting='unsafe')
# self.cube.dhat = self._kdir(direction)*1j*phi
self.cube.dhat[0,0,0] = 0
return self.cube.irfft()
self.cube.dhat[0,0,0] = 0
x = self.cube.irfft()
ret.append(x.copy() if copy else x)
return ret[0] if len(ret)==1 else ret
@ct.timeit_quiet
def lpt1(self, direction=0):
return self._gradient(self.dhat, direction)
@ -155,6 +169,7 @@ class LagrangianPerturbation(object):
self.cube.density = array
return self.cube.rfft()
@ct.timeit_quiet
def lpt2(self, direction=0):
# k2 = self._get_k2()
# k2[0,0,0] = 1
@ -170,12 +185,13 @@ class LagrangianPerturbation(object):
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={'div':div_phi2, '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) }
)
with ct.time_block("LPT2 elemental (%d,%d)" %(i,j)):
ne.evaluate('div + q * pot', out=div_phi2,
local_dict={'div':div_phi2, '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