187 lines
6.7 KiB
Python
187 lines
6.7 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).astype(np.float32)
|
|
|
|
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.reshape((x.size,)), y.reshape((y.size,)), z.reshape((z.size,))
|
|
|
|
def bin_power(P, L, bins=20, range=(0,1.), dev=False):
|
|
|
|
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)
|
|
|
|
if dev:
|
|
return Hw/(H-1), 0.5*(b[1:]+b[0:bins]), 1.0/np.sqrt(H)
|
|
else:
|
|
return Hw/(H-1), 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)
|
|
|
|
def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True, shiftPixel=False, needvel=True):
|
|
""" 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, L = ba.half_pixel_shift(borg_vol, doshift=shiftPixel)
|
|
|
|
lpt = LagrangianPerturbation(density, L, fourier=True, supersample=supersample)
|
|
|
|
# Generate grid
|
|
posq = gen_posgrid(N*supersample, 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)
|
|
|
|
D2 = -3./7 * D1_0**2
|
|
|
|
for j in xrange(3):
|
|
# Generate psi_j (displacement along j)
|
|
print("LPT1 axis=%d" % j)
|
|
psi = D1_0*lpt.lpt1(j)
|
|
psi = psi.reshape((psi.size,))
|
|
if do_lpt2:
|
|
print("LPT2 axis=%d" % j)
|
|
psi2 = lpt.lpt2(j)
|
|
psi += D2 * psi2.reshape((psi2.size,))
|
|
# Generate posx
|
|
posx.append(((posq[j] + psi)%L).astype(np.float32))
|
|
# Generate vel
|
|
if needvel:
|
|
vel.append((psi*velmul).astype(np.float32))
|
|
|
|
print("velmul=%lg" % (cosmo['h']*velmul))
|
|
|
|
lpt.cube.dhat = lpt.dhat
|
|
density = lpt.cube.irfft()
|
|
density *= (cgrowth.D(1)/cgrowth.D(a_borg))
|
|
|
|
return posx,vel,density,N*supersample,L,a_ic,cosmo
|
|
|
|
|
|
@ct.timeit_quiet
|
|
def whitify(density, L, cosmo, supergenerate=1, func='HU_WIGGLES'):
|
|
|
|
N = density.shape[0]
|
|
p = ct.CosmologyPower(**cosmo)
|
|
p.setFunction(func)
|
|
p.normalize(cosmo['SIGMA8'])
|
|
|
|
@ct.timeit_quiet
|
|
def build_Pk():
|
|
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)
|
|
return p.compute(k)*cosmo['h']**3*L**3
|
|
|
|
Pk = build_Pk()
|
|
Pk[0,0,0]=1
|
|
|
|
cube = CubeFT(N, L)
|
|
cube.density = density
|
|
density_hat = cube.rfft()
|
|
density_hat /= np.sqrt(Pk)
|
|
|
|
Ns = N*supergenerate
|
|
|
|
density_hat_super = np.zeros((Ns,Ns,Ns/2+1), dtype=np.complex128)
|
|
density_hat_super[:] = np.nan
|
|
# Copy density hat in place
|
|
|
|
hN = N/2
|
|
density_hat_super[:hN, :hN, :hN+1] = density_hat[:hN, :hN, :]
|
|
density_hat_super[:hN, (Ns-hN):Ns, :hN+1] = density_hat[:hN, hN:, :]
|
|
density_hat_super[(Ns-hN):Ns, (Ns-hN):Ns, :hN+1] = density_hat[hN:, hN:, :]
|
|
density_hat_super[(Ns-hN):Ns, :hN, :hN+1] = density_hat[hN:, :hN, :]
|
|
|
|
# The moved nyquist place is untouched (so loss of "noise") to keep the structure
|
|
# now we just add some noise term
|
|
if supergenerate > 1:
|
|
|
|
cond=np.isnan(density_hat_super)
|
|
print np.where(np.isnan(density_hat_super))[0].size
|
|
Nz = np.count_nonzero(cond)
|
|
density_hat_super.real[cond] = np.random.randn(Nz)
|
|
density_hat_super.imag[cond] = np.random.randn(Nz)
|
|
density_hat_super[cond] /= np.sqrt(2.0)
|
|
print np.where(np.isnan(density_hat_super))[0].size
|
|
|
|
# Now we have to fix the Nyquist plane
|
|
hNs = Ns/2
|
|
nyquist = density_hat_super[:, :, hNs]
|
|
Nplane = nyquist.size
|
|
nyquist.flat[:Nplane/2] = np.sqrt(2.0)*nyquist.flat[Nplane:Nplane/2:-1].conj()
|
|
|
|
print np.where(np.isnan(density_hat_super))[0].size
|
|
|
|
cube = CubeFT(Ns, L)
|
|
cube.dhat = density_hat_super
|
|
return np.fft.irfftn(density_hat_super)*Ns**1.5
|
|
|
|
|
|
|
|
def write_icfiles(*generated_ic, **kwargs):
|
|
"""Write the initial conditions from the tuple returned by run_generation"""
|
|
|
|
supergenerate=1
|
|
if 'supergenerate' in kwargs:
|
|
supergenerate=kwargs['supergenerate']
|
|
posx,vel,density,N,L,a_ic,cosmo = 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("ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
|
|
|
|
ct.writeGrafic("ic_deltab", density, L, a_ic, **cosmo)
|
|
|
|
ct.writeWhitePhase("white.dat", whitify(density, L, cosmo, supergenerate=supergenerate))
|
|
|
|
with file("white_params", mode="w") as f:
|
|
f.write("4\n%lg, %lg, %lg\n" % (cosmo['omega_M_0'], cosmo['omega_lambda_0'], 100*cosmo['h']))
|
|
f.write("%lg\n%lg\n-%lg\n0,0\n" % (cosmo['omega_B_0'],cosmo['ns'],cosmo['SIGMA8']))
|
|
f.write("-%lg\n1\n0\n\n\n\n\n" % L)
|
|
f.write("2\n\n0\nwhite.dat\n0\npadding_white.dat\n")
|
|
|