Add mask type methods, Add NoneType support to pyvisgen polarization argument

Author: tgross03

pyvisgrid/core/gridder.py

@@ -61,6 +61,38 @@ class GridData:
     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__(
@@ -242,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
@@ -275,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
@@ -304,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]
 

pyvisgrid/plotting/plotting.py

@@ -379,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",
@@ -389,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",