Merge pull request #29 from radionets-project/fix_ms_gridding

Fix Measurement Set Gridding

Author: tgross03

docs/changes/29.bugfix.rst

@@ -0,0 +1,4 @@
+- Downgraded the pinned version for ``numpy`` in ``pyproject.toml`` to avoid potential conflicts.
+- Updated the docstring for the ``vis_data`` attribute in `GridData` to specify its shape for better clarity.
+
+- Improved data selection and masking logic in ``GridData.from_ms`` to ensure correct row selection, flag masking, and shape alignment for measurement and channel data. This also fixes the handling of flagged data and ensures that the returned arrays are correctly filtered and shaped.

pyproject.toml

@@ -33,24 +33,22 @@ classifiers = [
 requires-python = ">=3.10"
 
 dependencies = [
-    "astropy>=6.1.0",
-    "click>=8.2.1",
-    "h5py>=3.14.0",
-    "matplotlib>=3.10.5",
-    "numpy>=2.2.6",
-    "pandas>=2.3.2",
-    "python-casacore>=3.7.1",
-    "scipy>=1.15.3",
-    "toml>=0.10.2",
-    "torch>=2.8.0",
-    "tqdm>=4.67.1",
+  "astropy>=6.1.0",
+  "click>=8.2.1",
+  "h5py>=3.14.0",
+  "matplotlib>=3.10.5",
+  "numpy>=1.26.0",
+  "pandas>=2.3.2",
+  "python-casacore>=3.7.1",
+  "scipy>=1.15.3",
+  "toml>=0.10.2",
+  "torch>=2.8.0",
+  "tqdm>=4.67.1",
 ]
 
 [project.optional-dependencies]
 pyvisgen = ["pyvisgen"]
-all = [
-  "pyvisgrid[pyvisgen]",
-]
+all = ["pyvisgrid[pyvisgen]"]
 
 
 [dependency-groups]
@@ -92,8 +90,8 @@ dev = [
   "ipython",
   "jupyter",
   "pre-commit",
-  {include-group = "tests"},
-  {include-group = "docs"},
+  { include-group = "tests" },
+  { include-group = "docs" },
 ]
 
 [project.urls]
@@ -121,11 +119,11 @@ extend-exclude = ["tests"]
 
 [tool.ruff.lint]
 extend-select = [
-  "I",  # isort
-  "E", # pycodestyle
-  "F", # Pyflakes
-  "UP", # pyupgrade
-  "B", # flake8-bugbear
+  "I",   # isort
+  "E",   # pycodestyle
+  "F",   # Pyflakes
+  "UP",  # pyupgrade
+  "B",   # flake8-bugbear
   "SIM", # flake8-simplify
 ]
 ignore = ["B905"]

src/pyvisgrid/core/gridder.py

@@ -29,7 +29,7 @@ class GridData:
     Attributes
     ----------
     vis_data : numpy.ndarray
-        The ungridded visibilities.
+        The ungridded visibilities of shape ``(N_MEASUREMENTS * N_CHANNELS,)``.
     fov : float
         The size of the Field Of View of the gridded data in arcseconds.
     mask : numpy.ndarray, optional
@@ -247,7 +247,7 @@ def grid(self, stokes_component: str = "I"):
             np.fft.ifft2(np.fft.fftshift(mask_real + 1j * mask_imag))
         )
 
-        return self[stokes_component]
+        return self.stokes[stokes_component]
 
     @classmethod
     def from_pyvisgen(
@@ -468,45 +468,41 @@ def from_ms(
 
         if desc_id is not None:
             mask = main_tab.getcol("DATA_DESC_ID") == desc_id
-            mask_idx = np.argwhere(mask).ravel()
 
-            main_tab = main_tab.selectrows(rownrs=mask_idx)
-
-            data = main_tab.getcol(data_colname)
-            uv = main_tab.getcol("UVW")[:, :2]
-            times = main_tab.getcol("TIME")
+            main_tab = main_tab.selectrows(rownrs=np.argwhere(mask).ravel())
 
             ref_frequency = spectral_tab.getcell("REF_FREQUENCY", desc_id)
             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
 
+        data = main_tab.getcol(data_colname)
+        uv = main_tab.getcol("UVW")[:, :2]
+        times = main_tab.getcol("TIME")
+
         if filter_flagged:
             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 = flag_mask.reshape((flag_mask.shape[0], -1)).astype(np.uint8)
+            flag_mask = np.prod(flag_mask, axis=1)
 
             flag_mask = np.logical_not(flag_mask.astype(bool))
 
         else:
-            flag_mask = np.ones(uv.shape[-1]).astype(bool)
+            flag_mask = np.ones(uv.shape[0]).astype(bool)
 
         uv = uv[flag_mask]
         data = data[flag_mask]
+        times = times[flag_mask]
 
         u_meter = uv[:, 0]
         v_meter = uv[:, 1]
 
         stokes_i = data[..., 0] + data[..., 1]
+        stokes_i = stokes_i.T  # ensure matching shape (N_CHANNELS, N_MEASUREMENTS)
 
         # FIXME: probably some kind of difference in normalization.
         # Factor 0.5 fixes this for now. Has to be investigated.
@@ -515,7 +511,7 @@ def from_ms(
         cls = cls(
             u_meter=u_meter,
             v_meter=v_meter,
-            times=Time(times[flag_mask] / 3600 / 24, format="mjd").mjd,
+            times=Time(times / 3600 / 24, format="mjd").mjd,
             img_size=img_size,
             fov=fov,
             ref_frequency=ref_frequency,

src/pyvisgrid/plotting/plotting.py

@@ -497,8 +497,7 @@ def plot_dirty_image(
         - ``abs``:      Plot the absolute value of the dirty image.
 
         Default is ``real``.
-
-    ax_unit: str | astropy.units.unit, optional
+    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!
@@ -508,7 +507,6 @@ def plot_dirty_image(
 
         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.