Updated icgen to support padded phases from BORG

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 = do_supergenerate(density_hat, mulfac=supergenerate)
-  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
 
   cube = CubeFT(L, Ns)
   cube.dhat = density_hat_super

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} )
+        inv_k2 = ne.evaluate('1/(kx**2+ky**2+kz**2)', {'kx':self._kdir(0),'ky':self._kdir(1),'kz':self._kdir(2)})
-        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[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) }

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
+supergen=4
-zstart=50
+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)

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)
+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,1.), 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