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 }