62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
|
import cosmotool as ct
|
||
|
|
import numpy as np
|
||
|
|
import cosmolopy as cpy
|
||
|
|
from cosmogrowth import *
|
||
|
|
import borgadaptor as ba
|
||
|
|
|
||
|
|
def gen_posgrid(N, L):
|
||
|
|
""" Generate an ordered lagrangian grid"""
|
||
|
|
|
||
|
|
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_borg, a_ic, **cosmo):
|
||
|
|
""" 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_hat, L = ba.half_pixel_shift(borg_vol)
|
||
|
|
|
||
|
|
|
||
|
|
lpt = LagrangianPerturbation(density_hat, L, fourier=True)
|
||
|
|
|
||
|
|
# Generate grid
|
||
|
|
posq = gen_posgrid(N, L)
|
||
|
|
vel= []
|
||
|
|
posx = []
|
||
|
|
|
||
|
|
# Compute LPT scaling coefficient
|
||
|
|
D1 = cgrowth.D(a_ic)
|
||
|
|
D1_0 = D1/cgrowth.D(a_borg)
|
||
|
|
velmul = cgrowth.compute_velmul(a_ic)*D1_0
|
||
|
|
|
||
|
|
D2 = 3./7 * D1**2
|
||
|
|
|
||
|
|
for j in xrange(3):
|
||
|
|
# Generate psi_j (displacement along j)
|
||
|
|
psi = D1_0*lpt.lpt1(j).flatten()*(N/L)**3
|
||
|
|
# Generate posx
|
||
|
|
posx.append(((posq[j] + psi)%L).astype(np.float32))
|
||
|
|
# Generate vel
|
||
|
|
vel.append((psi*velmul).astype(np.float32))
|
||
|
|
|
||
|
|
return posx,vel,N,L,a_ic
|
||
|
|
|
||
|
|
def write_icfiles(*generated_ic, **cosmo):
|
||
|
|
"""Write the initial conditions from the tuple returned by run_generation"""
|
||
|
|
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)
|
||
|
|
|