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

Added more timing. Fixed whitification.

Browse source
This commit is contained in:
Guilhem Lavaux 2014-06-12 09:45:46 +02:00
parent 401ddc8a8b
commit a6b08dfe8f
7 changed files with 95 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

66
python_sample/icgen/borgicgen.py
View file

@ -15,7 +15,7 @@ def gen_posgrid(N, L):
return x.reshape((x.size,)), y.reshape((y.size,)), z.reshape((z.size,))
def bin_power(P, L, bins=20, range=(0,1.)):
def bin_power(P, L, bins=20, range=(0,1.), dev=False):
N = P.shape[0]
ik = np.fft.fftfreq(N, d=L/N)*2*np.pi
@ -25,7 +25,10 @@ def bin_power(P, L, bins=20, range=(0,1.)):
H,b = np.histogram(k, bins=bins, range=range)
Hw,b = np.histogram(k, bins=bins, weights=P, range=range)
return Hw/H, 0.5*(b[1:]+b[0:bins])
if dev:
return Hw/(H-1), 0.5*(b[1:]+b[0:bins]), 1.0/np.sqrt(H)
else:
return Hw/(H-1), 0.5*(b[1:]+b[0:bins])
def compute_power_from_borg(input_borg, a_borg, cosmo, bins=10, range=(0,1)):
borg_vol = ct.read_borg_vol(input_borg)
@ -59,9 +62,9 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True,
cgrowth = CosmoGrowth(**cosmo)
density_hat, L = ba.half_pixel_shift(borg_vol, doshift=shiftPixel)
density, L = ba.half_pixel_shift(borg_vol, doshift=shiftPixel)
lpt = LagrangianPerturbation(density_hat, L, fourier=True, supersample=supersample)
lpt = LagrangianPerturbation(density, L, fourier=True, supersample=supersample)
# Generate grid
posq = gen_posgrid(N*supersample, L)
@ -73,7 +76,7 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True,
D1_0 = D1/cgrowth.D(a_borg)
velmul = cgrowth.compute_velmul(a_ic)
D2 = 3./7 * D1_0**2
D2 = -3./7 * D1_0**2
for j in xrange(3):
# Generate psi_j (displacement along j)
@ -91,23 +94,30 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True,
print("velmul=%lg" % (cosmo['h']*velmul))
density = np.fft.irfftn(lpt.dhat*D1_0)*(supersample*N/L)**3
density = cgrowth.D(1)/cgrowth.D(a_borg)*np.fft.irfftn(lpt.dhat)*(supersample*N/L)**3
return posx,vel,density,N*supersample,L,a_ic,cosmo
def whitify(density, L, cosmo, supergenerate=1, func='HU_WIGGLES'):
N = density.shape[0]
ik = np.fft.fftfreq(N, d=L/N)*2*np.pi
k = np.sqrt(ik[:,None,None]**2 + ik[None,:,None]**2 + ik[None,None,:(N/2+1)]**2)
@ct.timeit_quiet
def whitify(density, L, cosmo, supergenerate=1, func='HU_WIGGLES'):
N = density.shape[0]
p = ct.CosmologyPower(**cosmo)
p.setFunction(func)
p.normalize(cosmo['SIGMA8'])
Pk = p.compute(k)*cosmo['h']**3*L**3
@ct.timeit_quiet
def build_Pk():
ik = np.fft.fftfreq(N, d=L/N)*2*np.pi
k = np.sqrt(ik[:,None,None]**2 + ik[None,:,None]**2 + ik[None,None,:(N/2+1)]**2)
return p.compute(k)*cosmo['h']**3*L**3
Pk = build_Pk()
Pk[0,0,0]=1
density_hat = np.fft.rfftn(density)/N**3/np.sqrt(Pk)
density_hat = np.fft.rfftn(density)*(L/N)**3
density_hat /= np.sqrt(Pk)
Ns = N*supergenerate
@ -116,24 +126,28 @@ def whitify(density, L, cosmo, supergenerate=1, func='HU_WIGGLES'):
# Copy density hat in place
hN = N/2
density_hat_super[:hN, :hN, :hN+1] = density_hat[:hN, :hN, :]
density_hat_super[:hN, (Ns-hN):Ns, :hN+1] = density_hat[:hN, hN:, :]
density_hat_super[:hN, :hN, :hN+1] = density_hat[:hN, :hN, :]
density_hat_super[:hN, (Ns-hN):Ns, :hN+1] = density_hat[:hN, hN:, :]
density_hat_super[(Ns-hN):Ns, (Ns-hN):Ns, :hN+1] = density_hat[hN:, hN:, :]
density_hat_super[(Ns-hN):Ns, :hN, :hN+1] = density_hat[hN:, :hN, :]
density_hat_super[(Ns-hN):Ns, :hN, :hN+1] = density_hat[hN:, :hN, :]
# The moved nyquist place is untouched (so loss of "noise") to keep the structure
# now we just add some noise term
cond=np.isnan(density_hat_super)
x = np.random.randn(np.count_nonzero(cond),2)/np.sqrt(2.0)
density_hat_super[cond] = x[:,0] + 1j * x[:,1]
# Now we have to fix the Nyquist plane
hNs = Ns/2
nyquist = density_hat_super[:, :, :hNs]
Nplane = nyquist.size
nyquist.flat[:Nplane/2] = nyquist.flat[Nplane:Nplane/2:-1].conj()
if supergenerate > 1:
cond=np.isnan(density_hat_super)
x = np.random.randn(np.count_nonzero(cond),2)/np.sqrt(2.0)
density_hat_super[cond] = x[:,0] + 1j * x[:,1]
# Now we have to fix the Nyquist plane
hNs = Ns/2
nyquist = density_hat_super[:, :, hNs]
Nplane = nyquist.size
nyquist.flat[:Nplane/2] = np.sqrt(2.0)*nyquist.flat[Nplane:Nplane/2:-1].conj()
return np.fft.irfftn(density_hat_super)*Ns**1.5
return np.fft.irfftn(density_hat_super)
def write_icfiles(*generated_ic, **kwargs):
"""Write the initial conditions from the tuple returned by run_generation"""