def get_power(self, gridded_vis, kernel_weights, ps_dim=2):
"""
Determine the 2D Power Spectrum of the observation.
Parameters
----------
gridded_vis : complex (ngrid, ngrid, neta)-array
The gridded visibilities, fourier-transformed along the frequency axis. Units JyHz.
coords: list of 3 1D arrays.
The [u,v,eta] co-ordinates corresponding to the gridded fourier visibilities. u and v in 1/rad, and
eta in 1/Hz.
Returns
-------
PS : float (n_obs, n_eta, bins)-list
The cylindrical averaged (or 2D) Power Spectrum, with units JyHz**2.
"""
logger.info("Calculating the power spectrum")
PS = []
for vis in gridded_vis:
# The 3D power spectrum
power_3d = np.absolute(vis)**2
if ps_dim == 2:
P = angular_average_nd(
field=power_3d,
coords=[self.uvgrid, self.uvgrid, self.eta],
bins=self.u_edges,
n=ps_dim,
weights=np.sum(kernel_weights, axis=2), # weights,
bin_ave=False,
)[0]
elif ps_dim == 1:
P = angular_average_nd(
field=power_3d,
coords=[self.uvgrid, self.uvgrid, self.eta],
bins=self.u_edges,
weights=kernel_weights,
bin_ave=False,
)[0]
P[np.isnan(P)] = 0
PS.append(P)
return PS
def grid_weights(self):
"""The number of uv cells that go into a single u annulus (unrelated to baseline weights)"""
return angular_average_nd(
field=np.ones((len(self.uvgrid),) * 2),
coords=[self.uvgrid, self.uvgrid],
bins=self.u_edges, n=self.ps_dim, bin_ave=False,
average=False)[0]
def ps_3d_to_ps_2d(ps_3d, u, nu, bins=100):
"""
Take a 3D power spectrum and return a cylindrically-averaged 2D power spectrum.
Parameters
----------
ps_3d : 3D array
The power spectrum in 3D, with first axis corresponding to frequency.
u : 1D array
The grid-coordinates along a side of the `ps_3d` array. Assumes that the (u,v) grid is square.
nu : 1D array
The frequencies corresponding to the first dimension of `ps_3d`.
bins : int
Number of (regular linear) bins to form the average into.
Returns
-------
ps_2d : 2D array
The circularly-averaged PS, with first axis corresponding to frequency.
ubins : 1D array
Length `bins` array giving the average central-bin co-ordinate for u after averaging.
"""
# Perform cylindrical averaging.
ps_2d, ubins, _ = angular_average_nd(field=ps_3d.T,
coords=[u, u, nu],
bins=bins,
n=2)
return ps_2d.T, ubins
def compute_mps(lightcone, bins=None, nthreads=None):
# First get "visibilities"
vis, kperp = fft(lightcone.brightness_temp,
L=lightcone.user_params.BOX_LEN,
axes=(0, 1))
# vis has shape (HII_DIM, HII_DIM, lightcone_dim)
# Do wavelet transform
wvlts, kpar, _ = morlet_transform_c(vis.T,
lightcone.lightcone_coords,
nthreads=nthreads)
# wvlts has shape (len(kpar) + vis.T.shape) corresponding to (eta, nu_c, u,v)
# Now square it...
wvlts = np.abs(wvlts)**2
# Determine a nice number of bins.
if bins is None:
bins = int((np.product(kperp.shape) * len(kpar))**(1. / 3.) / 2.2)
# And angularly average
wvlts, k = angular_average_nd(wvlts.transpose((2, 3, 0, 1)),
list(kperp) +
[kpar, lightcone.lightcone_coords],
n=3,
bins=bins,
bin_ave=False,
get_variance=False)
return wvlts, k, lightcone.lightcone_coords
def test_angular_avg_nd_2_1_varnull():
x = np.linspace(-3,3,200)
P = np.ones((200,10))
coords = [x, np.linspace(-2,2,10)]
p_k, k_av_bins, var = angular_average_nd(P,coords, bins=20, n=1, get_variance=True)
assert np.all(var==0)
def test_angular_avg_nd_3():
x = np.linspace(-3,3,400)
X,Y = np.meshgrid(x,x)
r2 = X**2 + Y**2
P = r2**-1.
P = np.repeat(P,100).reshape(400,400,100)
freq = [x,x,np.linspace(-2,2,100)]
p_k, k_av_bins = angular_average_nd(P,freq,bins=50,n=2)
print(p_k[6:,0], k_av_bins[6:]**-2.)
assert np.max(np.abs((p_k[6:,0] - k_av_bins[6:]**-2.)/k_av_bins[6:]**-2.)) < 0.05