コード例 #1
0
ファイル: data.py プロジェクト: changhoonhahn/specmulator
def NeutHalos(mneut, nreal, nzbin):
    ''' Read halo catalogs generated by Paco
    
    parameters
    ----------
    mneut : float, 
        total neutrino mass 

    nreal : int,
        realization number 

    nzbin : int, 
        integer specifying the redshift of the snapshot. 
        nzbin = 0 --> z=3
        nzbin = 1 --> z=2
        nzbin = 2 --> z=1
        nzbin = 3 --> z=0.5
        nzbin = 4 --> z=0
    '''
    if mneut == 0.1:
        dir = ''.join([UT.dat_dir(), '0.10eV/', str(nreal)])
    else:
        dir = ''.join([UT.dat_dir(), str(mneut), 'eV/', str(nreal)])
    # read in Gadget header
    header = RS.read_gadget_header(''.join(
        [dir, '/snapdir_',
         str(nzbin).zfill(3), '/snap_',
         str(nzbin).zfill(3)]))

    # get cosmology from header
    Omega_b = 0.049  # fixed baryon
    cosmo = NBlab.cosmology.Planck15.clone(Omega_cdm=header['Omega_m'] -
                                           Omega_b,
                                           h=header['h'],
                                           Omega_b=Omega_b)

    Fof = readfof.FoF_catalog(dir,
                              nzbin,
                              long_ids=False,
                              swap=False,
                              SFR=False)
    group_data = {}
    group_data['Length'] = Fof.GroupLen
    group_data['Position'] = Fof.GroupPos / 1e3
    group_data['Velocity'] = Fof.GroupVel
    group_data['Mass'] = Fof.GroupMass * 1e10
    # calculate velocity offset
    rsd_factor = (1. + header['z']) / (100. * cosmo.efunc(header['z']))
    group_data['VelocityOffset'] = group_data['Velocity'] * rsd_factor
    #group_data['RSDPosition']       = group_data['Position'] + group_data['VelocityOffset'] * LOS

    # save to ArryCatalog for consistency
    cat = NBlab.ArrayCatalog(group_data,
                             BoxSize=np.array([1000., 1000., 1000.]))
    cat = NBlab.HaloCatalog(cat, cosmo=cosmo, redshift=header['z'], mdef='vir')
    return cat
コード例 #2
0
ファイル: qpm.py プロジェクト: yongjiguan/pySpectrum
def QPMspectra(rsd=False):
    ''' calculate the powerspectrum and bispectrum of the QPM 
    catalog.

    :param rsd: (default: False)
        if True calculate in redshift space. Otherwise, real-space 
    '''
    str_rsd = ''
    if rsd: str_rsd = '.rsd'
    f_halo = ''.join([dir_dat, 'halo_ascii.dat'])
    f_hdf5 = ''.join([dir_dat, 'halo.mlim1e13.Lbox1050.hdf5'])
    f_pell = ''.join([
        dir_dat, 'pySpec.Plk.halo.mlim1e13.Lbox1050', '.Ngrid360', str_rsd,
        '.dat'
    ])
    f_pnkt = ''.join([
        dir_dat, 'pySpec.Plk.halo.mlim1e13.Lbox1050', '.Ngrid360', '.nbodykit',
        str_rsd, '.dat'
    ])
    f_b123 = ''.join([
        dir_dat, 'pySpec.B123.halo.mlim1e13.Lbox1050', '.Ngrid360', '.Nmax40',
        '.Ncut3', '.step3', '.pyfftw', str_rsd, '.dat'
    ])

    Lbox = 1050.
    kf = 2. * np.pi / Lbox

    # 1. read in ascii file
    # 2. impose 10^13 halo mass limit
    # 3. calculate RSD positions
    # 4. write to hdf5 file
    if not os.path.isfile(f_hdf5):
        mh, x, y, z, vx, vy, vz = np.loadtxt(f_halo,
                                             unpack=True,
                                             skiprows=1,
                                             usecols=[0, 1, 2, 3, 4, 5, 6])
        xyz = np.zeros((len(x), 3))
        xyz[:, 0] = x
        xyz[:, 1] = y
        xyz[:, 2] = z

        vxyz = np.zeros((len(x), 3))
        vxyz[:, 0] = vx
        vxyz[:, 1] = vy
        vxyz[:, 2] = vz

        # RSD along the z axis
        xyz_s = pySpec.applyRSD(xyz.T,
                                vxyz.T,
                                0.55,
                                h=0.7,
                                omega0_m=0.340563,
                                LOS='z',
                                Lbox=Lbox)

        mlim = (mh > 1e13)

        mh = mh[mlim]
        xyz = xyz[mlim, :]
        vxyz = vxyz[mlim, :]
        xyz_s = xyz_s.T[mlim, :]

        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
        xyz_s = f['xyz_s'].value
        vxyz = f['vxyz'].value
        mh = f['mhalo'].value

    Nhalo = xyz.shape[0]
    print('# halos = %i in %.1f box' % (Nhalo, Lbox))
    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 powerspectrum monopole
        if not rsd:
            spec = pySpec.Pk_periodic(xyz.T,
                                      Lbox=Lbox,
                                      Ngrid=360,
                                      silent=False)
        else:
            spec = pySpec.Pk_periodic(xyz_s.T,
                                      Lbox=Lbox,
                                      Ngrid=360,
                                      silent=False)
        k = spec['k']
        p0k = spec['p0k']
        cnts = spec['counts']
        # save to file
        hdr = (
            'pyspectrum P_l=0(k) calculation. Lbox=%.1f, k_f=%.5e, SN=%.5e' %
            (Lbox, kf, 1. / nhalo))
        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/%.1f; P_shotnoise %f' % (
            Lbox, 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):
        # calculate bispectrum
        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')
    sub.plot(i_k * kf, p0k1, 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([dir_dat, 'qpm_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([dir_dat, 'qpm_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([dir_dat, 'qpm_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([dir_dat, 'qpm_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([dir_dat, 'qpm_B123', str_rsd, '.png']),
                bbox_inches='tight')
    return None
コード例 #3
0
ファイル: qpm.py プロジェクト: yongjiguan/pySpectrum
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
コード例 #4
0
ファイル: fastpm.py プロジェクト: yongjiguan/pySpectrum
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
コード例 #5
0
ファイル: quijote.py プロジェクト: emassara/boss-sbi
def Halos(halo_folder,
          z=0.5,
          Om=None,
          Ob=None,
          h=None,
          ns=None,
          s8=None,
          Mnu=0.):
    ''' read in Quijote halo catalog given the folder and snapshot # and store it as
    a nbodykit HaloCatalog object. The HaloCatalog object is convenient for 
    populating with galaxies and etc.


    Parameters
    ----------
    halo_folder : string
        directory that contains the halo catalogs e.g. on tiger it'd be
        something like: /projects/QUIJOTE/Halos/latin_hypercube/HR_0/

    Return 
    ------
    cat : nbodykit.lab.HaloCatalog 
        Quijote halo catalog  
    '''
    # if snapshot folder is not specified
    # then all values have to be specified in kwargs
    assert all([tt is not None for tt in [Om, Ol, z, h, Hz]])

    # redshift snapshot
    assert z in quijote_zsnap_dict.keys(
    ), 'snapshots are available at z=0, 0.5, 1, 2, 3'
    snapnum = quijote_zsnap_dict[z]

    # this cosmology is not used for anything. but it's included for nbodykit
    cosmo = NBlab.cosmology.Planck15.clone(h=h,
                                           Omega0_b=Ob,
                                           Omega0_cdm=Om - Ob,
                                           m_ncdm=Mnu,
                                           n_s=ns)
    Ol = 1. - Om
    Hz = 100.0 * np.sqrt(Om * (1. + z)**3 + Ol)  # km/s/(Mpc/h)

    # read FOF catalog (~90.6 ms)
    Fof = readfof.FoF_catalog(halo_folder,
                              snapnum,
                              read_IDs=False,
                              long_ids=False,
                              swap=False,
                              SFR=False)
    group_data = {}
    group_data['Length'] = Fof.GroupLen
    group_data['Position'] = Fof.GroupPos / 1e3
    group_data['Velocity'] = Fof.GroupVel
    group_data['Mass'] = Fof.GroupMass * 1e10
    # calculate velocity offset
    rsd_factor = (1. + z) / Hz
    group_data['VelocityOffset'] = group_data['Velocity'] * rsd_factor
    # save to ArryCatalog for consistency
    cat = NBlab.ArrayCatalog(group_data,
                             BoxSize=np.array([1000., 1000., 1000.]))
    cat = NBlab.HaloCatalog(cat, cosmo=cosmo, redshift=z, mdef='vir')

    cat.attrs['Om'] = Om
    cat.attrs['Ob'] = Ob
    cat.attrs['Ol'] = Ol
    cat.attrs['h'] = h
    cat.attrs['ns'] = ns
    cat.attrs['s8'] = s8
    cat.attrs['Hz'] = Hz  # km/s/(Mpc/h)
    cat.attrs['rsd_factor'] = rsd_factor
    return cat