updated types and coverage.

docs/changelog.md

@@ -4,6 +4,9 @@
 
 ## v0.3.0
 
+- Standardized treatment of numpy vs `torch.tensor`s, with preference for `torch.tensor` in many routines. This simplifies the internal logic of the routines and will make most operations run faster.
+- Standardized the input types of {class}:`mpol.fourier.NuFFT` and {class}:`mpol.fourier.NuFFTCached` to expect {class}`torch.Tensor`s (removed support for numpy arrays). This simplifies the internal logic of the routines and will make most operations run faster.
+- Changed {class}`mpol.fourier.make_fake_data` -> {class}`mpol.fourier.generate_fake_data`.
 - Changed base spatial frequency unit from k$\lambda$ to $\lambda$, closing issue [#223](https://github.com/MPoL-dev/MPoL/issues/223) and simplifying the internals of the codebase in numerous places. The following routines now expect inputs in units of $\lambda$:
   - {class}`mpol.coordinates.GridCoords`
   - {class}`mpol.coordinates.check_data_fit`

docs/ci-tutorials/fakedata.md

@@ -306,12 +306,12 @@ Thankfully, we see that we already chose a sufficiently small `cell_size`.
 
 ## Making the mock dataset
 
-With the {class}`~mpol.images.ImageCube`, $u,v$ and weight distributions now in hand, generating the mock visibilities is relatively straightforward using the {func}`mpol.fourier.make_fake_data` routine. This routine uses the {class}`~mpol.fourier.NuFFT` to produce loose visibilities at the $u,v$ locations and then adds random Gaussian noise to the visibilities, drawn from a probability distribution set by the value of the weights.
+With the {class}`~mpol.images.ImageCube`, $u,v$ and weight distributions now in hand, generating the mock visibilities is relatively straightforward using the {func}`mpol.fourier.generate_fake_data` routine. This routine uses the {class}`~mpol.fourier.NuFFT` to produce loose visibilities at the $u,v$ locations and then adds random Gaussian noise to the visibilities, drawn from a probability distribution set by the value of the weights.
 
 ```{code-cell} ipython3
 from mpol import fourier
 # will have the same shape as the uu, vv, and weight inputs
-data_noise, data_noiseless = fourier.make_fake_data(image, uu, vv, weight)
+data_noise, data_noiseless = fourier.generate_fake_data(image, uu, vv, weight)
 
 print(data_noise.shape)
 print(data_noiseless.shape)

pyproject.toml

@@ -99,9 +99,10 @@ module = [
     "MPoL.datasets",
     # "MPoL.fourier", # once we remove get_vis_residuals
     "MPoL.geometry",
-    # "MPoL.gridding", # once we sort check_data_inputs_2d
+    "MPoL.gridding", # once we sort check_data_inputs_2d
     "MPoL.images",
-    "MPoL.losses"
-    # "MPoL.utils", # once we fix check_baselines
+    "MPoL.losses",
+    "MPoL.precomposed",
+    "MPoL.utils"
     ]
 disallow_untyped_defs = true

src/mpol/coordinates.py

@@ -5,6 +5,7 @@
 import numpy as np
 import numpy.fft as np_fft
 import numpy.typing as npt
+import torch
 
 import mpol.constants as const
 from mpol.exceptions import CellSizeError
@@ -13,66 +14,66 @@
 
 class GridCoords:
     r"""
-    The GridCoords object uses desired image dimensions (via the ``cell_size`` and 
+    The GridCoords object uses desired image dimensions (via the ``cell_size`` and
     ``npix`` arguments) to define a corresponding Fourier plane grid.
 
     Args:
         cell_size (float): width of a single square pixel in [arcsec]
         npix (int): number of pixels in the width of the image
 
-    The Fourier grid is defined over the domain :math:`[-u,+u]`, :math:`[-v,+v]`, even 
+    The Fourier grid is defined over the domain :math:`[-u,+u]`, :math:`[-v,+v]`, even
     though for real images, technically we could use an RFFT grid from :math:`[0,+u]` to
-    :math:`[-v,+v]`. The reason we opt for a full FFT grid in this instance is 
+    :math:`[-v,+v]`. The reason we opt for a full FFT grid in this instance is
     implementation simplicity.
 
-    Images (and their corresponding Fourier transform quantities) are represented as 
-    two-dimensional arrays packed as ``[y, x]`` and ``[v, u]``.  This means that an 
-    image with dimensions ``(npix, npix)`` will also have a corresponding FFT Fourier 
-    grid with shape ``(npix, npix)``. The native :class:`~mpol.gridding.GridCoords` 
-    representation assumes the Fourier grid (and thus image) are laid out as one might 
+    Images (and their corresponding Fourier transform quantities) are represented as
+    two-dimensional arrays packed as ``[y, x]`` and ``[v, u]``.  This means that an
+    image with dimensions ``(npix, npix)`` will also have a corresponding FFT Fourier
+    grid with shape ``(npix, npix)``. The native :class:`~mpol.gridding.GridCoords`
+    representation assumes the Fourier grid (and thus image) are laid out as one might
     normally expect an image (i.e., no ``np.fft.fftshift`` has been applied).
 
     After the object is initialized, instance variables can be accessed, for example
 
     >>> myCoords = GridCoords(cell_size=0.005, 512)
     >>> myCoords.img_ext
 
-    :ivar dl: image-plane cell spacing in RA direction (assumed to be positive) 
+    :ivar dl: image-plane cell spacing in RA direction (assumed to be positive)
         [radians]
     :ivar dm: image-plane cell spacing in DEC direction [radians]
-    :ivar img_ext: The length-4 list of (left, right, bottom, top) expected by routines 
-        like ``matplotlib.pyplot.imshow`` in the ``extent`` parameter assuming 
+    :ivar img_ext: The length-4 list of (left, right, bottom, top) expected by routines
+        like ``matplotlib.pyplot.imshow`` in the ``extent`` parameter assuming
         ``origin='lower'``. Units of [arcsec]
-    :ivar packed_x_centers_2D: 2D array of l increasing, with fftshifted applied 
-        [arcseconds]. Useful for directly evaluating some function to create a 
+    :ivar packed_x_centers_2D: 2D array of l increasing, with fftshifted applied
+        [arcseconds]. Useful for directly evaluating some function to create a
         packed cube.
-    :ivar packed_y_centers_2D: 2D array of m increasing, with fftshifted applied 
-        [arcseconds]. Useful for directly evaluating some function to create a 
+    :ivar packed_y_centers_2D: 2D array of m increasing, with fftshifted applied
+        [arcseconds]. Useful for directly evaluating some function to create a
         packed cube.
-    :ivar sky_x_centers_2D: 2D array of l arranged for evaluating a sky image 
+    :ivar sky_x_centers_2D: 2D array of l arranged for evaluating a sky image
         [arcseconds]. l coordinate increases to the left (as on sky).
-    :ivar sky_y_centers_2D: 2D array of m arranged for evaluating a sky image 
-        [arcseconds]. 
-    :ivar du: Fourier-plane cell spacing in East-West direction 
+    :ivar sky_y_centers_2D: 2D array of m arranged for evaluating a sky image
+        [arcseconds].
+    :ivar du: Fourier-plane cell spacing in East-West direction
         [:math:`\mathrm{k}\lambda`]
-    :ivar dv: Fourier-plane cell spacing in North-South direction 
+    :ivar dv: Fourier-plane cell spacing in North-South direction
         [:math:`\mathrm{k}\lambda`]
-    :ivar u_centers: 1D array of cell centers in East-West direction 
+    :ivar u_centers: 1D array of cell centers in East-West direction
         [:math:`\mathrm{k}\lambda`].
-    :ivar v_centers: 1D array of cell centers in North-West direction 
+    :ivar v_centers: 1D array of cell centers in North-West direction
         [:math:`\mathrm{k}\lambda`].
-    :ivar u_edges: 1D array of cell edges in East-West direction 
+    :ivar u_edges: 1D array of cell edges in East-West direction
         [:math:`\mathrm{k}\lambda`].
-    :ivar v_edges: 1D array of cell edges in North-South direction 
+    :ivar v_edges: 1D array of cell edges in North-South direction
         [:math:`\mathrm{k}\lambda`].
     :ivar u_bin_min: minimum u edge [:math:`\mathrm{k}\lambda`]
     :ivar u_bin_max: maximum u edge [:math:`\mathrm{k}\lambda`]
     :ivar v_bin_min: minimum v edge [:math:`\mathrm{k}\lambda`]
     :ivar v_bin_max: maximum v edge [:math:`\mathrm{k}\lambda`]
-    :ivar max_grid: maximum spatial frequency enclosed by Fourier grid 
+    :ivar max_grid: maximum spatial frequency enclosed by Fourier grid
         [:math:`\mathrm{k}\lambda`]
-    :ivar vis_ext: length-4 list of (left, right, bottom, top) expected by routines 
-        like ``matplotlib.pyplot.imshow`` in the ``extent`` parameter assuming 
+    :ivar vis_ext: length-4 list of (left, right, bottom, top) expected by routines
+        like ``matplotlib.pyplot.imshow`` in the ``extent`` parameter assuming
         ``origin='lower'``. Units of [:math:`\mathrm{k}\lambda`]
     """
 
@@ -178,44 +179,62 @@ def __init__(self, cell_size: float, npix: int) -> None:
         self.sky_y_centers_2D = y_centers_2D  # [arcsec]
         self.sky_x_centers_2D = np.fliplr(x_centers_2D)  # [arcsec]
 
-    def check_data_fit(self, uu: npt.ArrayLike, vv: npt.ArrayLike) -> None:
+    def check_data_fit(
+        self,
+        uu: torch.Tensor | npt.NDArray[np.floating[Any]],
+        vv: torch.Tensor | npt.NDArray[np.floating[Any]],
+    ) -> bool:
         r"""
-        Test whether loose data visibilities fit within the Fourier grid defined by 
+        Test whether loose data visibilities fit within the Fourier grid defined by
         cell_size and npix.
 
-        Args:
-            uu (np.array): array of u spatial frequency coordinates. 
-                Units of [:math:`\mathrm{k}\lambda`]
-            vv (np.array): array of v spatial frequency coordinates. 
-                Units of [:math:`\mathrm{k}\lambda`]
-
-        Returns:
-            ``True`` if all visibilities fit within the Fourier grid defined by 
-            ``[u_bin_min, u_bin_max, v_bin_min, v_bin_max]``. Otherwise an 
-            ``AssertionError`` is raised on the first violated boundary.
+        Parameters
+        ----------
+        uu : :class:`torch.Tensor` of `torch.double`
+            u spatial frequency coordinates.
+            Units of [:math:`\mathrm{k}\lambda`]
+        vv : :class:`torch.Tensor` of `torch.double`
+            v spatial frequency coordinates.
+            Units of [:math:`\mathrm{k}\lambda`]
+
+        Returns
+        -------
+        bool
+            ``True`` if all visibilities fit within the Fourier grid defined by
+            ``[u_bin_min, u_bin_max, v_bin_min, v_bin_max]``. Otherwise a
+            :class:`mpol.exceptions.CellSizeError` is raised on the first violated
+            boundary.
         """
 
+        # we need this routine to work with both numpy.ndarray or torch.Tensor
+        # because it is called for DirtyImager setup (numpy only) 
+        # and torch layers
+        uu = torch.as_tensor(uu)
+        vv = torch.as_tensor(vv)
+
         # max freq in dataset
-        max_uu_vv = np.max(np.abs(np.concatenate([uu, vv])))
+        max_uu_vv = torch.max(torch.abs(torch.concatenate([uu, vv]))).item()
 
         # max freq needed to support dataset
         max_cell_size = get_maximum_cell_size(max_uu_vv)
 
-        if np.max(np.abs(uu)) > self.max_grid:
+        if torch.max(torch.abs(uu)) > self.max_grid:
             raise CellSizeError(
                 f"Dataset contains uu spatial frequency measurements larger than those in the proposed model image. Decrease cell_size below {max_cell_size} arcsec."
             )
 
-        if np.max(np.abs(vv)) > self.max_grid:
+        if torch.max(torch.abs(vv)) > self.max_grid:
             raise CellSizeError(
                 f"Dataset contains vv spatial frequency measurements larger than those in the proposed model image. Decrease cell_size below {max_cell_size} arcsec."
             )
 
+        return True
+
     def __eq__(self, other: Any) -> bool:
         if not isinstance(other, GridCoords):
             # don't attempt to compare against different types
             return NotImplemented
 
-        # GridCoords objects are considered equal if they have the same cell_size and 
+        # GridCoords objects are considered equal if they have the same cell_size and
         # npix, since all other attributes are derived from these two core properties.
         return self.cell_size == other.cell_size and self.npix == other.npix