diff --git a/python/cosmotool/__init__.py b/python/cosmotool/__init__.py index c4dcfcd..ebac8cb 100644 --- a/python/cosmotool/__init__.py +++ b/python/cosmotool/__init__.py @@ -1,4 +1,5 @@ from _cosmotool import * +from grafic import writeGrafic from borg import read_borg_vol diff --git a/python/cosmotool/grafic.py b/python/cosmotool/grafic.py new file mode 100644 index 0000000..468d373 --- /dev/null +++ b/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)) + diff --git a/python_sample/cosmogrowth.py b/python_sample/cosmogrowth.py new file mode 100644 index 0000000..8e378e5 --- /dev/null +++ b/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 + diff --git a/python_sample/gen_ic_from_borg.py b/python_sample/gen_ic_from_borg.py new file mode 100644 index 0000000..56b4330 --- /dev/null +++ b/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))