New python samples. New python function to write grafic files

python/cosmotool/__init__.py

@@ -1,4 +1,5 @@
 from _cosmotool import *
+from grafic import writeGrafic
 from borg import read_borg_vol
 
 

python/cosmotool/grafic.py

@@ -0,0 +1,24 @@
+import struct
+import numpy as np
+
+def writeGrafic(filename, field, BoxSize, scalefac, **cosmo):
+
+    with file(filename, mode="wb") as f:
+        checkPoint = 4*11
+        
+        Nx,Ny,Nz = field.shape
+        delta = BoxSize/Nx
+        bad = 0.0
+        
+        f.write(struct.pack("IIIIffffffffI", checkPoint, 
+                                          Nx, Ny, Nz, 
+                                          delta, 
+                                          bad, bad, bad, 
+                                          scalefac,
+                                          cosmo['omega_M_0'], cosmo['omega_lambda_0'], 100*cosmo['h'], checkPoint))
+        checkPoint = 4*Ny*Nz
+        for i in xrange(Nx):
+            f.write(struct.pack("I", checkPoint))
+            f.write(field[i].astype(np.float32).tostring())
+            f.write(struct.pack("I", checkPoint))
+        

python_sample/cosmogrowth.py

@@ -0,0 +1,37 @@
+import numpy as np
+import cosmolopy as cpy
+
+class CosmoGrowth:
+
+  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
+

python_sample/gen_ic_from_borg.py

@@ -0,0 +1,103 @@
+import cosmotool as ct
+import numpy as np
+import cosmolopy as cpy
+from cosmogrowth import *
+
+cosmo={'omega_M_0':0.3175, 'h':0.6711}
+cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
+cosmo['omega_k_0'] = 0
+
+a_start=0.001
+z_start=1/a_start-1
+
+
+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 half_pixel_shift(borg):
+
+    dhat,L,N = fourier_analysis(borg)
+
+    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)])
+    phase = np.cos(phi)+1j*np.sin(phi)
+    
+    return dhat*phase, L
+
+def new_shape(N, direction, q=3):
+    return ((1,)*direction) + (N,) + ((1,)*(q-1-direction))
+
+def build_dir(ik, direction, q=3):
+    if direction != q-1:
+        return ik.reshape(new_shape(ik.size, direction, q=q))
+    else:
+        N = ik.size/2+1
+        return ik[:N].reshape(new_shape(N, direction, q=q))
+
+def get_k2(ik, q=3):
+
+    N = ik.size
+    k2 = (ik.reshape(new_shape(N, 0, q=q))**2)
+    for d in xrange(1,q):
+        k2 = k2 + build_dir(ik, d, q=q)**2
+    return k2
+
+def get_displacement(dhat, L, direction=0):
+
+    N = dhat.shape[0]
+    ik = np.fft.fftfreq(N,d=1.0/N)*2*np.pi/L
+    
+    k2 = get_k2(ik)
+    k2[0,0,0] = 1
+    
+    return -build_dir(ik, direction)*1j*dhat / k2
+
+def gen_posgrid(N, L):
+
+    ix = np.arange(N)*L/N
+
+    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.flatten(), y.flatten(), z.flatten()
+
+
+def run_generation(input_borg, a_ic):
+
+    borg_vol = ct.read_borg_vol(input_borg)
+    N = borg_vol.density.shape[0]
+
+    cgrowth = CosmoGrowth(**cosmo)
+
+    density_hat, L = half_pixel_shift(borg_vol)
+
+    posq = gen_posgrid(N, L)
+    vel= []
+    posx = []
+    velmul = cgrowth.compute_velmul(a_start)
+
+    D1 = cgrowth.D(a_ic)
+    D1_0 = D1/cgrowth.D(a_start)
+    D2 = 3./7 * D1**2
+
+    for j in xrange(3):
+        psi = D1_0*np.fft.irfftn(get_displacement(density_hat, L, direction=j)).flatten()*(N/L)**3
+        posx.append(((posq[j] + psi)%L).astype(np.float32))
+        vel.append((psi*velmul).astype(np.float32)) 
+
+    return posx,vel,N,L,a_ic
+    
+def write_icfiles(*generated_ic):
+    posx,vel,N,L,a_ic = generated_ic
+
+    ct.simpleWriteGadget("borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic) 
+    for i,c in enumerate(['x','y','z']):
+        ct.writeGrafic("borg.ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
+        
+        
+if __name__=="__main__":
+    write_icfiles(*run_generation("initial_condition_borg.dat", 1.0))