Merge pull request #17 from radionets-project/update_ms_gridding

Update MS Gridding and Dirty Image Plotting

Author: tgross03

pyvisgrid/core/gridder.py

@@ -3,8 +3,9 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-import numpy
 import numpy as np
+from numpy.typing import ArrayLike
+
 from astropy.constants import c
 from astropy.io import fits
 from casatools.table import table
@@ -25,12 +26,15 @@ class GridData:
     DataClass to save the gridded and non-gridded visibilities for a
     specific Stokes component.
 
-    Parameters
+    Attributes
     ----------
 
     vis_data : numpy.ndarray
     The ungridded visibilities.
 
+    fov : float
+    The size of the Field Of View of the gridded data in arcseconds.
+
     mask : numpy.ndarray, optional
     The mask created from the given (u,v) coordinates. The mask contains
     the number of (u,v) coordinates per pixel.
@@ -47,25 +51,58 @@ class GridData:
 
     """
 
-    vis_data: numpy.ndarray
-    mask: numpy.ndarray | None = None
-    mask_real: numpy.ndarray | None = None
-    mask_imag: numpy.ndarray | None = None
-    dirty_image: numpy.ndarray | None = None
+    vis_data: np.ndarray
+    fov: float | None = None
+    mask: np.ndarray | None = None
+    mask_real: np.ndarray | None = None
+    mask_imag: np.ndarray | None = None
+    dirty_image: np.ndarray | None = None
 
     def __str__(self):
         return self.__dict__
 
+    def get_mask_complex(self):
+        """
+        Returns the gridded mask as a complex array with
+        the form ``mask.real + 1j * mask.imag``.
+
+        Returns
+        -------
+        numpy.ndarray: Complex mask
+        """
+        return self.mask_real + 1j * self.mask_imag
+
+    def get_mask_abs_phase(self):
+        """
+        Returns the gridded mask as amplitude and phase.
+
+        Returns
+        -------
+        tuple[numpy.ndarray]: Amplitude and phase of the complex mask.
+        """
+        mask_complex = self.get_mask_complex()
+        return np.abs(mask_complex), np.angle(mask_complex)
+
+    def get_mask_real_imag(self):
+        """
+        Returns the gridded mask as real and imaginary parts.
+
+        Returns
+        -------
+        tuple[numpy.ndarray]: Real and parts of the complex mask.
+        """
+        return self.mask_real, self.mask_imag
+
 
 class Gridder:
     def __init__(
         self,
-        u_meter: numpy.ndarray,
-        v_meter: numpy.ndarray,
+        u_meter: np.ndarray,
+        v_meter: np.ndarray,
         img_size: int,
         fov: float,
         ref_frequency: float,
-        frequency_offsets: numpy.typing.ArrayLike,
+        frequency_offsets: ArrayLike,
     ):
         """
 
@@ -191,27 +228,35 @@ def grid(self, stokes_component: str = "I"):
         )
 
         mask, *_ = np.histogram2d(
-            u_wave_full, v_wave_full, bins=[bins, bins], density=False
+            x=u_wave_full, y=v_wave_full, bins=[bins, bins], density=False
         )
+        mask = (
+            mask.T
+        )  # u (x) is histogrammed along the first dimension (rows) --> transpose
         mask[mask == 0] = 1
 
         mask_real, _, _ = np.histogram2d(
-            u_wave_full,
-            v_wave_full,
+            x=u_wave_full,
+            y=v_wave_full,
             bins=[bins, bins],
             weights=stokes_real_full,
             density=False,
         )
         mask_imag, _, _ = np.histogram2d(
-            u_wave_full,
-            v_wave_full,
+            x=u_wave_full,
+            y=v_wave_full,
             bins=[bins, bins],
             weights=stokes_imag_full,
             density=False,
         )
+
+        mask_real = mask_real.T  # see above
+        mask_imag = mask_imag.T  # see above
+
         mask_real /= mask
         mask_imag /= mask
 
+        grid_data.fov = self.fov
         grid_data.mask = mask
         grid_data.mask_real = mask_real
         grid_data.mask_imag = mask_imag
@@ -229,7 +274,7 @@ def from_pyvisgen(
         img_size: int,
         fov: float,
         stokes_components: list[str] | str = "I",
-        polarizations: list[str] | str = "",
+        polarizations: list[str] | str | None = None,
     ):
         """
         Initializes the gridder with the visibility data which is generated by the
@@ -262,8 +307,8 @@ def from_pyvisgen(
         This can either be a list of components (e.g. ``['I', 'V']``) or a single
         string. Default is ``'I'``.
 
-        polarizations : list[str] | str, optional
-        The polarization type. Default is ``''``.
+        polarizations : list[str] | str | None, optional
+        The polarization type. Default is ``None``.
         """
         u_meter = vis_data.u
         v_meter = vis_data.v
@@ -291,6 +336,9 @@ def from_pyvisgen(
         if isinstance(stokes_components, str):
             stokes_components = [stokes_components]
 
+        if polarizations is None:
+            polarizations = ""
+
         if isinstance(polarizations, str):
             polarizations = [polarizations]
 
@@ -460,26 +508,15 @@ def from_ms(
             mask = tab.getcol("DATA_DESC_ID") == desc_id
             mask_idx = np.argwhere(mask).ravel()
 
-            nrow = mask_idx[-1] - mask_idx[0] + 1
+            tab_subset = tab.selectrows(rownrs=mask_idx)
 
-            data = tab.getcolslice(
-                data_colname,
-                blc=[0, 0],
-                trc=[-1, -1],
-                incr=[1, 1],
-                startrow=mask_idx[0],
-                nrow=nrow,
-            )
-            data = data[..., mask_idx - mask_idx[0]]
+            data = tab_subset.getcol(data_colname)
+            uv = tab_subset.getcol("UVW")[:2]
 
-            uvw = tab.getcolslice(
-                "UVW", blc=[0], trc=[1], incr=[1], startrow=mask_idx[0], nrow=nrow
-            )
-            uvw = uvw[..., mask_idx - mask_idx[0]]
         else:
             mask = np.ones_like(tab.getcol("DATA_DESC_ID")).astype(bool)
             data = tab.getcol(data_colname)
-            uvw = tab.getcol("UVW")[:2]
+            uv = tab.getcol("UVW")[:2]
 
         spectral_tab = table(str(path / "SPECTRAL_WINDOW"))
 
@@ -502,13 +539,13 @@ def from_ms(
 
             flag_mask = np.logical_not(flag_mask.astype(bool))
         else:
-            flag_mask = np.ones(uvw.shape[-1]).astype(bool)
+            flag_mask = np.ones(uv.shape[-1]).astype(bool)
 
-        uvw = uvw[..., flag_mask]
+        uv = uv[..., flag_mask]
         data = data[..., flag_mask]
 
-        u_meter = uvw[0]
-        v_meter = uvw[1]
+        u_meter = uv[0]
+        v_meter = uv[1]
 
         stokes_i = data[0] + data[1]
 
@@ -736,14 +773,22 @@ def plot_dirty_image(self, stokes_component: str = "I", **kwargs):
 
         Default is ``real``.
 
-        crop : tuple[list[float | None]], optional
-        The crop of the image. This has to have the format
-        ``([x_left, x_right], [y_left, y_right])``, where the left and right
-        values for each axis are the upper and lower limits of the axes which
-        should be shown.
-        IMPORTANT: If one supplies the ``plt.imshow`` an ``extent`` parameter
-        via the ``plot_args`` parameter, this will be the scale in which one
-        has to give the crop! If not, the crop has to be in pixels.
+        ax_unit: str | astropy.units.unit, optional
+        The unit in which to show the ticks of the x and y-axes in.
+        The y-axis is the Declination (DEC) and the x-axis is the Right Ascension (RA).
+        The latter one is defined as increasing from left to right!
+        The unit has to be given as a string or an ``astropy.units.Unit``.
+        The string must correspond to the string representation of an ``astropy.units.Unit``.
+
+        Valid units are either ``pixel`` or angle units like ``arcsec``, ``degree`` etc..
+        Default is ``pixel``.
+
+        center_pos: tuple | None, optional
+        The coordinate center of the image. The coordinates have to
+        be given in the unit defined in the parameter ``ax_unit`` above.
+        If ``ax_unit`` is set to ``pixel`` this parameter is ignored.
+        Default is ``None``, meaning the coordinates of the axes will be
+        given as relative.
 
         norm : str | matplotlib.colors.Normalize | None, optional
         The name of the norm or a matplotlib norm.

pyvisgrid/plotting/plotting.py

@@ -2,6 +2,10 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
+import warnings
+
+import astropy.units as units
+
 
 def _configure_axes(
     fig: matplotlib.figure.Figure | None,
@@ -206,7 +210,7 @@ def plot_ungridded_uv(
 
     ax.scatter(x=np.append(-u, u), y=np.append(-v, v), s=marker_size, **plot_args)
 
-    ax.axis("equal")
+    ax.set_aspect("equal", "box")
 
     ax.set_xlabel(f"$u$ in {unit}")
     ax.set_ylabel(f"$v$ in {unit}")
@@ -375,8 +379,9 @@ def plot_mask(
                 im, ax=ax, shrink=colorbar_shrink, label="$(u,v)$ per frequel in 1/fq"
             )
         case "abs":
+            mask_abs, _ = grid_data.get_mask_abs_phase()
             im = ax.imshow(
-                np.abs(grid_data.mask_real + 1j * grid_data.mask_imag),
+                mask_abs,
                 norm=norm,
                 origin="lower",
                 interpolation="none",
@@ -385,8 +390,9 @@ def plot_mask(
             )
             fig.colorbar(im, ax=ax, shrink=colorbar_shrink, label="Amplitude in a.u.")
         case "phase":
+            _, mask_phase = grid_data.get_mask_abs_phase()
             im = ax.imshow(
-                np.angle(grid_data.mask_real + 1j * grid_data.mask_imag),
+                mask_phase,
                 norm=norm,
                 origin="lower",
                 interpolation="none",
@@ -439,7 +445,8 @@ def plot_mask(
 def plot_dirty_image(
     grid_data,
     mode: str = "real",
-    crop: tuple[list[float | None]] = ([None, None], [None, None]),
+    ax_unit: str | units.Unit = "pixel",
+    center_pos: tuple[float] | None = None,
     norm: str | matplotlib.colors.Normalize = None,
     colorbar_shrink: float = 1,
     cmap: str | matplotlib.colors.Colormap | None = "inferno",
@@ -464,7 +471,6 @@ def plot_dirty_image(
 
     mode : str, optional
     The mode specifying which values of the mask should be plotted.
-
     Possible values are:
 
     - ``real``:     Plots the real part of the dirty image.
@@ -475,14 +481,22 @@ def plot_dirty_image(
 
     Default is ``real``.
 
-    crop : tuple[list[float | None]], optional
-    The crop of the image. This has to have the format
-    ``([x_left, x_right], [y_left, y_right])``, where the left and right
-    values for each axis are the upper and lower limits of the axes which
-    should be shown.
-    IMPORTANT: If one supplies the ``plt.imshow`` an ``extent`` parameter
-    via the ``plot_args`` parameter, this will be the scale in which one
-    has to give the crop! If not, the crop has to be in pixels.
+    ax_unit: str | astropy.units.unit, optional
+    The unit in which to show the ticks of the x and y-axes in.
+    The y-axis is the Declination (DEC) and the x-axis is the Right Ascension (RA).
+    The latter one is defined as increasing from left to right!
+    The unit has to be given as a string or an ``astropy.units.Unit``.
+    The string must correspond to the string representation of an ``astropy.units.Unit``.
+
+    Valid units are either ``pixel`` or angle units like ``arcsec``, ``degree`` etc..
+    Default is ``pixel``.
+
+    center_pos: tuple | None, optional
+    The coordinate center of the image. The coordinates have to
+    be given in the unit defined in the parameter ``ax_unit`` above.
+    If ``ax_unit`` is set to ``pixel`` this parameter is ignored.
+    Default is ``None``, meaning the coordinates of the axes will be
+    given as relative.
 
     norm : str | matplotlib.colors.Normalize | None, optional
     The name of the norm or a matplotlib norm.
@@ -571,19 +585,51 @@ def plot_dirty_image(
                 "The given mode does not exist! Valid modes are: real, imag, abs"
             )
 
+    unit = units.Unit(ax_unit)
+
+    if unit.physical_type == "angle":
+        img_size = dirty_image.shape[0]
+        cell_size = grid_data.fov / img_size
+
+        extent = (
+            np.array([-img_size / 2, img_size / 2] * 2) * cell_size * units.arcsecond
+        ).to(unit)
+
+        if center_pos is not None:
+            center_pos = np.array(center_pos) * unit
+            extent[:2] += center_pos[0]
+            extent[2:] += center_pos[1]
+            label_prefix = ""
+        else:
+            label_prefix = "Relative "
+
+        ax.set_xlabel(f"{label_prefix}RA in {unit}")
+        ax.set_ylabel(f"{label_prefix}DEC in {unit}")
+
+        extent = extent.value
+
+    else:
+        if unit != units.pixel:
+            warnings.warn(
+                f"The given unit {unit} is no angle unit! Using pixels instead."
+            )
+
+        extent = None
+
+        ax.set_xlabel("Pixels")
+        ax.set_ylabel("Pixels")
+
     im = ax.imshow(
         dirty_image,
         norm=norm,
         origin="lower",
         interpolation="none",
         cmap=cmap,
+        extent=extent,
         **plot_args,
     )
     fig.colorbar(im, ax=ax, shrink=colorbar_shrink, label="Flux Density in Jy/px")
 
-    ax.set_xlabel("Pixels")
-    ax.set_ylabel("Pixels")
-
     if save_to is not None:
         fig.savefig(save_to, **save_args)