# grid size for figure array
ngrid = 2 * 512
# physical size of figure array
cutsize = 256.0 #Mpc/h
# offset vector [Mpc/h]:
### don't know what htis is - perhaps it could be zero?
com = [1.30208333, 1.10677083, 0.944]
com[0] += offset_zoom[0]+cellsize_zoom * \
(BBox_out_zoom[0,1]-BBox_out_zoom[0,0])/2.0-cutsize/2.0
com[1] += offset_zoom[1]+cellsize_zoom * \
(BBox_out_zoom[1,1]-BBox_out_zoom[1,0])/2.0-cutsize/2.0
com[2] += offset_zoom[2]+cellsize_zoom * \
(BBox_out_zoom[2,1]-BBox_out_zoom[2,0])/2.0-cutsize/2.0
density, den_k, den_fft, _ = initialize_density(ngrid, ngrid, ngrid)
density.fill(0.0)
"""
# Lay down fine particles on density array with CiC:
CICDeposit_3(
px_zoom-com[0],
py_zoom-com[1],
pz_zoom-com[2],
px_zoom,py_zoom,pz_zoom, #dummies
density,
cellsize_zoom,
cutsize/float(ngrid),
0,
0, # dummy
0, # dummy
def evolve(cellsize,
sx_full,
sy_full,
sz_full,
sx2_full,
sy2_full,
sz2_full,
FULL=False,
cellsize_zoom=0,
sx_full_zoom=None,
sy_full_zoom=None,
sz_full_zoom=None,
sx2_full_zoom=None,
sy2_full_zoom=None,
sz2_full_zoom=None,
offset_zoom=None,
BBox_in=None,
ngrid_x=None,
ngrid_y=None,
ngrid_z=None,
gridcellsize=None,
ngrid_x_lpt=None,
ngrid_y_lpt=None,
ngrid_z_lpt=None,
gridcellsize_lpt=None,
Om=0.274,
Ol=1.0 - 0.274,
a_initial=1. / 15.,
a_final=1.0,
n_steps=15,
nCola=-2.5,
save_to_file=False,
file_npz_out='tmp.npz'):
"""
:math:`\vspace{-1mm}`
Evolve a set of initial conditions forward in time using the COLA
method in both the spatial and temporal domains.
**Arguments**:
* ``cellsize`` -- a float. The inter-particle spacing in Lagrangian space.
* ``sx_full,sy_full,sz_full`` -- 3-dim NumPy float32 arrays containing the
components of the particle displacements today as calculated in
the ZA in the full box. These particles should cover the COLA
volume only. If a refined subvolume is provided, these crude
particles which reside inside that subvolume are discarded and
replaced with the fine particles. There arrays are overwritten.
* ``sx2_full,sy2_full,sz2_full`` -- same as above but for the second
order displacement field.
* ``FULL`` -- a boolean (default: ``False``). If True, it indicates to
the code that the COLA volume covers the full box. In that case,
LPT in the COLA volume is not calculated, as that matches the LPT
in the full box.
* ``cellsize_zoom`` -- a float (default: ``0``). The inter-particle
spacing in Lagrangian space for the refined subvolume, if such is
provided. If not, ``cellsize_zoom`` must be set to zero
(default), as that is used as a check for the presence of that
subvolume.
* ``s*_full_zoom,s*2_full_zoom`` -- same as without ``_zoom``
above, but for the refined region (default: ``None``).
* ``offset_zoom`` -- a 3-vector of floats (default: ``None``). Gives the
physical coordinates of the origin of the refinement region
relative to the the origin of the full box.
* ``BBox_in`` -- a 3x2 array of integers (default: ``None``). It has the
form ``[[i0,i1],[j0,j1],[k0,k1]]``, which gives the bounding box
for the refinement region in units of the crude particles
Lagrangian index. Thus, the particles with displacements
``sx_full|sy_full|sz_full[i0:i1,j0:j1,k0:k1]`` are replaced with the fine
particles with displacements ``sx_full_zoom|sy_full_zoom|sz_full_zoom``.
* ``ngrid_x,ngrid_y,ngrid_z`` -- integers (default: ``None``). Provide the size of the
PM grid, which the algorithm uses to calculate the forces for the
Kicks.
* ``gridcellsize`` --float (default: ``None``). Provide the grid spacing of the PM
grid, which the algorithm uses to calculate the forces for the
Kicks.
* ``ngrid_x_lpt,ngrid_y_lpt,ngrid_z_lpt,gridcellsize_lpt`` -- the
same as without ``_lpt`` above but for calculating the LPT
displacements in the COLA volume. These better match their
counterparts above for the force calculation, as mismatches often
lead to unexpected non-cancellations and artifacts.
* ``Om`` -- a float (default: ``0.274``), giving the matter density, :math:`\Omega_m`, today.
* ``Ol`` -- a float (default: ``1.-0.274``), giving the vacuum density, :math:`\Omega_\Lambda`, today.
* ``a_initial`` -- a float (default: ``1./15.``). The initial scale
factor from which to start the COLA evolution. This better be
near ``1/n_steps``.
* ``a_final`` -- a float (default: ``1.0``). The final scale
factor for the COLA evolution.
* ``n_steps`` -- an integer (default: ``15``). The total number of
timesteps which the COLA algorithm should make.
* ``nCola`` -- a float (default: ``-2.5``). The spectral index for the time-domain COLA.
Sane values lie in the range ``(-4,3.5)``. Cannot be ``0``, but of course can be near ``0`` (say ``0.001``).
See Section A.3 of [temporalCOLA]_.
* ``save_to_file`` -- a boolean (default: ``False``). Whether to save the final snapshot to file.
* ``file_npz_out`` -- a string (default: ``'tmp.npz'``), giving the
filename for the ``.npz`` `SciPy
file <http://docs.scipy.org/doc/numpy/reference/generated/numpy.savez.html>`_,
in which to save the snapshot. See the source file for what is actually saved.
**Return**:
* ``px,py,pz`` -- 3-dim float32 arrays containing the components of the particle positions inside the COLA volume.
* ``vx,vy,vz`` -- 3-dim float32 arrays containing the components of
the particle velocities, :math:`\bm{v}`. Velocities are in units
of :math:`\mathrm{Mpc}/h` and are calculated according to:
.. math::
:nowrap:
\begin{eqnarray}
\bm{v}\equiv \frac{1}{a\,H(a)}\frac{d\bm{x}}{d\eta}
\end{eqnarray}
where :math:`\eta` is conformal time; :math:`a` is the final scale
factor ``a_final``; :math:`H(a)` is the Hubble parameter;
:math:`\bm{x}` is the comoving position. This definition allows
calculating redshift-space positions trivial: one simply has to
add the line-of-sight velocity to the particle position.
"""
from numpy import array, zeros
from ic import initial_positions, ic_2lpt_engine
from growth import _vel_coef, _displ_coef, growth_factor_solution, growth_2lpt, d_growth2
from potential import get_phi, initialize_density
from acceleration import grad_phi_engine
from scipy import interpolate
from cic import CICDeposit_3
print "*** FULL:", FULL
print "*** cellsize_zoom:", cellsize_zoom
if (cellsize_zoom == 0):
BBox_in = array([[0, 0], [0, 0], [0, 0]], dtype='int32')
else:
offset_zoom = offset_zoom.astype('float32')
offset = array([0.0, 0.0, 0.0], dtype='float32')
# time-related stuff
da = (a_final - a_initial) / float(n_steps)
d = growth_factor_solution(Om, Ol)
growth = interpolate.interp1d(d[:, 0].tolist(),
d[:, 1].tolist(),
kind='linear')
d_growth = interpolate.interp1d(d[:, 0].tolist(),
d[:, 2].tolist(),
kind='linear')
initial_growth_factor = growth(a_initial)
initial_growth2_factor = growth_2lpt(a_initial, initial_growth_factor, Om)
final_d_growth = d_growth(a_final)
final_d_growth2 = d_growth2(a_final, final_d_growth, Om, Ol)
initial_d_growth = d_growth(a_initial)
initial_d_growth2 = d_growth2(a_initial, initial_d_growth, Om, Ol)
del d_growth
#############
npart_x, npart_y, npart_z = sx_full.shape
npart_x_zoom = None
npart_y_zoom = None
npart_z_zoom = None
if (cellsize_zoom != 0):
npart_x_zoom, npart_y_zoom, npart_z_zoom = sx_full_zoom.shape
#############
import time
start = time.time()
#####################
# Do LPT in COLA box
#####################
if (FULL):
# if (COLA box)=(full box), then their lpt's match:
sx = sx_full
sy = sy_full
sz = sz_full
sx2 = sx2_full
sy2 = sy2_full
sz2 = sz2_full
if (cellsize_zoom != 0):
sx_zoom = sx_full_zoom
sy_zoom = sy_full_zoom
sz_zoom = sz_full_zoom
sx2_zoom = sx2_full_zoom
sy2_zoom = sy2_full_zoom
sz2_zoom = sz2_full_zoom
else:
# if (COLA box) != (full box), then we need the lpt in the COLA box:
print "Calculate LPT in the COLA box..."
sx, sy, sz, sx2, sy2, sz2, sx_zoom, sy_zoom, sz_zoom, sx2_zoom, sy2_zoom, sz2_zoom = ic_2lpt_engine(
sx_full,
sy_full,
sz_full,
cellsize,
ngrid_x_lpt,
ngrid_y_lpt,
ngrid_z_lpt,
gridcellsize_lpt,
with_2lpt=True,
sx2_full=sx2_full,
sy2_full=sy2_full,
sz2_full=sz2_full,
cellsize_zoom=cellsize_zoom,
BBox_in=BBox_in,
sx_full_zoom=sx_full_zoom,
sy_full_zoom=sy_full_zoom,
sz_full_zoom=sz_full_zoom,
sx2_full_zoom=sx2_full_zoom,
sy2_full_zoom=sy2_full_zoom,
sz2_full_zoom=sz2_full_zoom,
offset_zoom=offset_zoom)
print "... done"
#######################
# Some initializations:
#######################
# density:
density, den_k, den_fft, phi_fft = initialize_density(
ngrid_x, ngrid_y, ngrid_z)
# positions:
print "Initializing particle positions..."
px, py, pz = initial_positions(sx_full, sy_full, sz_full, sx2_full,
sy2_full, sz2_full, cellsize,
initial_growth_factor,
initial_growth2_factor, ngrid_x, ngrid_y,
ngrid_z, gridcellsize)
if (cellsize_zoom != 0):
px_zoom, py_zoom, pz_zoom = initial_positions(sx_full_zoom,
sy_full_zoom,
sz_full_zoom,
sx2_full_zoom,
sy2_full_zoom,
sz2_full_zoom,
cellsize_zoom,
initial_growth_factor,
initial_growth2_factor,
ngrid_x,
ngrid_y,
ngrid_z,
gridcellsize,
offset=offset_zoom)
print "...done"
# velocities:
# Initial residual velocities are zero in COLA.
# This corresponds to the L_- operator in 1301.0322.
# But to avoid short-scale
# noise, we do the smoothing trick explained in the new paper.
# However, that smoothing should not affect the IC velocities!
# So, first add the full vel, then further down subtract the same but smoothed.
#
# This smoothing is not really needed if FULL=True. But that case is
# not very interesting here, so we do it just the same.
vx = initial_d_growth * (sx_full) + initial_d_growth2 * (sx2_full)
vy = initial_d_growth * (sy_full) + initial_d_growth2 * (sy2_full)
vz = initial_d_growth * (sz_full) + initial_d_growth2 * (sz2_full)
if (cellsize_zoom != 0):
vx_zoom = initial_d_growth * (sx_full_zoom) + initial_d_growth2 * (
sx2_full_zoom)
vy_zoom = initial_d_growth * (sy_full_zoom) + initial_d_growth2 * (
sy2_full_zoom)
vz_zoom = initial_d_growth * (sz_full_zoom) + initial_d_growth2 * (
sz2_full_zoom)
### DJB: idiot necessary type casting
from numpy import float64, float32
sx = sx.astype(float32)
sy = sy.astype(float32)
sz = sz.astype(float32)
sx2 = sx2.astype(float32)
sy2 = sy2.astype(float32)
sz2 = sz2.astype(float32)
sx_zoom = sx_zoom.astype(float32)
sy_zoom = sy_zoom.astype(float32)
sz_zoom = sz_zoom.astype(float32)
sx2_zoom = sx2_zoom.astype(float32)
sy2_zoom = sy2_zoom.astype(float32)
sz2_zoom = sz2_zoom.astype(float32)
sx_full = sx_full.astype(float32)
sy_full = sy_full.astype(float32)
sz_full = sz_full.astype(float32)
sx_full_zoom = sx_full_zoom.astype(float32)
sy_full_zoom = sy_full_zoom.astype(float32)
sz_full_zoom = sz_full_zoom.astype(float32)
print "Smoothing arrays for the COLA game ..."
from box_smooth import box_smooth
tmp = zeros(sx_full.shape, dtype='float32')
box_smooth(sx_full, tmp)
sx_full[:] = tmp[:]
box_smooth(sy_full, tmp)
sy_full[:] = tmp[:]
box_smooth(sz_full, tmp)
sz_full[:] = tmp[:]
box_smooth(sx2_full, tmp)
sx2_full[:] = tmp[:]
box_smooth(sy2_full, tmp)
sy2_full[:] = tmp[:]
box_smooth(sz2_full, tmp)
sz2_full[:] = tmp[:]
#
box_smooth(sx, tmp)
sx[:] = tmp[:]
box_smooth(sy, tmp)
sy[:] = tmp[:]
box_smooth(sz, tmp)
sz[:] = tmp[:]
box_smooth(sx2, tmp)
sx2[:] = tmp[:]
box_smooth(sy2, tmp)
sy2[:] = tmp[:]
box_smooth(sz2, tmp)
sz2[:] = tmp[:]
del tmp
if (cellsize_zoom != 0):
tmp = zeros(sx_full_zoom.shape, dtype='float32')
box_smooth(sx_full_zoom, tmp)
sx_full_zoom[:] = tmp[:]
box_smooth(sy_full_zoom, tmp)
sy_full_zoom[:] = tmp[:]
box_smooth(sz_full_zoom, tmp)
sz_full_zoom[:] = tmp[:]
box_smooth(sx2_full_zoom, tmp)
sx2_full_zoom[:] = tmp[:]
box_smooth(sy2_full_zoom, tmp)
sy2_full_zoom[:] = tmp[:]
box_smooth(sz2_full_zoom, tmp)
sz2_full_zoom[:] = tmp[:]
#
box_smooth(sx_zoom, tmp)
sx_zoom[:] = tmp[:]
box_smooth(sy_zoom, tmp)
sy_zoom[:] = tmp[:]
box_smooth(sz_zoom, tmp)
sz_zoom[:] = tmp[:]
box_smooth(sx2_zoom, tmp)
sx2_zoom[:] = tmp[:]
box_smooth(sy2_zoom, tmp)
sy2_zoom[:] = tmp[:]
box_smooth(sz2_zoom, tmp)
sz2_zoom[:] = tmp[:]
del tmp
print "... done"
#All s* arrays are now smoothed!
#Next subtract smoothed vels as prescribed above.
vx -= initial_d_growth * (sx_full) + initial_d_growth2 * (sx2_full)
vy -= initial_d_growth * (sy_full) + initial_d_growth2 * (sy2_full)
vz -= initial_d_growth * (sz_full) + initial_d_growth2 * (sz2_full)
if (cellsize_zoom != 0):
vx_zoom -= initial_d_growth * (sx_full_zoom) + initial_d_growth2 * (
sx2_full_zoom)
vy_zoom -= initial_d_growth * (sy_full_zoom) + initial_d_growth2 * (
sy2_full_zoom)
vz_zoom -= initial_d_growth * (sz_full_zoom) + initial_d_growth2 * (
sz2_full_zoom)
#vx = zeros(sx.shape,dtype='float32')
#vy = zeros(sx.shape,dtype='float32')
#vz = zeros(sx.shape,dtype='float32')
#if (cellsize_zoom!=0):
# vx_zoom = zeros(sx_zoom.shape,dtype='float32')
# vy_zoom = zeros(sx_zoom.shape,dtype='float32')
# vz_zoom = zeros(sx_zoom.shape,dtype='float32')
#scale factors:
# initialize scale factor
aiKick = a_initial
aKick = a_initial
aiDrift = a_initial
aDrift = a_initial
#dummy values, to initialize as global
afKick = 0
afDrift = 0
dummy = 0.0 # yet another dummy
####################
#DO THE TIMESTEPS!!!
####################
for i in range(n_steps + 1):
if i == 0 or i == n_steps:
afKick = aiKick + da / 2.0
else:
afKick = aiKick + da
################
# FORCES!
################
# Calculate PM density:
density.fill(0.0)
CICDeposit_3(
px,
py,
pz,
px,
py,
pz, #dummies
density,
cellsize,
gridcellsize,
0,
dummy,
dummy,
BBox_in,
offset,
1)
if (cellsize_zoom != 0):
CICDeposit_3(
px_zoom,
py_zoom,
pz_zoom,
px_zoom,
py_zoom,
pz_zoom, #dummies
density,
cellsize_zoom,
gridcellsize,
0,
dummy,
dummy,
array([[0, 0], [0, 0], [0, 0]], dtype='int32'),
offset_zoom,
1)
density -= 1.0
#print "ave den",density.mean(dtype=float64)
# Calculate potential
get_phi(density, den_k, den_fft, phi_fft, ngrid_x, ngrid_y, ngrid_z,
gridcellsize)
phi = density # density now holds phi, so rename it
################
# KICK!
################
beta = -1.5 * aDrift * Om * _vel_coef(aiKick, afKick, aDrift, nCola,
Om, Ol)
d = growth(aDrift)
Om143 = (Om / (Om + (1.0 - Om) * aDrift * aDrift * aDrift))**(1. /
143.)
# Note that grad_phi_engine() will subtract the lpt forces in the COLA volume
# before doing the kick.
### DJB: idiot necessary type casting
beta = float32(beta)
vx = vx.astype(float32)
vy = vy.astype(float32)
vz = vz.astype(float32)
vx_zoom = vx_zoom.astype(float32)
vy_zoom = vy_zoom.astype(float32)
vz_zoom = vz_zoom.astype(float32)
grad_phi_engine(px, py, pz, vx, vy, vz, sx, sy, sz, sx2, sy2, sz2,
beta, beta, npart_x, npart_y, npart_z, phi, ngrid_x,
ngrid_y, ngrid_z, cellsize, gridcellsize, d,
d * d * (1.0 + 7. / 3. * Om143),
array([0.0, 0.0, 0.0], dtype='float32'), 0)
if (cellsize_zoom != 0):
grad_phi_engine(px_zoom, py_zoom, pz_zoom, vx_zoom, vy_zoom,
vz_zoom, sx_zoom, sy_zoom, sz_zoom, sx2_zoom,
sy2_zoom, sz2_zoom, beta, beta, npart_x_zoom,
npart_y_zoom, npart_z_zoom, phi, ngrid_x, ngrid_y,
ngrid_z, cellsize_zoom, gridcellsize, d,
d * d * (1.0 + 7. / 3. * Om143),
array([0.0, 0.0, 0.0], dtype='float32'), 0)
del phi
aKick = afKick
aiKick = afKick
print "Kicked to a = " + str(aKick)
################
# DRIFT!
################
if i < n_steps:
afDrift = aiDrift + da
alpha = _displ_coef(aiDrift, afDrift, aKick, nCola, Om, Ol)
gamma = 1.0 * (growth(afDrift) - growth(aiDrift))
gamma2 = 1.0 * (growth_2lpt(afDrift, growth(afDrift), Om) -
growth_2lpt(aiDrift, growth(aiDrift), Om))
# Drift, but also add the lpt displacement from the full volume:
px += vx * alpha + sx_full * gamma + sx2_full * gamma2 + float(
ngrid_x) * gridcellsize
py += vy * alpha + sy_full * gamma + sy2_full * gamma2 + float(
ngrid_y) * gridcellsize
pz += vz * alpha + sz_full * gamma + sz2_full * gamma2 + float(
ngrid_z) * gridcellsize
px %= float(ngrid_x) * gridcellsize
py %= float(ngrid_y) * gridcellsize
pz %= float(ngrid_z) * gridcellsize
if (cellsize_zoom != 0):
px_zoom += vx_zoom * alpha + sx_full_zoom * gamma + sx2_full_zoom * gamma2 + float(
ngrid_x) * gridcellsize
py_zoom += vy_zoom * alpha + sy_full_zoom * gamma + sy2_full_zoom * gamma2 + float(
ngrid_y) * gridcellsize
pz_zoom += vz_zoom * alpha + sz_full_zoom * gamma + sz2_full_zoom * gamma2 + float(
ngrid_z) * gridcellsize
px_zoom %= float(ngrid_x) * gridcellsize
py_zoom %= float(ngrid_y) * gridcellsize
pz_zoom %= float(ngrid_z) * gridcellsize
aiDrift = afDrift
aDrift = afDrift
print "Drifted to a = " + str(aDrift)
del den_k, den_fft, phi_fft, density
# Add back LPT velocity to velocity residual
# This corresponds to the L_+ operator in 1301.0322.
vx += final_d_growth * (sx_full) + final_d_growth2 * (sx2_full)
vy += final_d_growth * (sy_full) + final_d_growth2 * (sy2_full)
vz += final_d_growth * (sz_full) + final_d_growth2 * (sz2_full)
if (cellsize_zoom != 0):
vx_zoom += final_d_growth * (sx_full_zoom) + final_d_growth2 * (
sx2_full_zoom)
vy_zoom += final_d_growth * (sy_full_zoom) + final_d_growth2 * (
sy2_full_zoom)
vz_zoom += final_d_growth * (sz_full_zoom) + final_d_growth2 * (
sz2_full_zoom)
# Now convert velocities to
# v_{rsd}\equiv (ds/d\eta)/(a H(a)):
from growth import _q_factor
rsd_fac = a_final / _q_factor(a_final, Om, Ol)
vx *= rsd_fac
vy *= rsd_fac
vz *= rsd_fac
if (cellsize_zoom != 0):
vx_zoom *= rsd_fac
vy_zoom *= rsd_fac
vz_zoom *= rsd_fac
end = time.time()
print "Time elapsed on small box (including IC): " + str(
end - start) + " seconds."
#### Set all these params to zero if cellsize_zoom==0 for ease.
##### these prob shoudl have been initialized elsewhere.
px_zoom = 0
py_zoom = 0
pz_zoom = 0
vx_zoom = 0
vy_zoom = 0
vz_zoom = 0
if save_to_file:
from numpy import savez
savez(
file_npz_out,
#px_zoom=px_zoom,py_zoom=py_zoom,pz_zoom=pz_zoom,
#vx_zoom=vx_zoom,vy_zoom=vy_zoom,vz_zoom=vz_zoom,
#cellsize_zoom=cellsize_zoom,
px=px,
py=py,
pz=pz
#vx=vx,vy=vy,vz=vz,
#cellsize=cellsize,
#ngrid_x=ngrid_x,ngrid_y=ngrid_y,ngrid_z=ngrid_z,
#z_final=1.0/(aDrift)-1.0,z_init=1.0/(a_initial)-1.0,
#n_steps=n_steps,Om=Om,Ol=Ol,nCola=nCola,
#ngrid_x_lpt=ngrid_x_lpt,ngrid_y_lpt=ngrid_y_lpt,ngrid_z_lpt=ngrid_z_lpt,
#gridcellsize=gridcellsize,
#gridcellsize_lpt=gridcellsize_lpt
)
return px, py, pz, vx, vy, vz, px_zoom, py_zoom, pz_zoom, vx_zoom, vy_zoom, vz_zoom
def ic_2lpt_engine(sx_full,
sy_full,
sz_full,
cellsize,
ngrid_x,
ngrid_y,
ngrid_z,
gridcellsize,
growth_2pt_calc=0.05,
with_4pt_rule=False,
factor_4pt=2.0,
with_2lpt=False,
sx2_full=None,
sy2_full=None,
sz2_full=None,
cellsize_zoom=0,
BBox_in=None,
sx_full_zoom=None,
sy_full_zoom=None,
sz_full_zoom=None,
sx2_full_zoom=None,
sy2_full_zoom=None,
sz2_full_zoom=None,
offset_zoom=None):
r"""
:math:`\vspace{-1mm}`
The same as :func:`ic.ic_2lpt` above, but calculates the 2LPT displacements for the particles in the
COLA volume as generated by same particles displaced
according to the 2LPT of the full box. (todo: *expand this!*) In fact, :func:`ic.ic_2lpt` works
by making a call to this function.
**Arguments**:
* ``sx_full,sy_full,sz_full`` -- 3-dim NumPy arrays containing the
components of the particle displacements today as calculated in
the ZA of the full box. These particles should cover the whole box. If a refined
subvolume is provided, the crude particles which reside inside
that subvolume are discarded and replaced with the fine
particles.
* ``cellsize`` -- a float. The inter-particle spacing in Lagrangian space.
* ``ngrid_x,ngrid_y,ngrid_z`` -- integers. Provide the size of the
PM grid, which the algorithm
uses to calculate the 2LPT displacements.
* ``gridcellsize`` --float. Provide the grid spacing of the PM
grid, which the algorithm
uses to calculate the 2LPT displacements.
* ``growth_2pt_calc`` -- a float (default: ``0.05``). The
linear growth factor used internally in the 2LPT calculation. A
value of 0.05 gives excellent cross-correlation between the 2LPT
field returned by this function, and the 2LPT returned using the
usual fft tricks for a 100:math:`\mathrm{Mpc}/h` box. Yet, some
irrelevant short-scale noise is present, which one may decide to
filter out. That noise is probably due to lack of force accuracy
for too low ``growth_2pt_calc``. Experiment with this value as
needed.
* ``with_4pt_rule`` -- a boolean (default: ``False``). Whether to use
the 4-point force rule to evaluate the ZA and 2LPT displacements
in the COLA region. See the Algorithm section below. If set to
False, it uses the 2-point force rule.
* ``factor_4pt`` -- a float, different from ``1.0`` (default:
``2.0``). Used for the 4-point
force rule. See the Algorithm section below.
* ``with_2lpt`` -- a boolean (default: ``False``). Whether the second
order displacement field over the full box is provided. One must
provide it if the COLA volume is different from the full box.
Only if they are the same (as in the case of ``ic_2lpt()``) can
one set ``with_2lpt=False``.
* ``sx2_full,sy2_full,sz2_full`` -- 3-dim NumPy float arrays giving the second
order displacement field over the full box. Needs ``with_2lpt=True``.
* The rest of the input is as in :func:`ic.ic_2lpt`, with all LPT
quantities provided for the whole box.
**Return**:
* If no refined subregion is supplied (indicated by
``cellsize_zoom=0``), then return:
``(sx,sy,sz,sx2,sy2,sz2)`` -- 3-dim NumPy
arrays containing the components of the first and second (``s``:sub:`i`\ ``2``)
order particle displacements today as calculated in 2LPT in the
COLA volume.
* If a refined subregion is supplied (indicated by
``cellsize_zoom>0``), then return:
``(sx,sy,sz,sx2,sy2,sz2,sx_zoom,sy_zoom,sz_zoom,sx2_zoom,sy2_zoom,sz2_zoom)``
The first 6 arrays are as
above. The last 6 give the second order displacements today
for the particles in the refined subvolume of the COLA volume.
**Algorithm**:
The first-order and second-order displacements,
:math:`\bm{s}^{(1)}_{\mathrm{COLA}}` and
:math:`\bm{s}^{(2)}_{\mathrm{COLA}}`, in the COLA volume at
redshift zero are calculated according to the following 2-pt or
4-pt (denoted by subscript) equations:
.. math::
:nowrap:
\begin{eqnarray}
\bm{s}_{\mathrm{COLA},\mathrm{2pt}}^{(1)} & = & - \frac{1}{2g} \left[\bm{F}(g,\beta g^2)-\bm{F}(-g,\beta g^2)\right] \\
\bm{s}_{\mathrm{COLA},\mathrm{2pt}}^{(2)} & = & - \frac{\alpha}{2g^2} \left[\bm{F}(g,\beta g^2)+\bm{F}(-g,\beta g^2)\right] \\
\bm{s}_{\mathrm{COLA},\mathrm{4pt}}^{(1)} & = & - \frac{1}{2g} \frac{a^2}{a^2-1}\bigg[\bm{F}(g,\beta g^2)-\bm{F}(-g,\beta g^2)-\\
& & \quad \quad \quad \quad \quad \quad - \frac{1}{a^3}\bigg(\bm{F}\left(a g,\beta a^2 g^2\right)-\bm{F}\left(-a g,\beta a^2 g^2\right)\bigg)\bigg] \\
\bm{s}_{\mathrm{COLA},\mathrm{4pt}}^{(2)} & = & - \frac{\alpha}{2g^2}\frac{a^2}{a^2-1}\bigg[\bm{F}(g,\beta g^2)+\bm{F}(-g,\beta g^2)-\\
& & \quad \quad \quad \quad \quad \quad - \frac{1}{a^4}\bigg(\bm{F}\left(a g,\beta a^2 g^2\right)+\bm{F}\left(-a g,\beta a^2 g^2\right)\bigg)\bigg]
\end{eqnarray}
where:
:math:`a=` ``factor_4pt``
:math:`g=` ``growth_2pt_calc``
if ``with_2lpt`` then:
:math:`\beta=1` and :math:`\alpha=(3/10)\Omega_{m}^{1/143}`
else:
:math:`\beta=0` and :math:`\alpha=(3/7)\Omega_{m}^{1/143}`
The factors of :math:`\Omega_{m}^{1/143}` (:math:`\Omega_m` being
the matter density today) are needed to rescale the second order
displacements to matter domination and are correct to
:math:`\mathcal{O}(\max(10^{-4},g^3/143))` in
:math:`\Lambda\mathrm{CDM}`. The force :math:`\bm{F}(g_1,g_2)` is
given by:
.. math::
:nowrap:
\begin{eqnarray}
\bm{F}(g_1,g_2) = \bm{\nabla}\nabla^{-2}\delta\left[g_1\bm{s}_{\mathrm{full}}^{(1)}+g_2\Omega_{m}^{-1/143}\bm{s}_{\mathrm{full}}^{(2)}\right]
\end{eqnarray}
where :math:`\delta[\bm{s}]` is the cloud-in-cell fractional
overdensity calculated from a grid of particles displaced by the
input displacement vector field :math:`\bm{s}`. Above,
:math:`\bm{s}_{\mathrm{full}}^{(1)}/\bm{s}_{\mathrm{full}}^{(2)}` are
the input first/second-order input displacement fields calculated
in the full box at redshift zero.
It is important to note that implicitly for each particle at
Lagrangian position :math:`\bm{q}`, the force
:math:`\bm{F}(g_1,g_2)` is evaluated at the corresponding Eulerian position:
:math:`\bm{q}+g_1\bm{s}_{\mathrm{full}}^{(1)}+g_2\Omega_{m}^{-1/143}\bm{s}_{\mathrm{full}}^{(2)}`.
As noted above, ``with_2lpt=False`` is only allowed if the COLA
volume covers the full box volume. In that case,
:math:`\bm{s}_{\mathrm{full}}^{(2)}` is not needed as input since
:math:`\beta=0`. Instead, the output
:math:`\bm{s}_{\mathrm{COLA}}^{(2)}` can be used as a good
approximation to :math:`\bm{s}_{\mathrm{full}}^{(2)}`. This fact
is used in :func:`ic.ic_2lpt` to calculate
:math:`\bm{s}_{\mathrm{full}}^{(2)}` from
:math:`\bm{s}_{\mathrm{full}}^{(1)}`.
.. note:: If ``with_4pt_rule=False``, then the first/second order
displacements receive corrections at third/fourth order. If
``with_4pt_rule=True``, then those corrections are fifth/sixth
order. However, when using the 4-point rule instead of the
2-point rule, one must make two more force evaluations at a
slightly different growth factor given by
``growth_2pt_calc*factor_4pt``. Since the code is single
precision and is using a simple PM grid to evaluate forces, one
cannot make ``factor_4pt`` and ``growth_2pt_calc`` too small due
to noise issues. Thus, when comparing the 2-pt and 4-pt rule, we
should assume ``factor_4pt>1``. And again due to numerical
precision issues, one cannot choose ``factor_4pt`` to be too
close to one; hence, the default value of ``2.0``.
Therefore, as the higher order corrections for the 4-pt rule are
proportional to powers of ``growth_2pt_calc*factor_4pt``, one
may be better off using the 2-pt rule (the default) in this
particular implementation. Yet for codes where force accuracy is
not an issue, one may consider using the 4-pt rule. Thus, its
inclusion in this code is mostly done as an illustration.
"""
from numpy import float64, float32
if (cellsize_zoom != 0):
cellsize_zoom = float32(cellsize_zoom)
offset_zoom = offset_zoom.astype('float32')
npart_x, npart_y, npart_z = sx_full.shape
if (cellsize_zoom != 0):
npart_x_zoom, npart_y_zoom, npart_z_zoom = sx_full_zoom.shape
from numpy import zeros, array
sx = zeros((npart_x, npart_y, npart_z), dtype='float32')
sy = zeros((npart_x, npart_y, npart_z), dtype='float32')
sz = zeros((npart_x, npart_y, npart_z), dtype='float32')
sx_minus = zeros((npart_x, npart_y, npart_z), dtype='float32')
sy_minus = zeros((npart_x, npart_y, npart_z), dtype='float32')
sz_minus = zeros((npart_x, npart_y, npart_z), dtype='float32')
if (cellsize_zoom != 0):
sx_zoom = zeros((npart_x_zoom, npart_y_zoom, npart_z_zoom),
dtype='float32')
sy_zoom = zeros((npart_x_zoom, npart_y_zoom, npart_z_zoom),
dtype='float32')
sz_zoom = zeros((npart_x_zoom, npart_y_zoom, npart_z_zoom),
dtype='float32')
sx_minus_zoom = zeros((npart_x_zoom, npart_y_zoom, npart_z_zoom),
dtype='float32')
sy_minus_zoom = zeros((npart_x_zoom, npart_y_zoom, npart_z_zoom),
dtype='float32')
sz_minus_zoom = zeros((npart_x_zoom, npart_y_zoom, npart_z_zoom),
dtype='float32')
if (with_4pt_rule):
sx_4pt_minus = zeros((npart_x, npart_y, npart_z), dtype='float32')
sy_4pt_minus = zeros((npart_x, npart_y, npart_z), dtype='float32')
sz_4pt_minus = zeros((npart_x, npart_y, npart_z), dtype='float32')
sx_4pt = zeros((npart_x, npart_y, npart_z), dtype='float32')
sy_4pt = zeros((npart_x, npart_y, npart_z), dtype='float32')
sz_4pt = zeros((npart_x, npart_y, npart_z), dtype='float32')
if (cellsize_zoom != 0):
sx_4pt_minus_zoom = zeros(
(npart_x_zoom, npart_y_zoom, npart_z_zoom), dtype='float32')
sy_4pt_minus_zoom = zeros(
(npart_x_zoom, npart_y_zoom, npart_z_zoom), dtype='float32')
sz_4pt_minus_zoom = zeros(
(npart_x_zoom, npart_y_zoom, npart_z_zoom), dtype='float32')
sx_4pt_zoom = zeros((npart_x_zoom, npart_y_zoom, npart_z_zoom),
dtype='float32')
sy_4pt_zoom = zeros((npart_x_zoom, npart_y_zoom, npart_z_zoom),
dtype='float32')
sz_4pt_zoom = zeros((npart_x_zoom, npart_y_zoom, npart_z_zoom),
dtype='float32')
else:
sx_4pt_minus = 0.0
sy_4pt_minus = 0.0
sz_4pt_minus = 0.0
sx_4pt = 0.0
sy_4pt = 0.0
sz_4pt = 0.0
if (cellsize_zoom != 0):
sx_4pt_minus_zoom = 0.0
sy_4pt_minus_zoom = 0.0
sz_4pt_minus_zoom = 0.0
sx_4pt_zoom = 0.0
sy_4pt_zoom = 0.0
sz_4pt_zoom = 0.0
from potential import get_phi, initialize_density
from cic import CICDeposit_3
from acceleration import grad_phi
###
density, den_k, den_fft, phi_fft = initialize_density(
ngrid_x, ngrid_y, ngrid_z)
density.fill(0.0)
Om = 0.274
dd = Om**(
1. / 143.
) # this is a good enough approximation at early times and is ~0.95
if with_2lpt:
cc = 3. / 10. * dd
L2 = growth_2pt_calc * growth_2pt_calc / dd
else:
cc = 3. / 7. * dd
L2 = 0.0
gridcellsize = float32(gridcellsize)
growth_2pt_calc = float32(growth_2pt_calc)
L2 = float32(L2)
offset = array([0.0, 0.0, 0.0], dtype='float32')
if (cellsize_zoom == 0):
BBox_in = array([[0, 0], [0, 0], [0, 0]], dtype='int32')
if not with_2lpt:
sx2_full = zeros((0, 0, 0), dtype='float32')
sy2_full = zeros((0, 0, 0), dtype='float32')
sz2_full = zeros((0, 0, 0), dtype='float32')
CICDeposit_3(sx_full, sy_full, sz_full, sx2_full, sy2_full, sz2_full,
density, cellsize, gridcellsize, 1, growth_2pt_calc, L2,
BBox_in, offset, 1)
if (cellsize_zoom != 0):
CICDeposit_3(sx_full_zoom, sy_full_zoom, sz_full_zoom, sx2_full_zoom,
sy2_full_zoom, sz2_full_zoom, density, cellsize_zoom,
gridcellsize, 1, growth_2pt_calc, L2,
array([[0, 0], [0, 0], [0, 0]],
dtype='int32'), offset_zoom, 1)
density -= 1.0
#print "den ic",density.mean(dtype=float64)
if (with_4pt_rule):
density *= -0.5 / growth_2pt_calc / (1.0 -
1.0 / factor_4pt / factor_4pt)
else:
density *= -0.5 / growth_2pt_calc
get_phi(density, den_k, den_fft, phi_fft, ngrid_x, ngrid_y, ngrid_z,
gridcellsize)
phi = density # density now holds phi, so rename it
grad_phi(sx_full, sy_full, sz_full, sx2_full, sy2_full, sz2_full, sx, sy,
sz, npart_x, npart_y, npart_z, phi, ngrid_x, ngrid_y, ngrid_z,
cellsize, gridcellsize, growth_2pt_calc, L2, offset)
if (cellsize_zoom != 0):
grad_phi(sx_full_zoom, sy_full_zoom, sz_full_zoom, sx2_full_zoom,
sy2_full_zoom, sz2_full_zoom, sx_zoom, sy_zoom, sz_zoom,
npart_x_zoom, npart_y_zoom, npart_z_zoom, phi, ngrid_x,
ngrid_y, ngrid_z, cellsize_zoom, gridcellsize,
growth_2pt_calc, L2, offset_zoom)
#######
density.fill(0.0)
CICDeposit_3(sx_full, sy_full, sz_full, sx2_full, sy2_full, sz2_full,
density, cellsize, gridcellsize, 1, -growth_2pt_calc, L2,
BBox_in, offset, 1)
if (cellsize_zoom != 0):
CICDeposit_3(sx_full_zoom, sy_full_zoom, sz_full_zoom, sx2_full_zoom,
sy2_full_zoom, sz2_full_zoom, density, cellsize_zoom,
gridcellsize, 1, -growth_2pt_calc, L2,
array([[0, 0], [0, 0], [0, 0]],
dtype='int32'), offset_zoom, 1)
density -= 1.0
#print "den ic",density.mean(dtype=float64)
if (with_4pt_rule):
density *= -0.5 / growth_2pt_calc / (1.0 -
1.0 / factor_4pt / factor_4pt)
else:
density *= -0.5 / growth_2pt_calc
get_phi(density, den_k, den_fft, phi_fft, ngrid_x, ngrid_y, ngrid_z,
gridcellsize)
phi = density # density now holds phi, so rename it
grad_phi(sx_full, sy_full, sz_full, sx2_full, sy2_full, sz2_full, sx_minus,
sy_minus, sz_minus, npart_x, npart_y, npart_z, phi, ngrid_x,
ngrid_y, ngrid_z, cellsize, gridcellsize, -growth_2pt_calc, L2,
offset)
if (cellsize_zoom != 0):
grad_phi(sx_full_zoom, sy_full_zoom, sz_full_zoom, sx2_full_zoom,
sy2_full_zoom, sz2_full_zoom, sx_minus_zoom, sy_minus_zoom,
sz_minus_zoom, npart_x_zoom, npart_y_zoom, npart_z_zoom, phi,
ngrid_x, ngrid_y, ngrid_z, cellsize_zoom, gridcellsize,
-growth_2pt_calc, L2, offset_zoom)
######
######
###### Two more force evaluations in the case of a 4pt rule.
######
######
if (with_4pt_rule):
density.fill(0.0)
CICDeposit_3(sx_full, sy_full, sz_full, sx2_full, sy2_full, sz2_full,
density, cellsize, gridcellsize, 1,
growth_2pt_calc * factor_4pt,
L2 * factor_4pt * factor_4pt, BBox_in, offset, 1)
if (cellsize_zoom != 0):
CICDeposit_3(sx_full_zoom, sy_full_zoom, sz_full_zoom,
sx2_full_zoom, sy2_full_zoom, sz2_full_zoom, density,
cellsize_zoom, gridcellsize, 1,
growth_2pt_calc * factor_4pt,
L2 * factor_4pt * factor_4pt,
array([[0, 0], [0, 0], [0, 0]],
dtype='int32'), offset_zoom, 1)
density -= 1.0
#print "den ic",density.mean(dtype=float64)
density *= -0.5 / growth_2pt_calc / (
1.0 - 1.0 / factor_4pt / factor_4pt) * (-1.0 / factor_4pt**3)
get_phi(density, den_k, den_fft, phi_fft, ngrid_x, ngrid_y, ngrid_z,
gridcellsize)
phi = density # density now holds phi, so rename it
grad_phi(sx_full, sy_full, sz_full, sx2_full, sy2_full, sz2_full,
sx_4pt, sy_4pt, sz_4pt, npart_x, npart_y, npart_z, phi,
ngrid_x, ngrid_y, ngrid_z, cellsize, gridcellsize,
growth_2pt_calc * factor_4pt, L2 * factor_4pt * factor_4pt,
offset)
if (cellsize_zoom != 0):
grad_phi(sx_full_zoom, sy_full_zoom, sz_full_zoom, sx2_full_zoom,
sy2_full_zoom, sz2_full_zoom, sx_4pt_zoom, sy_4pt_zoom,
sz_4pt_zoom, npart_x_zoom, npart_y_zoom, npart_z_zoom,
phi, ngrid_x, ngrid_y, ngrid_z, cellsize_zoom,
gridcellsize, growth_2pt_calc * factor_4pt,
L2 * factor_4pt * factor_4pt, offset_zoom)
#######
density.fill(0.0)
CICDeposit_3(sx_full, sy_full, sz_full, sx2_full, sy2_full, sz2_full,
density, cellsize, gridcellsize, 1,
-growth_2pt_calc * factor_4pt,
L2 * factor_4pt * factor_4pt, BBox_in, offset, 1)
if (cellsize_zoom != 0):
CICDeposit_3(sx_full_zoom, sy_full_zoom, sz_full_zoom,
sx2_full_zoom, sy2_full_zoom, sz2_full_zoom, density,
cellsize_zoom, gridcellsize, 1,
-growth_2pt_calc * factor_4pt,
L2 * factor_4pt * factor_4pt,
array([[0, 0], [0, 0], [0, 0]],
dtype='int32'), offset_zoom, 1)
density -= 1.0
density *= -0.5 / growth_2pt_calc / (
1.0 - 1.0 / factor_4pt / factor_4pt) * (-1.0 / factor_4pt**3)
get_phi(density, den_k, den_fft, phi_fft, ngrid_x, ngrid_y, ngrid_z,
gridcellsize)
phi = density # density now holds phi, so rename it
grad_phi(sx_full, sy_full, sz_full, sx2_full, sy2_full, sz2_full,
sx_4pt_minus, sy_4pt_minus, sz_4pt_minus, npart_x, npart_y,
npart_z, phi, ngrid_x, ngrid_y, ngrid_z, cellsize,
gridcellsize, -growth_2pt_calc * factor_4pt,
L2 * factor_4pt * factor_4pt, offset)
if (cellsize_zoom != 0):
grad_phi(sx_full_zoom, sy_full_zoom, sz_full_zoom, sx2_full_zoom,
sy2_full_zoom, sz2_full_zoom, sx_4pt_minus_zoom,
sy_4pt_minus_zoom, sz_4pt_minus_zoom, npart_x_zoom,
npart_y_zoom, npart_z_zoom, phi, ngrid_x, ngrid_y,
ngrid_z, cellsize_zoom, gridcellsize,
-growth_2pt_calc * factor_4pt,
L2 * factor_4pt * factor_4pt, offset_zoom)
######
######
###### Done with the two more force evaluations in the case of a 4pt rule.
######
######
del density, den_k, den_fft, phi, phi_fft
if (
cellsize_zoom != 0
): # the variables *_4pt* (except factor_4pt) are init'd to zero if 4pt rule is not requested
sx2_zoom = (sx_zoom + sx_minus_zoom +
(sx_4pt_zoom + sx_4pt_minus_zoom) /
factor_4pt) * cc / growth_2pt_calc
sy2_zoom = (sy_zoom + sy_minus_zoom +
(sy_4pt_zoom + sy_4pt_minus_zoom) /
factor_4pt) * cc / growth_2pt_calc
sz2_zoom = (sz_zoom + sz_minus_zoom +
(sz_4pt_zoom + sz_4pt_minus_zoom) /
factor_4pt) * cc / growth_2pt_calc
sx_zoom += sx_4pt_zoom - (sx_minus_zoom + sx_4pt_minus_zoom)
sy_zoom += sy_4pt_zoom - (sy_minus_zoom + sy_4pt_minus_zoom)
sz_zoom += sz_4pt_zoom - (sz_minus_zoom + sz_4pt_minus_zoom)
sx2 = (sx + sx_minus +
(sx_4pt + sx_4pt_minus) / factor_4pt) * cc / growth_2pt_calc
sy2 = (sy + sy_minus +
(sy_4pt + sy_4pt_minus) / factor_4pt) * cc / growth_2pt_calc
sz2 = (sz + sz_minus +
(sz_4pt + sz_4pt_minus) / factor_4pt) * cc / growth_2pt_calc
sx += sx_4pt - (sx_minus + sx_4pt_minus)
sy += sy_4pt - (sy_minus + sy_4pt_minus)
sz += sz_4pt - (sz_minus + sz_4pt_minus)
del sx_minus, sy_minus, sz_minus
if (cellsize_zoom != 0):
del sx_minus_zoom, sy_minus_zoom, sz_minus_zoom
if (with_4pt_rule):
del sx_4pt_minus, sy_4pt_minus, sz_4pt_minus
if (cellsize_zoom != 0):
del sx_4pt_minus_zoom, sy_4pt_minus_zoom, sz_4pt_minus_zoom
del sx_4pt, sy_4pt, sz_4pt
if (cellsize_zoom != 0):
del sx_4pt_zoom, sy_4pt_zoom, sz_4pt_zoom
# The lines below fix the box-size irrotational condition for periodic boxes
#Mx=sx.mean(axis=0,dtype=float64)
#My=sy.mean(axis=1,dtype=float64)
#Mz=sz.mean(axis=2,dtype=float64)
#for i in range(npart_x):
#sx[i,:,:]-=Mx
#for i in range(npart_y):
#sy[:,i,:]-=My
#for i in range(npart_z):
#sz[:,:,i]-=Mz
#del Mx,My,Mz
#Mx=sx2.mean(axis=0,dtype=float64)
#My=sy2.mean(axis=1,dtype=float64)
#Mz=sz2.mean(axis=2,dtype=float64)
#for i in range(npart_x):
#sx2[i,:,:]-=Mx
#for i in range(npart_y):
#sy2[:,i,:]-=My
#for i in range(npart_z):
#sz2[:,:,i]-=Mz
#del Mx,My,Mz
#Mx=sx_zoom.mean(axis=0,dtype=float64)
#My=sy_zoom.mean(axis=1,dtype=float64)
#Mz=sz_zoom.mean(axis=2,dtype=float64)
#for i in range(npart_x_zoom):
#sx_zoom[i,:,:]-=Mx
#for i in range(npart_y_zoom):
#sy_zoom[:,i,:]-=My
#for i in range(npart_z_zoom):
#sz_zoom[:,:,i]-=Mz
#del Mx,My,Mz
#Mx=sx2_zoom.mean(axis=0,dtype=float64)
#My=sy2_zoom.mean(axis=1,dtype=float64)
#Mz=sz2_zoom.mean(axis=2,dtype=float64)
#for i in range(npart_x_zoom):
#sx2_zoom[i,:,:]-=Mx
#for i in range(npart_y_zoom):
#sy2_zoom[:,i,:]-=My
#for i in range(npart_z_zoom):
#sz2_zoom[:,:,i]-=Mz
#del Mx,My,Mz
if (cellsize_zoom != 0):
return sx, sy, sz, sx2, sy2, sz2, sx_zoom, sy_zoom, sz_zoom, sx2_zoom, sy2_zoom, sz2_zoom
else:
return sx, sy, sz, sx2, sy2, sz2
def run(music_file):
import numpy as np
import matplotlib.pyplot as plt
from aux import boundaries
from ic import ic_2lpt, import_music_snapshot
from evolve import evolve
from cic import CICDeposit_3
from potential import initialize_density
# Set up the parameters from the MUSIC ic snapshot:
# Set up according to instructions for
# aux.boundaries()
boxsize = 100.0 # in Mpc/h
level = 7
level_zoom = 8
gridscale = 1
# Set up according to instructions for
# ic.import_music_snapshot()
level0 = '07' # should match level above
level1 = '08' # should match level_zoom above
# Set how much to cut from the sides of the full box.
# This makes the COLA box to be of the following size in Mpc/h:
# (2.**level-(cut_from_sides[0]+cut_from_sides[1]))/2.**level*boxsize
# This is the full box. Set FULL=True in evolve() below
#cut_from_sides=[0,0]# 100Mpc/h.
#
# These are the interesting cases:
#cut_from_sides=[64,64]# 75Mpc/h
#cut_from_sides=[128,128] # 50Mpc/h
# cut_from_sides=[192,192] # 25Mpc/h
cut_from_sides = [192 / 4, 192 / 4] # 25Mpc/h with level/zoom = 7/8
sx_full1, sy_full1, sz_full1, sx_full_zoom1, sy_full_zoom1, \
sz_full_zoom1, offset_from_code \
= import_music_snapshot(music_file, \
boxsize,level0=level0,level1=level1)
sx_full1 = sx_full1.astype('float32')
sy_full1 = sy_full1.astype('float32')
sz_full1 = sz_full1.astype('float32')
sx_full_zoom1 = sx_full_zoom1.astype('float32')
sy_full_zoom1 = sy_full_zoom1.astype('float32')
sz_full_zoom1 = sz_full_zoom1.astype('float32')
NPART_zoom = list(sx_full_zoom1.shape)
print "Starting 2LPT on full box."
#Get bounding boxes for full box with 1 refinement level for MUSIC.
BBox_in, offset_zoom, cellsize, cellsize_zoom, \
offset_index, BBox_out, BBox_out_zoom, \
ngrid_x, ngrid_y, ngrid_z, gridcellsize \
= boundaries(boxsize, level, level_zoom, \
NPART_zoom, offset_from_code, [0,0], gridscale)
sx2_full1, sy2_full1, sz2_full1, sx2_full_zoom1, \
sy2_full_zoom1, sz2_full_zoom1 \
= ic_2lpt(
cellsize,
sx_full1 ,
sy_full1 ,
sz_full1 ,
cellsize_zoom=cellsize_zoom,
sx_zoom = sx_full_zoom1,
sy_zoom = sy_full_zoom1,
sz_zoom = sz_full_zoom1,
boxsize=100.00,
ngrid_x_lpt=ngrid_x,ngrid_y_lpt=ngrid_y,ngrid_z_lpt=ngrid_z,
offset_zoom=offset_zoom,BBox_in=BBox_in)
#Get bounding boxes for the COLA box with 1 refinement level for MUSIC.
BBox_in, offset_zoom, cellsize, cellsize_zoom, \
offset_index, BBox_out, BBox_out_zoom, \
ngrid_x, ngrid_y, ngrid_z, gridcellsize \
= boundaries(
boxsize, level, level_zoom, \
NPART_zoom, offset_from_code, cut_from_sides, gridscale)
# Trim full-box displacement fields down to COLA volume.
sx_full = sx_full1[BBox_out[0,0]:BBox_out[0,1], \
BBox_out[1,0]:BBox_out[1,1], \
BBox_out[2,0]:BBox_out[2,1]]
sy_full = sy_full1[BBox_out[0,0]:BBox_out[0,1],
BBox_out[1,0]:BBox_out[1,1], \
BBox_out[2,0]:BBox_out[2,1]]
sz_full = sz_full1[BBox_out[0,0]:BBox_out[0,1], \
BBox_out[1,0]:BBox_out[1,1], \
BBox_out[2,0]:BBox_out[2,1]]
sx_full_zoom = sx_full_zoom1[BBox_out_zoom[0,0]:BBox_out_zoom[0,1], \
BBox_out_zoom[1,0]:BBox_out_zoom[1,1], \
BBox_out_zoom[2,0]:BBox_out_zoom[2,1]]
sy_full_zoom = sy_full_zoom1[BBox_out_zoom[0,0]:BBox_out_zoom[0,1], \
BBox_out_zoom[1,0]:BBox_out_zoom[1,1], \
BBox_out_zoom[2,0]:BBox_out_zoom[2,1]]
sz_full_zoom = sz_full_zoom1[BBox_out_zoom[0,0]:BBox_out_zoom[0,1], \
BBox_out_zoom[1,0]:BBox_out_zoom[1,1], \
BBox_out_zoom[2,0]:BBox_out_zoom[2,1]]
del sx_full1, sy_full1, sz_full1, sx_full_zoom1, sy_full_zoom1, sz_full_zoom1
sx2_full = sx2_full1[BBox_out[0,0]:BBox_out[0,1], \
BBox_out[1,0]:BBox_out[1,1], \
BBox_out[2,0]:BBox_out[2,1]]
sy2_full = sy2_full1[BBox_out[0,0]:BBox_out[0,1], \
BBox_out[1,0]:BBox_out[1,1], \
BBox_out[2,0]:BBox_out[2,1]]
sz2_full = sz2_full1[BBox_out[0,0]:BBox_out[0,1], \
BBox_out[1,0]:BBox_out[1,1], \
BBox_out[2,0]:BBox_out[2,1]]
sx2_full_zoom = sx2_full_zoom1[BBox_out_zoom[0,0]:BBox_out_zoom[0,1], \
BBox_out_zoom[1,0]:BBox_out_zoom[1,1], \
BBox_out_zoom[2,0]:BBox_out_zoom[2,1]]
sy2_full_zoom = sy2_full_zoom1[BBox_out_zoom[0,0]:BBox_out_zoom[0,1], \
BBox_out_zoom[1,0]:BBox_out_zoom[1,1], \
BBox_out_zoom[2,0]:BBox_out_zoom[2,1]]
sz2_full_zoom = sz2_full_zoom1[BBox_out_zoom[0,0]:BBox_out_zoom[0,1], \
BBox_out_zoom[1,0]:BBox_out_zoom[1,1], \
BBox_out_zoom[2,0]:BBox_out_zoom[2,1]]
del sx2_full1, sy2_full1, sz2_full1, sx2_full_zoom1, sy2_full_zoom1, sz2_full_zoom1
print "2LPT on full box is done."
print "Starting COLA!"
px, py, pz, vx, vy, vz, \
px_zoom, py_zoom, pz_zoom, vx_zoom, vy_zoom, vz_zoom \
= evolve(
cellsize,
sx_full, sy_full, sz_full,
sx2_full, sy2_full, sz2_full,
FULL=False,
cellsize_zoom=cellsize_zoom,
sx_full_zoom = sx_full_zoom ,
sy_full_zoom = sy_full_zoom ,
sz_full_zoom = sz_full_zoom ,
sx2_full_zoom = sx2_full_zoom,
sy2_full_zoom = sy2_full_zoom,
sz2_full_zoom = sz2_full_zoom,
offset_zoom=offset_zoom,
BBox_in=BBox_in,
ngrid_x=ngrid_x,
ngrid_y=ngrid_y,
ngrid_z=ngrid_z,
gridcellsize=gridcellsize,
ngrid_x_lpt=ngrid_x,
ngrid_y_lpt=ngrid_y,
ngrid_z_lpt=ngrid_z,
gridcellsize_lpt=gridcellsize,
a_final=1.,
a_initial=1./10.,
n_steps=10,
save_to_file=False, # set this to True to output the snapshot to a file
file_npz_out='tmp.npz',
)
del vx_zoom, vy_zoom, vz_zoom
del vx, vy, vz
print "Making a figure ..."
# grid size for figure array
ngrid = 2 * 128
# physical size of figure array
cutsize = 12.0 #Mpc/h
# offset vector [Mpc/h]:
com = [1.30208333, 1.10677083, 0.944]
com[0] += offset_zoom[0]+cellsize_zoom * \
(BBox_out_zoom[0,1]-BBox_out_zoom[0,0])/2.0-cutsize/2.0
com[1] += offset_zoom[1]+cellsize_zoom * \
(BBox_out_zoom[1,1]-BBox_out_zoom[1,0])/2.0-cutsize/2.0
com[2] += offset_zoom[2]+cellsize_zoom * \
(BBox_out_zoom[2,1]-BBox_out_zoom[2,0])/2.0-cutsize/2.0
density, den_k, den_fft, _ = initialize_density(ngrid, ngrid, ngrid)
density.fill(0.0)
# Lay down fine particles on density array with CiC:
CICDeposit_3(
px_zoom - com[0],
py_zoom - com[1],
pz_zoom - com[2],
px_zoom,
py_zoom,
pz_zoom, #dummies
density,
cellsize_zoom,
cutsize / float(ngrid),
0,
0, # dummy
0, # dummy
np.array([[0, 0], [0, 0], [0, 0]], dtype='int32'),
np.array([0.0, 0.0, 0.0], dtype='float32'),
0)
# Lay down any present crude particles on density array with CiC:
CICDeposit_3(
px - com[0],
py - com[1],
pz - com[2],
px,
py,
pz, #dummies
density,
cellsize,
cutsize / float(ngrid),
0,
0, # dummy
0, # dummy
BBox_in,
np.array([0.0, 0.0, 0.0], dtype='float32'),
0)
# make the figure:
plt.imshow(np.arcsinh(
(density.mean(axis=2)) * np.sinh(1.0) / 10.0)**(1. / 3.),
vmin=0.0,
vmax=1.75,
interpolation='bicubic',
cmap='CMRmap_r')
plt.axis('off')
plt.show()