def box_pyramid_pspec(cube, angular_extent, r_mpc, Zs, kz = None, Nkbins=100, error=False, cosmo=False, return_dV=False):
"""
Estimate the power spectrum in a "pyramid" appoximation:
angular_extent = in radians, the width of the box
For each r_mpc, calculate the width of the projected region from the angular diameter distance.
In the radial direction, calculate discrete fourier transform. In tangential directions, do FFT.
"""
Zs.sort()
r_mpc.sort()
Nx, Ny, Nz = cube.shape
D_mpc = WMAP9.comoving_transverse_distance(Zs).to("Mpc").value
L = D_mpc * angular_extent # Transverse lengths in Mpc
if cosmo:
kfact=2*np.pi
# dV = (Nz) = (L/Nx) * (L/Ny) * np.diff(
pixel_size = angular_extent**2/(Nx * Ny)
dz_redshift = np.diff(Zs)[0] # Assume uniform redshift interval
dV = WMAP9.differential_comoving_volume(Zs).to("Mpc3/sr").value * dz_redshift * pixel_size
Lz = r_mpc[-1] - r_mpc[0]
pfact = 1/(L**2 * Lz)
else:
kfact=1
dV = 1
pfact = 1/float(Nx*Ny*Nz)
if kz is None:
dz = np.abs(r_mpc[-1] - r_mpc[0])/float(Nz) #Mean spacing
kz = np.fft.fftfreq(Nz,d=dz) * kfact
else:
assert kz.size == r_mpc.size
## DFT in radial direction
M = np.exp(np.pi * 2 * (-1j) * np.outer(kz,r_mpc) )
_c = np.apply_along_axis(lambda x: np.dot(M,x),2, cube)
## 2D FFT in the transverse directions
_c = np.fft.fft2(_c,axes=(0,1))
_c = _c * dV
## kx and ky
kx = np.array([np.fft.fftfreq(Nx,d=l/Nx) for l in L])*kfact
ky = np.array([np.fft.fftfreq(Ny,d=l/Ny) for l in L])*kfact
kx = np.moveaxis(np.tile(kx[:,:,np.newaxis],Ny),[1,2],[0,1])
ky = np.moveaxis(np.tile(ky[:,:,np.newaxis],Nx),[1,2],[1,0])
kz = np.tile(kz[np.newaxis,:],Nx*Ny).reshape(Nx,Ny,Nz)
pk3d = np.abs(_c)**2 * pfact
results = bin_1d(pk3d,(kx,ky,kz),Nkbins=Nkbins,error=error)
if return_dV:
return results, dV
return results
def shell_project_pspec(shell, nside, radius, dims=None,r_mpc=None, hpx_inds = None, N_sections=None,
freqs = None, kz=None, method='fft', Nkbins='auto',cosmo=False, error=False):
""" Estimate the power spectrum of a healpix shell by projecting regions to Cartesian space.
Shell = (Npix, Nfreq) Healpix shell, or section of one.
radius = (analogous to beam width)
N_sections = Compute the power spectrum for this many regions (like snapshots in observing)
dims (3) = Box dimensions in Mpc. If not available, check for freqs and use cosmology module
method = Choose estimator method (lomb-scargle, dft, pyramid)
"""
## For now, pick a random point within the (assumed full) spherical shell, and get the pixels within distance radius around it.
if hpx_inds is None:
Npix = hp.nside2npix(nside)
hpx_inds = np.arange(Npix)
else:
Npix = hpx_inds.size
if N_sections is None: N_sections=1
radius *= np.pi/180. #Deg2Rad
if freqs is None and method=='lomb_scargle': raise ValueError(" Frequency data must be provided to use the Lomb-Scargle method ")
if dims is None:
if freqs is None:
if dist is None: raise ValueError("Comoving distance or frequencies must be defined")
Lx = 2*radius*dist #Dist = comoving distance in Mpc
Ly = Lx; Lz = Lx
else:
Zs = 1420./freqs - 1.
rcomov = WMAP9.comoving_distance(Zs).to("Mpc").value
Lz = np.max(rcomov) - np.min(rcomov)
dist = WMAP9.comoving_transverse_distance(np.mean(Zs)).value #Identical to comoving if Omega_k = 0
Lx = 2*radius*dist
Ly = Lx
dims = [Lx,Ly,Lz]
if method=='lomb_scargle' or method=='pyramid':
try:
r_mpc
Zs
except:
Zs = 1420./freqs - 1.
r_mpc = WMAP9.comoving_distance(Zs).to("Mpc").value
print 'Dimensions: ',dims
pk = None
kbins = None
errs = None
## Loop over sections, choosing a different center each time.
## Average power spectrum estimates from different parts of the sky.
for s in range(N_sections):
cent = hp.pix2vec(nside,np.random.choice(hpx_inds))
inds = hp.query_disc(nside,cent,radius)
print "Section, pixels: ",s, inds.shape
cube = orthoslant_project(shell, cent, radius)
#if s==0:
# np.savez('projected_cube',cube=cube,dims=dims)
if method=='lomb_scargle':
if kz is not None:
results_i = box_ls_pspec(cube,dims,r_mpc,Nkbins=Nkbins,error=error,kz=kz,cosmo=cosmo)
else:
results_i = box_ls_pspec(cube,dims,r_mpc,Nkbins=Nkbins,error=error,cosmo=cosmo)
elif method=='pyramid':
results_i = box_pyramid_pspec(cube, radius, r_mpc, Zs, Nkbins=Nkbins,error=error,cosmo=cosmo)
else:
results_i = box_dft_pspec(cube,dims,Nkbins=Nkbins,r_mpc=r_mpc,error=error,cosmo=cosmo)
if pk is None:
pk = results_i[1]
kbins = results_i[0]
Nkbins = pk.size
if error:
errs = results_i[2]
else:
pk += results_i[1]
kbins += results_i[0]
if error:
errs += results_i[2]
print s, np.sum(np.isnan(results_i[1]))
pk /= N_sections
kbins /= N_sections
if error:
errs /= N_sections
return kbins, pk, errs
else:
return kbins, pk
