Merge pull request #27 from radionets-project/add_times

Change table reading package to casacore, Add times in gridder and uv plotting

Author: tgross03

.github/workflows/ci.yml

@@ -67,11 +67,6 @@ jobs:
         if: matrix.environment-type == 'mamba'
         run: micromamba list
 
-      - name: Initialize CASAData
-        run: |
-          mkdir -p ~/.casa/data/
-          python -c "import casatools;"
-
       - name: Tests
         run: |
           pytest -vv --cov --cov-branch --cov-report=xml

pyproject.toml

@@ -43,7 +43,7 @@ dependencies = [
   "toml",
   "torch",
   "tqdm",
-  "casatools",
+  "python-casacore",
 ]
 
 [project.optional-dependencies]

pyvisgrid/core/gridder.py

@@ -7,7 +7,8 @@
 import numpy as np
 from astropy.constants import c
 from astropy.io import fits
-from casatools.table import table
+from astropy.time import Time
+from casacore.tables import table
 from numpy.exceptions import AxisError
 from numpy.typing import ArrayLike
 
@@ -89,6 +90,7 @@ def __init__(
         self,
         u_meter: np.ndarray,
         v_meter: np.ndarray,
+        times: np.ndarray,
         img_size: int,
         fov: float,
         ref_frequency: float,
@@ -107,6 +109,8 @@ def __init__(
             The u coordinates in meter.
         v_meter : numpy.ndarray
             The v coordinates in meter.
+        times : numpy.ndarray
+            The times (in MJD) at which the visibilities were measured.
         img_size : int
             The size of the image in pixels.
         fov : float
@@ -131,6 +135,8 @@ def __init__(
         self.u_meter = u_meter
         self.v_meter = v_meter
 
+        self.times = Time(np.tile(times, reps=self.frequencies.size), format="mjd")
+
         self.img_size = img_size
 
         self.fov = np.deg2rad(fov / 3600)  # convert from asec to rad
@@ -281,6 +287,8 @@ def from_pyvisgen(
         u_meter = vis_data.u
         v_meter = vis_data.v
 
+        times = obs.baselines.time / 3600 / 24
+
         vis_data = vis_data.get_values()
 
         if vis_data.ndim != 7:
@@ -295,6 +303,7 @@ def from_pyvisgen(
         cls = cls(
             u_meter=u_meter.cpu().numpy(),
             v_meter=v_meter.cpu().numpy(),
+            times=times,
             img_size=img_size,
             fov=fov,
             ref_frequency=obs.ref_frequency.cpu().numpy(),
@@ -363,12 +372,6 @@ def from_fits(
             uv_colnames = dict(u=None, v=None)
 
         path = Path(path)
-
-        if not path.is_file() or path.suffix.lower() != ".fits":
-            raise FileNotFoundError(
-                f"The file {path} is not valid! You have to select a valid .fits file!"
-            )
-
         file = fits.open(path)
 
         data = file[0].data.T
@@ -391,6 +394,8 @@ def from_fits(
             u_meter = data[uv_colnames[0]].T * c.value
             v_meter = data[uv_colnames[1]].T * c.value
 
+        times = Time(data["DATE"], format="jd").mjd
+
         vis = file[0].data["DATA"]
         stokes_i = (
             (vis[..., 0, 0] + 1j * vis[..., 0, 1])
@@ -400,6 +405,7 @@ def from_fits(
         cls = cls(
             u_meter=u_meter,
             v_meter=v_meter,
+            times=times,
             img_size=img_size,
             fov=fov,
             ref_frequency=file[0].header["CRVAL4"],
@@ -454,54 +460,52 @@ def from_ms(
                 f"This measurement set does not exist under the path {path}"
             )
 
-        tab = table(str(path))
+        main_tab = table(str(path), ack=False)
+        spectral_tab = table(str(path / "SPECTRAL_WINDOW"), ack=False)
 
         data_colname = "DATA" if not use_calibrated else "MODEL_DATA"
 
         if desc_id is not None:
-            mask = tab.getcol("DATA_DESC_ID") == desc_id
+            mask = main_tab.getcol("DATA_DESC_ID") == desc_id
             mask_idx = np.argwhere(mask).ravel()
 
-            tab_subset = tab.selectrows(rownrs=mask_idx)
+            main_tab = main_tab.selectrows(rownrs=mask_idx)
 
-            data = tab_subset.getcol(data_colname)
-            uv = tab_subset.getcol("UVW")[:2]
+            data = main_tab.getcol(data_colname)
+            uv = main_tab.getcol("UVW")[:, :2]
+            times = main_tab.getcol("TIME")
 
-        else:
-            mask = np.ones_like(tab.getcol("DATA_DESC_ID")).astype(bool)
-            data = tab.getcol(data_colname)
-            uv = tab.getcol("UVW")[:2]
-
-        spectral_tab = table(str(path / "SPECTRAL_WINDOW"))
-
-        if desc_id is not None:
             ref_frequency = spectral_tab.getcell("REF_FREQUENCY", desc_id)
-        else:
-            ref_frequency = spectral_tab.getcell("REF_FREQUENCY", 0)
-
-        if desc_id is not None:
             frequency_offsets = (
                 spectral_tab.getcell("CHAN_FREQ", desc_id) - ref_frequency
             )
+
         else:
+            mask = np.ones_like(main_tab.getcol("DATA_DESC_ID")).astype(bool)
+            data = main_tab.getcol(data_colname)
+            uv = main_tab.getcol("UVW")[:, :2]
+            times = main_tab.getcol("TIME")
+
+            ref_frequency = spectral_tab.getcell("REF_FREQUENCY", 0)
             frequency_offsets = spectral_tab.getcell("CHAN_FREQ", 0) - ref_frequency
 
         if filter_flagged:
-            flag_mask = tab.getcol("FLAG")[..., mask]
-            flag_mask = flag_mask.reshape((-1, flag_mask.shape[-1])).astype(np.uint8)
+            flag_mask = main_tab.getcol("FLAG")
+            flag_mask = flag_mask.reshape((-1, flag_mask.shape[0])).astype(np.uint8)
             flag_mask = np.prod(flag_mask, axis=0)
 
             flag_mask = np.logical_not(flag_mask.astype(bool))
+
         else:
             flag_mask = np.ones(uv.shape[-1]).astype(bool)
 
-        uv = uv[..., flag_mask]
-        data = data[..., flag_mask]
+        uv = uv[flag_mask]
+        data = data[flag_mask]
 
-        u_meter = uv[0]
-        v_meter = uv[1]
+        u_meter = uv[:, 0]
+        v_meter = uv[:, 1]
 
-        stokes_i = data[0] + data[1]
+        stokes_i = data[..., 0] + data[..., 1]
 
         # FIXME: probably some kind of difference in normalization.
         # Factor 0.5 fixes this for now. Has to be investigated.
@@ -510,6 +514,7 @@ def from_ms(
         cls = cls(
             u_meter=u_meter,
             v_meter=v_meter,
+            times=Time(times[flag_mask] / 3600 / 24, format="mjd").mjd,
             img_size=img_size,
             fov=fov,
             ref_frequency=ref_frequency,
@@ -525,12 +530,28 @@ def plot_ungridded_uv(self, **kwargs):
 
         Parameters
         ----------
+        gridder : pyvisgrid.Gridder
+            The gridder from which to take the (u,v) coordinates.
         mode : str, optional
             The mode specifying the scale of the (u,v) coordinates.
             This can be either ``wave``, meaning the coordinates are
             plotted in units of the reference wavelength, or ``meter``,
             meaning the (u,v) coordinates will be plotted in meter.
             Default is ``wave``.
+        show_times : bool, optional
+            Whether to show the timestamps of the measured visibilities
+            as a colormap. Default is ``True``.
+        use_relative_time : bool, optional
+            Whether to show the times relative to the timestamp of the
+            first measurement in hours.
+            Default is ``True``.
+        times_cmap: str | matplotlib.colors.Colormap, optional
+            The colormap to be used for the time component of the plot.
+            Default is ``'inferno'``.
+        colorbar_shrink: float, optional
+            The shrink parameter of the colorbar. This can be needed if the plot is
+            included as a subplot to adjust the size of the colorbar.
+            Default is ``1``, meaning original scale.
         marker_size : float | None, optional
             The size of the scatter markers in points**2.
             Default is ``None``, meaning the default value supplied by