Update cosmotool 2nd part

external/cosmotool/python_sample/icgen/borgicgen.py

@@ -0,0 +1,228 @@
+import cosmotool as ct
+import numpy as np
+import cosmolopy as cpy
+from cosmogrowth import *
+import borgadaptor as ba
+
+@ct.timeit
+def gen_posgrid(N, L, delta=1, dtype=np.float32):
+    """ Generate an ordered lagrangian grid"""
+
+    ix = (np.arange(N)*(L/N*delta)).astype(dtype)
+
+    x = ix[:,None,None].repeat(N, axis=1).repeat(N, axis=2)
+    y = ix[None,:,None].repeat(N, axis=0).repeat(N, axis=2)
+    z = ix[None,None,:].repeat(N, axis=0).repeat(N, axis=1)
+
+    return x.reshape((x.size,)), y.reshape((y.size,)), z.reshape((z.size,))
+
+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
+
+  k = np.sqrt(ik[:,None,None]**2 + ik[None,:,None]**2 + ik[None,None,:(N/2+1)]**2)
+
+  H,b = np.histogram(k, bins=bins, range=range)
+  Hw,b = np.histogram(k, bins=bins, weights=P, range=range)
+  
+  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)
+  N = borg_vol.density.shape[0]
+
+  cgrowth = CosmoGrowth(**cosmo)
+  D1 = cgrowth.D(1)
+  D1_0 = D1/cgrowth.D(a_borg)
+  print("D1_0=%lg" % D1_0)
+
+  density_hat, L = ba.half_pixel_shift(borg_vol)
+
+  return bin_power(D1_0**2*np.abs(density_hat)**2/L**3, L, bins=bins, range=range)
+  
+def compute_ref_power(L, N, cosmo, bins=10, range=(0,1), func='HU_WIGGLES'):
+  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)
+  p = ct.CosmologyPower(**cosmo)
+  p.setFunction(func)
+  p.normalize(cosmo['SIGMA8'])
+
+  return bin_power(p.compute(k)*cosmo['h']**3, L, bins=bins, range=range)
+
+
+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)*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)."""
+
+    borg_vol = ct.read_borg_vol(input_borg)
+    N = borg_vol.density.shape[0]
+
+    cgrowth = CosmoGrowth(**cosmo)
+
+    density, L = ba.half_pixel_shift(borg_vol, doshift=shiftPixel)
+
+
+    # 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*supersample, L)
+    vel= []
+    posx = []
+    
+    velmul = cgrowth.compute_velmul(a_ic) if not psi_instead else 1
+    
+    D2 = -3./7 * D1_0**2
+
+    if do_lpt2:
+      psi2 = lpt.lpt2('all')
+    for j in xrange(3):
+        # Generate psi_j (displacement along j)
+        print("LPT1 axis=%d" % j)
+        psi = D1_0*lpt.lpt1(j)
+        psi = psi.reshape((psi.size,))
+        if do_lpt2:
+          print("LPT2 axis=%d" % j)
+          psi += D2 * psi2[j].reshape((psi2[j].size,))
+        # Generate posx
+        posx.append(((posq[j] + psi)%L).astype(np.float32))
+        # Generate vel
+        if needvel:
+          vel.append((psi*velmul).astype(np.float32)) 
+
+    print("velmul=%lg" % (cosmo['h']*velmul))
+
+    lpt.cube.dhat = lpt.dhat
+    density = lpt.cube.irfft()
+    density *= (cgrowth.D(1)/cgrowth.D(a_borg))
+
+    return posx,vel,density,N*supergenerate*supersample,L,a_ic,cosmo
+
+
+@ct.timeit_quiet
+def whitify(density, L, cosmo, supergenerate=1, zero_fill=False, func='HU_WIGGLES'):
+
+  N = density.shape[0]
+  p = ct.CosmologyPower(**cosmo)
+  p.setFunction(func)
+  p.normalize(cosmo['SIGMA8'])
+
+  @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)*L**3
+    
+  Pk = build_Pk()
+  Pk[0,0,0]=1
+
+  cube = CubeFT(L, N)
+  cube.density = density
+  density_hat = cube.rfft()
+  density_hat /= np.sqrt(Pk)
+
+  Ns = N*supergenerate
+
+  density_hat_super = do_supergenerate(density_hat, mulfac=supergenerate)
+
+  cube = CubeFT(L, Ns)
+  cube.dhat = density_hat_super
+  return np.fft.irfftn(density_hat_super)*Ns**1.5
+
+
+    
+def write_icfiles(*generated_ic, **kwargs):
+  """Write the initial conditions from the tuple returned by run_generation"""
+  
+  supergenerate=kwargs.get('supergenerate', 1)
+  zero_fill=kwargs.get('zero_fill', False)
+  posx,vel,density,N,L,a_ic,cosmo = generated_ic
+
+  ct.simpleWriteGadget("Data/borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic) 
+  for i,c in enumerate(["z","y","x"]):
+    ct.writeGrafic("Data/ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
+
+  ct.writeGrafic("Data/ic_deltab", density, L, a_ic, **cosmo)
+  
+  ct.writeWhitePhase("Data/white.dat", whitify(density, L, cosmo, supergenerate=supergenerate,zero_fill=zero_fill))
+  
+  with file("Data/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")
+