Added parallelization to _project.pyx code

This commit is contained in:
Guilhem Lavaux 2014-06-11 11:14:33 +02:00
parent 5b133a9ac8
commit 401ddc8a8b
9 changed files with 173 additions and 122 deletions

View file

@ -93,17 +93,67 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True,
density = np.fft.irfftn(lpt.dhat*D1_0)*(supersample*N/L)**3
return posx,vel,density,N*supersample,L,a_ic
return posx,vel,density,N*supersample,L,a_ic,cosmo
def write_icfiles(*generated_ic, **cosmo):
"""Write the initial conditions from the tuple returned by run_generation"""
posx,vel,density,N,L,a_ic = generated_ic
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
ct.simpleWriteGadget("borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic)
for i,c in enumerate(["x","y","z"]):
ct.writeGrafic("ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
# This used to be necessary. However this has been fixed in writeGrafic now
# ct.writeGrafic("ic_velc%s" % c, vel[i].reshape((N,N,N)).transpose(), L, a_ic, **cosmo)
k = np.sqrt(ik[:,None,None]**2 + ik[None,:,None]**2 + ik[None,None,:(N/2+1)]**2)
p = ct.CosmologyPower(**cosmo)
p.setFunction(func)
p.normalize(cosmo['SIGMA8'])
Pk = p.compute(k)*cosmo['h']**3*L**3
Pk[0,0,0]=1
density_hat = np.fft.rfftn(density)/N**3/np.sqrt(Pk)
Ns = N*supergenerate
density_hat_super = np.zeros((Ns,Ns,Ns/2+1), dtype=np.complex128)
density_hat_super[:] = np.nan
# 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[(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, :]
# 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()
return np.fft.irfftn(density_hat_super)
def write_icfiles(*generated_ic, **kwargs):
"""Write the initial conditions from the tuple returned by run_generation"""
supergenerate=1
if 'supergenerate' in kwargs:
supergenerate=kwargs['supergenerate']
posx,vel,density,N,L,a_ic,cosmo = generated_ic
ct.simpleWriteGadget("borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic)
for i,c in enumerate(["x","y","z"]):
ct.writeGrafic("ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
ct.writeGrafic("ic_deltab", density, L, a_ic, **cosmo)
ct.writeWhitePhase("white.dat", whitify(density, L, cosmo, supergenerate=supergenerate))
with file("white_params", mode="w") as f:
f.write("4\n%lg, %lg, %lg\n" % (cosmo['omega_M_0'], cosmo['omega_lambda_0'], 100*cosmo['h']))
f.write("%lg\n%lg\n-%lg\n0,0\n" % (cosmo['omega_B_0'],cosmo['ns'],cosmo['SIGMA8']))
f.write("-%lg\n1\n0\n\n\n\n\n" % L)
f.write("2\n\n0\nwhite.dat\n0\npadding_white.dat\n")
ct.writeGrafic("ic_deltab", density, L, a_ic, **cosmo)
ct.writeWhitePhase("white.dat", density)