Added parallelization to _project.pyx code
This commit is contained in:
parent
5b133a9ac8
commit
401ddc8a8b
9 changed files with 173 additions and 122 deletions
|
@ -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)
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue