cosmotool/python_sample/icgen/borgicgen.py

110 lines
3.6 KiB
Python
Raw Normal View History

2014-06-01 18:07:44 +02:00
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)
2014-06-04 13:57:17 +02:00
return x.reshape((x.size,)), y.reshape((y.size,)), z.reshape((z.size,))
2014-06-01 18:07:44 +02:00
def bin_power(P, L, bins=20, range=(0,1.)):
N = P.shape[0]
ik = np.fft.fftfreq(N, d=L/N)*2*np.pi
k = np.sqrt(ik[:,None,None]**2 + ik[None,:,None]**2 + ik[None,None,:(N/2+1)]**2)
H,b = np.histogram(k, bins=bins, range=range)
Hw,b = np.histogram(k, bins=bins, weights=P, range=range)
return Hw/H, 0.5*(b[1:]+b[0:bins])
def compute_power_from_borg(input_borg, a_borg, cosmo, bins=10, range=(0,1)):
borg_vol = ct.read_borg_vol(input_borg)
N = borg_vol.density.shape[0]
cgrowth = CosmoGrowth(**cosmo)
D1 = cgrowth.D(1)
D1_0 = D1/cgrowth.D(a_borg)
print("D1_0=%lg" % D1_0)
density_hat, L = ba.half_pixel_shift(borg_vol)
return bin_power(D1_0**2*np.abs(density_hat)**2/L**3, L, bins=bins, range=range)
def compute_ref_power(L, N, cosmo, bins=10, range=(0,1), func='HU_WIGGLES'):
ik = np.fft.fftfreq(N, d=L/N)*2*np.pi
k = np.sqrt(ik[:,None,None]**2 + ik[None,:,None]**2 + ik[None,None,:(N/2+1)]**2)
p = ct.CosmologyPower(**cosmo)
p.setFunction(func)
p.normalize(cosmo['SIGMA8'])
return bin_power(p.compute(k)*cosmo['h']**3, L, bins=bins, range=range)
2014-06-04 09:13:34 +02:00
def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True, shiftPixel=False):
2014-06-01 18:07:44 +02:00
""" 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)
2014-06-04 09:13:34 +02:00
density_hat, L = ba.half_pixel_shift(borg_vol, doshift=shiftPixel)
2014-06-01 18:07:44 +02:00
lpt = LagrangianPerturbation(density_hat, L, fourier=True, supersample=supersample)
2014-06-01 18:07:44 +02:00
# Generate grid
posq = gen_posgrid(N*supersample, L)
2014-06-01 18:07:44 +02:00
vel= []
posx = []
# Compute LPT scaling coefficient
D1 = cgrowth.D(a_ic)
D1_0 = D1/cgrowth.D(a_borg)
2014-06-03 18:50:04 +02:00
velmul = cgrowth.compute_velmul(a_ic)
2014-06-01 18:07:44 +02:00
2014-06-04 13:57:17 +02:00
D2 = 3./7 * D1_0**2
2014-06-01 18:07:44 +02:00
for j in xrange(3):
# Generate psi_j (displacement along j)
print("LPT1 axis=%d" % j)
2014-06-04 13:57:17 +02:00
psi = D1_0*lpt.lpt1(j)
psi = psi.reshape((psi.size,))
if do_lpt2:
print("LPT2 axis=%d" % j)
2014-06-04 13:57:17 +02:00
psi2 = lpt.lpt2(j)
psi += D2 * psi2.reshape((psi2.size,))
2014-06-01 18:07:44 +02:00
# Generate posx
posx.append(((posq[j] + psi)%L).astype(np.float32))
# Generate vel
vel.append((psi*velmul).astype(np.float32))
2014-06-03 18:50:04 +02:00
print("velmul=%lg" % (cosmo['h']*velmul))
2014-06-04 09:13:34 +02:00
density = np.fft.irfftn(lpt.dhat*D1_0)*(supersample*N/L)**3
2014-06-04 09:13:34 +02:00
return posx,vel,density,N*supersample,L,a_ic
2014-06-01 18:07:44 +02:00
def write_icfiles(*generated_ic, **cosmo):
"""Write the initial conditions from the tuple returned by run_generation"""
posx,vel,density,N,L,a_ic = generated_ic
2014-06-01 18:07:44 +02:00
ct.simpleWriteGadget("borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic)
2014-06-04 09:13:34 +02:00
for i,c in enumerate(["x","y","z"]):
2014-06-04 13:57:17 +02:00
ct.writeGrafic("ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
# This used to be necessary. However this has been fixed in writeGrafic now
# ct.writeGrafic("ic_velc%s" % c, vel[i].reshape((N,N,N)).transpose(), L, a_ic, **cosmo)
2014-06-01 18:07:44 +02:00
ct.writeGrafic("ic_deltab", density, L, a_ic, **cosmo)
ct.writeWhitePhase("white.dat", density)