diff --git a/docs/changelog.md b/docs/changelog.md index 26bdc713..b76a9692 100644 --- a/docs/changelog.md +++ b/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` diff --git a/docs/ci-tutorials/fakedata.md b/docs/ci-tutorials/fakedata.md index 926fca3a..d4e1cbb6 100644 --- a/docs/ci-tutorials/fakedata.md +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index dbdc7009..fcf8ffa3 100644 --- a/pyproject.toml +++ b/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 \ No newline at end of file diff --git a/src/mpol/coordinates.py b/src/mpol/coordinates.py index e52b6b6f..f939999b 100644 --- a/src/mpol/coordinates.py +++ b/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,23 +14,23 @@ 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 @@ -37,42 +38,42 @@ class GridCoords: >>> 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 diff --git a/src/mpol/fourier.py b/src/mpol/fourier.py index c3404051..091cd8fc 100644 --- a/src/mpol/fourier.py +++ b/src/mpol/fourier.py @@ -84,7 +84,7 @@ def ground_vis(self) -> torch.Tensor: Returns ------- :class:`torch.Tensor` of :class:`torch.complex128` of shape ``(nchan, npix, npix)`` - complex-valued FFT of the image cube (i.e., the visibility cube), in + complex-valued FFT of the image cube (i.e., the visibility cube), in 'ground' format. """ @@ -116,6 +116,7 @@ def ground_phase(self) -> torch.Tensor: """ return torch.angle(self.ground_vis) + class NuFFT(nn.Module): r""" This layer translates input from an :class:`mpol.images.ImageCube` to loose, @@ -271,8 +272,8 @@ def _assemble_ktraj(self, uu: torch.Tensor, vv: torch.Tensor) -> torch.Tensor: def forward( self, cube: torch.Tensor, - uu, - vv, + uu: torch.Tensor, + vv: torch.Tensor, sparse_matrices: bool = False, ) -> torch.Tensor: r""" @@ -282,22 +283,29 @@ def forward( points. In general, you probably do not want to provide baselines that include Hermitian pairs. - Args: - cube (torch.double tensor): of shape ``(nchan, npix, npix)``). The cube - should be a "prepacked" image cube, for example, - from :meth:`mpol.images.ImageCube.forward` - uu (array-like): array of the u (East-West) spatial frequency coordinate - [klambda]. - vv (array-like): array of the v (North-South) spatial frequency coordinate - [klambda] (must be the same shape as uu) - sparse_matrices (bool): If False, use the default table-based interpolation - of TorchKbNufft.If True, use TorchKbNuFFT sparse matrices (generally - slower but more accurate). Note that sparse matrices are incompatible - with multi-channel `uu` and `vv` arrays (see below). + Parameters + ---------- + cube : :class:`torch.Tensor` of :class:`torch.double` + shape ``(nchan, npix, npix)``). The cube + should be a "prepacked" image cube, for example, + from :meth:`mpol.images.ImageCube.forward` + uu : :class:`torch.Tensor` of :class:`torch.double` + 2D array of the u (East-West) spatial frequency coordinate + [klambda] of shape ``(nchan, npix)`` + vv : :class:`torch.Tensor` of :class:`torch.double` + 2D array of the v (North-South) spatial frequency coordinate + [klambda] (must be the same shape as uu) + sparse_matrices : bool + If False, use the default table-based interpolation + of TorchKbNufft.If True, use TorchKbNuFFT sparse matrices (generally + slower but more accurate). Note that sparse matrices are incompatible + with multi-channel `uu` and `vv` arrays (see below). - Returns: - torch.complex tensor: Fourier samples of shape ``(nchan, nvis)``, evaluated - at the ``uu``, ``vv`` points + Returns + ------- + :class:`torch.Tensor` of :class:`torch.complex128` + Fourier samples of shape ``(nchan, nvis)``, evaluated at the ``uu``, + ``vv`` points **Dimensionality**: You should consider the dimensionality of your image and your visibility samples when using this method. If your image has multiple @@ -358,10 +366,6 @@ def forward( depending on your problem. """ - # permit numpy, but prefer tensor - uu = torch.as_tensor(uu) - vv = torch.as_tensor(vv) - # make sure that the nchan assumptions for the ImageCube and the NuFFT # setup are the same if cube.size(0) != self.nchan: @@ -503,8 +507,6 @@ class NuFFTCached(NuFFT): this will result in a warning. Args: - cell_size (float): the width of an image-plane pixel [arcseconds] - npix (int): the number of pixels per image side coords (GridCoords): an object already instantiated from the GridCoords class. If providing this, cannot provide ``cell_size`` or ``npix``. nchan (int): the number of channels in the :class:`mpol.images.ImageCube`. @@ -519,8 +521,8 @@ class NuFFTCached(NuFFT): def __init__( self, coords: GridCoords, - uu, - vv, + uu: torch.Tensor, + vv: torch.Tensor, nchan: int = 1, sparse_matrices: bool = True, ): @@ -543,10 +545,6 @@ def __init__( self.sparse_matrices = sparse_matrices self.same_uv = same_uv - # permit numpy, but prefer tensor - uu = torch.as_tensor(uu) - vv = torch.as_tensor(vv) - self.register_buffer("k_traj", self._assemble_ktraj(uu, vv)) self.k_traj: torch.Tensor @@ -629,55 +627,73 @@ def forward(self, cube): return output -def make_fake_data( - image_cube: ImageCube, - uu: NDArray[floating[Any]], - vv: NDArray[floating[Any]], - weight: NDArray[floating[Any]], -) -> tuple[NDArray[complexfloating[Any, Any]], ...]: +def generate_fake_data( + packed_cube: torch.Tensor, + coords: GridCoords, + uu: torch.Tensor, + vv: torch.Tensor, + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: r""" - Create a fake dataset from a supplied :class:`mpol.images.ImageCube`. See - :ref:`mock-dataset-label` for more details on how to prepare a generic image for - use in an :class:`~mpol.images.ImageCube`. + Create a fake dataset from a supplied packed tensor cube using + :class:`mpol.fourier.NuFFT`. See :ref:`mock-dataset-label` for more details on how + to prepare a generic image to a `packed_cube`. - The provided visibilities can be 1d for a single continuum channel, or 2d for - image cube. If 1d, visibilities will be converted to 2d arrays of shape - ``(1, nvis)``. + The ``uu`` and ``vv`` baselines can either be 1D or 2D, depending on the desired + broadcasting behavior from the :class:`mpol.fourier.NuFFT`. - Args: - imageCube (:class:`~mpol.images.ImageCube`): the image layer to put into a - fake dataset - uu (numpy array): array of u spatial frequency coordinates, not including - Hermitian pairs. Units of [:math:`\mathrm{k}\lambda`] - vv (numpy array): array of v spatial frequency coordinates, not including - Hermitian pairs. Units of [:math:`\mathrm{k}\lambda`] - weight (2d numpy array): length array of thermal weights - :math:`w_i = 1/\sigma_i^2`. Units of [:math:`1/\mathrm{Jy}^2`] - - Returns: - (2-tuple): a two tuple of the fake data. The first array is the mock dataset + If the ``weight`` array is 1D, the routine assumes the weights will be broadcasted + to all ``nchan``. Otherwise, provide a 2D weight array. + + Parameters + ---------- + packed_cube : :class:`torch.Tensor` of `class:`torch.double` + the image in "packed" format with shape (`nchan`, `npix`, `npix`) + coords : :class:`mpol.coordinates.GridCoords` + uu : :class:`torch.Tensor` of `class:`torch.double` + array of u spatial frequency coordinates, + not including Hermitian pairs. Units of [:math:`\mathrm{k}\lambda`] + vv : :class:`torch.Tensor` of `class:`torch.double` + array of v spatial frequency coordinates, + not including Hermitian pairs. Units of [:math:`\mathrm{k}\lambda`] + weight : :class:`torch.Tensor` of `class:`torch.double` + shape ``(nchan, nvis)`` array of thermal weights + :math:`w_i = 1/\sigma_i^2`. Units of [:math:`1/\mathrm{Jy}^2`] + + Returns + ------- + :class:`torch.Tensor` of `class:`torch.complex128` + A 2-tuple of the fake data. The first array is the mock dataset including noise, the second array is the mock dataset without added noise. + Each array is shape ``(nchan, npix, npix)``. """ - # instantiate a NuFFT object based on the ImageCube - # OK if uu shape (nvis,) - nufft = NuFFT(coords=image_cube.coords, nchan=image_cube.nchan) - - # make into a multi-channel dataset, even if only a single-channel provided - if uu.ndim == 1: - uu = np.atleast_2d(uu) - vv = np.atleast_2d(vv) - weight = np.atleast_2d(weight) + # instantiate a NuFFT object based on the + nchan, npix, _ = packed_cube.size() + assert coords.npix == npix, "npix for packed_cube {:} and coords {:} differ".format( + nchan, coords.npix + ) + nufft = NuFFT(coords=coords, nchan=nchan) # carry it forward to the visibilities, which will be (nchan, nvis) - vis_noiseless: NDArray[complexfloating[Any, Any]] - vis_noiseless = nufft(image_cube.cube, uu, vv).detach().numpy() + vis_noiseless: torch.Tensor + vis_noiseless = nufft(packed_cube, uu, vv) # generate complex noise - sigma = 1 / np.sqrt(weight) - noise = np.random.normal( - loc=0, scale=sigma, size=uu.shape - ) + 1.0j * np.random.normal(loc=0, scale=sigma, size=uu.shape) + # we could use torch.normal with complex quantities directly, but they treat + # a complex normal with std of 1 as having an amplitude of 1. + # wheras ALMA "weights" have the definition that their scatter is for the reals + # (or imaginaries). So there is a factor of sqrt(2) between them. + # we split things up, work with real quantities, and then combine later + mean = torch.zeros(vis_noiseless.size()) + + # broadcast weight to (nchan, nvis) if it isn't already + weight = torch.broadcast_to(weight, vis_noiseless.size()) + sigma = torch.sqrt(1 / weight) + + noise_re = torch.normal(mean, sigma) + noise_im = torch.normal(mean, sigma) + noise = torch.complex(noise_re, noise_im) # add to data vis_noise = vis_noiseless + noise @@ -713,7 +729,9 @@ def get_vis_residuals(model, u_true, v_true, V_true, return_Vmod=False, channel= """ nufft = NuFFT(coords=model.coords, nchan=model.nchan) - vis_model = nufft(model.icube.cube.to("cpu"), u_true, v_true) # TODO: remove 'to' call + vis_model = nufft( + model.icube.cube.to("cpu"), u_true, v_true + ) # TODO: remove 'to' call # convert to numpy, select channel vis_model = vis_model.detach().numpy()[channel] diff --git a/src/mpol/precomposed.py b/src/mpol/precomposed.py index 227b3e3d..365e99dc 100644 --- a/src/mpol/precomposed.py +++ b/src/mpol/precomposed.py @@ -37,10 +37,10 @@ class SimpleNet(torch.nn.Module): def __init__( self, - coords=None, - nchan=1, - base_cube=None, - ): + coords: GridCoords, + nchan: int = 1, + base_cube: torch.Tensor | None = None, + ) -> None: super().__init__() self.coords = coords @@ -55,7 +55,7 @@ def __init__( self.icube = images.ImageCube(coords=self.coords, nchan=self.nchan) self.fcube = fourier.FourierCube(coords=self.coords) - def forward(self): + def forward(self) -> torch.Tensor: r""" Feed forward to calculate the model visibilities. In this step, a :class:`~mpol.images.BaseCube` is fed to a :class:`~mpol.images.HannConvCube` @@ -67,5 +67,5 @@ def forward(self): x = self.bcube() x = self.conv_layer(x) x = self.icube(x) - vis = self.fcube(x) + vis: torch.Tensor = self.fcube(x) return vis diff --git a/src/mpol/utils.py b/src/mpol/utils.py index f8ddb7cf..0bf33840 100644 --- a/src/mpol/utils.py +++ b/src/mpol/utils.py @@ -5,7 +5,7 @@ from typing import Any import numpy.typing as npt -from mpol.constants import arcsec, c_ms, cc, deg, kB +from mpol.constants import arcsec, cc, deg, kB def torch2npy(t: torch.Tensor) -> npt.NDArray: @@ -234,6 +234,7 @@ def check_baselines(q, min_feasible_q=1e0, max_feasible_q=1e5): "[lambda].".format(min(q), min_feasible_q * 1e3) ) + def get_max_spatial_freq(cell_size: float, npix: int) -> float: r""" Calculate the maximum spatial frequency that the image can represent and still @@ -272,8 +273,8 @@ def get_maximum_cell_size(uu_vv_point: float) -> float: def get_optimal_image_properties( image_width: float, - u: npt.NDArray[np.floating[Any]], - v: npt.NDArray[np.floating[Any]], + u: torch.Tensor, + v: torch.Tensor, ) -> tuple[float, int]: r""" For an image of desired width, determine the maximum pixel size that @@ -285,7 +286,7 @@ def get_optimal_image_properties( image_width : float, unit = arcsec Desired width of the image (for a square image of size `image_width` :math:`\times` `image_width`). - u, v : np.ndarray of np.float, unit = :math:`k\lambda` + u, v : :class:`torch.Tensor` of :class:`torch.double`, unit = :math:`k\lambda` `u` and `v` baselines. Returns @@ -299,9 +300,9 @@ def get_optimal_image_properties( ----- Assumes baselines are as-observed. """ - max_freq = max(max(abs(u)), max(abs(v))) + max_freq = torch.max(torch.abs(torch.concat([u,v]))) - cell_size = get_maximum_cell_size(max_freq) + cell_size = get_maximum_cell_size(max_freq.item()) # round npix up to nearest integer npix = math.ceil(image_width / cell_size) diff --git a/test/conftest.py b/test/conftest.py index d46bdca3..cbf0d1dc 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -56,9 +56,11 @@ def baselines_1D(baselines_m): uu, vv = baselines_m # klambda for now - return 1e-3 * visread.process.convert_baselines( - uu, 230.0e9 - ), 1e-3 * visread.process.convert_baselines(vv, 230.0e9) + uu = 1e-3 * visread.process.convert_baselines(uu, 230.0e9) + vv = 1e-3 * visread.process.convert_baselines(vv, 230.0e9) + + # convert to torch + return torch.as_tensor(uu), torch.as_tensor(vv) @pytest.fixture @@ -70,7 +72,10 @@ def baselines_2D(baselines_m): ) # klambda for now - return 1e-3 * u_lam, 1e-3 * v_lam + u_klam = 1e-3 * torch.as_tensor(u_lam) + v_klam = 1e-3 * torch.as_tensor(v_lam) + + return u_klam, v_klam # to replace everything with the mock dataset (and pass), we need to replace diff --git a/test/coordinates_test.py b/test/coordinates_test.py index 7b1b4db2..cfddf554 100644 --- a/test/coordinates_test.py +++ b/test/coordinates_test.py @@ -1,5 +1,5 @@ import matplotlib.pyplot as plt -import numpy as np +import torch import pytest from mpol import coordinates @@ -97,7 +97,7 @@ def test_grid_coords_fail(baselines_2D): coords = coordinates.GridCoords(cell_size=0.05, npix=800) - print("max u data", np.max(uu)) + print("max u data", torch.max(uu)) print("max u grid", coords.max_grid) with pytest.raises(CellSizeError): diff --git a/test/fourier_test.py b/test/fourier_test.py index fb7b5084..613c0f1a 100644 --- a/test/fourier_test.py +++ b/test/fourier_test.py @@ -1,9 +1,9 @@ import matplotlib.pyplot as plt +import numpy as np import torch from pytest import approx from mpol import fourier, images, utils -from mpol.constants import * def test_fourier_cube(coords, tmp_path): @@ -258,21 +258,23 @@ def test_nufft_accuracy_single_chan(coords, baselines_1D, tmp_path): img_packed_tensor = torch.tensor(img_packed[np.newaxis, :, :], requires_grad=True) # use the NuFFT to predict the values of the cube at the u,v locations - num_output = ( - layer(img_packed_tensor, uu, vv)[0].detach().numpy() - ) # take the channel dim out + num_output = layer(img_packed_tensor, uu, vv)[0] # take the channel dim out # calculate the values analytically an_output = utils.fourier_gaussian_klambda_arcsec(uu, vv, **kw) # find max difference diff = num_output - an_output - max_diff = np.max(np.abs(diff)) - max = np.max(np.abs(num_output)) + max_diff = torch.max(torch.abs(diff)) + max = torch.max(torch.abs(num_output)) print(max_diff, max) # collapse the function into 1D by doing q - qq = np.hypot(uu, vv) + qq = utils.torch2npy(torch.hypot(uu, vv)) + + # convert to numpy for plotting + num_output = utils.torch2npy(num_output) + diff = utils.torch2npy(diff) fig, ax = plt.subplots(nrows=4, sharex=True) ax[0].scatter(qq, an_output.real, s=3, label="analytic") @@ -327,7 +329,7 @@ def test_nufft_cached_accuracy_single_chan(coords, baselines_1D, tmp_path): # use the NuFFT to predict the values of the cube at the u,v locations num_output = ( - layer(img_packed_tensor)[0].detach().numpy() + layer(img_packed_tensor)[0] ) # take the channel dim out # calculate the values analytically @@ -335,12 +337,17 @@ def test_nufft_cached_accuracy_single_chan(coords, baselines_1D, tmp_path): # find max difference diff = num_output - an_output - max_diff = np.max(np.abs(diff)) - max = np.max(np.abs(num_output)) + max_diff = torch.max(torch.abs(diff)) + max = torch.max(torch.abs(num_output)) print(max_diff, max) + # collapse the function into 1D by doing q - qq = np.hypot(uu, vv) + qq = utils.torch2npy(torch.hypot(uu, vv)) + + # convert to numpy for plotting + num_output = utils.torch2npy(num_output) + diff = utils.torch2npy(diff) fig, ax = plt.subplots(nrows=4, sharex=True) ax[0].scatter(qq, an_output.real, s=3, label="analytic") @@ -449,15 +456,20 @@ def test_nufft_cached_accuracy_batch_broadcast(coords, baselines_2D, tmp_path): ) # use the NuFFT to predict the values of the cube at the u,v locations - num_output = layer(img_packed_tensor).detach().numpy() + num_output = layer(img_packed_tensor) # plot a single channel, to check ichan = 1 - qq = np.hypot(uu[ichan], vv[ichan]) an_output = utils.fourier_gaussian_klambda_arcsec(uu[ichan], vv[ichan], **kw) + diff = num_output[ichan] - an_output + # convert for plotting + qq = utils.torch2npy(torch.hypot(uu[ichan], vv[ichan])) + num_output = utils.torch2npy(num_output) + diff = utils.torch2npy(diff) + fig, ax = plt.subplots(nrows=4, sharex=True) ax[0].scatter(qq, an_output.real, s=3, label="analytic") ax[0].scatter(qq, num_output[ichan].real, s=1, label="NuFFT")