Difference in 21cm power spectrum with and without inner structure

paper_plot.py

from prettyfigure.style import *
define_figure_style()

data = np.load('temp.npz')
pkdata = data['pkdata']
pkdata_corrected = data['pkdata_corrected']
pkdata_other = data['pkdata_other']
pkdata_b = data['pkdata_b']
slice1 = data['slice1']
slice2 = data['slice2']

plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata, label='P_k(tot)')
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata_corrected, '--', label='P_k(nf)')
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata_other, ':', label='P_k(f)')
temp=pkdata-pkdata_corrected-pkdata_other
temp[temp<0]=np.nan
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*np.abs(temp),'k-',lw=0.5, label='|P_k(tot)-P_k(nf)-P_k(f)|')
temp=pkdata-pkdata_corrected-pkdata_other
temp[temp>0]=np.nan
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*np.abs(temp),'k:',lw=0.5)
plt.xscale('log')
plt.yscale('log')
plt.legend(loc=2)
plt.ylabel(r'$k^3/(2\pi^2)\;P_k(x)$')
plt.xlabel(r'$k$')

plot_difference.py

import matplotlib
matplotlib.use('PS')

def pk(data, box_size, k_list_phys, mode=0, usefftw=False):
    N = data.shape[0]
    k_list = k_list_phys*box_size/N
    if usefftw: # Runs, but does not work !!!!!!!!!!!!!
        import pyfftw
        input_size = (N, N, N)
        padded_size = (input_size[0],input_size[1], 2*(input_size[2]//2 + 1))
        full_array = pyfftw.n_byte_align_empty(padded_size, pyfftw.simd_alignment, 
                    'float32')
        real_input = full_array[:, :, :input_size[2]]
        complex_output = full_array.view('complex64')
        fftw_obj = pyfftw.FFTW(real_input, complex_output, axes=(0, 1, 2)) # Plus whatever other args you want
        real_input[:] = data.copy() # with junk
        fftw_obj() # returns complex_output and overwrites the input
    else:
        data=np.fft.rfftn(data)
    kx, ky, kz = np.mgrid[:N, :N, :(N/2+1)]
    kx[kx > N/2-1] = kx[kx > N/2-1]-N
    ky[ky > N/2-1] = ky[ky > N/2-1]-N
    kz[kz > N/2-1] = kz[kz > N/2-1]-N
    k=2.0*np.pi*np.sqrt(kx**2+ky**2+kz**2)/N
    if mode == 1:
        kf = 2.0*np.pi/N
        res = np.zeros(len(k_list)-1)
        for i in range(len(k_list)-1):
            if np.sum((k >= k_list[i]) & (k < k_list[i+1]))>0:
                res[i] = np.mean(np.abs(data[(k >= k_list[i]) & (k < k_list[i+1])])**2)
        return res*box_size**3 / N**6
    if mode==0:
        kf=2.0*np.pi/N
        h1, dump = np.histogram(k.flat,weights=np.abs(data.flat)**2,bins=k_list)
        h2, dump = np.histogram(k.flat,bins=k_list)
        h2[h2==0] = 1.0
        res = h1/h2
        return res*box_size**3/N**6
    if mode==2:
        kf=2.0*np.pi/N
        res=np.zeros(len(k_list))
        for i in range(len(k_list)):
            res[i]=np.mean(np.abs(data[(k>=k_list[i]-kf) & (k<k_list[i]+kf)])**2)
        return res*box_size**3/N**6
    if mode==2:
        return k,data


def readifrit(path, nvar=[0], moden=2, skipmoden=2, dtype='i', dtypef='f4'):
    import struct
    import numpy as np
    openfile=open(path, "rb")
    dump = np.fromfile(openfile, dtype=dtype, count=moden)
    N1,N2,N3=np.fromfile(openfile, dtype=dtype, count=3)
    dump = np.fromfile(openfile, dtype=dtype, count=moden)
    print N1,N2,N3
    data=np.zeros([N1,N2,N3,np.size(nvar)])
    j=0
    for i in range(np.max(nvar)+1):
        dump = np.fromfile(openfile, dtype='i', count=moden)
        if i==nvar[j]:
            data[:,:,:,j]=np.reshape(np.fromfile(openfile, dtype=dtypef, count=N1*N2*N3),[N1,N2,N3])
            j=j+1
        else:
            dump=np.fromfile(openfile, dtype=dtypef, count=N1*N2*N3)
        dump = np.fromfile(openfile, dtype='i', count=moden)
    openfile.close()
    return N1,N2,N3,data

import glob
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

#path = glob.glob('/lustre/gnedin/REI/D/L/SET.B40.N256L2.sf=1_uv=0.1_bw=10/A/IFRIT/*')
#path.sort()
#n1,n2,n3,data = readifrit('/lustre/gnedin/REI/D/L/SET.B40.N256L2.sf=1_uv=0.1_bw=10/A/IFRIT/ifrit-box-a=0.1230.bin', nvar=[0,1,2])
n1,n2,n3,data = readifrit('ifrit-box-a=0.1230.bin', nvar=[0,1,2])
N=512
box_size=40

n1,n2,n3,data = readifrit('/lustre/gnedin/REI/D/H/SET2.B80.N512L2.sf=1_uv=0.15_bw=10_res=200/A/IFRIT/ifrit-a=0.1135.bin', nvar=[0,1], dtype='>i4', dtypef='>f4')
N=1024
box_size=80

kappa = (data[:,:,:,1])*(1.0-data[:,:,:,0])**2/data[:,:,:,0]

#plt.imshow(kappa[:,:,0])
#plt.show()

def cut_func(d):
    d0 = [-1,    0,      0.25,      0.5,    0.75,   1.0,    1.5,    5.0]
    k0 = [2.14,  2.14,   1.8,       1.4,   0.6,    0.3,    0.0,    0.0]
    return np.interp(d, d0, k0)

plt.hist2d(np.log10(data[:,:,:,1].flatten()), np.log10(kappa.flatten()), 100, norm=LogNorm())
d_list=np.linspace(-1,5,100)
plt.plot(d_list, cut_func(d_list))
plt.show()

k=2.0*np.pi*np.logspace(0, np.log10(N), 20)/N
k_phys_list=k*N/box_size
k_phys = np.sqrt(k_phys_list[1:]*k_phys_list[:-1])

rhoH = data[:,:,:,1]*data[:,:,:,0]
temp = data[:,:,:,0].copy()
filt = np.log10(kappa)>cut_func(np.log10(data[:,:,:,1]))
# temp[(kappa>1) & (data[:,:,:,1]>6)] = 0.001
temp[filt] = temp[filt]*0.001
rhoH_corrected = temp*data[:,:,:,1]
rhoH_other = rhoH-rhoH_corrected

pkdata = pk(rhoH, box_size, k_phys_list, mode=0)
pkdata_corrected = pk(rhoH_corrected, box_size, k_phys_list, mode=0)
pkdata_other = pk(rhoH_other, box_size, k_phys_list, mode=0)
pkdata_b = pk(data[:,:,:,1], box_size, k_phys_list, mode=0)


ax = plt.subplot(131)
plt.imshow(np.log10(data[:,:,0,0]), vmin=-3, vmax=0.0)
plt.colorbar()
plt.subplot(132, sharex = ax, sharey = ax)
plt.imshow(np.log10(temp[:,:,0]), vmin=-3, vmax=0.0)
plt.colorbar()
plt.subplot(133)
# plt.plot(k_phys, pkdata_corrected/pkdata)
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata)
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata_corrected, '--')
plt.xscale('log')
plt.yscale('log')
plt.ylabel(r'$\Delta$')
plt.ylabel(r'$P_{without}/P_{with}$')
plt.xlabel(r'$k$')
plt.show()



mpl_fig_obj= plt.figure()
ax = plt.subplot(121)
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata, label='P_k(tot)')
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata_corrected, '--', label='P_k(nf)')
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata_other, ':', label='P_k(f)')
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata_b,'-',lw=3, label='P_k(baryons)')
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*(pkdata_corrected*pkdata_other)**0.5,'k-',lw=0.5, label='(P_k(f)*P_k(nf))^0.5')
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*np.abs(pkdata-pkdata_corrected-pkdata_other),'k--',lw=0.5, label='|P_k(tot)-P_k(nf)-P_k(f)|')
plt.xscale('log')
plt.yscale('log')
plt.legend(loc=2)
plt.ylabel(r'$k^3/(2\pi^2)\;P_k(x)$')
plt.xlabel(r'$k$')
plt.subplot(122)
plt.plot(k_phys, pkdata/pkdata_b)
plt.plot(k_phys, pkdata_corrected/pkdata_b, '--')
plt.plot(k_phys, pkdata_other/pkdata_b, ':')
plt.xscale('log')
plt.yscale('log')
plt.ylabel(r'$P_{x}/P_{baryons}$')
plt.xlabel(r'$k$')
# plt.show()



np.savez('temp.npz', k_phys=k_phys, 
                     pkdata=pkdata, 
                     pkdata_corrected=pkdata_corrected,
                     pkdata_other=pkdata_other, 
                     pkdata_b=pkdata_b,
                     slice1=data[:,:,0,0],
                     slice2=temp[:,:,0])

plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata, label='P_k(tot)')
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata_corrected, '--', label='P_k(nf)')
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*pkdata_other, ':', label='P_k(f)')
temp=pkdata-pkdata_corrected-pkdata_other
temp[temp<0]=np.nan
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*np.abs(temp),'k-',lw=0.5, label='|P_k(tot)-P_k(nf)-P_k(f)|')
temp=pkdata-pkdata_corrected-pkdata_other
temp[temp>0]=np.nan
plt.plot(k_phys, k_phys**3/(2*np.pi**2)*np.abs(temp),'k:',lw=0.5)
plt.xscale('log')
plt.yscale('log')
plt.legend(loc=2)
plt.ylabel(r'$k^3/(2\pi^2)\;P_k(x)$')
plt.xlabel(r'$k$')