Use PYFFTW for ICgen

This commit is contained in:
Guilhem Lavaux 2014-07-03 15:36:58 +02:00
parent f2b81d863f
commit 33806a690f
4 changed files with 101 additions and 38 deletions

View file

@ -7,7 +7,7 @@ import borgadaptor as ba
def gen_posgrid(N, L):
""" Generate an ordered lagrangian grid"""
ix = np.arange(N)*L/N
ix = (np.arange(N)*L/N).astype(np.float32)
x = ix[:,None,None].repeat(N, axis=1).repeat(N, axis=2)
y = ix[None,:,None].repeat(N, axis=0).repeat(N, axis=2)
@ -52,7 +52,7 @@ def compute_ref_power(L, N, cosmo, bins=10, range=(0,1), func='HU_WIGGLES'):
p.normalize(cosmo['SIGMA8'])
return bin_power(p.compute(k)*cosmo['h']**3, L, bins=bins, range=range)
def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True, shiftPixel=False, needvel=True):
""" Generate particles and velocities from a BORG snapshot. Returns a tuple of
(positions,velocities,N,BoxSize,scale_factor)."""
@ -95,7 +95,9 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True,
print("velmul=%lg" % (cosmo['h']*velmul))
density = cgrowth.D(1)/cgrowth.D(a_borg)*np.fft.irfftn(lpt.dhat)*(supersample*N/L)**3
lpt.cube.dhat = lpt.dhat
density = lpt.cube.irfft()
density *= (cgrowth.D(1)/cgrowth.D(a_borg))
return posx,vel,density,N*supersample,L,a_ic,cosmo
@ -117,7 +119,9 @@ def whitify(density, L, cosmo, supergenerate=1, func='HU_WIGGLES'):
Pk = build_Pk()
Pk[0,0,0]=1
density_hat = np.fft.rfftn(density)*(L/N)**3
cube = CubeFT(N, L)
cube.density = density
density_hat = cube.rfft()
density_hat /= np.sqrt(Pk)
Ns = N*supergenerate
@ -152,6 +156,8 @@ def whitify(density, L, cosmo, supergenerate=1, func='HU_WIGGLES'):
print np.where(np.isnan(density_hat_super))[0].size
cube = CubeFT(Ns, L)
cube.dhat = density_hat_super
return np.fft.irfftn(density_hat_super)*Ns**1.5