guilhem_lavaux/cosmotool
0
0
Fork 0
Code Issues Pull requests Projects Releases 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

10
python_sample/icgen/borgicgen.py
View file

@ -4,10 +4,11 @@ import cosmolopy as cpy
from cosmogrowth import *
import borgadaptor as ba
def gen_posgrid(N, L):
@ct.timeit
def gen_posgrid(N, L, delta=1, dtype=np.float32):
""" Generate an ordered lagrangian grid"""
ix = (np.arange(N)*L/N).astype(np.float32)
ix = (np.arange(N)*(L/N*delta)).astype(dtype)
x = ix[:,None,None].repeat(N, axis=1).repeat(N, axis=2)
y = ix[None,:,None].repeat(N, axis=0).repeat(N, axis=2)
@ -147,6 +148,8 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, supergenerate
D2 = -3./7 * D1_0**2
if do_lpt2:
psi2 = lpt.lpt2('all')
for j in xrange(3):
# Generate psi_j (displacement along j)
print("LPT1 axis=%d" % j)
@ -154,8 +157,7 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, supergenerate
psi = psi.reshape((psi.size,))
if do_lpt2:
print("LPT2 axis=%d" % j)
psi2 = lpt.lpt2(j)
psi += D2 * psi2.reshape((psi2.size,))
psi += D2 * psi2[j].reshape((psi2[j].size,))
# Generate posx
posx.append(((posq[j] + psi)%L).astype(np.float32))
# Generate vel

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