diff --git a/python_sample/icgen/borgicgen.py b/python_sample/icgen/borgicgen.py index 9aab639..fcfee2c 100644 --- a/python_sample/icgen/borgicgen.py +++ b/python_sample/icgen/borgicgen.py @@ -53,7 +53,65 @@ def compute_ref_power(L, N, cosmo, bins=10, range=(0,1), func='HU_WIGGLES'): 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, psi_instead=False, needvel=True): + +def do_supergenerate(density, density_out=None, mulfac=None,zero_fill=False,Pk=None,L=None,h=None): + + N = density.shape[0] + + if density_out is None: + assert mulfac is not None + Ns = mulfac*N + density_out = np.zeros((Ns,Ns,Ns/2+1), dtype=np.complex128) + density_out[:] = np.nan + elif mulfac is None: + mulfac = density_out.shape[0] / N + Ns = density_out.shape[0] + assert (density_out.shape[0] % N) == 0 + + assert len(density_out.shape) == 3 + assert density_out.shape[0] == density_out.shape[1] + assert density_out.shape[2] == (density_out.shape[0]/2+1) + + hN = N/2 + density_out[:hN, :hN, :hN+1] = density[:hN, :hN, :] + density_out[:hN, (Ns-hN):Ns, :hN+1] = density[:hN, hN:, :] + density_out[(Ns-hN):Ns, (Ns-hN):Ns, :hN+1] = density[hN:, hN:, :] + density_out[(Ns-hN):Ns, :hN, :hN+1] = density[hN:, :hN, :] + + if mulfac > 1: + + cond=np.isnan(density_out) + if zero_fill: + density_out[cond] = 0 + else: + + if Pk is not None: + assert L is not None and h is not None + + @ct.timeit_quiet + def build_Pk(): + ik = np.fft.fftfreq(Ns, d=L/Ns)*2*np.pi + k = ne.evaluate('sqrt(kx**2 + ky**2 + kz**2)', {'kx':ik[:,None,None], 'ky':ik[None,:,None], 'kz':ik[None,None,:(Ns/2+1)]}) + return Pk.compute(k)*(h*L)**3 + + print np.where(np.isnan(density_out))[0].size + Nz = np.count_nonzero(cond) + amplitude = np.sqrt(build_Pk()[cond]/2) if Pk is not None else (1.0/np.sqrt(2)) + density_out.real[cond] = np.random.randn(Nz) * amplitude + density_out.imag[cond] = np.random.randn(Nz) * amplitude + print np.where(np.isnan(density_out))[0].size + + # Now we have to fix the Nyquist plane + hNs = Ns/2 + nyquist = density_out[:, :, hNs] + Nplane = nyquist.size + nyquist.flat[:Nplane/2] = np.sqrt(2.0)*nyquist.flat[Nplane:Nplane/2:-1].conj() + + + return density_out + +@ct.timeit_quiet +def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, supergenerate=1, do_lpt2=True, shiftPixel=False, psi_instead=False, needvel=True, func='HU_WIGGLES'): """ Generate particles and velocities from a BORG snapshot. Returns a tuple of (positions,velocities,N,BoxSize,scale_factor).""" @@ -64,17 +122,27 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True, density, L = ba.half_pixel_shift(borg_vol, doshift=shiftPixel) - lpt = LagrangianPerturbation(density, L, fourier=True, supersample=supersample) - # Generate grid - posq = gen_posgrid(N*supersample, L) - vel= [] - posx = [] - # Compute LPT scaling coefficient D1 = cgrowth.D(a_ic) D1_0 = D1/cgrowth.D(a_borg) + Dborg = cgrowth.D(a_borg)/cgrowth.D(1.0) print "D1_0=%lg" % D1_0 + + if supergenerate>1: + print("Doing supergeneration (factor=%d)" % supergenerate) + p = ct.CosmologyPower(**cosmo) + p.setFunction(func) + p.normalize(cosmo['SIGMA8']*Dborg) + density = do_supergenerate(density,mulfac=supergenerate,Pk=p,L=L,h=cosmo['h']) + + lpt = LagrangianPerturbation(density, L, fourier=True, supersample=supersample) + + # Generate grid + posq = gen_posgrid(N*supergenerate, L) + vel= [] + posx = [] + velmul = cgrowth.compute_velmul(a_ic) if not psi_instead else 1 D2 = -3./7 * D1_0**2 @@ -100,7 +168,7 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True, density = lpt.cube.irfft() density *= (cgrowth.D(1)/cgrowth.D(a_borg)) - return posx,vel,density,N*supersample,L,a_ic,cosmo + return posx,vel,density,N*supergenerate*supersample,L,a_ic,cosmo @ct.timeit_quiet @@ -127,39 +195,7 @@ def whitify(density, L, cosmo, supergenerate=1, zero_fill=False, func='HU_WIGGLE 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 - if supergenerate > 1: - - cond=np.isnan(density_hat_super) - if zero_fill: - density_hat_super[cond] = 0 - else: - - print np.where(np.isnan(density_hat_super))[0].size - Nz = np.count_nonzero(cond) - density_hat_super.real[cond] = np.random.randn(Nz) - density_hat_super.imag[cond] = np.random.randn(Nz) - density_hat_super[cond] /= np.sqrt(2.0) - print np.where(np.isnan(density_hat_super))[0].size - - # 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() - - print np.where(np.isnan(density_hat_super))[0].size + density_hat_super = do_supergenerate(density_hat, mulfac=supergenerate) cube = CubeFT(L, Ns) cube.dhat = density_hat_super diff --git a/python_sample/icgen/cosmogrowth.py b/python_sample/icgen/cosmogrowth.py index 4e9bacb..234524c 100644 --- a/python_sample/icgen/cosmogrowth.py +++ b/python_sample/icgen/cosmogrowth.py @@ -116,7 +116,7 @@ class LagrangianPerturbation(object): def _gradient(self, phi, direction): ne.evaluate('phi_hat * i * kv / (kx**2 + ky**2 + kz**2)', out=self.cube.dhat, local_dict={'i':-1j, 'phi_hat':phi, 'kv':self._kdir(direction), - 'kx':self._kx, 'ky':self._ky, 'kz':self._kz} + 'kx':self._kx, 'ky':self._ky, 'kz':self._kz},casting='unsafe') # self.cube.dhat = self._kdir(direction)*1j*phi self.cube.dhat[0,0,0] = 0 return self.cube.irfft() @@ -159,8 +159,10 @@ class LagrangianPerturbation(object): # k2 = self._get_k2() # k2[0,0,0] = 1 - potgen0 = lambda i: ne.evaluate('kdir**2*d/k2',out=self.cube.dhat,local_dict={'kdir':self._kdir(i),'d':self.dhat,'k2':k2} ) - potgen = lambda i,j: ne.evaluate('kdir0*kdir1*d/k2',out=self.cube.dhat,local_dict={'kdir0':self._kdir(i),'kdir1':self._kdir(j),'d':self.dhat,'k2':k2} ) + inv_k2 = ne.evaluate('1/(kx**2+ky**2+kz**2)', {'kx':self._kdir(0),'ky':self._kdir(1),'kz':self._kdir(2)}) + inv_k2[0,0,0]=0 + potgen0 = lambda i: ne.evaluate('kdir**2*d*ik2',out=self.cube.dhat,local_dict={'kdir':self._kdir(i),'d':self.dhat,'ik2':inv_k2}, casting='unsafe' ) + potgen = lambda i,j: ne.evaluate('kdir0*kdir1*d*ik2',out=self.cube.dhat,local_dict={'kdir0':self._kdir(i),'kdir1':self._kdir(j),'d':self.dhat,'ik2':inv_k2}, casting='unsafe' ) if 'lpt2_potential' not in self.cache: print("Rebuilding potential...") @@ -169,7 +171,7 @@ class LagrangianPerturbation(object): q = self._do_irfft( potgen0(j) ).copy() for i in xrange(j+1, 3): ne.evaluate('div + q * pot', out=div_phi2, - local_dict={'q':q,'pot':self._do_irfft( potgen0(i), copy=False ) } + local_dict={'div':div_phi2, 'q':q,'pot':self._do_irfft( potgen0(i), copy=False ) } ) ne.evaluate('div - pot**2',out=div_phi2, local_dict={'div':div_phi2,'pot':self._do_irfft(potgen(i,j), copy=False) } diff --git a/python_sample/icgen/gen_ic_from_borg.py b/python_sample/icgen/gen_ic_from_borg.py index ab5024c..8716b2f 100644 --- a/python_sample/icgen/gen_ic_from_borg.py +++ b/python_sample/icgen/gen_ic_from_borg.py @@ -14,11 +14,11 @@ cosmo['omega_B_0']=0.049 cosmo['SIGMA8']=0.8344 cosmo['ns']=0.9624 -supergen=8 -zstart=50 +supergen=4 +zstart=99 astart=1/(1.+zstart) halfPixelShift=False zero_fill=False if __name__=="__main__": - bic.write_icfiles(*bic.run_generation("initial_density_1872.dat", 0.001, astart, cosmo, supersample=1, shiftPixel=halfPixelShift, do_lpt2=False), supergenerate=supergen, zero_fill=zero_fill) + bic.write_icfiles(*bic.run_generation("initial_density_1872.dat", 0.001, astart, cosmo, supersample=1, shiftPixel=halfPixelShift, do_lpt2=False, supergenerate=supergen), supergenerate=1, zero_fill=zero_fill) diff --git a/python_sample/icgen/test_ic_from_borg.py b/python_sample/icgen/test_ic_from_borg.py index 7f1355a..566af86 100644 --- a/python_sample/icgen/test_ic_from_borg.py +++ b/python_sample/icgen/test_ic_from_borg.py @@ -37,7 +37,7 @@ if doSimulation: dsim_hat = np.fft.rfftn(dsim)*(L/N0)**3 Psim, bsim = bic.bin_power(np.abs(dsim_hat)**2/L**3, L, range=(0,1.), bins=150) -pos,_,density,N,L,_,_ = bic.run_generation("initial_density_988.dat", 0.001, astart, cosmo, supersample=1, do_lpt2=True) +pos,_,density,N,L,_,_ = bic.run_generation("initial_density_1872.dat", 0.001, astart, cosmo, supersample=1, do_lpt2=False, supergenerate=2) dcic = ct.cicParticles(pos, L, N0) dcic /= np.average(np.average(np.average(dcic, axis=0), axis=0), axis=0) @@ -46,17 +46,16 @@ dcic -= 1 dcic_hat = np.fft.rfftn(dcic)*(L/N0)**3 dens_hat = np.fft.rfftn(density)*(L/N0)**3 -Pcic, bcic = bic.bin_power(np.abs(dcic_hat)**2/L**3, L, range=(0,1.), bins=150) -Pdens, bdens = bic.bin_power(np.abs(dens_hat)**2/L**3, L, range=(0,1.), bins=150) +Pcic, bcic = bic.bin_power(np.abs(dcic_hat)**2/L**3, L, range=(0,4.), bins=150) +Pdens, bdens = bic.bin_power(np.abs(dens_hat)**2/L**3, L, range=(0,4.), bins=150) cgrowth = cg.CosmoGrowth(**cosmo) D1 = cgrowth.D(astart) D1_0 = D1/cgrowth.D(1)#0.001) -Pref, bref = bic.compute_ref_power(L, N0, cosmo, range=(0,1.), bins=150) +Pref, bref = bic.compute_ref_power(L, N0, cosmo, range=(0,4.), bins=150) 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