diff --git a/python/cosmotool/__init__.py b/python/cosmotool/__init__.py index 3db5697..6c1d376 100644 --- a/python/cosmotool/__init__.py +++ b/python/cosmotool/__init__.py @@ -4,4 +4,4 @@ from grafic import writeGrafic, writeWhitePhase, readGrafic, readWhitePhase from borg import read_borg_vol from cic import cicParticles from simu import loadRamsesAll, simpleWriteGadget, SimulationBare -from timing import timeit +from timing import timeit, timeit_quiet diff --git a/python/cosmotool/timing.py b/python/cosmotool/timing.py index 58e319d..f04087e 100644 --- a/python/cosmotool/timing.py +++ b/python/cosmotool/timing.py @@ -13,3 +13,16 @@ def timeit(method): return timed +def timeit_quiet(method): + + def timed(*args, **kw): + ts = time.time() + result = method(*args, **kw) + te = time.time() + + print '%r %2.2f sec' % \ + (method.__name__, te-ts) + return result + + return timed + diff --git a/python_sample/icgen/borgadaptor.py b/python_sample/icgen/borgadaptor.py index 2ee6018..07637df 100644 --- a/python_sample/icgen/borgadaptor.py +++ b/python_sample/icgen/borgadaptor.py @@ -6,15 +6,50 @@ def fourier_analysis(borg_vol): return np.fft.rfftn(borg_vol.density)*(L/N)**3, L, N +def borg_upgrade_sampling(dhat, supersample): + N = dhat.shape[0] + N2 = N * supersample + dhat_new = np.zeros((N2, N2, N2/2+1), dtype=np.complex128) + + hN = N/2 + dhat_new[:hN, :hN, :hN+1] = dhat[:hN, :hN, :] + dhat_new[:hN, (N2-hN):N2, :hN+1] = dhat[:hN, hN:, :] + dhat_new[(N2-hN):N2, (N2-hN):N2, :hN+1] = dhat[hN:, hN:, :] + dhat_new[(N2-hN):N2, :hN, :hN+1] = dhat[hN:, :hN, :] + + return dhat_new, N2 + def half_pixel_shift(borg, doshift=False): dhat,L,N = fourier_analysis(borg) if not doshift: return dhat, L + return bare_half_pixel_shift(dhat, L, N) + +def bare_half_pixel_shift(dhat, L, N, doshift=False): + +# dhat_new,N2 = borg_upgrade_sampling(dhat, 2) +# d = (np.fft.irfftn(dhat_new)*(N2/L)**3)[1::2,1::2,1::2] +# del dhat_new +# dhat = np.fft.rfftn(d)*(L/N)**3 +# return dhat, L + +# dhat2 = np.zeros((N,N,N),dtype=np.complex128) +# dhat2[:,:,:N/2+1] = dhat +# dhat2[N:0:-1, N:0:-1, N:N/2:-1] = np.conj(dhat[1:,1:,1:N/2]) +# dhat2[0, N:0:-1, N:N/2:-1] = np.conj(dhat[0, 1:, 1:N/2]) +# dhat2[N:0:-1, 0, N:N/2:-1] = np.conj(dhat[1:, 0, 1:N/2]) +# dhat2[0,0,N:N/2:-1] = np.conj(dhat[0, 0, 1:N/2]) + ik = np.fft.fftfreq(N,d=L/N)*2*np.pi phi = 0.5*L/N*(ik[:,None,None]+ik[None,:,None]+ik[None,None,:(N/2+1)]) +# phi %= 2*np.pi phase = np.cos(phi)+1j*np.sin(phi) + dhat = dhat*phase + dhat[N/2,:,:] = 0 + dhat[:,N/2,:] = 0 + dhat[:,:,N/2] = 0 - return dhat*phase, L + return dhat, L diff --git a/python_sample/icgen/borgicgen.py b/python_sample/icgen/borgicgen.py index 111c125..f2d173f 100644 --- a/python_sample/icgen/borgicgen.py +++ b/python_sample/icgen/borgicgen.py @@ -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""" diff --git a/python_sample/icgen/cosmogrowth.py b/python_sample/icgen/cosmogrowth.py index 6263126..4a1d399 100644 --- a/python_sample/icgen/cosmogrowth.py +++ b/python_sample/icgen/cosmogrowth.py @@ -110,7 +110,7 @@ class LagrangianPerturbation(object): div_phi2 -= (np.fft.irfftn( self._kdir(j)*self._kdir(i)*self.dhat / k2 ))**2 div_phi2 *= (self.N/self.L)**6 - phi2_hat = np.fft.rfftn(div_phi2) * ((self.L/self.N)**3) / k2 + phi2_hat = -np.fft.rfftn(div_phi2) * ((self.L/self.N)**3) / k2 self.cache['lpt2_potential'] = phi2_hat del div_phi2 else: diff --git a/python_sample/icgen/gen_ic_from_borg.py b/python_sample/icgen/gen_ic_from_borg.py index aa6509d..a3b48cd 100644 --- a/python_sample/icgen/gen_ic_from_borg.py +++ b/python_sample/icgen/gen_ic_from_borg.py @@ -9,9 +9,10 @@ cosmo['omega_B_0']=0.049 cosmo['SIGMA8']=0.8344 cosmo['ns']=0.9624 +supergen=4 zstart=50 astart=1/(1.+zstart) halfPixelShift=True if __name__=="__main__": - bic.write_icfiles(*bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=2, shiftPixel=halfPixelShift, do_lpt2=False)) + bic.write_icfiles(*bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=1, shiftPixel=halfPixelShift, do_lpt2=False), supergenerate=supergen) diff --git a/python_sample/icgen/test_ic_from_borg.py b/python_sample/icgen/test_ic_from_borg.py index dc0265d..2287340 100644 --- a/python_sample/icgen/test_ic_from_borg.py +++ b/python_sample/icgen/test_ic_from_borg.py @@ -51,4 +51,6 @@ Pcic /= D1_0**2 Pdens /= D1_0**2 borg_evolved = ct.read_borg_vol("final_density_1380.dat") +dborg_hat = np.fft.rfftn(borg_evolved.density)*L**3/borg_evolved.density.size +Pborg, bborg = bic.bin_power(np.abs(dborg_hat)**2/L**3, L, range=(0,1.),bins=150)