diff --git a/.coveragerc b/.coveragerc index 24bf59e2..e5fbc3d2 100644 --- a/.coveragerc +++ b/.coveragerc @@ -1,7 +1,6 @@ [run] omit = venv/* - src/mpol/spheroidal_gridding.py relative_files = True diff --git a/docs/changelog.md b/docs/changelog.md index ce21054d..65fd36af 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -4,6 +4,7 @@ ## v0.2.1 +- Removed custom `spheroidal_gridding` routines, tests, and the `UVDataset` object that used them. These have been superseded by the TorchKbNuFFT package. For reference, the old routines (including the tricky `corrfun` math) is preserved in a Gist [here](https://gist.github.com/iancze/f3d2769005a9e2c6731ee6977f166a83). - Changed API of {class}`~mpol.fourier.NuFFT`. Previous signature took `uu` and `vv` points at initialization (`__init__`), and the `.forward` method took only an image cube. This behaviour is preserved in a new class {class}`~mpol.fourier.NuFFTCached`. The updated signature of {class}`~mpol.fourier.NuFFT` *does not* take `uu` and `vv` at initialization. Rather, its `forward` method is modified to take an image cube and the `uu` and `vv` points. This allows an instance of this class to be used with new `uu` and `vv` points in each forward call. This follows the standard expectation of a layer (e.g., a linear regression function predicting at new `x`) and the pattern of the TorchKbNuFFT package itself. It is expected that the new `NuFFT` will be the default routine and `NuFFTCached` will only be used in specialized circumstances (and possibly deprecated/removed in future updates). Changes implemented by [#232](https://github.com/MPoL-dev/MPoL/pull/232). - Moved "Releasing a new version of MPoL" from the wiki to the Developer Documentation ({ref}`releasing-new-version-label`). diff --git a/src/mpol/datasets.py b/src/mpol/datasets.py index 3a3a07bb..6eaf2577 100644 --- a/src/mpol/datasets.py +++ b/src/mpol/datasets.py @@ -12,7 +12,7 @@ from mpol.coordinates import GridCoords from mpol.exceptions import WrongDimensionError -from . import spheroidal_gridding, utils +from . import utils from .constants import * from .utils import loglinspace @@ -160,106 +160,6 @@ def ground_mask(self) -> torch.Tensor: """ return utils.packed_cube_to_ground_cube(self.mask) - -# custom dataset loader -class UVDataset(torch_ud.Dataset): - r""" - Container for loose interferometric visibilities. - - Args: - uu (2d numpy array): (nchan, nvis) length array of u spatial frequency coordinates. Units of [:math:`\mathrm{k}\lambda`] - vv (2d numpy array): (nchan, nvis) length array of v spatial frequency coordinates. Units of [:math:`\mathrm{k}\lambda`] - data_re (2d numpy array): (nchan, nvis) length array of the real part of the visibility measurements. Units of [:math:`\mathrm{Jy}`] - data_im (2d numpy array): (nchan, nvis) length array of the imaginary part of the visibility measurements. Units of [:math:`\mathrm{Jy}`] - weights (2d numpy array): (nchan, nvis) length array of thermal weights. Units of [:math:`1/\mathrm{Jy}^2`] - cell_size (float): the image pixel size in arcsec. Defaults to None, but if both `cell_size` and `npix` are set, the visibilities will be pre-gridded to the RFFT output dimensions. - npix (int): the number of pixels per image side (square images only). Defaults to None, but if both `cell_size` and `npix` are set, the visibilities will be pre-gridded to the RFFT output dimensions. - device (torch.device) : the desired device of the dataset. If ``None``, defalts to current device. - - If both `cell_size` and `npix` are set, the dataset will be automatically pre-gridded to the RFFT output grid. This will greatly speed up performance. - - If you have just a single channel, you can pass 1D numpy arrays for `uu`, `vv`, `weights`, `data_re`, and `data_im` and they will automatically be promoted to 2D with a leading dimension of 1 (i.e., ``nchan=1``). - """ - - def __init__( - self, - uu: NDArray[floating[Any]], - vv: NDArray[floating[Any]], - weights: NDArray[floating[Any]], - data_re: NDArray[floating[Any]], - data_im: NDArray[floating[Any]], - cell_size: float | None = None, - npix: int | None = None, - device: torch.device | str | None = None, - **kwargs, - ): - # ensure that all vectors are the same shape - if not all( - array.shape == uu.shape for array in [vv, weights, data_re, data_im] - ): - raise WrongDimensionError("All dataset inputs must be the same shape.") - - if uu.ndim == 1: - uu = np.atleast_2d(uu) - vv = np.atleast_2d(vv) - data_re = np.atleast_2d(data_re) - data_im = np.atleast_2d(data_im) - weights = np.atleast_2d(weights) - - if np.any(weights <= 0.0): - raise ValueError("Not all thermal weights are positive, check inputs.") - - self.nchan = uu.shape[0] - self.gridded = False - - if cell_size is not None and npix is not None: - ( - uu, - vv, - grid_mask, - weights, - data_re, - data_im, - ) = spheroidal_gridding.grid_dataset( - uu, - vv, - weights, - data_re, - data_im, - cell_size, - npix, - ) - - # grid_mask (nchan, npix, npix//2 + 1) bool: a boolean array the same size as the output of the RFFT - # designed to directly index into the output to evaluate against pre-gridded visibilities. - self.grid_mask = torch.tensor(grid_mask, dtype=torch.bool, device=device) - self.cell_size = cell_size * arcsec # [radians] - self.npix = npix - self.gridded = True - - self.uu = torch.tensor(uu, dtype=torch.double, device=device) # klambda - self.vv = torch.tensor(vv, dtype=torch.double, device=device) # klambda - self.weights = torch.tensor( - weights, dtype=torch.double, device=device - ) # 1/Jy^2 - self.re = torch.tensor(data_re, dtype=torch.double, device=device) # Jy - self.im = torch.tensor(data_im, dtype=torch.double, device=device) # Jy - - # TODO: store kwargs to do something for antenna self-cal - - def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]: - return ( - self.uu[index], - self.vv[index], - self.weights[index], - self.re[index], - self.im[index], - ) - - def __len__(self) -> int: - return len(self.uu) - - class Dartboard: r""" A polar coordinate grid relative to a :class:`~mpol.coordinates.GridCoords` object, reminiscent of a dartboard layout. The main utility of this object is to support splitting a dataset along radial and azimuthal bins for k-fold cross validation. diff --git a/src/mpol/spheroidal_gridding.py b/src/mpol/spheroidal_gridding.py deleted file mode 100644 index a8553fb9..00000000 --- a/src/mpol/spheroidal_gridding.py +++ /dev/null @@ -1,521 +0,0 @@ -r""" -The gridding module contains routines to manipulate visibility data between uniform and non-uniform samples in the :math:`(u,v)` plane. -You probably wont need to use these routines in the normal workflow of MPoL, but they are documented here for reference. -""" - -import numpy as np -from scipy.sparse import lil_matrix - -from . import utils -from .constants import arcsec - - -def horner(x, a): - r""" - Use `Horner's method `_ to compute and return the polynomial - - .. math:: - - f(x) = a_0 + a_1 x^1 + a_2 x^2 + \ldots + a_{n-1} x^{(n-1)} - - Args: - x (float): input - a (list): list of polynomial coefficients - - Returns: - float: the polynomial evaluated at `x` - """ - result = 0 - for i in range(len(a) - 1, -1, -1): - result = a[i] + (x * result) - return result - - -@np.vectorize -def spheroid(eta): - r""" - Prolate spheroidal wavefunction function assuming :math:`\alpha = 1` and :math:`m=6` for speed." - - Args: - eta (float) : the value between ``[0, 1]`` - - Returns: - float : the value of the spheroid at :math:`\eta` - """ - - # Since the function is symmetric, overwrite eta - eta = np.abs(eta) - - if eta <= 0.75: - nn = eta ** 2 - 0.75 ** 2 - - return ( - horner( - nn, - np.array( - [8.203343e-2, -3.644705e-1, 6.278660e-1, -5.335581e-1, 2.312756e-1] - ), - ) - / horner(nn, np.array([1.0, 8.212018e-1, 2.078043e-1])) - ) - - elif eta <= 1.0: - nn = eta ** 2 - 1.0 - - return ( - horner( - nn, - np.array( - [4.028559e-3, -3.697768e-2, 1.021332e-1, -1.201436e-1, 6.412774e-2] - ), - ) - / horner(nn, np.array([1.0, 9.599102e-1, 2.918724e-1])) - ) - - elif eta <= 1.0 + 1e-7: - # case to allow some floating point error - return 0.0 - - else: - # Now you're really outside of the bounds - print( - "The spheroid is only defined on the domain -1.0 <= eta <= 1.0. (modulo machine precision.)" - ) - raise ValueError - - -def corrfun(eta): - r""" - The gridding *correction* function is applied to the image plane to counter-act the effects of the convolutional interpolation in the Fourier plane. For more information see `Schwab 1984 `_. - - Args: - eta (float): the value in [0, 1]. Accepts floats or vectors of float. - - Returns: - float: the correction function evaluated at :math:`\eta` - """ - return spheroid(eta) - - -def corrfun_mat(alphas, deltas): - """ - Calculate the image correction function over deminsionalities corresponding to the image. Apply to the image using a a multiply operation. The input coordinates must be fftshifted the same way as the image. - - Args: - alphas (1D array): RA list - deltas (1D array): DEC list - - Returns: - 2D array: correction function matrix evaluated over alphas and deltas - - """ - - ny = len(deltas) - nx = len(alphas) - - mat = np.empty((ny, nx), dtype="float") - - # The size of one half-of the image. - # sometimes ra and dec will be symmetric about 0, othertimes they won't - # so this is a more robust way to determine image half-size - maxra = np.abs(alphas[2] - alphas[1]) * nx / 2 - maxdec = np.abs(deltas[2] - deltas[1]) * ny / 2 - - for i in range(nx): - for j in range(ny): - etax = (alphas[i]) / maxra - etay = (deltas[j]) / maxdec - - if (np.abs(etax) > 1.0) or (np.abs(etay) > 1.0): - # We would be querying outside the shifted image - # bounds, so set this emission to 0.0 - mat[j, i] = 0.0 - else: - mat[j, i] = 1 / (corrfun(etax) * corrfun(etay)) - - return mat - - -def gcffun(eta): - r""" - The gridding *convolution* function for the convolution and interpolation of the visibilities in - the Fourier domain. This is also the Fourier transform of ``corrfun``. For more information see `Schwab 1984 `_. - - Args: - eta (float): in the domain of [0,1] - - Returns: - float : the gridding convolution function evaluated at :math:`\eta` - """ - - return np.abs(1.0 - eta ** 2) * spheroid(eta) - - -def calc_matrices(u_data, v_data, u_model, v_model): - r""" - Calculate real :math:`C_\Re` and imaginary :math:`C_\Im` sparse interpolation matrices for RFFT output, using spheroidal wave functions. - - Let :math:`W_\Re` and :math:`W_\Im` represent the real and imaginary output from the RFFT operatation carried out on an image that was pre-multiplied by the gridding correction function. To interpolate from the RFFT grid to the locations of `u_data` and `v_data`, one uses these matrices with a sparse matrix multiply to carry out - - .. math:: - - V_\Re = C_\Re W_\Re - - and - - .. math:: - - V_\Im = C_\Im W_\Im - - such that :math:`V_\Re` and :math:`V_\Im` are the real and imaginary visibilities interpolated to `u_data` and `v_data`. - - Args: - - u_data (1D numpy array): the :math:`u` coordinates of the dataset (in [:math:`k\lambda`]) - v_data (1D numpy array): the :math:`v` coordinates of the dataset (in [:math:`k\lambda`]) - u_model: the :math:`u` axis delivered by the rfft (unflattened, in [:math:`k\lambda`]). Assuming this is trailing dimension, which is the one over which the RFFT was carried out. - v_model: the :math:`v` axis delivered by the rfft (unflattened, in [:math:`k\lambda`]). Assuming this is leading dimension, which is the one over which the FFT was carried out. - - Returns: - 2-tuple : `coo` format sparse matrices :math:`C_\Re` and :math:`C_\Im` - - Begin with an image packed like ``Img[j, i]``, where ``i`` is the :math:`l` index and ``j`` is the :math:`m` index. - Then the RFFT output will look like ``RFFT[j, i]``, where ``i`` is the u index and ``j`` is the v index. - - This assumes the `u_model` array is like the output from `np.fft.rfftfreq `_ :: - - f = [0, 1, ..., n/2-1, n/2] / (d*n) if n is even - - and that the `v_model` array is like the output from `np.fft.fftfreq `_ :: - - f = [0, 1, ..., n/2-1, -n/2, ..., -1] / (d*n) if n is even - """ - - # calculate the stride needed to advance one v point in the flattened array = the length of the u row - vstride = len(u_model) - Npix = len(v_model) - - data_points = np.array([u_data, v_data]).T - # assert that the maximum baseline is contained within the model grid, which is the same as saying - # that the image-plane pixels are small enough. - assert ( - np.max(np.abs(u_data)) < u_model[-4] - ) # choose -4, because we'll need 3 pixels to do the interpolation - assert ( - np.max(np.abs(v_data)) < v_model[Npix // 2 - 4] - ) # choose -4, because we'll need 3 pixels to do the interpolation - - # number of visibility points in the dataset - N_vis = len(data_points) - - # initialize two sparse lil matrices for the instantiation - # convert to csc at the end - C_real = lil_matrix((N_vis, (Npix * vstride)), dtype="float") - C_imag = lil_matrix((N_vis, (Npix * vstride)), dtype="float") - - # determine model grid spacing - du = np.abs(u_model[1] - u_model[0]) - dv = np.abs(v_model[1] - v_model[0]) - - # for each data_point within the grid, calculate the row and insert it into the matrix - for row_index, (u, v) in enumerate(data_points): - - # assuming the grid stretched for -/+ values easily - i0 = int(np.ceil(u / du)) - j0 = int(np.ceil(v / dv)) - - i_indices = np.arange(i0 - 3, i0 + 3) - j_indices = np.arange(j0 - 3, j0 + 3) - - # calculate the etas and weights - u_etas = (u / du - i_indices) / 3 - v_etas = (v / dv - j_indices) / 3 - - uw = gcffun(u_etas) - vw = gcffun(v_etas) - - w = np.sum(uw) * np.sum(vw) - - l_indices = np.zeros(36, dtype="int") - weights_real = np.zeros(36, dtype="float") - weights_imag = np.zeros(36, dtype="float") - - # loop through all 36 points and calculate the matrix element - # do it in this order because u indices change the quickest - for j in range(6): - for i in range(6): - k = j * 6 + i - - i_index = i_indices[i] - - if i_index >= 0: - j_index = j_indices[j] - imag_prefactor = 1.0 - - else: # i_index < 0: - # map negative i index to positive i_index - i_index = -i_index - - # map j index to opposite sign - j_index = -j_indices[j] - - # take the complex conjugate for the imaginary weight - imag_prefactor = -1.0 - - # shrink j to fit in the range of [0, Npix] - if j_index < 0: - j_index += Npix - - l_indices[k] = i_index + j_index * vstride - weights_real[k] = uw[i] * vw[j] / w - weights_imag[k] = imag_prefactor * uw[i] * vw[j] / w - - # TODO: at the end, decide whether there is overlap and consolidate - l_sorted, unique_indices, unique_inverse, unique_counts = np.unique( - l_indices, return_index=True, return_inverse=True, return_counts=True - ) - if len(unique_indices) < 36: - # Some indices are querying the same point in the RFFT grid, and their weights need to be - # consolidated - - N_unique = len(l_sorted) - condensed_weights_real = np.zeros(N_unique) - condensed_weights_imag = np.zeros(N_unique) - - # find where unique_counts > 1 - ind_multiple = unique_counts > 1 - - # take the single weights from the array - ind_single = ~ind_multiple - - # these are the indices back to the original weights of sorted singles - # unique_indices[ind_single] - - # get the weights of the values that only occur once - condensed_weights_real[ind_single] = weights_real[ - unique_indices[ind_single] - ] - condensed_weights_imag[ind_single] = weights_imag[ - unique_indices[ind_single] - ] - - # the indices that occur multiple times - # l_sorted[ind_multiple] - - # the indices of indices that occur multiple times - ind_arg = np.where(ind_multiple)[0] - - for repeated_index in ind_arg: - # figure out the indices of the repeated indices in the original flattened index array - repeats = np.where(l_sorted[repeated_index] == l_indices) - - # stuff the sum of these weights into the condensed array - condensed_weights_real[repeated_index] = np.sum(weights_real[repeats]) - condensed_weights_imag[repeated_index] = np.sum(weights_imag[repeats]) - - # remap these variables to the shortened arrays - l_indices = l_sorted - weights_real = condensed_weights_real - weights_imag = condensed_weights_imag - - C_real[row_index, l_indices] = weights_real - C_imag[row_index, l_indices] = weights_imag - - return C_real.tocoo(), C_imag.tocoo() - - -def grid_datachannel(uu, vv, weights, re, im, cell_size, npix, debug=False, **kwargs): - r""" - Rather than interpolating the complex model visibilities from these grid points to the non-uniform :math:`(u,v)` points, pre-average the data visibilities to the nearest grid point. This saves time by eliminating an interpolation operation after every new model evaluation, since the model visibilities correspond to the locations of the gridded visibilities. - - Args: - uu (list): the uu points (in [:math:`k\lambda`]) - vv (list): the vv points (in [:math:`k\lambda`]) - weights (list): the thermal weights (in [:math:`\mathrm{Jy}^{-2}`]) - re (list): the real component of the visibilities (in [:math:`\mathrm{Jy}`]) - im (list): the imaginary component of the visibilities (in [:math:`\mathrm{Jy}`]) - cell_size (float): the image cell size (in arcsec) - npix (int): the number of pixels in each dimension of the square image - debug (bool): provide the full 2D grid output of the gridded reals, imaginaries, and weights. Default False. - - Returns: - 6-tuple: (`u_grid`, `v_grid`, `grid_mask`, `avg_weights`, `avg_re`, `avg_im`) tuple of arrays. `grid_mask` has shape (npix, npix//2 + 1) corresponding to the RFFT output of an image with `cell_size` and dimensions (npix, npix). The remaining arrays are 1D and have length corresponding to the number of true elements in `ind`. They correspond to the model visibilities when the RFFT output is indexed with `grid_mask`. If debug is true, return a 7-tuple which includes the number of visibilities per cell. - - An image `cell_size` and `npix` correspond to particular `u_grid` and `v_grid` values from the RFFT. - - This pre-gridding procedure is similar to "uniform" weighting of visibilities for imaging, but not exactly the same in practice. This is because we are still evaluating a visibility likelihood function which incorporates the uncertainties of the measured spatial frequencies, whereas imaging routines use the weights to adjust the contributions of the measured visibility to the synthesize image. Evaluating a model against these gridded visibilities should be equivalent to the full interpolated calculation so long as it is true that - - 1) the visibility function is approximately constant over a :math:`(u,v)` cell - 2) the measurement uncertainties on the real and imaginary components of individual visibilities are correct and described by Gaussian noise - - If (1) is violated, you can always increase the width and npix of the image (keeping `cell_size` constant) to shrink the size of the :math:`(u,v)` cells (i.e., see https://docs.scipy.org/doc/numpy/reference/generated/numpy.fft.rfftfreq.html). If you need to do this, this probably indicates that you weren't using an image size appropriate for your dataset. - - If (2) is violated, then it's not a good idea to procede with this faster routine. Instead, revisit the calibration of your dataset, or build in self-calibration adjustments based upon antenna and time metadata of the visibilities. One downside of the pre-gridding operation is that it eliminates the possibility of using self-calibration type loss functions. - """ - assert npix % 2 == 0, "Image must have an even number of pixels" - - assert np.max(np.abs(uu)) < utils.get_max_spatial_freq( - cell_size, npix - ), "Dataset contains uu spatial frequency measurements larger than those in the proposed model image." - assert np.max(np.abs(vv)) < utils.get_max_spatial_freq( - cell_size, npix - ), "Dataset contains vv spatial frequency measurements larger than those in the proposed model image." - - assert np.all(weights > 0), "some weights are negative, check input args" - - # calculate the grid spacings - cell_size = cell_size * arcsec # [radians] - # cell_size is also the differential change in sky angles - # dll = dmm = cell_size #[radians] - - # the output spatial frequencies of the RFFT routine - uu_grid = np.fft.rfftfreq(npix, d=cell_size) * 1e-3 # convert to [kλ] - vv_grid = np.fft.fftfreq(npix, d=cell_size) * 1e-3 # convert to [kλ] - - nu = len(uu_grid) - nv = len(vv_grid) - - du = np.abs(uu_grid[1] - uu_grid[0]) - dv = np.abs(vv_grid[1] - vv_grid[0]) - - # The RFFT outputs u in the range [0, +] and v in the range [-, +], - # but the dataset contains measurements at u [-,+] and v [-, +]. - # Find all the u < 0 points and convert them via complex conj - ind_u_neg = uu < 0.0 - uu[ind_u_neg] *= -1.0 # swap axes so all u > 0 - vv[ind_u_neg] *= -1.0 # swap axes - im[ind_u_neg] *= -1.0 # complex conjugate - - # calculate the sum of the weights within each cell - # create the cells as edges around the existing points - # note that at this stage, the bins are strictly increasing - # when in fact, later on, we'll need to put this into fftshift format for the RFFT - weight_cell, v_edges, u_edges = np.histogram2d( - vv, - uu, - bins=[nv, nu], - range=[ - (np.min(vv_grid) - dv / 2, np.max(vv_grid) + dv / 2), - (np.min(uu_grid) - du / 2, np.max(uu_grid) + du / 2), - ], - weights=weights, - ) - - if debug: - nvis, v_edges, u_edges = np.histogram2d( - vv, - uu, - bins=[nv, nu], - range=[ - (np.min(vv_grid) - dv / 2, np.max(vv_grid) + dv / 2), - (np.min(uu_grid) - du / 2, np.max(uu_grid) + du / 2), - ], - ) - - # calculate the weighted average and weighted variance for each cell - # https://en.wikipedia.org/wiki/Weighted_arithmetic_mean - # also Bevington, "Data Reduction in the Physical Sciences", pg 57, 3rd ed. - # where weight = 1/sigma^2 - - # blank out the cells that have zero counts - weight_cell[(weight_cell == 0.0)] = np.nan - ind_ok = ~np.isnan(weight_cell) - - # weighted_mean = np.sum(x_i * weight_i) / weight_cell - real_part, v_edges, u_edges = np.histogram2d( - vv, - uu, - bins=[nv, nu], - range=[ - (np.min(vv_grid) - dv / 2, np.max(vv_grid) + dv / 2), - (np.min(uu_grid) - du / 2, np.max(uu_grid) + du / 2), - ], - weights=re * weights, - ) - - imag_part, v_edges, u_edges = np.histogram2d( - vv, - uu, - bins=[nv, nu], - range=[ - (np.min(vv_grid) - dv / 2, np.max(vv_grid) + dv / 2), - (np.min(uu_grid) - du / 2, np.max(uu_grid) + du / 2), - ], - weights=im * weights, - ) - - weighted_mean_real = real_part / weight_cell - weighted_mean_imag = imag_part / weight_cell - - # do an fftshift on weighted_means and sigmas to get this into a 2D grid - # that matches the RFFT output directly - - ind = np.fft.fftshift(ind_ok, axes=0) # RFFT indices that are not nan - # return ind as a 2D grid, so we can index the model RFFT output to directly - # compare to these values - - # list of gridded visibilities - if debug: - avg_weights = np.fft.fftshift(weight_cell, axes=0) - avg_re = np.fft.fftshift(weighted_mean_real, axes=0) - avg_im = np.fft.fftshift(weighted_mean_imag, axes=0) - nvis = np.fft.fftshift(nvis, axes=0) - - return (uu_grid, vv_grid, ind, avg_weights, avg_re, avg_im, nvis) - - else: - avg_weights = np.fft.fftshift(weight_cell, axes=0)[ind] - avg_re = np.fft.fftshift(weighted_mean_real, axes=0)[ind] - avg_im = np.fft.fftshift(weighted_mean_imag, axes=0)[ind] - - return (uu_grid, vv_grid, ind, avg_weights, avg_re, avg_im) - - -def grid_dataset(uus, vvs, weights, res, ims, cell_size, npix, **kwargs): - """ - Pre-grid a dataset containing multiple channels to the expected `u_grid` and `v_grid` points from the RFFT routine. - - Note that `nvis` need not be the same for each channel (i.e., `uus`, `vvs`, etc... can be ragged arrays, as long as each is iterable across the channel dimension). This routine iterates through the channels, calling :meth:`mpol.gridding.grid_datachannel` for each one. - - Args: - uus (nchan, nvis) list: the uu points (in klambda) - vvs (nchan, nvis) list: the vv points (in klambda) - weights (nchan, nvis) list: the thermal weights - res (nchan, nvis) list: the real component of the visibilities - ims (nchan, nvis) list: the imaginary component of the visibilities - cell_size: the image cell size (in arcsec) - npix: the number of pixels in each dimension of the square image - - Returns: - 6-tuple: (`u_grid`, `v_grid`, `grid_mask`, `avg_weights`, `avg_re`, `avg_im`) tuple of arrays. `u_grid` and `v_grid` are 1D arrays representing the coordinates of the RFFT grid, and are provided for reference. `grid_mask` has shape (npix, npix//2 + 1) corresponding to the RFFT output of an image with `cell_size` and dimensions (npix, npix). The remaining arrays are 1D and have length corresponding to the number of true elements in `ind`. They correspond to the model visibilities when the RFFT output is indexed with `grid_mask`. - """ - - nchan = uus.shape[0] - - # pre-allocate the index array - ind = np.zeros((nchan, npix, npix // 2 + 1), dtype="bool") - - avg_weights = [] - avg_re = [] - avg_im = [] - - for i in range(nchan): - uu_temp, vv_temp, ind_temp, w_temp, re_temp, im_temp = grid_datachannel( - uus[i], vvs[i], weights[i], res[i], ims[i], cell_size, npix - ) - ind[i] = ind_temp - avg_weights.append(w_temp) - avg_re.append(re_temp) - avg_im.append(im_temp) - - # flatten all visibilities to a single vector - # these are done because, once gridded, it's not certain that there will - # be the same number of nvis per channel, depending on where the visibilities - # fall in u,v space and to which cell they are gridded. - avg_weights = np.concatenate(avg_weights) - avg_re = np.concatenate(avg_re) - avg_im = np.concatenate(avg_im) - - # just take the last uu and vv values, since they are the same for all channels - - return (uu_temp, vv_temp, ind, avg_weights, avg_re, avg_im) diff --git a/test/spheroidal_gridding_test.py b/test/spheroidal_gridding_test.py deleted file mode 100644 index caa99cbb..00000000 --- a/test/spheroidal_gridding_test.py +++ /dev/null @@ -1,394 +0,0 @@ -import matplotlib.pyplot as plt -import numpy as np -import pytest - -from mpol import spheroidal_gridding, utils -from mpol.constants import * - - -def sky_plane( - alpha, - dec, - a=1, - delta_alpha=1.0 * arcsec, - delta_delta=1.0 * arcsec, - sigma_alpha=1.0 * arcsec, - sigma_delta=1.0 * arcsec, - Omega=0.0, -): - """ - Calculates a Gaussian on the sky plane - - Args: - alpha: ra (in radians) - delta: dec (in radians) - a : amplitude - delta_alpha : offset (in radians) - delta_dec : offset (in radians) - sigma_alpha : width (in radians) - sigma_dec : width (in radians) - Omega : position angle of ascending node (in degrees east of north) - - Returns: - Gaussian evaluated at input args - """ - return a * np.exp( - -( - (alpha - delta_alpha) ** 2 / (2 * sigma_alpha ** 2) - + (dec - delta_delta) ** 2 / (2 * sigma_delta ** 2) - ) - ) - - -def fourier_plane( - u, - v, - a=1, - delta_alpha=1.0 * arcsec, - delta_delta=1.0 * arcsec, - sigma_alpha=1.0 * arcsec, - sigma_delta=1.0 * arcsec, - Omega=0.0, -): - """ - Calculates the Analytic Fourier transform of the sky plane Gaussian. - - Args: - u: spatial freq (in kλ) - v: spatial freq (in kλ) - a : amplitude - delta_alpha : offset (in radians) - delta_dec : offset (in radians) - sigma_alpha : width (in radians) - sigma_dec : width (in radians) - Omega : position angle of ascending node (in degrees east of north) - - Returns: - FT Gaussian evaluated at u, v points - """ - - # convert back to lambda - u = u * 1e3 - v = v * 1e3 - - return ( - 2 - * np.pi - * a - * sigma_alpha - * sigma_delta - * np.exp( - -2 * np.pi ** 2 * (sigma_alpha ** 2 * u ** 2 + sigma_delta ** 2 * v ** 2) - - 2 * np.pi * 1.0j * (delta_alpha * u + delta_delta * v) - ) - ) - - -@pytest.fixture(scope="module") -def image_dict(): - """ - A useful fixture containing the analytic Gaussian evaluated on the sky plane. - """ - N_alpha = 128 - N_dec = 128 - img_radius = 15.0 * arcsec - - # full span of the image - ra = utils.fftspace(img_radius, N_alpha) # [arcsec] - dec = utils.fftspace(img_radius, N_dec) # [arcsec] - - # fill out an image - img = np.empty((N_dec, N_alpha), "float") - - for i, delta in enumerate(dec): - for j, alpha in enumerate(ra): - img[i, j] = sky_plane(alpha, delta) - - return { - "N_alpha": N_alpha, - "N_dec": N_dec, - "img_radius": img_radius, - "ra": ra, - "dec": dec, - "img": img, - } - - -@pytest.fixture(scope="module") -def corrfun_mat(image_dict): - """ - A useful fixture containing the correction function evaluated over the range of the sky plane test image. - """ - ra = image_dict["ra"] - dec = image_dict["dec"] - - # pre-multiply the image by the correction function - return spheroidal_gridding.corrfun_mat(np.fft.fftshift(ra), np.fft.fftshift(dec)) - - -def test_plot_input(image_dict, tmp_path): - fig, ax = plt.subplots(nrows=1) - ax.imshow(image_dict["img"], origin="upper", interpolation="none", aspect="equal") - ax.set_xlabel(r"$\Delta \alpha \cos \delta$") - ax.set_ylabel(r"$\Delta \delta$") - ax.set_title("Input image") - fig.savefig(tmp_path / "input.png", dpi=300) - - -def test_fill_corrfun_matrix(corrfun_mat, tmp_path): - - fig, ax = plt.subplots(nrows=1) - im = ax.imshow(corrfun_mat, origin="upper", interpolation="none", aspect="equal") - plt.colorbar(im) - ax.set_xlabel(r"$\Delta \alpha \cos \delta$") - ax.set_ylabel(r"$\Delta \delta$") - ax.set_title("Correction function") - fig.savefig(tmp_path / "corrfun.png", dpi=300) - - -@pytest.fixture(scope="module") -def vis_dict(image_dict, corrfun_mat): - """ - A useful fixture containing the corrected on non-corrected outputs from the FFT of the sky plane image. - """ - - img = image_dict["img"] - - # calculate the Fourier coordinates - dalpha = (2 * image_dict["img_radius"]) / image_dict["N_alpha"] - ddelta = (2 * image_dict["img_radius"]) / image_dict["N_dec"] - - us = np.fft.rfftfreq(image_dict["N_alpha"], d=dalpha) * 1e-3 # convert to [kλ] - vs = np.fft.fftfreq(image_dict["N_dec"], d=ddelta) * 1e-3 # convert to [kλ] - - # calculate the FFT, but first shift all axes. - # normalize output properly - vis = ( - dalpha * ddelta * np.fft.rfftn(corrfun_mat * np.fft.fftshift(img), axes=(0, 1)) - ) - vis_no_cor = dalpha * ddelta * np.fft.rfftn(np.fft.fftshift(img), axes=(0, 1)) - - return {"us": us, "vs": vs, "vis": vis, "vis_no_cor": vis_no_cor} - - -@pytest.fixture(scope="module") -def vis_analytical_full(vis_dict): - vs = vis_dict["vs"] - - XX_full, YY_full = np.meshgrid(vs, vs) - return fourier_plane(XX_full, YY_full) - - -@pytest.fixture(scope="module") -def vis_analytical_half(vis_dict): - us = vis_dict["us"] - vs = vis_dict["vs"] - - XX_full, YY_full = np.meshgrid(us, vs) - return fourier_plane(XX_full, YY_full) - - -def test_plot_full_analytical(tmp_path, vis_dict, vis_analytical_full): - vs_limit = np.fft.fftshift(vis_dict["vs"]) - - ext_full = [vs_limit[-1], vs_limit[0], vs_limit[-1], vs_limit[0]] - fig, ax = plt.subplots(nrows=2) - ax[0].imshow( - np.real(np.fft.fftshift(vis_analytical_full)), origin="upper", extent=ext_full - ) - ax[1].imshow( - np.imag(np.fft.fftshift(vis_analytical_full)), origin="upper", extent=ext_full - ) - fig.savefig(tmp_path / "analytical_full.png", dpi=300) - - -@pytest.fixture(scope="module") -def vis_diff(vis_dict, vis_analytical_half): - # compare to the analytic version - vis_no_cor = vis_dict["vis_no_cor"] - - return np.fft.fftshift(vis_no_cor - vis_analytical_half, axes=0) - - -def test_vis_prolate_diff(vis_diff): - # create a numerical test to make sure the difference is small - # i.e., we're actually comparing the right numerical grid - # relative to the expected analytical function - assert np.sum(np.abs(vis_diff)) < 1e-10 - - -def test_plot_analytical_prolate(tmp_path, vis_dict, vis_analytical_half, vis_diff): - vis_no_cor = vis_dict["vis_no_cor"] - - us = vis_dict["us"] - vs_limit = np.fft.fftshift(vis_dict["vs"]) - ext = [us[0], us[-1], vs_limit[-1], vs_limit[0]] - - fig, ax = plt.subplots(nrows=2, ncols=3, figsize=(7, 5)) - fig.suptitle("no interpolation") - ax[0, 0].set_title("numerical") - ax[0, 0].imshow( - np.real(np.fft.fftshift(vis_no_cor, axes=0)), - origin="upper", - interpolation="none", - aspect="equal", - extent=ext, - ) - ax[1, 0].imshow( - np.imag(np.fft.fftshift(vis_no_cor, axes=0)), - origin="upper", - interpolation="none", - aspect="equal", - extent=ext, - ) - - ax[0, 1].set_title("analytical") - im_ral = ax[0, 1].imshow( - np.real(np.fft.fftshift(vis_analytical_half, axes=0)), - origin="upper", - interpolation="none", - aspect="equal", - extent=ext, - ) - plt.colorbar(im_ral, ax=ax[0, 1]) - ax[1, 1].imshow( - np.imag(np.fft.fftshift(vis_analytical_half, axes=0)), - origin="upper", - interpolation="none", - aspect="equal", - extent=ext, - ) - - np.imag(np.fft.fftshift(vis_no_cor - vis_analytical_half, axes=0)) - ax[0, 2].set_title("difference") - im_real = ax[0, 2].imshow( - np.real(vis_diff), - origin="upper", - interpolation="none", - aspect="equal", - extent=ext, - ) - plt.colorbar(im_real, ax=ax[0, 2]) - im_imag = ax[1, 2].imshow( - np.imag(vis_diff), - origin="upper", - interpolation="none", - aspect="equal", - extent=ext, - ) - plt.colorbar(im_imag, ax=ax[1, 2]) - fig.subplots_adjust(wspace=0.05) - - fig.savefig(tmp_path / "comparison.png", dpi=300) - - -@pytest.fixture(scope="module") -def baselines(): - # Make sure we get enough samples in different quadrants. - return np.array( - [ - [50.0, 10.1], - [50.0, 0.0], - [39.5, -1.0], - [-50.1, 10.0], - [5.1, 1.0], - [-5.3, 1.0], - [5.8, 20.0], - [-5.0, -21.3], - ] - ) - - -@pytest.fixture(scope="module") -def analytic_samples(baselines): - u_data = baselines[:, 0] - v_data = baselines[:, 1] - return fourier_plane(u_data, v_data) - - -@pytest.fixture(scope="module") -def interpolation_matrices(vis_dict, baselines): - u_data = baselines[:, 0] - v_data = baselines[:, 1] - us = vis_dict["us"] - vs = vis_dict["vs"] - - # calculate and visualize the C_real and C_imag matrices - # these are scipy csc sparse matrices - return spheroidal_gridding.calc_matrices(u_data, v_data, us, vs) - - -def test_plot_interpolation_matrices(tmp_path, interpolation_matrices): - C_real, C_imag = interpolation_matrices - - C_real = C_real.toarray() - C_imag = C_imag.toarray() - - fig, ax = plt.subplots(nrows=1, figsize=(6, 6)) - vvmax = np.max(np.abs(C_real[:, 0:300])) - ax.imshow( - C_real[:, 0:300], - interpolation="none", - origin="upper", - cmap="RdBu", - aspect="auto", - vmin=-vvmax, - vmax=vvmax, - ) - fig.savefig(tmp_path / "C_real.png", dpi=300) - - fig, ax = plt.subplots(ncols=1, figsize=(6, 6)) - vvmax = np.max(np.abs(C_imag[:, 0:300])) - ax.imshow( - C_imag[:, 0:300], - interpolation="none", - origin="upper", - cmap="RdBu", - aspect="auto", - vmin=-vvmax, - vmax=vvmax, - ) - fig.savefig(tmp_path / "C_imag.png", dpi=300) - - -@pytest.fixture(scope="module") -def interpolated_prolate(vis_dict, interpolation_matrices): - vis = vis_dict["vis"] - C_real, C_imag = interpolation_matrices - return C_real.dot(np.real(vis.flatten())), C_imag.dot(np.imag(vis.flatten())) - - -def test_interpolate_points_prolate(analytic_samples, interpolated_prolate, vis_dict): - vis = vis_dict["vis_no_cor"] - vis_range = np.max(np.abs(vis)) - np.min(np.abs(vis)) - - interp_real, interp_imag = interpolated_prolate - interp = interp_real + interp_imag * 1.0j - - # the relative difference is using the full visibility range - assert np.all(np.abs((np.real(analytic_samples) - interp_real) / vis_range) < 1e-3) - assert np.all(np.abs((np.imag(analytic_samples) - interp_imag) / vis_range) < 1e-3) - assert np.all(np.abs(np.abs(analytic_samples - interp) / vis_range) < 1e-3) - - -def test_plot_points_prolate( - tmp_path, analytic_samples, interpolated_prolate, vis_dict -): - vis = vis_dict["vis_no_cor"] - vis_range = np.max(np.abs(vis)) - np.min(np.abs(vis)) - - interp_real, interp_imag = interpolated_prolate - fig, ax = plt.subplots(nrows=4, figsize=(4, 5)) - fig.suptitle("relative across full vis range") - ax[0].plot(np.real(analytic_samples), ".", ms=4) - ax[0].plot(interp_real, ".", ms=3) - ax[0].set_ylabel("real") - ax[1].plot((np.real(analytic_samples) - interp_real) / vis_range, ".") - ax[1].set_ylabel("rel difference") - ax[2].plot(np.imag(analytic_samples), ".", ms=4) - ax[2].plot(interp_imag, ".", ms=3) - ax[2].set_ylabel("imag") - ax[3].plot((np.imag(analytic_samples) - interp_imag) / vis_range, ".") - ax[3].set_ylabel("rel difference") - fig.subplots_adjust(hspace=0.4, left=0.2) - fig.savefig(tmp_path / "interpolated_comparison.png", dpi=300)