diff --git a/python_sample/icgen/borgadaptor.py b/python_sample/icgen/borgadaptor.py index 93fcd9b..6d906a5 100644 --- a/python_sample/icgen/borgadaptor.py +++ b/python_sample/icgen/borgadaptor.py @@ -9,6 +9,7 @@ def fourier_analysis(borg_vol): def half_pixel_shift(borg): dhat,L,N = fourier_analysis(borg) + return dhat, L 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)]) diff --git a/python_sample/icgen/borgicgen.py b/python_sample/icgen/borgicgen.py index f3a51ce..d76ed0c 100644 --- a/python_sample/icgen/borgicgen.py +++ b/python_sample/icgen/borgicgen.py @@ -50,7 +50,7 @@ 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): +def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True): """ Generate particles and velocities from a BORG snapshot. Returns a tuple of (positions,velocities,N,BoxSize,scale_factor).""" @@ -61,10 +61,10 @@ def run_generation(input_borg, a_borg, a_ic, **cosmo): density_hat, L = ba.half_pixel_shift(borg_vol) - lpt = LagrangianPerturbation(density_hat, L, fourier=True) + lpt = LagrangianPerturbation(density_hat, L, fourier=True, supersample=supersample) # Generate grid - posq = gen_posgrid(N, L) + posq = gen_posgrid(N*supersample, L) vel= [] posx = [] @@ -73,11 +73,15 @@ def run_generation(input_borg, a_borg, a_ic, **cosmo): D1_0 = D1/cgrowth.D(a_borg) velmul = cgrowth.compute_velmul(a_ic)*D1_0 - D2 = 3./7 * D1**2 + D2 = -3./7 * D1_0**2 for j in xrange(3): # Generate psi_j (displacement along j) + print("LPT1 axis=%d" % j) psi = D1_0*lpt.lpt1(j).flatten() + if do_lpt2: + print("LPT2 axis=%d" % j) + psi += D2 * lpt.lpt2(j).flatten() # Generate posx posx.append(((posq[j] + psi)%L).astype(np.float32)) # Generate vel diff --git a/python_sample/icgen/cosmogrowth.py b/python_sample/icgen/cosmogrowth.py index bee7fd3..6263126 100644 --- a/python_sample/icgen/cosmogrowth.py +++ b/python_sample/icgen/cosmogrowth.py @@ -42,13 +42,29 @@ class CosmoGrowth(object): class LagrangianPerturbation(object): - def __init__(self,density,L, fourier=False): + def __init__(self,density,L, fourier=False, supersample=1): self.L = L self.N = density.shape[0] self.dhat = np.fft.rfftn(density)*(L/self.N)**3 if not fourier else density + if supersample > 1: + self.upgrade_sampling(supersample) self.ik = np.fft.fftfreq(self.N, d=L/self.N)*2*np.pi - self.cache = weakref.WeakValueDictionary() + self.cache = {}#weakref.WeakValueDictionary() + + def upgrade_sampling(self, supersample): + N2 = self.N * supersample + N = self.N + dhat_new = np.zeros((N2, N2, N2/2+1), dtype=np.complex128) + + hN = N/2 + dhat_new[:hN, :hN, :hN+1] = self.dhat[:hN, :hN, :] + dhat_new[:hN, (N2-hN):N2, :hN+1] = self.dhat[:hN, hN:, :] + dhat_new[(N2-hN):N2, (N2-hN):N2, :hN+1] = self.dhat[hN:, hN:, :] + dhat_new[(N2-hN):N2, :hN, :hN+1] = self.dhat[hN:, :hN, :] + + self.dhat = dhat_new + self.N = N2 def _gradient(self, phi, direction): return np.fft.irfftn(self._kdir(direction)*1j*phi)*(self.N/self.L)**3 @@ -85,15 +101,16 @@ class LagrangianPerturbation(object): k2[0,0,0] = 1 if 'lpt2_potential' not in self.cache: - div_phi2 = np.zeros((N,N,N), dtype=np.float64) + print("Rebuilding potential...") + div_phi2 = np.zeros((self.N,self.N,self.N), dtype=np.float64) for j in xrange(3): - q = np.fft.irfftn( build_dir(ik, j)**2*self.dhat / k2 ) + q = np.fft.irfftn( self._kdir(j)**2*self.dhat / k2 ) for i in xrange(j+1, 3): - div_phi2 += q * np.fft.irfftn( build_dir(ik, i)**2*self.dhat / k2 ) - div_phi2 -= (np.fft.irfftn( build_dir(ik, j)*build_dir(ik, i)*self.dhat / k2 ))**2 + div_phi2 += q * np.fft.irfftn( self._kdir(i)**2*self.dhat / k2 ) + div_phi2 -= (np.fft.irfftn( self._kdir(j)*self._kdir(i)*self.dhat / k2 ))**2 - div_phi2 *= (self.N/self.L)**3 - phi2_hat = np.fft.rfftn(div_phi2) * ((L/N)**3) / k2 + div_phi2 *= (self.N/self.L)**6 + 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 5f9d210..81d5d35 100644 --- a/python_sample/icgen/gen_ic_from_borg.py +++ b/python_sample/icgen/gen_ic_from_borg.py @@ -1,3 +1,4 @@ +import numpy as np import cosmotool as ct import borgicgen as bic @@ -7,12 +8,23 @@ cosmo['omega_k_0'] = 0 cosmo['omega_B_0']=0.049 cosmo['SIGMA8']=0.8344 -zstart=0 +TestCase=True +zstart=10 astart=1/(1.+zstart) -pos,_,density,N,L,_ = bic.run_generation("initial_condition_borg.dat", 0.001, astart, **cosmo) +if TestCase: + pos,_,density,N,L,_ = bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=1, do_lpt2=True) -dcic = ct.cicParticles(pos, L, N) + dcic = ct.cicParticles(pos, L, N) + dcic /= np.average(np.average(np.average(dcic, axis=0), axis=0), axis=0) + dcic -= 1 -#if __name__=="__main__": -# bic.write_icfiles(*bic.run_generation("initial_condition_borg.dat", 0.001, astart, **cosmo), **cosmo) + dcic_hat = np.fft.rfftn(dcic)*(L/N)**3 + + Pcic, bcic = bic.bin_power(np.abs(dcic_hat)**2/L**3, L, bins=50) + + borg_evolved = ct.read_borg_vol("final_density_1380.dat") + +if __name__=="__main__": + if not TestCase: + bic.write_icfiles(*bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, do_lpt2=True), **cosmo)