104 lines
2.7 KiB
Python
104 lines
2.7 KiB
Python
|
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))
|