Update cosmotool 2nd part

external/cosmotool/python_sample/icgen/__init__.py

@@ -0,0 +1,2 @@
+from borgicgen import *
+import cosmogrowth

external/cosmotool/python_sample/icgen/borgadaptor.py

@@ -0,0 +1,55 @@
+import numpy as np
+
+def fourier_analysis(borg_vol):
+  L = (borg_vol.ranges[1]-borg_vol.ranges[0])
+  N = borg_vol.density.shape[0]
+    
+  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, L
+

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")
+

external/cosmotool/python_sample/icgen/cosmogrowth.py

@@ -0,0 +1,202 @@
+import numexpr as ne
+import multiprocessing
+import pyfftw
+import weakref
+import numpy as np
+import cosmolopy as cpy
+import cosmotool as ct
+
+class CubeFT(object):
+  def __init__(self, L, N, max_cpu=-1):
+  
+      self.N = N
+      self.align = pyfftw.simd_alignment
+      self.L = L
+      self.max_cpu = multiprocessing.cpu_count() if max_cpu < 0 else max_cpu
+      self._dhat = pyfftw.n_byte_align_empty((self.N,self.N,self.N/2+1), self.align, dtype='complex64')
+      self._density = pyfftw.n_byte_align_empty((self.N,self.N,self.N), self.align, dtype='float32')
+      self._irfft = pyfftw.FFTW(self._dhat, self._density, axes=(0,1,2), direction='FFTW_BACKWARD', threads=self.max_cpu, normalize_idft=False)
+      self._rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False)
+
+  def rfft(self):
+      return ne.evaluate('c*a', local_dict={'c':self._rfft(normalise_idft=False),'a':(self.L/self.N)**3})
+      
+  def irfft(self):
+      return ne.evaluate('c*a', local_dict={'c':self._irfft(normalise_idft=False),'a':(1/self.L)**3})
+   
+  def get_dhat(self):
+      return self._dhat
+  def set_dhat(self, in_dhat):
+      self._dhat[:] = in_dhat 
+  dhat = property(get_dhat, set_dhat, None)
+      
+  def get_density(self):
+      return self._density      
+  def set_density(self, d):
+      self._density[:] = d 
+  density = property(get_density, set_density, None)
+      
+
+class CosmoGrowth(object):
+
+  def __init__(self, **cosmo):
+    self.cosmo = cosmo
+
+  def D(self, a):
+    return cpy.perturbation.fgrowth(1/a-1, self.cosmo['omega_M_0'], unnormed=True)
+
+  def compute_E(self, a):
+    om = self.cosmo['omega_M_0']
+    ol = self.cosmo['omega_lambda_0']
+    ok = self.cosmo['omega_k_0']
+
+    E = np.sqrt(om/a**3 + ol + ok/a**2)
+  
+    H2 = -3*om/a**4  - 2*ok/a**3
+
+    Eprime = 0.5*H2/E
+
+    return E,Eprime
+
+  def Ddot(self, a):
+    E,Eprime = self.compute_E(a)
+    D = self.D(a)
+    Ddot_D = Eprime/E + 2.5 * self.cosmo['omega_M_0']/(a**3*E**2*D)
+    Ddot_D *= a
+    return Ddot_D
+
+  def compute_velmul(self, a):
+    E,_ = self.compute_E(a)
+    velmul = self.Ddot(a)
+    velmul *= 100 * a * E
+    return velmul
+
+
+
+
+
+class LagrangianPerturbation(object):
+
+    def __init__(self,density,L, fourier=False, supersample=1, max_cpu=-1):
+    
+        self.L = L
+        self.N = density.shape[0]
+        
+        self.max_cpu = max_cpu
+        self.cube = CubeFT(self.L, self.N, max_cpu=max_cpu)
+        
+        if not fourier:
+          self.cube.density = density
+          self.dhat = self.cube.rfft().copy()
+        else:
+          self.dhat = density.copy()
+
+        if supersample > 1:
+          self.upgrade_sampling(supersample)
+        self.ik = np.fft.fftfreq(self.N, d=L/self.N)*2*np.pi
+        self._kx = self.ik[:,None,None]
+        self._ky = self.ik[None,:,None]
+        self._kz = self.ik[None,None,:(self.N/2+1)]
+        self.cache = {}#weakref.WeakValueDictionary()
+
+    @ct.timeit_quiet
+    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
+        self.cube = CubeFT(self.L, self.N, max_cpu=self.max_cpu)
+        
+    @ct.timeit_quiet
+    def _gradient(self, phi, direction):
+        if direction == 'all':
+          dirs = [0,1,2]
+          copy = True
+        else:
+          dirs = [direction]
+          copy = False
+        ret=[]
+        for dir in dirs:
+          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(dir),
+                          '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
+          x = self.cube.irfft()
+          ret.append(x.copy() if copy else x)
+        return ret[0] if len(ret)==1 else ret
+            
+    @ct.timeit_quiet
+    def lpt1(self, direction=0):
+        return self._gradient(self.dhat, direction)
+        
+    def new_shape(self,direction, q=3, half=False):
+        N0 = (self.N/2+1) if half else self.N
+        return ((1,)*direction) + (N0,) + ((1,)*(q-1-direction))
+
+    def _kdir(self, direction, q=3):
+        if direction != q-1:
+            return self.ik.reshape(self.new_shape(direction, q=q))
+        else:
+            return self.ik[:self.N/2+1].reshape(self.new_shape(direction, q=q, half=True))
+
+    def _get_k2(self, q=3):
+        if 'k2' in self.cache:
+            return self.cache['k2']
+            
+        k2 = self._kdir(0, q=q)**2
+        for d in xrange(1,q):
+            k2 = k2 + self._kdir(d, q=q)**2
+            
+        self.cache['k2'] = k2
+        return k2
+        
+    def _do_irfft(self, array, copy=True):
+        if copy:
+          self.cube.dhat = array
+        return self.cube.irfft()
+        
+    def _do_rfft(self, array, copy=True):
+        if copy:
+          self.cube.density = array
+        return self.cube.rfft()
+
+    @ct.timeit_quiet
+    def lpt2(self, direction=0):
+#        k2 = self._get_k2()
+#        k2[0,0,0] = 1
+
+        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...")
+            div_phi2 = np.zeros((self.N,self.N,self.N), dtype=np.float64)
+            for j in xrange(3):
+                q = self._do_irfft( potgen0(j) ).copy()
+                for i in xrange(j+1, 3):
+                    with ct.time_block("LPT2 elemental (%d,%d)" %(i,j)):
+                      ne.evaluate('div + q * pot', out=div_phi2, 
+                            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) }
+                            )
+
+            phi2_hat = self._do_rfft(div_phi2)
+            #self.cache['lpt2_potential'] = phi2_hat
+            del div_phi2
+        else:
+            phi2_hat = self.cache['lpt2_potential']
+            
+        return self._gradient(phi2_hat, direction)

external/cosmotool/python_sample/icgen/gen_ic_from_borg.py

@@ -0,0 +1,24 @@
+import pyfftw
+import numpy as np
+import cosmotool as ct
+import borgicgen as bic
+import pickle
+
+with file("wisdom") as f:
+  pyfftw.import_wisdom(pickle.load(f))
+
+cosmo={'omega_M_0':0.3175, 'h':0.6711}
+cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
+cosmo['omega_k_0'] = 0
+cosmo['omega_B_0']=0.049
+cosmo['SIGMA8']=0.8344
+cosmo['ns']=0.9624
+
+supergen=1
+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), supergenerate=1, zero_fill=zero_fill)

external/cosmotool/python_sample/icgen/test_ic_from_borg.py

@@ -0,0 +1,63 @@
+import numpy as np
+import cosmotool as ct
+import borgicgen as bic
+import cosmogrowth as cg
+import sys
+
+cosmo={'omega_M_0':0.3175, 'h':0.6711}
+cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
+cosmo['omega_k_0'] = 0
+cosmo['omega_B_0']=0.049
+cosmo['SIGMA8']=0.8344
+cosmo['ns']=0.9624
+N0=256
+
+doSimulation=False
+simShift=False
+
+snap_id=int(sys.argv[1])
+astart=1/100.
+
+if doSimulation:
+  s = ct.loadRamsesAll("/nethome/lavaux/remote2/borgsim3/", snap_id, doublePrecision=True)
+  astart=s.getTime()
+  L = s.getBoxsize()
+
+  p = s.getPositions()
+  Nsim = int( np.round( p[0].size**(1./3)) )
+  print("Nsim = %d" % Nsim)
+
+  if simShift:
+    p = [(q-0.5*L/Nsim)%L for q in p]
+
+  dsim = ct.cicParticles(p[::-1], L, N0)
+  dsim /= np.average(np.average(np.average(dsim, axis=0), axis=0), axis=0)
+  dsim -= 1
+
+  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_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)
+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,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,4.), bins=150)
+
+Pcic /= 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)

external/cosmotool/python_sample/icgen/test_whitify.py

@@ -0,0 +1,41 @@
+import numpy as np
+import cosmotool as ct
+import borgicgen as bic
+from matplotlib import pyplot as plt
+
+cosmo={'omega_M_0':0.3175, 'h':0.6711}
+cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
+cosmo['omega_k_0'] = 0
+cosmo['omega_B_0']=0.049
+cosmo['SIGMA8']=0.8344
+cosmo['ns']=0.9624
+
+zstart=50
+astart=1/(1.+zstart)
+halfPixelShift=False
+
+posx,vel,density,N,L,a_ic,cosmo = bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=1, shiftPixel=halfPixelShift, do_lpt2=False)
+
+w1 = bic.whitify(density, L, cosmo, supergenerate=1)
+w2 = bic.whitify(density, L, cosmo, supergenerate=2)
+
+N = w1.shape[0]
+Ns = w2.shape[0]
+
+w1_hat = np.fft.rfftn(w1)*(L/N)**3
+w2_hat = np.fft.rfftn(w2)*(L/Ns)**3
+
+P1, b1, dev1 = bic.bin_power(np.abs(w1_hat)**2, L, range=(0,3),bins=150,dev=True)
+P2, b2, dev2 = bic.bin_power(np.abs(w2_hat)**2, L, range=(0,3),bins=150,dev=True)
+
+fig = plt.figure(1)
+fig.clf()
+plt.fill_between(b1, P1*(1-dev1), P1*(1+dev1), label='Supergen=1', color='b')
+plt.fill_between(b2, P2*(1-dev2), P2*(1+dev2), label='Supergen=2', color='g', alpha=0.5)
+ax = plt.gca()
+ax.set_xscale('log')
+plt.ylim(0.5,1.5)
+plt.xlim(1e-2,4)
+plt.axhline(1.0, color='red', lw=4.0)
+plt.legend()
+plt.show()