New python samples. New python function to write grafic files

This commit is contained in:
Guilhem Lavaux 2014-06-01 15:32:15 +02:00
parent 5d3b56bfcb
commit 8258bd0fe9
4 changed files with 165 additions and 0 deletions

View File

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

View File

@ -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))

View File

@ -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

View File

@ -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))