def test_roundtrip_fb(g2d, a, b):
Fx, freq = fft(g2d['fx'], L=g2d['L'], a=a, b=b, left_edge=-g2d['L'] / 2)
Lk = -2 * np.min(freq)
fx, x = ifft(Fx, Lk=Lk, a=a, b=b)
assert np.max((fx.real - g2d['fx'])) < 1e-10 # Test FT result
assert np.max(x[0] - g2d['x']) < 1e-10 # Test x-grid
def run_roundtrip(a, b):
Fx, freq = fft(fx, L=L, a=a, b=b)
Lk = -2 * np.min(freq)
fx_, x_ = ifft(Fx, Lk=Lk, a=a, b=b)
assert np.max(np.real((fx_ - fx))) < 1e-10 # Test FT result
assert np.max(x_[0] - x) < 1e-10 # Test x-grid
def test_roundtrip_bf(g2d, a, b):
fx, freq = ifft(g2d['fx'], Lk=g2d['L'], a=a, b=b)
L = -2 * np.min(freq)
Fk, k = fft(fx, L=L, a=a, b=b)
assert np.max((Fk.real - g2d['fx'])) < 1e-10 # Test FT result
assert np.max(k[0] - g2d['x']) < 1e-10 # Test x-grid
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 visibility(self):
"""
The visibilities of :meth:`visible_sky` over a uniform (u,v)-grid, defined as :attr:`ugrid_raw0`.
"""
# Note that base fft in numpy is equivalent to \int f(x) e^{-2pi i kx} dx
vsky = self.visible_sky()
vsky[np.isnan(vsky)] = 0.0
return ((fft(vsky, axes=(-2, -1), a=0, b=2 * np.pi)[0]).T *
self.spatial_dist.cell_area).T
def test_forward_unitary_ordinary():
a, b = 0, 2 * np.pi
Fx, freq, grid = fft(fx, L=L, a=a, b=b, ret_cubegrid=True)
Fx_anl_fc = lambda k: (1. / a_squared) * np.exp(-np.pi * k**2 / a_squared)
# Analytic transform
Fx_anl = Fx_anl_fc(grid)
Fx_circ = angular_average(np.abs(Fx), grid, int(N / 2.2))[0]
Fx_anl_circ = angular_average(Fx_anl, grid, int(N / 2.2))[0]
assert np.max(np.abs(Fx_circ - Fx_anl_circ)) < 1e-10
def test_mixed_unitary_ordinary_unitary_angular_1d_bf():
N = 1000
Lk = 10.
dk = Lk / N
k = np.arange(-Lk / 2, Lk / 2, dk)[:N]
alpha = np.pi
Fk = np.exp(-alpha * k**2)
fx, x = ifft(Fk, Lk=Lk, a=0, b=1)
Fk_, k_ = fft(fx, -2 * np.min(x), a=0, b=2 * np.pi)
assert np.max(
np.abs(Fk_ - np.sqrt(2 * np.pi) * np.exp(-alpha * (
(2 * np.pi) * k_[0])**2))) < 1e-10
def vis_to_3d_ps(vis, dnu, taper=None):
"""
Apply a taper and perform a FT over the frequency dimension of a grid of visibilities to get the 3D power spectrum.
Parameters
----------
vis : 3D array
Array with first axis corresponding to nu and final two axes corresponding to u,v.
dnu : float
The (regular) interval between frequency bins.
taper : callable, optional
A taper/filter function to apply over the frequency axis.
Returns
-------
ps_3d : 3D array
3D Power Spectrum, with first axis corresponding to frequency.
eta : 1D array
The Fourier-dual of input frequencies.
"""
if taper is not None:
taper = taper(len(vis))
else:
taper = 1
# Do the DFT to eta-space
vistot, eta = dft.fft((taper * vis.value.T).T,
L=dnu.value,
a=0,
b=2 * np.pi,
axes=(0, ))
ps_3d = np.abs(vistot)**2 # Form the power spectrum
# ps_3d = np.zeros(vis.shape)
# for i in range(vis.shape[1]):
# for j in range(vis.shape[1]):
# vistot, eta = dft.fft(taper * vis[:,i, j].value,
# L=dnu.value, a=0, b=2 * np.pi)
#
# ps_3d[:,i, j] = np.abs(vistot)**2
return ps_3d, eta[
0] # eta is a list of arrays, so take first (and only) entry
def frequency_fft(vis, freq, dim, taper=np.ones_like, n_obs=1):
"""
Fourier-transform a gridded visibility along the frequency axis.
Parameters
----------
vis : complex (ncells, ncells, nfreq)-array
The gridded visibilities.
freq : (nfreq)-array
The linearly-spaced frequencies of the observation.
taper : callable, optional
A function which computes a taper function on an nfreq-array. Default is to have no taper. Callable should
take single argument, N.
n_obs : int, optional
Number of observations used to separate the visibilities into different bandwidths.
Returns
-------
ft : (ncells, ncells, nfreq/2)-array
The fourier-transformed signal, with negative eta removed.
eta : (nfreq/2)-array
The eta-coordinates, without negative values.
"""
ft = []
W = (freq.max() - freq.min()) / n_obs
L = int(len(freq) / n_obs)
for ii in range(n_obs):
ft.append(
fft(vis[:, :, ii * L:(ii + 1) * L] * taper(L),
W,
axes=(2, ),
a=0,
b=2 * np.pi)[0][:, :,
int(L / 2):]) # return the positive part)
ft = np.array(ft)
return ft
def test_mixed_unitary_ordinary_unitary_angular_2d_bf():
N = 1000
Lk = 10.
dk = Lk / N
k = np.arange(-Lk / 2, Lk / 2, dk)[:N]
KX, KY = np.meshgrid(k, k)
alpha = np.pi
Fk = np.exp(-alpha * (KX**2 + KY**2))
fx, x = ifft(Fk, Lk=Lk, a=0, b=1)
Fk_, k_, kgrid = fft(fx,
-2 * np.min(x),
a=0,
b=2 * np.pi,
ret_cubegrid=True)
Fk_, bins = angular_average(Fk_, kgrid, 200)
assert np.max(
np.abs(Fk_ - 2 * np.pi * np.exp(-alpha * ((2 * np.pi) * bins)**2)))
def test_mixed_unitary_ordinary_unitary_angular_1d_fb():
N = 1000
L = 10.
dx = L / N
x = np.arange(-L / 2, L / 2, dx)[:N]
alpha = np.pi
fx = np.exp(-alpha * x**2)
Fk, freq = fft(
fx,
L=L,
a=0,
b=2 * np.pi,
)
Lk = -2 * np.min(freq)
fx_, x_ = ifft(Fk, Lk=Lk, a=0, b=1)
assert np.max(
np.abs(fx_ - np.exp(-np.pi *
(x_[0] /
(2 * np.pi))**2) / np.sqrt(2 * np.pi))) < 1e-10
def test_mixed_unitary_ordinary_unitary_angular_2d_fb():
N = 1000
L = 10.
dx = L / N
x = np.arange(-L / 2, L / 2, dx)[:N]
X, Y = np.meshgrid(x, x)
a_squared = 1
fx = np.exp(-np.pi * a_squared * (X**2 + Y**2))
Fk, freq = fft(fx, L=L, a=0, b=2 * np.pi)
Lk = -2 * np.min(freq)
fx_, x_, xgrid = ifft(Fk, Lk=Lk, a=0, b=1, ret_cubegrid=True)
fx_, bins = angular_average(fx_, xgrid, 200)
print(
np.max(
np.abs(fx_ - np.exp(-np.pi * (bins / (2 * np.pi))**2) /
(2 * np.pi))))
assert np.max(
np.abs(fx_ - np.exp(-np.pi * (bins / (2 * np.pi))**2) /
(2 * np.pi))) < 1e-4
def image_to_uv(sky, L):
"""
Transform a box from image plan to UV plane.
Parameters
----------
sky : (ncells, ncells, nfreq)-array
The frequency-dependent sky brightness (in arbitrary units)
L : float
The size of the box in radians.
Returns
-------
uvsky : (ncells, ncells, nfreq)-array
The UV-plane representation of the sky. Units are units of the sky times radians.
uv_scale : list of two arrays.
The u and v co-ordinates of the uvsky, respectively. Units are inverse of L.
"""
logger.info("Converting to UV space...")
ft, uv_scale = fft(sky, L, axes=(0, 1), a=0, b=2 * np.pi)
return ft, uv_scale
def power_from_vis(vis, f, taper=None):
"""
Get the power spectrum at a bunch of vectors u, from the (gridded) visibilities (as functions of frequency)
there.
Parameters
----------
vis : (Nf,N)-array
Array of N visibilities at Nf frequencies.
f : (Nf)-array
Frequencies of observation, normalised to reference frequency.
taper : (Nf)-array or None
A frequency taper to apply
Returns
-------
power : (Nf/2-1)-array
The power spectrum at each visibility
omega : (Nf/2-1)-array
The fourier-dual of frequency corresponding to `f`. Note that these correspond to `f` *not* `nu`.
To yield a power spectrum which corresponds to nu, multiply the power by nu0**2 and eta by 1/nu0
"""
if taper is None:
taper = 1
res, omega = fft((taper * vis.T).T,
L=(f[-1] - f[0]),
axes=(0, ),
b=2 * np.pi,
a=0)
omega = omega[0]
res = np.abs(res)**2
res = res[len(f) // 2 + 1:]
omega = omega[len(f) // 2 + 1:]
return res, omega
def test_mixed_2d_bf(g2d, a, b, ainv, binv):
Fk, freq = ifft(g2d['fx'], Lk=g2d['L'], a=ainv, b=binv)
L = -2 * np.min(freq)
fx, x, xgrid = fft(Fk, L=L, a=a, b=b, left_edge=-L / 2, ret_cubegrid=True)
assert np.max(
np.abs(fx.real - analytic_mix(xgrid, a, binv, ainv, b))) < 1e-10
def imaging(chain=None, cores=None, lk=None, freq_ind=0):
"""
Create a plot of the imaging capability of the current setup.
Uses the loaded cores to create a simulated sky, then "observe" this with given baselines. Then uses an
Instrumental2D likelihood to grid those baselines and transform back to the image plane. Every step of
this process is output as a panel in a plot to be compared.
Parameters
----------
chain : :class:`~py21cmmc.mcmc.cosmoHammer.LikelihoodComputationChain.LikelihoodComputationChain` instance, optional
A computation chain which contains loaded likelihoods and cores.
cores : list of :class:`~py21cmmc.mcmc.core.CoreBase` instances, optional
A list of cores defining the sky and instrument. Only required if `chain` is not given.
lk : :class:`~likelihood.LikelihoodInstrumental2D` class, optional
An instrumental likelihood, required if `chain` not given.
freq_ind : int, optional
The index of the frequency to actually show plots for (default is the first frequency channel).
Returns
-------
fig :
A matplotlib figure object.
"""
if chain is None and (cores is None or lk is None):
raise ValueError("Either chain or both cores and likelihood must be given.")
# Create a likelihood computation chain.
if chain is None:
chain = build_computation_chain(cores, lk)
chain.setup()
else:
lk = chain.getLikelihoodModules()[0]
cores = chain.getCoreModules()
if not isinstance(lk, LikelihoodInstrumental2D):
raise ValueError("likelihood needs to be a Instrumental2D likelihood")
if not hasattr(lk, "LikelihoodComputationChain"):
chain.setup()
# Call all core simulators.
ctx = chain.build_model_data()
visgrid = lk.grid_visibilities(ctx.get("visibilities"))
# Do a direct FT there and back, rather than baselines.
print(ctx.get("new_sky"))
direct_vis, direct_u = fft(ctx.get("new_sky")[:, :, freq_ind], L=lk._instr_core.sky_size, a=0, b=2 * np.pi)
direct_img, direct_l = ifft(direct_vis, Lk=(lk.uvgrid[1] - lk.uvgrid[0]) * len(lk.uvgrid), a=0, b=2 * np.pi)
# Get the reconstructed image
image_plane, image_grid = ifft(visgrid[:, :, freq_ind], Lk=(lk.uvgrid[1] - lk.uvgrid[0]) * len(lk.uvgrid), a=0, b=2 * np.pi)
# Make a figure.
if len(cores) == 2:
fig, ax = plt.subplots(2, 4, figsize=(12, 6))
mid_row = 0
else:
fig, ax = plt.subplots(3, max((4, len(cores))), figsize=(3*max((4, len(cores))), 9))
mid_row = 1
# Show original sky(s) (before Beam)
i = 0
for core in cores:
if isinstance(core, ForegroundsBase):
# TODO: frequency plotted here does not necessarily match the frequency in other plots.
mp = ax[0, i].imshow(ctx.get("foregrounds")[i][:, :, freq_ind].T, origin='lower',
extent=(-core.sky_size / 2, core.sky_size / 2) * 2)
ax[0, i].set_title("Orig. %s FG" % core.__class__.__name__)
cbar = plt.colorbar(mp, ax=ax[0, i])
ax[0, i].set_xlabel("l")
ax[0, i].set_ylabel("m")
cbar.set_label("Brightness Temp. [Jy/sr]")
i += 1
# # TODO: add lightcone plot
# if isinstance(core, CoreLightConeModule):
# mp = ax[0, i].imshow(ctx.get("foregrounds")[i][:, :, -freq_ind].T, origin='lower',
# extent=(-core.sky_size / 2, core.sky_size / 2) * 2)
# ax[0, i].set_title("Original %s foregrounds" % core.__class__.__name__)
# cbar = plt.colorbar(mp, ax=ax[0, i])
# ax[0, i].set_xlabel("l")
# ax[0, i].set_ylabel("m")
# cbar.set_label("Brightness Temp. [K]")
# i += 1
# Show tiled (if applicable) and attenuated sky
mp = ax[mid_row, 1].imshow(
ctx.get("new_sky")[:, :, freq_ind].T, origin='lower',
extent=(-lk._instr_core.sky_size / 2, lk._instr_core.sky_size / 2) * 2
)
ax[mid_row, 1].set_title("Tiled+Beam FG")
cbar = plt.colorbar(mp, ax=ax[mid_row, 1])
ax[mid_row, 1].set_xlabel("l")
ax[mid_row, 1].set_ylabel("m")
cbar.set_label("Brightness Temp. [K]")
# Show UV weights
mp = ax[mid_row, 2].imshow(
lk.nbl_uvnu[:, :, freq_ind].T, origin='lower',
extent=(lk.uvgrid.min(), lk.uvgrid.max()) * 2
)
ax[mid_row, 2].set_title("UV weights")
cbar = plt.colorbar(mp, ax=ax[mid_row, 2])
ax[mid_row, 2].set_xlabel("u")
ax[mid_row, 2].set_ylabel("v")
cbar.set_label("Weight")
# Show raw visibilities
wvlength = 3e8 / ctx.get("frequencies")[freq_ind]
mp = ax[mid_row, 3].scatter(ctx.get("baselines")[:, 0] / wvlength, ctx.get("baselines")[:, 1] / wvlength,
c=np.real(ctx.get("visibilities")[:, freq_ind]))
ax[mid_row, 3].set_title("Raw Vis.")
cbar = plt.colorbar(mp, ax=ax[mid_row, 3])
ax[mid_row, 3].set_xlabel("u")
ax[mid_row, 3].set_xlabel("v")
cbar.set_label("Re[Vis] [Jy?]")
# Show Gridded Visibilities
mp = ax[mid_row+1, 3].imshow(
np.real(visgrid[:, :, freq_ind].T), origin='lower',
extent=(lk.uvgrid.min(), lk.uvgrid.max()) * 2
)
ax[mid_row+1, 3].set_title("Gridded Vis")
cbar = plt.colorbar(mp, ax=ax[mid_row+1, 3])
ax[mid_row+1, 3].set_xlabel("u")
ax[mid_row+1, 3].set_ylabel("v")
cbar.set_label("Jy")
# Show directly-calculated UV plane
mp = ax[mid_row+1, 2].imshow(
np.real(direct_vis), origin='lower',
extent=(direct_u[0].min(), direct_u[0].max()) * 2
)
ax[mid_row+1, 2].set_title("Direct Vis")
cbar = plt.colorbar(mp, ax=ax[mid_row+1, 2])
ax[mid_row+1, 2].set_xlabel("u")
ax[mid_row+1, 2].set_ylabel("v")
cbar.set_label("Jy")
# Show final "image"
mp = ax[mid_row+1, 1].imshow(np.abs(image_plane).T, origin='lower',
extent=(image_grid[0].min(), image_grid[0].max(),) * 2)
ax[mid_row+1, 1].set_title("Recon. FG")
cbar = plt.colorbar(mp, ax=ax[mid_row+1, 1])
ax[mid_row+1, 1].set_xlabel("l")
ax[mid_row+1, 1].set_ylabel("m")
cbar.set_label("Flux Density. [Jy]")
# Show direct reconstruction
mp = ax[mid_row+1, 0].imshow(np.abs(direct_img).T, origin='lower',
extent=(direct_l[0].min(), direct_l[0].max(),) * 2)
ax[mid_row+1, 0].set_title("Recon. direct FG")
cbar = plt.colorbar(mp, ax=ax[mid_row+1, 0])
ax[mid_row+1, 0].set_xlabel("l")
ax[mid_row+1, 0].set_ylabel("m")
cbar.set_label("Flux Density. [Jy]")
plt.tight_layout()
return fig
def test_forward_only(g1d, a, b):
Fx, freq = fft(g1d['fx'], L=g1d['L'], a=a, b=b, left_edge=-g1d['L'] / 2)
assert np.max(np.abs(Fx.real - gauss_ft(freq[0], a, b, n=1))) < 1e-10
def test_mixed_1d_fb(g1d, a, b, ainv, binv):
Fk, freq = fft(g1d['fx'], L=g1d['L'], a=a, b=b, left_edge=-g1d['L'] / 2)
Lk = -2 * np.min(freq)
fx, x = ifft(Fk, Lk=Lk, a=ainv, b=binv)
assert np.max(
np.abs(fx.real - analytic_mix(x[0], a, b, ainv, binv, n=1))) < 1e-10
def test_mixed_1d_bf(g1d, a, b, ainv, binv):
Fk, freq = ifft(g1d['fx'], Lk=g1d['L'], a=ainv, b=binv)
L = -2 * np.min(freq)
fx, x = fft(Fk, L=L, a=a, b=b, left_edge=-L / 2)
assert np.max(
np.abs(fx.real - analytic_mix(x[0], a, binv, ainv, b, n=1))) < 1e-10
def test_mixed_2d_fb(g2d, a, b, ainv, binv):
Fk, freq = fft(g2d['fx'], L=g2d['L'], a=a, b=b, left_edge=-g2d['L'] / 2)
Lk = -2 * np.min(freq)
fx, x, xgrid = ifft(Fk, Lk=Lk, a=ainv, b=binv, ret_cubegrid=True)
assert np.max(
np.abs(fx.real - analytic_mix(xgrid, a, b, ainv, binv))) < 1e-10
def visibility(self):
# Note that base fft in numpy is equivalent to \int f(x) e^{-2pi i kx} dx
return (
(fft(self.visible_sky(), axes=(-2, -1), a=0, b=2 * np.pi)[0]).T *
self.cell_area).T
def sim_to_vis(box,
antenna_pos,
numin=150.,
numax=180.,
cosmo=Planck15,
beam=CircularGaussian):
"""
Parameters
----------
box : :class:`tocm_tools.Box` instance or str
Either a 21cmFAST simulation box, or a string specifying a filename to one.
antenna_pos : array
2D array of (x,y) positions of antennae in meters (shape (Nantennate,2)).
numin : float
Minimum frequency (in MHz) to include in the "observation"
numax : float
Maximum frequency (in MHz) to include in the "observation"
cosmo : :class:`astropy.cosmology.FLRW` instance
The cosmology to use for all calculations
beam : :class:`spore.model.beam.Beam` instance
The telescope beam model to use.
Returns
-------
uvsample : array
2D complex array in which the first dimension has length of Nbaselines, and the second has Nnu. This is the
Fourier Transform of the sky at the baselines, has units of Jy.
baselines : array
The baseline vectors of the observation (shape (Nbaseline,2)). Units m.
nu : array
1D array of frequencies of observation. Units MHz.
"""
# READ THE BOX
box, Nbox, L, d, nu, z = get_cut_box(box, numin, numax)
lam = 3e8 / (nu * 1e6) # in m
# Convert to specific intensity (Jy/sr)
#Jy/(sr?) mK->K to Jy (J/K) /m^2
box *= 1e-3 * 1e26 * 2 * k_B / lam**2
# INITIALISE A BEAM MODEL
beam = beam(nu.min(), np.linspace(1, nu.max() / nu.min(), len(nu)))
# Box width at different redshifts
width = L / cosmo.angular_diameter_distance(z).value
# Minimum umax available in simulation
maxl = 2 * np.sin(width.max() / 2 - width.max() / Nbox / 2)
maxu = Nbox / 2 / maxl
# GENERATE BASELINE VECTORS
baselines = pos_to_baselines(antenna_pos)
u0, baselines = baselines_to_u0(baselines, nu[0], maxu, ret_baselines=True)
uvsample = np.zeros((len(baselines), len(nu)), dtype="complex128")
# ATTENUATE BOX BY THE BEAM
for i in range(len(nu)):
dl = width[i] / Nbox
l = np.sin(
np.linspace(-width[i] / 2 + dl / 2, width[i] / 2 - dl / 2, Nbox))
L, M = np.meshgrid(l, l)
slice = box[:, :, i] * np.exp(-(L**2 + M**2) / (2 * beam.sigma[i]**2))
# Interpolate onto regular l,m grid
spl_lm = RectBivariateSpline(l, l, slice)
l = np.linspace(l.min(), l.max(), len(l))
slice = spl_lm(l, l, grid=True)
FT, freq = dft.fft(slice, L=l.max() - l.min(), a=0, b=2 * np.pi)
uvsample[:, i] = interpolate_visibility_onto_baselines(
FT, freq[0], nu[i] / nu[0], u0)
return uvsample, u0, nu, {
'slice': slice,
'box': box,
"FT": FT,
'freq': freq,
"l": l
}
