Merge pull request #239 from MPoL-dev/spheroidal-remove

Remove spheroidal gridding routines since deprecated by the NuFFT

.coveragerc

@@ -1,7 +1,6 @@
 [run]
 omit =
     venv/*
-    src/mpol/spheroidal_gridding.py
 relative_files = True
 
 

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`).
 

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.