def egBk():
''' calculate example Bk
'''
Lbox = 2600. # box size
Ngrid = 360 # fft grid size
fnbox = h5py.File(os.path.join(UT.dat_dir(), 'BoxN1.hdf5'), 'r')
xyz = fnbox['xyz'].value
vxyz = fnbox['vxyz'].value
xyz_s = pySpec.applyRSD(xyz.T, vxyz.T, 0.5, h=0.7, omega0_m=0.3, LOS='z', Lbox=2600.)
bisp = pySpec.Bk_periodic(xyz_s, Lbox=Lbox, Ngrid=Ngrid, step=3, Ncut=3, Nmax=40, fft='pyfftw')
pickle.dump(bisp, open(os.path.join(UT.dat_dir(), 'egBk.p'), 'wb'))
return None
def demoBk():
'''
'''
kmax = 0.5
kf = 2.*np.pi/2600.
# read in Bk
bisp = pickle.load(open(os.path.join(UT.dat_dir(), 'egBk.p'), 'rb'))
i_k, j_k, l_k, bk = bisp['i_k1'], bisp['i_k2'], bisp['i_k3'], bisp['b123']
klim = ((i_k*kf <= kmax) & (j_k*kf <= kmax) & (l_k*kf <= kmax))
i_k, j_k, l_k = i_k[klim], j_k[klim], l_k[klim]
ijl = UT.ijl_order(i_k, j_k, l_k, typ='GM') # order of triangles
for itri in range(1500):
fig = plt.figure(figsize=(15,4.8))
sub = fig.add_subplot(111)
sub.plot(range(1500), bk[klim][ijl][:1500], c='k', lw=0.5)
sub.scatter(range(1500), bk[klim][ijl][:1500], c='k', s=2)
sub.scatter(itri, bk[klim][ijl][itri], c='C1', s=30, zorder=10)
itri_i, itri_j, itri_l = i_k[ijl][itri], j_k[ijl][itri], l_k[ijl][itri]
trixy = np.zeros((3,2))
trixy[0,:] = np.array([1400, 2e10])
trixy[1,:] = np.array([1400 - 3.33*itri_i, 2e10])
theta23 = np.arccos(-0.5*(itri_l**2 - itri_i**2 - itri_j**2)/itri_i/itri_j)[0]
trixy[2,:] = np.array([1400-3.33*itri_j[0]*np.cos(theta23),
10**(10.301 - 0.022*itri_j[0]*np.sin(theta23))])
tri = plt.Polygon(trixy, fill=None, edgecolor='k')
fig.gca().add_patch(tri)
sub.text(0.5*(trixy[0,0]+trixy[1,0]), 0.5*(trixy[0,1]+trixy[1,1]), '$k_3$',
fontsize=8, ha='center', va='bottom')
sub.text(0.5*(trixy[0,0]+trixy[2,0]), 0.5*(trixy[0,1]+trixy[2,1]), '$k_2$',
fontsize=8, ha='left', va='top')
sub.text(0.5*(trixy[1,0]+trixy[2,0]), 0.5*(trixy[1,1]+trixy[2,1]), '$k_1$',
fontsize=8, ha='right', va='top')
sub.set_xlabel('triangle configurations', fontsize=25)
sub.set_xlim([0, 1500])
sub.set_ylabel(r'$B(k_1, k_2, k_3)$', fontsize=25)
sub.set_yscale('log')
sub.set_ylim([1e7, 3e10])
fig.savefig(os.path.join(UT.dat_dir(), 'demo', 'demoBk%i.png' % itri), bbox_inches='tight')
plt.close()
return None
def AEM_rzspace():
''' comparison between real and redshift space
'''
f_b123 = lambda sim, rsd: os.path.join(
UT.dat_dir(), sim,
'pySpec.B123.halo.mlim1e13.Ngrid360.Nmax40.Ncut3.step3.pyfftw%s.dat' %
['', '.rsd'][rsd])
Lbox_aem = 1050.
kf_aem = 2. * np.pi / Lbox_aem
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts, _ = np.loadtxt(
f_b123('aemulus', False), skiprows=1, unpack=True, usecols=range(10))
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123_s, q123_s, counts, _ = np.loadtxt(
f_b123('aemulus', True), skiprows=1, unpack=True, usecols=range(10))
klim = ((i_k * kf_aem <= 0.3) & (j_k * kf_aem <= 0.3) &
(l_k * kf_aem <= 0.3))
ijl = UT.ijl_order(i_k[klim], j_k[klim], l_k[klim],
typ='GM') # order of triangles
sub.scatter(range(np.sum(klim)),
b123[klim][ijl],
c='k',
s=5,
label='real-space')
sub.plot(range(np.sum(klim)), b123[klim][ijl], c='k')
sub.scatter(range(np.sum(klim)),
b123_s[klim][ijl],
c='C1',
s=5,
label='z-space')
sub.plot(range(np.sum(klim)), b123_s[klim][ijl], c='C1')
print(b123_s[klim][ijl] / b123[klim][ijl]).flatten()
sub.legend(loc='upper right', markerscale=4, handletextpad=0., fontsize=20)
sub.set_ylabel('$B(k_1, k_2, k_3)$', fontsize=25)
sub.set_yscale('log')
sub.set_xlabel(r'$k_1 \le k_2 \le k_3$ triangle indices', fontsize=25)
sub.set_xlim([0, np.sum(klim)])
fig.savefig(''.join([UT.dat_dir(), 'qpm/B123_aemulus_rzspace.png']),
bbox_inches='tight')
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
sub.scatter(range(np.sum(klim)),
q123[klim][ijl],
c='k',
s=2,
label='real-space')
sub.plot(range(np.sum(klim)), q123[klim][ijl], c='k', lw=1)
sub.scatter(range(np.sum(klim)),
q123_s[klim][ijl],
c='C1',
s=2,
label='redshift-space')
sub.plot(range(np.sum(klim)), q123_s[klim][ijl], c='C1', lw=1)
sub.legend(loc='lower right',
handletextpad=0.2,
markerscale=10,
fontsize=20)
sub.set_ylabel('$Q(k_1, k_2, k_3)$', fontsize=25)
sub.set_xlabel('$k_1 > k_2 > k_3$ triangle indices', fontsize=25)
sub.set_xlim([0, np.sum(klim)])
sub.set_ylim([0., 1.])
fig.savefig(''.join([UT.dat_dir(), 'qpm/Q123_aemulus_rzspace.png']),
bbox_inches='tight')
return None
def QPM_AEM(rsd=False):
str_rsd = ''
if rsd: str_rsd = '.rsd'
f_pell = lambda sim: ''.join([
UT.dat_dir(), sim, '/pySpec.Plk.halo.mlim1e13.Ngrid360', str_rsd,
'.dat'
])
f_pnkt = lambda sim: ''.join([
UT.dat_dir(), sim, '/pySpec.Plk.halo.mlim1e13.Ngrid360.nbodykit',
str_rsd, '.dat'
])
f_b123 = lambda sim: ''.join([
UT.dat_dir(), sim,
'/pySpec.B123.halo.mlim1e13.Ngrid360.Nmax40.Ncut3.step3.pyfftw',
str_rsd, '.dat'
])
Lbox_qpm = 1050.
Lbox_aem = 1050.
kf_qpm = 2. * np.pi / Lbox_qpm
kf_aem = 2. * np.pi / Lbox_aem
fig = plt.figure(figsize=(5, 5))
sub = fig.add_subplot(111)
k, p0k = np.loadtxt(f_pell('aemulus'),
unpack=True,
skiprows=1,
usecols=[0, 1])
klim = (k < 0.3)
sub.plot(k[klim], p0k[klim], c='k', lw=1, label='Aemulus')
k, p0k = np.loadtxt(f_pnkt('aemulus'),
unpack=True,
skiprows=1,
usecols=[0, 1])
klim = (k < 0.3)
sub.plot(k[klim], p0k[klim], c='k', ls='--', lw=1)
k, p0k = np.loadtxt(f_pell('qpm'), unpack=True, skiprows=1, usecols=[0, 1])
klim = (k < 0.3)
sub.plot(k[klim], p0k[klim], c='C1', lw=1, label='QPM')
k, p0k = np.loadtxt(f_pnkt('qpm'), unpack=True, skiprows=1, usecols=[0, 1])
klim = (k < 0.3)
sub.plot(k[klim], p0k[klim], c='C1', ls='--', lw=1)
sub.legend(loc='upper right', fontsize=20)
sub.set_ylabel('$P_0(k)$', fontsize=25)
sub.set_yscale('log')
sub.set_xlabel('$k$', fontsize=25)
sub.set_xlim([8e-3, 0.5])
sub.set_xscale('log')
fig.savefig(''.join([UT.dat_dir(), 'qpm/p0k_qpm_aemulus', str_rsd,
'.png']),
bbox_inches='tight')
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts, b123_sn = np.loadtxt(
f_b123('aemulus'), skiprows=1, unpack=True, usecols=range(10))
klim = ((i_k * kf_aem <= 0.22) & (i_k * kf_aem >= 0.03) &
(j_k * kf_aem <= 0.22) & (j_k * kf_aem >= 0.03) &
(l_k * kf_aem <= 0.22) & (l_k * kf_aem >= 0.03))
i_k, j_k, l_k = i_k[klim], j_k[klim], l_k[klim]
ijl = UT.ijl_order(i_k, j_k, l_k, typ='GM') # order of triangles
sub.scatter(range(np.sum(klim)),
b123[klim][ijl],
c='k',
s=5,
label='Aemulus')
sub.plot(range(np.sum(klim)), b123[klim][ijl], c='k')
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts, b123_sn = np.loadtxt(
f_b123('qpm'), skiprows=1, unpack=True, usecols=range(10))
klim = ((i_k * kf_qpm <= 0.22) & (i_k * kf_qpm >= 0.03) &
(j_k * kf_qpm <= 0.22) & (j_k * kf_qpm >= 0.03) &
(l_k * kf_qpm <= 0.22) & (l_k * kf_qpm >= 0.03))
i_k, j_k, l_k, = i_k[klim], j_k[klim], l_k[klim]
ijl = UT.ijl_order(i_k, j_k, l_k, typ='GM') # order of triangles
sub.scatter(range(np.sum(klim)), b123[klim][ijl], c='C1', s=5, label='QPM')
sub.plot(range(np.sum(klim)), b123[klim][ijl], c='C1')
sub.legend(loc='upper right', markerscale=4, handletextpad=0., fontsize=20)
sub.set_ylabel('$B(k_1, k_2, k_3)$', fontsize=25)
sub.set_yscale('log')
sub.set_ylim([1e8, 6e9])
sub.set_xlabel(r'$k_1 \le k_2 \le k_3$ triangle indices', fontsize=25)
sub.set_xlim([0, np.sum(klim)])
fig.savefig(''.join(
[UT.dat_dir(), 'qpm/B123_qpm_aemulus', str_rsd, '.png']),
bbox_inches='tight')
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts, b123_sn = np.loadtxt(
f_b123('aemulus'), skiprows=1, unpack=True, usecols=range(10))
sub.scatter(range(len(q123)), q123, c='k', s=2, label='Aemulus')
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts, b123_sn = np.loadtxt(
f_b123('qpm'), skiprows=1, unpack=True, usecols=range(10))
sub.scatter(range(len(q123)), q123, c='C1', s=2, label='QPM')
sub.legend(loc='upper right', fontsize=20)
sub.set_ylabel('$Q(k_1, k_2, k_3)$', fontsize=25)
sub.set_xlabel('$k_1 > k_2 > k_3$ triangle indices', fontsize=25)
sub.set_xlim([0, len(q123)])
sub.set_ylim([0., 1.])
fig.savefig(''.join(
[UT.dat_dir(), 'qpm/Q123_qpm_aemulus', str_rsd, '.png']),
bbox_inches='tight')
return None
def AEMspectra(rsd=False):
''' calculate the powerspectrum and bispectrum of the Aemulus simulation box
'''
str_rsd = ''
if rsd: str_rsd = '.rsd'
f_halo = ''.join(
[UT.dat_dir(), 'aemulus/aemulus_test002_halos.mlim1e13.hdf5'])
f_hdf5 = ''.join(
[UT.dat_dir(), 'aemulus/aemulus_test002_halos.mlim1e13.hdf5'])
f_pell = ''.join([
UT.dat_dir(), 'aemulus/pySpec.Plk.halo.mlim1e13.Ngrid360', str_rsd,
'.dat'
])
f_pnkt = ''.join([
UT.dat_dir(), 'aemulus/pySpec.Plk.halo.mlim1e13.Ngrid360.nbodykit',
str_rsd, '.dat'
])
f_b123 = ''.join([
UT.dat_dir(),
'aemulus/pySpec.B123.halo.mlim1e13.Ngrid360.Nmax40.Ncut3.step3.pyfftw',
str_rsd, '.dat'
])
Lbox = 1050.
kf = 2. * np.pi / Lbox
if not os.path.isfile(f_hdf5):
f = h5py.File(f_halo, 'r')
xyz = f['xyz'].value
vxyz = f['vxyz'].value
mh = f['mhalo'].value
xyz_s = pySpec.applyRSD(xyz.T,
vxyz.T,
0.55,
h=0.7,
omega0_m=0.340563,
LOS='z',
Lbox=Lbox)
xyz_s = xyz_s.T
f = h5py.File(f_hdf5, 'w')
f.create_dataset('xyz', data=xyz)
f.create_dataset('vxyz', data=vxyz)
f.create_dataset('xyz_s', data=xyz_s)
f.create_dataset('mhalo', data=mh)
f.close()
else:
f = h5py.File(f_hdf5, 'r')
xyz = f['xyz'].value
vxyz = f['vxyz'].value
xyz_s = f['xyz_s'].value
mh = f['mhalo'].value
f.close()
Nhalo = xyz.shape[0]
print('# halos = %i' % Nhalo)
nhalo = float(Nhalo) / Lbox**3
print('number density = %f' % nhalo)
print('1/nbar = %f' % (1. / nhalo))
# calculate powerspectrum
if not os.path.isfile(f_pell):
# calculate FFTs
if not rsd:
delta = pySpec.FFTperiodic(xyz.T,
fft='fortran',
Lbox=Lbox,
Ngrid=360,
silent=False)
else:
delta = pySpec.FFTperiodic(xyz_s.T,
fft='fortran',
Lbox=Lbox,
Ngrid=360,
silent=False)
delta_fft = pySpec.reflect_delta(delta, Ngrid=360)
# calculate powerspectrum monopole
k, p0k, cnts = pySpec.Pk_periodic(delta_fft)
k = k * kf
p0k = p0k / kf**3 - 1. / nhalo
# save to file
hdr = 'pyspectrum P_l=0(k) calculation. k_f = 2pi/1050.'
np.savetxt(f_pell,
np.array([k, p0k, cnts]).T,
fmt='%.5e %.5e %.5e',
delimiter='\t',
header=hdr)
else:
k, p0k, cnts = np.loadtxt(f_pell,
skiprows=1,
unpack=True,
usecols=[0, 1, 2])
# calculate P(k) using nbodykit for santiy check
if not os.path.isfile(f_pnkt):
# get cosmology from header
Omega_m = 0.3175
Omega_b = 0.049 # fixed baryon
h = 0.6711
cosmo = NBlab.cosmology.Planck15.clone(Omega_cdm=Omega_m - Omega_b,
h=h,
Omega_b=Omega_b)
halo_data = {}
if not rsd: halo_data['Position'] = xyz
else: halo_data['Position'] = xyz_s
halo_data['Velocity'] = vxyz
halo_data['Mass'] = mh
print("putting it into array catalog")
halos = NBlab.ArrayCatalog(halo_data,
BoxSize=np.array([Lbox, Lbox, Lbox]))
print("putting it into halo catalog")
halos = NBlab.HaloCatalog(halos, cosmo=cosmo, redshift=0., mdef='vir')
print("putting it into mesh")
mesh = halos.to_mesh(window='tsc',
Nmesh=360,
compensated=True,
position='Position')
print("calculating powerspectrum")
r = NBlab.FFTPower(mesh, mode='1d', dk=kf, kmin=kf, poles=[0, 2, 4])
poles = r.poles
plk = {'k': poles['k']}
for ell in [0, 2, 4]:
P = (poles['power_%d' % ell].real)
if ell == 0:
P = P - poles.attrs[
'shotnoise'] # subtract shotnoise from monopole
plk['p%dk' % ell] = P
plk['shotnoise'] = poles.attrs['shotnoise'] # save shot noise term
# header
hdr = 'pyspectrum P_l(k) calculation. k_f = 2pi/1050; P_shotnoise ' + str(
plk['shotnoise'])
# write to file
np.savetxt(f_pnkt,
np.array([plk['k'], plk['p0k'], plk['p2k'], plk['p4k']]).T,
header=hdr)
else:
_k, _p0k, _p2k, _p4k = np.loadtxt(f_pnkt,
skiprows=1,
unpack=True,
usecols=[0, 1, 2, 3])
plk = {}
plk['k'] = _k
plk['p0k'] = _p0k
plk['p2k'] = _p2k
plk['p4k'] = _p4k
# calculate bispectrum
if not os.path.isfile(f_b123):
if not rsd:
bispec = pySpec.Bk_periodic(xyz.T,
Lbox=Lbox,
Ngrid=360,
Nmax=40,
Ncut=3,
step=3,
fft='pyfftw',
nthreads=1,
silent=False)
else:
bispec = pySpec.Bk_periodic(xyz_s.T,
Lbox=Lbox,
Ngrid=360,
Nmax=40,
Ncut=3,
step=3,
fft='pyfftw',
nthreads=1,
silent=False)
i_k = bispec['i_k1']
j_k = bispec['i_k2']
l_k = bispec['i_k3']
p0k1 = bispec['p0k1']
p0k2 = bispec['p0k2']
p0k3 = bispec['p0k3']
b123 = bispec['b123']
b123_sn = bispec['b123_sn']
q123 = bispec['q123']
counts = bispec['counts']
# save to file
hdr = 'pyspectrum bispectrum calculation test. k_f = 2pi/%.1f' % Lbox
np.savetxt(f_b123,
np.array([
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts,
b123_sn
]).T,
fmt='%i %i %i %.5e %.5e %.5e %.5e %.5e %.5e %.5e',
delimiter='\t',
header=hdr)
else:
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts, b123_sn = np.loadtxt(
f_b123, skiprows=1, unpack=True, usecols=range(10))
# plot powerspecrtrum shape triangle plot
fig = plt.figure(figsize=(5, 5))
sub = fig.add_subplot(111)
sub.plot(k, p0k, c='k', lw=1, label='pySpectrum')
sub.plot(plk['k'], plk['p0k'], c='C1', lw=1, label='nbodykit')
iksort = np.argsort(i_k)
sub.plot(i_k[iksort] * kf,
p0k1[iksort],
c='k',
lw=1,
ls='--',
label='bispectrum code')
sub.legend(loc='lower left', fontsize=20)
sub.set_ylabel('$P_0(k)$', fontsize=25)
#sub.set_ylim([1e2, 3e4])
sub.set_yscale('log')
sub.set_xlabel('$k$', fontsize=25)
sub.set_xlim([3e-3, 1.])
sub.set_xscale('log')
fig.savefig(''.join([UT.dat_dir(), 'aemulus/aemulus_p0k', str_rsd,
'.png']),
bbox_inches='tight')
# plot bispectrum shape triangle plot
nbin = 31
x_bins = np.linspace(0., 1., nbin + 1)
y_bins = np.linspace(0.5, 1., (nbin // 2) + 1)
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
Bgrid = Plots._BorQgrid(
l_k.astype(float) / i_k.astype(float),
j_k.astype(float) / i_k.astype(float), q123, counts, x_bins, y_bins)
bplot = plt.pcolormesh(x_bins,
y_bins,
Bgrid.T,
vmin=0,
vmax=1,
cmap='RdBu')
cbar = plt.colorbar(bplot, orientation='vertical')
sub.set_title(r'$Q(k_1, k_2, k_3)$ QPM halo catalog', fontsize=25)
sub.set_xlabel('$k_3/k_1$', fontsize=25)
sub.set_ylabel('$k_2/k_1$', fontsize=25)
fig.savefig(''.join(
[UT.dat_dir(), 'aemulus/aemulus_Q123_shape', str_rsd, '.png']),
bbox_inches='tight')
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
Bgrid = Plots._BorQgrid(
l_k.astype(float) / i_k.astype(float),
j_k.astype(float) / i_k.astype(float), b123, counts, x_bins, y_bins)
bplot = plt.pcolormesh(x_bins,
y_bins,
Bgrid.T,
norm=LogNorm(vmin=1e6, vmax=1e8),
cmap='RdBu')
cbar = plt.colorbar(bplot, orientation='vertical')
sub.set_title(r'$B(k_1, k_2, k_3)$ QPM halo catalog', fontsize=25)
sub.set_xlabel('$k_3/k_1$', fontsize=25)
sub.set_ylabel('$k_2/k_1$', fontsize=25)
fig.savefig(''.join(
[UT.dat_dir(), 'aemulus/aemulus_B123_shape', str_rsd, '.png']),
bbox_inches='tight')
# plot bispectrum amplitude
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
sub.scatter(range(len(b123)), q123, c='k', s=1)
sub.set_xlabel(r'$k_1 > k_2 > k_3$ triangle index', fontsize=25)
sub.set_xlim([0, len(b123)])
sub.set_ylabel(r'$Q(k_1, k_2, k_3)$', fontsize=25)
sub.set_ylim([0., 1.])
fig.savefig(''.join(
[UT.dat_dir(), 'aemulus/aemulus_Q123', str_rsd, '.png']),
bbox_inches='tight')
# plot bispectrum amplitude
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
sub.scatter(range(len(b123)), b123, c='k', s=1)
sub.set_xlabel(r'$k_1 > k_2 > k_3$ triangle index', fontsize=25)
sub.set_xlim([0, len(b123)])
sub.set_ylabel(r'$B(k_1, k_2, k_3)$', fontsize=25)
sub.set_yscale('log')
fig.savefig(''.join(
[UT.dat_dir(), 'aemulus/aemulus_B123', str_rsd, '.png']),
bbox_inches='tight')
return None
'''
calculate bispectrum
'''
import time
import numpy as np
# -- pyspectrum --
from pyspectrum import util as UT
from pyspectrum import pyspectrum as pySpec
if __name__ == "__main__":
x, y, z = np.loadtxt(''.join([UT.dat_dir(), 'BoxN1.mock']),
unpack=True,
usecols=[0, 1, 2])
xyz = np.zeros((3, len(x)))
xyz[0, :] = x
xyz[1, :] = y
xyz[2, :] = z
delta = pySpec.FFTperiodic(xyz,
fft='fortran',
Lbox=2600,
Ngrid=360,
silent=False)
delta_fft = pySpec.reflect_delta(delta, Ngrid=360)
t0 = time.time()
# calculate bispectrum
i_k, j_k, l_k, b123, q123 = pySpec.Bk123_periodic(delta_fft,
Nmax=40,
import numpy as np
from scipy.io import FortranFile
# -- pyspectrum --
from pyspectrum import util as UT
from pyspectrum import pyspectrum as pySpec
if __name__ == "__main__":
counts = pySpec._counts_Bk123(Ngrid=360,
Nmax=40,
Ncut=3,
step=3,
fft_method='pyfftw',
silent=False)
counts_f77 = pySpec._counts_Bk123_f77(Ngrid=360,
Nmax=40,
Ncut=3,
step=3,
silent=False)
# check counts
f = FortranFile(''.join([UT.dat_dir(), 'counts2_n360_nmax40_ncut3_s3']),
'r')
_counts = f.read_reals(dtype=np.float64)
_counts = np.reshape(_counts, (40, 40, 40), order='F')
_counts_swap = np.swapaxes(_counts, 0,
2) # axes are swaped for whatever reason
assert np.allclose(_counts_swap, counts_f77)
assert np.allclose(_counts_swap, counts)
# both the python and fortran wrapped implementations work well
mpl.rcParams['axes.linewidth'] = 1.5
mpl.rcParams['axes.xmargin'] = 1
mpl.rcParams['xtick.labelsize'] = 'x-large'
mpl.rcParams['xtick.major.size'] = 5
mpl.rcParams['xtick.major.width'] = 1.5
mpl.rcParams['ytick.labelsize'] = 'x-large'
mpl.rcParams['ytick.major.size'] = 5
mpl.rcParams['ytick.major.width'] = 1.5
mpl.rcParams['legend.frameon'] = False
Lbox = 2600. # box size
Ngrid = 360 # fft grid size
kf = 2 * np.pi / Lbox # fundament mode
# read in Nseries box data
fnbox = h5py.File(os.path.join(UT.dat_dir(), 'BoxN1.hdf5'), 'r')
xyz = fnbox['xyz'].value
vxyz = fnbox['vxyz'].value
nhalo = xyz.shape[0]
nbar = float(nhalo) / Lbox**3
# real-space power/bispectrum with pySpectrum
t0 = time.time()
pspec = pySpec.Pk_periodic(xyz.T, Lbox=Lbox, Ngrid=Ngrid, silent=False)
print('--pySpec.Pk_periodic: %f sec' % ((time.time() - t0) / 60.))
t0 = time.time()
bispec = pySpec.Bk_periodic(xyz.T, Lbox=Lbox, Ngrid=Ngrid, silent=False)
print('--pySpec.Bk_periodic: %f sec' % ((time.time() - t0) / 60.))
# compare real-space bispectrum with Roman's output
f_bk_rs = os.path.join(UT.dat_dir(), 'bk.BoxN1.mock')
mpl.rcParams['axes.linewidth'] = 1.5
mpl.rcParams['axes.xmargin'] = 1
mpl.rcParams['xtick.labelsize'] = 'x-large'
mpl.rcParams['xtick.major.size'] = 5
mpl.rcParams['xtick.major.width'] = 1.5
mpl.rcParams['ytick.labelsize'] = 'x-large'
mpl.rcParams['ytick.major.size'] = 5
mpl.rcParams['ytick.major.width'] = 1.5
mpl.rcParams['legend.frameon'] = False
Lbox = 2600. # box size
Ngrid = 360 # fft grid size
kf = 2 * np.pi / Lbox # fundament mode
# read in Nseries box data
fnbox = h5py.File(os.path.join(UT.dat_dir(), 'BoxN1.hdf5'), 'r')
xyz = fnbox['xyz'][...]
vxyz = fnbox['vxyz'][...]
nhalo = xyz.shape[0]
nbar = float(nhalo) / Lbox**3
N = xyz.shape[1] # number of positions
kf = 2 * np.pi / Lbox
# Pk from real-space powerspectrum function
t0 = time.time()
pkout = pySpec.Pk_periodic(xyz.T,
Lbox=Lbox,
Ngrid=Ngrid,
fft='pyfftw',
silent=False)
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
mpl.rcParams['text.usetex'] = True
mpl.rcParams['font.family'] = 'serif'
mpl.rcParams['axes.linewidth'] = 1.5
mpl.rcParams['axes.xmargin'] = 1
mpl.rcParams['xtick.labelsize'] = 'x-large'
mpl.rcParams['xtick.major.size'] = 5
mpl.rcParams['xtick.major.width'] = 1.5
mpl.rcParams['ytick.labelsize'] = 'x-large'
mpl.rcParams['ytick.major.size'] = 5
mpl.rcParams['ytick.major.width'] = 1.5
mpl.rcParams['legend.frameon'] = False
dir_mc = os.path.join(UT.dat_dir(), 'gqc_mock_challenge')
def stage1(name):
''' stage 1 of mock challenge. run power spectrum for periodic box
'''
# read in mock
if name == 'unit':
_fname = 'UNIT_DESI_Shadab_HOD_snap97_ELG_v0.txt'
Lbox = 1000.
elif name == 'pmill':
_fname = 'DESI_ELG_z0.76_catalogue.dat'
Lbox = 542.16
x, y, z, z_rsd = np.loadtxt(os.path.join(dir_mc, _fname),
unpack=True,
'''
calculate powerspectrum
'''
import time
import numpy as np
import matplotlib.pyplot as plt
from pyspectrum import util as UT
from pyspectrum import pyspectrum as pySpec
if __name__ == "__main__":
# read in FFTed density grid
f_fft = ''.join([UT.dat_dir(), 'FFT.BoxN1.mock.Ngrid360'])
delta = pySpec.read_fortFFT(file=f_fft)
# calculate powerspectrum monopole
k, p0k = pySpec.Pk_periodic(delta)
kf = 2. * np.pi / 2600.
_k, _p0k = np.loadtxt(''.join([UT.dat_dir(), 'PK.BoxN1.mock.Ngrid360']),
unpack=True,
usecols=[0, 1])
plt.plot(kf * k, p0k / (2. * np.pi)**3 / kf**3)
plt.plot(_k, _p0k, ls='--')
plt.xscale("log")
plt.yscale("log")
plt.show()
def fastpm_v_tng(z,
str_flag='',
mh_lim=15.,
Lbox=205.,
Nmax=40,
Ncut=3,
step=3):
''' make plots that compare the powerspectrum and bispectrum of
fastpm and illustris tng
'''
dir_fpm = os.path.join(UT.dat_dir(), 'fastpm')
f_fpm = ('halocat_FastPM_40step_N250_IC500_B2_z%.2f%s.mlim%.fe10' %
(z, str_flag, mh_lim))
f_tng = ('halocat_TNG300Dark_z%.2f.mlim%.fe10' % (z, mh_lim))
print('FastPM file: %s' % f_fpm)
print('Illustris TNG file: %s' % f_tng)
f_pell = lambda f_halo: ('%s/pySpec.Plk.%s.Lbox%.f.Ngrid360.dat' %
(dir_fpm, f_halo, Lbox))
f_pnkt = lambda f_halo: ('%s/pySpec.Plk.%s.Lbox%.f.Ngrid360.nbodykit.dat' %
(dir_fpm, f_halo, Lbox))
f_b123 = lambda f_halo: (
'%s/pySpec.Bk.%s.Lbox%.f.Ngrid360.step%i.Ncut%i.Nmax%i.dat' %
(dir_fpm, f_halo, Lbox, step, Ncut, Nmax))
kf = 2. * np.pi / Lbox
# P(k) comparison
fig = plt.figure(figsize=(5, 5))
sub = fig.add_subplot(111)
k, p0k = np.loadtxt(f_pell(f_fpm), unpack=True, skiprows=1, usecols=[0, 1])
sub.plot(k, p0k, c='k', lw=1, label='FastPM')
k, p0k_fpm = np.loadtxt(f_pnkt(f_fpm),
unpack=True,
skiprows=1,
usecols=[0, 1])
sub.plot(k, p0k_fpm, c='k', ls='--', lw=1)
k, p0k = np.loadtxt(f_pell(f_tng), unpack=True, skiprows=1, usecols=[0, 1])
sub.plot(k, p0k, c='C1', lw=1, label='Illustris TNG')
k, p0k_tng = np.loadtxt(f_pnkt(f_tng),
unpack=True,
skiprows=1,
usecols=[0, 1])
sub.plot(k, p0k_tng, c='C1', ls='--', lw=1)
print(p0k_fpm / p0k_tng)[k < 0.2] - 1.
sub.legend(loc='lower left', fontsize=20)
sub.set_ylabel('$P_0(k)$', fontsize=25)
sub.set_yscale('log')
sub.set_ylim([1e0, 1e4])
sub.set_xlabel('$k$', fontsize=25)
sub.set_xscale('log')
sub.set_xlim([1e-2, 10.])
fig.savefig(os.path.join(dir_fpm,
'p0k_fpm_tng_z%.2f%s.png' % (z, str_flag)),
bbox_inches='tight')
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
for fh, lbl, c in zip([f_fpm, f_tng], ['FastPM', 'Illustris TNG'],
['k', 'C1']):
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts, b123_sn = np.loadtxt(
f_b123(fh), skiprows=1, unpack=True, usecols=range(10))
klim = ((i_k * kf <= 1.) & (i_k * kf >= 0.01) & (j_k * kf <= 1.) &
(j_k * kf >= 0.01) & (l_k * kf <= 1.) & (l_k * kf >= 0.01))
i_k, j_k, l_k = i_k[klim], j_k[klim], l_k[klim]
ijl = UT.ijl_order(i_k, j_k, l_k, typ='GM') # order of triangles
sub.scatter(range(np.sum(klim)), b123[klim][ijl], c=c, s=5, label=lbl)
sub.plot(range(np.sum(klim)), b123[klim][ijl], c=c)
sub.legend(loc='upper right', markerscale=4, handletextpad=0., fontsize=20)
sub.set_ylabel('$B(k_1, k_2, k_3)$', fontsize=25)
sub.set_yscale('log')
sub.set_ylim([1e3, 5e6])
sub.set_xlabel(r'$k_1 \le k_2 \le k_3$ triangle indices', fontsize=25)
sub.set_xlim([0, np.sum(klim)])
fig.savefig(os.path.join(dir_fpm,
'bk_fpm_tng_z%.2f%s.png' % (z, str_flag)),
bbox_inches='tight')
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
for fh, lbl, c in zip([f_fpm, f_tng], ['FastPM', 'Illustris TNG'],
['k', 'C1']):
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts, b123_sn = np.loadtxt(
f_b123(fh), skiprows=1, unpack=True, usecols=range(10))
klim = ((i_k * kf >= 0.01) & (j_k * kf >= 0.01) & (l_k * kf >= 0.01))
i_k, j_k, l_k = i_k[klim], j_k[klim], l_k[klim]
ijl = UT.ijl_order(i_k, j_k, l_k, typ='GM') # order of triangles
sub.scatter(range(np.sum(klim)), q123[klim][ijl], c=c, s=5, label=lbl)
sub.plot(range(np.sum(klim)), q123[klim][ijl], c=c)
sub.legend(loc='upper right', markerscale=4, handletextpad=0., fontsize=20)
sub.set_ylabel('$Q(k_1, k_2, k_3)$', fontsize=25)
sub.set_ylim([0., 1.])
sub.set_xlabel(r'$k_1 \le k_2 \le k_3$ triangle indices', fontsize=25)
sub.set_xlim([0, np.sum(klim)])
fig.savefig(os.path.join(dir_fpm,
'qk_fpm_tng_z%.2f%s.png' % (z, str_flag)),
bbox_inches='tight')
return None
def fastPM(z, str_flag='', mh_lim=15., Lbox=205., Nmax=40, Ncut=3, step=3):
''' calculate the powerspectrum and bispectrum of the fastPM catalog.
'''
dir_fpm = os.path.join(UT.dat_dir(), 'fastpm')
f_halo = ('halocat_FastPM_40step_N250_IC500_B2_z%.2f%s.txt' %
(z, str_flag))
f_mlim = ('halocat_FastPM_40step_N250_IC500_B2_z%.2f%s.mlim%.fe10' %
(z, str_flag, mh_lim))
f_hdf5 = ('%s/%s.hdf5' % (dir_fpm, f_mlim))
f_pell = ('%s/pySpec.Plk.%s.Lbox%.f.Ngrid360.dat' %
(dir_fpm, f_mlim, Lbox))
f_pnkt = ('%s/pySpec.Plk.%s.Lbox%.f.Ngrid360.nbodykit.dat' %
(dir_fpm, f_mlim, Lbox))
f_b123 = ('%s/pySpec.Bk.%s.Lbox%.f.Ngrid360.step%i.Ncut%i.Nmax%i.dat' %
(dir_fpm, f_mlim, Lbox, step, Ncut, Nmax))
kf = 2. * np.pi / Lbox
if not os.path.isfile(f_hdf5):
# read in halo catalog
dat_halo = np.loadtxt(os.path.join(dir_fpm, f_halo),
unpack=True,
usecols=[0, 1, 2, 3, 7, 8, 9])
mh = dat_halo[0]
Nhalo = len(mh)
print('%i halos in %.f Mpc/h box' % (len(mh), Lbox))
print('%f < M_h/10^10Msun < %f' % (mh.min(), mh.max()))
xyz = np.zeros((Nhalo, 3))
xyz[:, 0] = dat_halo[1]
xyz[:, 1] = dat_halo[2]
xyz[:, 2] = dat_halo[3]
print('%f < x < %f' % (xyz[:, 0].min(), xyz[:, 0].max()))
print('%f < y < %f' % (xyz[:, 1].min(), xyz[:, 1].max()))
print('%f < z < %f' % (xyz[:, 2].min(), xyz[:, 2].max()))
vxyz = np.zeros((Nhalo, 3))
vxyz[:, 0] = dat_halo[4]
vxyz[:, 1] = dat_halo[5]
vxyz[:, 2] = dat_halo[6]
mlim = (mh > 15.)
Nhalo = np.sum(mlim)
print('%i halos in %.f Mpc/h box with Mh > %f' % (Nhalo, Lbox, mh_lim))
mh = mh[mlim]
xyz = xyz[mlim, :]
vxyz = vxyz[mlim, :]
f = h5py.File(f_hdf5, 'w')
f.create_dataset('xyz', data=xyz)
f.create_dataset('vxyz', data=vxyz)
f.create_dataset('mhalo', data=mh)
f.close()
else:
f = h5py.File(f_hdf5, 'r')
xyz = f['xyz'].value
vxyz = f['vxyz'].value
mh = f['mhalo'].value
Nhalo = xyz.shape[0]
print('%i halos in %.f Mpc/h box with Mh > %f' %
(len(mh), Lbox, mh_lim))
nhalo = float(Nhalo) / Lbox**3
print('number density = %f' % nhalo)
print('1/nbar = %f' % (1. / nhalo))
# calculate powerspectrum
if not os.path.isfile(f_pell):
delta = pySpec.FFTperiodic(xyz.T,
fft='fortran',
Lbox=Lbox,
Ngrid=360,
silent=False)
delta_fft = pySpec.reflect_delta(delta, Ngrid=360)
# calculate powerspectrum monopole
k, p0k, cnts = pySpec.Pk_periodic(delta_fft)
k *= kf
p0k = p0k / (kf**3) - 1. / nhalo
# save to file
hdr = ('pySpectrum P_l=0(k). Nhalo=%i, Lbox=%.f, k_f=%.5e, SN=%.5e' %
(Nhalo, Lbox, kf, 1. / nhalo))
hdr += '\n k, p0k, counts'
np.savetxt(f_pell,
np.array([k, p0k, cnts]).T,
fmt='%.5e %.5e %.5e',
delimiter='\t',
header=hdr)
else:
k, p0k, cnts = np.loadtxt(f_pell,
skiprows=1,
unpack=True,
usecols=[0, 1, 2])
# calculate P(k) using nbodykit for santiy check
if not os.path.isfile(f_pnkt):
cosmo = NBlab.cosmology.Planck15
halo_data = {}
halo_data['Position'] = xyz
halo_data['Velocity'] = vxyz
halo_data['Mass'] = mh
print("putting it into array catalog")
halos = NBlab.ArrayCatalog(halo_data,
BoxSize=np.array([Lbox, Lbox, Lbox]))
print("putting it into halo catalog")
halos = NBlab.HaloCatalog(halos, cosmo=cosmo, redshift=z, mdef='vir')
print("putting it into mesh")
mesh = halos.to_mesh(window='tsc',
Nmesh=360,
compensated=True,
position='Position')
print("calculating powerspectrum")
r = NBlab.FFTPower(mesh, mode='1d', dk=kf, kmin=kf, poles=[0, 2, 4])
poles = r.poles
plk = {'k': poles['k']}
for ell in [0, 2, 4]:
P = (poles['power_%d' % ell].real)
if ell == 0:
P = P - poles.attrs[
'shotnoise'] # subtract shotnoise from monopole
plk['p%dk' % ell] = P
plk['shotnoise'] = poles.attrs['shotnoise'] # save shot noise term
# header
hdr = ('pySpectrum P_l(k). Nhalo=%i, Lbox=%.f, k_f=%.5e, SN=%.5e' %
(Nhalo, Lbox, kf, plk['shotnoise']))
hdr += '\n k, p0k, p2k, p4k'
# save to file
np.savetxt(f_pnkt,
np.array([plk['k'], plk['p0k'], plk['p2k'], plk['p4k']]).T,
header=hdr)
else:
_k, _p0k, _p2k, _p4k = np.loadtxt(f_pnkt,
skiprows=1,
unpack=True,
usecols=[0, 1, 2, 3])
plk = {}
plk['k'] = _k
plk['p0k'] = _p0k
plk['p2k'] = _p2k
plk['p4k'] = _p4k
# calculate bispectrum
if not os.path.isfile(f_b123):
# calculate bispectrum
bispec = pySpec.Bk_periodic(xyz.T,
Lbox=Lbox,
Ngrid=360,
Nmax=40,
Ncut=3,
step=3,
fft='pyfftw',
nthreads=1,
silent=False)
i_k = bispec['i_k1']
j_k = bispec['i_k2']
l_k = bispec['i_k3']
p0k1 = bispec['p0k1']
p0k2 = bispec['p0k2']
p0k3 = bispec['p0k3']
b123 = bispec['b123']
b123_sn = bispec['b123_sn']
q123 = bispec['q123']
counts = bispec['counts']
# save to file
hdr = 'pyspectrum bispectrum calculation test. k_f = 2pi/%.1f' % Lbox
hdr += '\n i_k1, i_k2, i_k3, p0k1, p0k2, p0k3, bk, qk, counts, bk_shotnoise'
np.savetxt(f_b123,
np.array([
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts,
b123_sn
]).T,
fmt='%i %i %i %.5e %.5e %.5e %.5e %.5e %.5e %.5e',
delimiter='\t',
header=hdr)
else:
i_k, j_k, l_k, p0k1, p0k2, p0k3, b123, q123, counts, b123_sn = np.loadtxt(
f_b123, skiprows=1, unpack=True, usecols=range(10))
# plot powerspecrtrum shape triangle plot
fig = plt.figure(figsize=(5, 5))
sub = fig.add_subplot(111)
sub.plot(k, p0k, c='k', lw=1, label='pySpectrum')
sub.plot(plk['k'], plk['p0k'], c='C1', lw=1, label='nbodykit')
iksort = np.argsort(i_k)
sub.plot(i_k[iksort] * kf,
p0k1[iksort],
c='k',
lw=1,
ls='--',
label='bispectrum code')
sub.legend(loc='lower left', fontsize=20)
sub.set_ylabel('$P_0(k)$', fontsize=25)
sub.set_ylim([1e0, 1e4])
sub.set_yscale('log')
sub.set_xlabel('$k$', fontsize=25)
sub.set_xlim([1e-2, 10.])
sub.set_xscale('log')
fig.savefig(f_pell.replace('.dat', '.png'), bbox_inches='tight')
# plot bispectrum shape triangle plot
nbin = 31
x_bins = np.linspace(0., 1., nbin + 1)
y_bins = np.linspace(0.5, 1., (nbin // 2) + 1)
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
Bgrid = Plots._BorQgrid(
l_k.astype(float) / i_k.astype(float),
j_k.astype(float) / i_k.astype(float), q123, counts, x_bins, y_bins)
bplot = plt.pcolormesh(x_bins,
y_bins,
Bgrid.T,
vmin=0,
vmax=1,
cmap='RdBu')
cbar = plt.colorbar(bplot, orientation='vertical')
sub.set_title(r'$Q(k_1, k_2, k_3)$ FastPM halo catalog', fontsize=25)
sub.set_xlabel('$k_3/k_1$', fontsize=25)
sub.set_ylabel('$k_2/k_1$', fontsize=25)
fig.savefig(f_b123.replace('.dat', '.Qk_shape.png'), bbox_inches='tight')
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
Bgrid = Plots._BorQgrid(
l_k.astype(float) / i_k.astype(float),
j_k.astype(float) / i_k.astype(float), b123, counts, x_bins, y_bins)
bplot = plt.pcolormesh(x_bins,
y_bins,
Bgrid.T,
norm=LogNorm(vmin=1e6, vmax=1e8),
cmap='RdBu')
cbar = plt.colorbar(bplot, orientation='vertical')
sub.set_title(r'$B(k_1, k_2, k_3)$ FastPM halo catalog', fontsize=25)
sub.set_xlabel('$k_3/k_1$', fontsize=25)
sub.set_ylabel('$k_2/k_1$', fontsize=25)
fig.savefig(f_b123.replace('.dat', '.Bk_shape.png'), bbox_inches='tight')
# plot bispectrum amplitude
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
sub.scatter(range(len(b123)), q123, c='k', s=1)
sub.set_xlabel(r'$k_1 > k_2 > k_3$ triangle index', fontsize=25)
sub.set_xlim([0, len(b123)])
sub.set_ylabel(r'$Q(k_1, k_2, k_3)$', fontsize=25)
sub.set_ylim([0., 1.])
fig.savefig(f_b123.replace('.dat', '.Qk.png'), bbox_inches='tight')
fig = plt.figure(figsize=(10, 5))
sub = fig.add_subplot(111)
sub.scatter(range(len(b123)), b123, c='k', s=1)
sub.set_xlabel(r'$k_1 > k_2 > k_3$ triangle index', fontsize=25)
sub.set_xlim([0, len(b123)])
sub.set_ylabel(r'$B(k_1, k_2, k_3)$', fontsize=25)
sub.set_yscale('log')
fig.savefig(f_b123.replace('.dat', '.Bk.png'), bbox_inches='tight')
return None
mpl.rcParams['text.usetex'] = True
mpl.rcParams['font.family'] = 'serif'
mpl.rcParams['axes.linewidth'] = 1.5
mpl.rcParams['axes.xmargin'] = 1
mpl.rcParams['xtick.labelsize'] = 'x-large'
mpl.rcParams['xtick.major.size'] = 5
mpl.rcParams['xtick.major.width'] = 1.5
mpl.rcParams['ytick.labelsize'] = 'x-large'
mpl.rcParams['ytick.major.size'] = 5
mpl.rcParams['ytick.major.width'] = 1.5
mpl.rcParams['legend.frameon'] = False
if __name__ == "__main__":
kf = 2. * np.pi / 2600.
x, y, z, vx, vy, vz = np.loadtxt(os.path.join(UT.dat_dir(), 'BoxN1.mock'),
unpack=True,
usecols=[0, 1, 2, 3, 4, 5])
xyz = np.zeros((3, len(x)))
xyz[0, :] = x
xyz[1, :] = y
xyz[2, :] = z
vxyz = np.zeros((3, len(x)))
vxyz[0, :] = vx
vxyz[1, :] = vy
vxyz[2, :] = vz
s_xyz = pySpec.applyRSD(xyz,
vxyz,
0.5,