xradio 0.0.34__py3-none-any.whl → 0.0.36__py3-none-any.whl

xradio/vis/_processing_set.py

@@ -1,19 +1,57 @@
 import pandas as pd
+from xradio._utils.list_and_array import to_list
+import numbers
 
 
 class processing_set(dict):
+    """
+    A dictionary subclass representing a Processing Set (PS) that is a set of Measurement Sets v4 (MS).
+
+    This class extends the built-in `dict` class and provides additional methods for manipulating and selecting subsets of the Processing Set.
+
+    Attributes:
+        meta (dict): A dictionary containing metadata information about the Processing Set.
+
+    Methods:
+        summary(data_group="base"): Returns a summary of the Processing Set as a Pandas table.
+        get_ps_max_dims(): Returns the maximum dimension of all the MSs in the Processing Set.
+        get_ps_freq_axis(): Combines the frequency axis of all MSs.
+        sel(query:str=None, **kwargs): Selects a subset of the Processing Set based on column names and values or a Pandas query.
+        ms_sel(**kwargs): Selects a subset of the Processing Set by applying the `sel` method to each individual MS.
+        ms_isel(**kwargs): Selects a subset of the Processing Set by applying the `isel` method to each individual MS.
+    """
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.meta = {"summary": {}}
 
     def summary(self, data_group="base"):
+        """
+        Returns a summary of the Processing Set as a Pandas table.
+
+        Args:
+            data_group (str): The data group to summarize. Default is "base".
+
+        Returns:
+            pandas.DataFrame: A DataFrame containing the summary information.
+        """
         if data_group in self.meta["summary"]:
             return self.meta["summary"][data_group]
         else:
-            self.meta["summary"][data_group] = self._summary(data_group)
+            self.meta["summary"][data_group] = self._summary(data_group).sort_values(
+                by=["name"], ascending=True
+            )
             return self.meta["summary"][data_group]
 
     def get_ps_max_dims(self):
+        """
+        Returns the maximum dimension of all the MSs in the Processing Set.
+
+        For example, if the Processing Set contains two MSs with dimensions (50, 20, 30) and (10, 30, 40), the maximum dimensions will be (50, 30, 40).
+
+        Returns:
+            dict: A dictionary containing the maximum dimensions of the Processing Set.
+        """
         if "max_dims" in self.meta:
             return self.meta["max_dims"]
         else:
@@ -21,6 +59,12 @@ class processing_set(dict):
             return self.meta["max_dims"]
 
     def get_ps_freq_axis(self):
+        """
+        Combines the frequency axis of all MSs.
+
+        Returns:
+            xarray.DataArray: The frequency axis of the Processing Set.
+        """
         if "freq_axis" in self.meta:
             return self.meta["freq_axis"]
         else:
@@ -33,11 +77,14 @@ class processing_set(dict):
             "obs_mode": [],
             "shape": [],
             "polarization": [],
+            "scan_number": [],
             "spw_name": [],
             # "field_id": [],
             "field_name": [],
             # "source_id": [],
             "source_name": [],
+            # "num_lines": [],
+            "line_name": [],
             "field_coords": [],
             "start_frequency": [],
             "end_frequency": [],
@@ -52,6 +99,9 @@ class processing_set(dict):
                 value.attrs["partition_info"]["spectral_window_name"]
             )
             summary_data["polarization"].append(value.polarization.values)
+            summary_data["scan_number"].append(
+                value.attrs["partition_info"]["scan_number"]
+            )
 
             if "visibility" in value.attrs["data_groups"][data_group]:
                 data_name = value.attrs["data_groups"][data_group]["visibility"]
@@ -72,8 +122,14 @@ class processing_set(dict):
             summary_data["source_name"].append(
                 value.attrs["partition_info"]["source_name"]
             )
-            summary_data["start_frequency"].append(value["frequency"].values[0])
-            summary_data["end_frequency"].append(value["frequency"].values[-1])
+
+            summary_data["line_name"].append(value.attrs["partition_info"]["line_name"])
+
+            # summary_data["num_lines"].append(value.attrs["partition_info"]["num_lines"])
+            summary_data["start_frequency"].append(
+                to_list(value["frequency"].values)[0]
+            )
+            summary_data["end_frequency"].append(to_list(value["frequency"].values)[-1])
 
             if value[data_name].attrs["field_and_source_xds"].is_ephemeris:
                 summary_data["field_coords"].append("Ephemeris")
@@ -117,7 +173,7 @@ class processing_set(dict):
         for ms_xds in self.values():
             assert (
                 frame == ms_xds.frequency.attrs["frame"]
-            ), "Frequency reference frame not consistent in processing set."
+            ), "Frequency reference frame not consistent in Processing Set."
             if ms_xds.frequency.attrs["spectral_window_id"] not in spw_ids:
                 spw_ids.append(ms_xds.frequency.attrs["spectral_window_id"])
                 freq_axis_list.append(ms_xds.frequency)
@@ -142,19 +198,71 @@ class processing_set(dict):
     def get(self, id):
         return self[list(self.keys())[id]]
 
-    def sel(self, **kwargs):
+    def sel(self, string_exact_match: bool = True, query: str = None, **kwargs):
+        """
+        Selects a subset of the Processing Set based on column names and values or a Pandas query.
+
+        The following columns are supported: name, obs_mode, polarization, spw_name, field_name, source_name, field_coords, start_frequency, end_frequency.
+
+        This function will not apply any selection on the MS data so data will not be dropped for example if a MS has field_name=['field_0','field_10','field_08'] and ps.sel(field_name='field_0') is done the resulting MS will still have field_name=['field_0','field_10','field_08'].
+
+        Examples:
+            ps.sel(obs_mode='OBSERVE_TARGET#ON_SOURCE', polarization=['RR', 'LL']) # Select all MSs with obs_mode 'OBSERVE_TARGET#ON_SOURCE' and polarization 'RR' or 'LL'.
+            ps.sel(query='start_frequency > 100e9 AND end_frequency < 200e9') # Select all MSs with start_frequency greater than 100 GHz and less than 200 GHz.
+
+        Args:
+            query (str): A Pandas query string. Default is None.
+            string_exact_match (bool): If True, the selection will be an exact match for string and string list columns. Default is True.
+            **kwargs: Keyword arguments representing column names and values to filter the Processing Set.
+
+        Returns:
+            processing_set: The subset of the Processing Set.
+        """
         import numpy as np
 
+        # def select_rows(df, col, input_strings):
+        #     return df[df[col].apply(lambda x: any(i in x for i in input_strings))]
+
+        # def select_rows(df, col, sel, string_exact_match):
+        #     def check_selection(row_val):
+        #         if isinstance(row_val, numbers.Number) or string_exact_match:
+        #             return any(i == row_val for i in sel) #If values are numbers
+        #         return any(i in row_val for i in sel) #If values are strings
+        #     return df[df[col].apply(check_selection)]
+
+        def select_rows(df, col, sel_vals, string_exact_match):
+            def check_selection(row_val):
+                row_val = to_list(
+                    row_val
+                )  # make sure that it is a list so that we can iterate over it.
+
+                for rw in row_val:
+                    for s in sel_vals:
+                        if string_exact_match:
+                            if rw == s:
+                                return True
+                        else:
+                            if s in rw:
+                                return True
+                return False
+
+            return df[df[col].apply(check_selection)]
+
         summary_table = self.summary()
         for key, value in kwargs.items():
-            if isinstance(value, list) or isinstance(value, np.ndarray):
-                summary_table = summary_table[summary_table[key].isin(value)]
-            elif isinstance(value, slice):
+            value = to_list(value)  # make sure value is a list.
+
+            if len(value) == 1 and isinstance(value[0], slice):
                 summary_table = summary_table[
-                    summary_table[key].between(value.start, value.stop)
+                    summary_table[key].between(value[0].start, value[0].stop)
                 ]
             else:
-                summary_table = summary_table[summary_table[key] == value]
+                summary_table = select_rows(
+                    summary_table, key, value, string_exact_match
+                )
+
+        if query is not None:
+            summary_table = summary_table.query(query)
 
         sub_ps = processing_set()
         for key, val in self.items():
@@ -164,12 +272,30 @@ class processing_set(dict):
         return sub_ps
 
     def ms_sel(self, **kwargs):
+        """
+        Selects a subset of the Processing Set by applying the `sel` method to each MS.
+
+        Args:
+            **kwargs: Keyword arguments representing column names and values to filter the Processing Set.
+
+        Returns:
+            processing_set: The subset of the Processing Set.
+        """
         sub_ps = processing_set()
         for key, val in self.items():
             sub_ps[key] = val.sel(kwargs)
         return sub_ps
 
     def ms_isel(self, **kwargs):
+        """
+        Selects a subset of the Processing Set by applying the `isel` method to each MS.
+
+        Args:
+            **kwargs: Keyword arguments representing dimension names and indices to select from the Processing Set.
+
+        Returns:
+            processing_set: The subset of the Processing Set.
+        """
         sub_ps = processing_set()
         for key, val in self.items():
             sub_ps[key] = val.isel(kwargs)

xradio/vis/_vis_utils/_ms/_tables/read.py

@@ -455,6 +455,9 @@ def redimension_ms_subtable(xds: xr.Dataset, subt_name: str) -> xr.Dataset:
         "SOURCE": ["SOURCE_ID", "TIME", "SPECTRAL_WINDOW_ID"],
         "SYSCAL": ["ANTENNA_ID", "FEED_ID", "SPECTRAL_WINDOW_ID", "TIME"],
         "WEATHER": ["ANTENNA_ID", "TIME"],
+        "PHASE_CAL": ["ANTENNA_ID", "TIME", "SPECTRAL_WINDOW_ID"],
+        "GAIN_CURVE": ["ANTENNA_ID", "TIME", "SPECTRAL_WINDOW_ID"],
+        "FEED": ["ANTENNA_ID", "SPECTRAL_WINDOW_ID"],
         # added tables (MSv3 but not preent in MSv2). Build it from "EPHEMi_... tables
         # Not clear what to do about 'time' var/dim:  , "time"],
         "EPHEMERIDES": ["ephemeris_row_id", "ephemeris_id"],
@@ -645,6 +648,9 @@ def load_generic_table(
         "SOURCE",
         "SYSCAL",
         "WEATHER",
+        "PHASE_CAL",
+        "GAIN_CURVE",
+        "FEED",
     ]:
         xds = redimension_ms_subtable(xds, tname)
 
@@ -944,6 +950,9 @@ def raw_col_data_to_coords_vars(
         "SOURCE",
         "SYSCAL",
         "WEATHER",
+        "PHASE_CAL",
+        "GAIN_CURVE",
+        "FEED",
     )
     dim_prefix = "dim"
 

xradio/vis/_vis_utils/_ms/conversion.py

@@ -4,17 +4,22 @@ from .._zarr.encoding import add_encoding
 from typing import Dict, Union
 import graphviper.utils.logger as logger
 import os
+import pathlib
 
 import numpy as np
 import xarray as xr
 
 from casacore import tables
-from .msv4_sub_xdss import create_ant_xds, create_pointing_xds, create_weather_xds
+from xradio.vis._vis_utils._ms.msv4_sub_xdss import (
+    create_pointing_xds,
+    create_weather_xds,
+)
+from xradio.vis._vis_utils._ms.create_antenna_xds import create_antenna_xds
 from xradio.vis._vis_utils._ms.create_field_and_source_xds import (
     create_field_and_source_xds,
 )
+from xradio._utils.schema import column_description_casacore_to_msv4_measure
 from .msv2_to_msv4_meta import (
-    column_description_casacore_to_msv4_measure,
     create_attribute_metadata,
     col_to_data_variable_names,
     col_dims,
@@ -400,7 +405,7 @@ def calc_indx_for_row_split(tb_tool, taql_where):
 
 
 def create_coordinates(
-    xds, in_file, ddi, utime, interval, baseline_ant1_id, baseline_ant2_id
+    xds, in_file, ddi, utime, interval, baseline_ant1_id, baseline_ant2_id, scan_id
 ):
     coords = {
         "time": utime,
@@ -408,6 +413,7 @@ def create_coordinates(
         "baseline_antenna2_id": ("baseline_id", baseline_ant2_id),
         "uvw_label": ["u", "v", "w"],
         "baseline_id": np.arange(len(baseline_ant1_id)),
+        "scan_number": ("time", scan_id),
     }
 
     ddi_xds = load_generic_table(in_file, "DATA_DESCRIPTION").sel(row=ddi)
@@ -446,12 +452,12 @@ def create_coordinates(
     )
     xds.frequency.attrs.update(msv4_measure)
 
-    if (spectral_window_xds.NAME.values.item() is None) or (
-        spectral_window_xds.NAME.values.item() == "none"
-    ):
+    spw_name = spectral_window_xds.NAME.values.item()
+    if (spw_name is None) or (spw_name == "none") or (spw_name == ""):
         spw_name = "spw_" + str(spectral_window_id)
     else:
-        spw_name = spectral_window_xds.NAME.values.item()
+        # spw_name = spectral_window_xds.NAME.values.item()
+        spw_name = spw_name + "_" + str(spectral_window_id)
 
     xds.frequency.attrs["spectral_window_name"] = spw_name
     msv4_measure = column_description_casacore_to_msv4_measure(
@@ -588,7 +594,7 @@ def create_data_variables(
                 logger.debug(
                     "Time to read column " + str(col) + " : " + str(time.time() - start)
                 )
-            except:
+            except Exception as e:
                 logger.debug("Could not load column", col)
 
                 if ("WEIGHT_SPECTRUM" == col) and (
@@ -606,6 +612,18 @@ def create_data_variables(
                     )
 
 
+def add_missing_data_var_attrs(xds):
+    """Adds in attributes expected metadata that cannot be found
+    in the input MSv2. For now specifically for missing
+    single-dish/SPECTRUM metadata"""
+    data_var_names = ["SPECTRUM", "SPECTRUM_CORRECTED"]
+    for var_name in data_var_names:
+        if var_name in xds.data_vars:
+            xds.data_vars[var_name].attrs["units"] = ["Jy"]
+
+    return xds
+
+
 def get_weight(
     xds,
     col,
@@ -662,7 +680,7 @@ def create_taql_query(partition_info):
 def convert_and_write_partition(
     in_file: str,
     out_file: str,
-    ms_v4_id: int,
+    ms_v4_id: Union[int, str],
     partition_info: Dict,
     use_table_iter: bool,
     partition_scheme: str = "ddi_intent_field",
@@ -671,6 +689,7 @@ def convert_and_write_partition(
     pointing_chunksize: Union[Dict, float, None] = None,
     pointing_interpolate: bool = False,
     ephemeris_interpolate: bool = False,
+    phase_cal_interpolate: bool = False,
     compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
     storage_backend="zarr",
     overwrite: bool = False,
@@ -775,8 +794,19 @@ def convert_and_write_partition(
             else:
                 interval = interval_unique[0]
 
+            scan_id = np.full(time_baseline_shape, -42, dtype=int)
+            scan_id[tidxs, bidxs] = tb_tool.getcol("SCAN_NUMBER")
+            scan_id = np.max(scan_id, axis=1)
+
             xds = create_coordinates(
-                xds, in_file, ddi, utime, interval, baseline_ant1_id, baseline_ant2_id
+                xds,
+                in_file,
+                ddi,
+                utime,
+                interval,
+                baseline_ant1_id,
+                baseline_ant2_id,
+                scan_id,
             )
             logger.debug("Time create coordinates " + str(time.time() - start))
 
@@ -792,9 +822,11 @@ def convert_and_write_partition(
                 use_table_iter,
             )
 
-            # Add data_groups and field_info
+            # Add data_groups
             xds, is_single_dish = add_data_groups(xds)
 
+            xds = add_missing_data_var_attrs(xds)
+
             if (
                 "WEIGHT" not in xds.data_vars
             ):  # Some single dish datasets don't have WEIGHT.
@@ -811,6 +843,9 @@ def convert_and_write_partition(
 
             logger.debug("Time create data variables " + str(time.time() - start))
 
+            # To constrain the time range to load (in pointing, ephemerides, phase_cal data_vars)
+            time_min_max = find_min_max_times(tb_tool, taql_where)
+
             # Create ant_xds
             start = time.time()
             feed_id = unique_1d(
@@ -826,17 +861,26 @@ def convert_and_write_partition(
                     [xds["baseline_antenna1_id"].data, xds["baseline_antenna2_id"].data]
                 )
             )
+            if phase_cal_interpolate:
+                phase_cal_interp_time = xds.time.values
+            else:
+                phase_cal_interp_time = None
 
-            ant_xds = create_ant_xds(
+            ant_xds = create_antenna_xds(
                 in_file,
                 xds.frequency.attrs["spectral_window_id"],
                 antenna_id,
                 feed_id,
                 telescope_name,
+                time_min_max,
+                phase_cal_interp_time,
             )
 
             # Change antenna_ids to antenna_names
             xds = antenna_ids_to_names(xds, ant_xds)
+            ant_xds = ant_xds.drop_vars(
+                "antenna_id"
+            )  # No longer needed after converting to name.
 
             logger.debug("Time ant xds  " + str(time.time() - start))
 
@@ -845,9 +889,7 @@ def convert_and_write_partition(
             weather_xds = create_weather_xds(in_file)
             logger.debug("Time weather " + str(time.time() - start))
 
-            # To constrain the time range to load (in pointing, ephemerides data_vars)
-            time_min_max = find_min_max_times(tb_tool, taql_where)
-
+            # Create pointing_xds
             if with_pointing:
                 start = time.time()
                 if pointing_interpolate:
@@ -869,6 +911,7 @@ def convert_and_write_partition(
                 )
 
             start = time.time()
+            xds.attrs["type"] = "visibility"
 
             # Time and frequency should always be increasing
             if len(xds.frequency) > 1 and xds.frequency[1] - xds.frequency[0] < 0:
@@ -884,10 +927,6 @@ def convert_and_write_partition(
             else:
                 ephemeris_interp_time = None
 
-            scan_id = np.full(time_baseline_shape, -42, dtype=int)
-            scan_id[tidxs, bidxs] = tb_tool.getcol("SCAN_NUMBER")
-            scan_id = np.max(scan_id, axis=1)
-
             if "FIELD_ID" not in partition_scheme:
                 field_id = np.full(time_baseline_shape, -42, dtype=int)
                 field_id[tidxs, bidxs] = tb_tool.getcol("FIELD_ID")
@@ -901,7 +940,7 @@ def convert_and_write_partition(
             # assert len(col_unique) == 1, col_name + " is not consistent."
             # return col_unique[0]
 
-            field_and_source_xds, source_id = create_field_and_source_xds(
+            field_and_source_xds, source_id, num_lines = create_field_and_source_xds(
                 in_file,
                 field_id,
                 xds.frequency.attrs["spectral_window_id"],
@@ -935,11 +974,20 @@ def convert_and_write_partition(
 
             file_name = os.path.join(
                 out_file,
-                out_file.replace(".vis.zarr", "").replace(".zarr", "").split("/")[-1]
+                pathlib.Path(out_file)
+                .name.replace(".vis.zarr", "")
+                .replace(".zarr", "")
                 + "_"
                 + str(ms_v4_id),
             )
 
+            if "line_name" in field_and_source_xds.coords:
+                line_name = to_list(
+                    unique_1d(np.ravel(field_and_source_xds.line_name.values))
+                )
+            else:
+                line_name = []
+
             xds.attrs["partition_info"] = {
                 # "spectral_window_id": xds.frequency.attrs["spectral_window_id"],
                 "spectral_window_name": xds.frequency.attrs["spectral_window_name"],
@@ -948,11 +996,14 @@ def convert_and_write_partition(
                     np.unique(field_and_source_xds.field_name.values)
                 ),
                 # "source_id": to_list(unique_1d(source_id)),
+                "line_name": line_name,
+                "scan_number": to_list(np.unique(scan_id)),
                 "source_name": to_list(
                     np.unique(field_and_source_xds.source_name.values)
                 ),
                 "polarization_setup": to_list(xds.polarization.values),
-                "obs_mode": obs_mode,
+                "num_lines": num_lines,
+                "obs_mode": obs_mode.split(","),
                 "taql": taql_where,
             }
 
@@ -968,7 +1019,7 @@ def convert_and_write_partition(
                         mode=mode,
                     )
 
-                if with_pointing:
+                if with_pointing and len(pointing_xds.data_vars) > 1:
                     pointing_xds.to_zarr(
                         store=os.path.join(file_name, "POINTING"), mode=mode
                     )
@@ -987,69 +1038,26 @@ def convert_and_write_partition(
 
 
 def antenna_ids_to_names(xds, ant_xds):
+    ant_xds = ant_xds.set_xindex(
+        "antenna_id"
+    )  # Allows for non-dimension coordinate selection.
 
-    if ant_xds.attrs["overall_telescope_name"] in ["ALMA", "VLA", "NOEMA", "EVLA"]:
-        moving_antennas = True
-    else:
-        moving_antennas = False
-
-    if moving_antennas:
-        if "baseline_antenna1_id" in xds:  # Interferometer
-
-            baseline_ant1_name = np.core.defchararray.add(
-                ant_xds["name"].sel(antenna_id=xds["baseline_antenna1_id"]).values, "_"
-            )
-            baseline_ant1_name = np.core.defchararray.add(
-                baseline_ant1_name,
-                ant_xds["station"].sel(antenna_id=xds["baseline_antenna1_id"]).values,
-            )
-            baseline_ant2_name = np.core.defchararray.add(
-                ant_xds["name"].sel(antenna_id=xds["baseline_antenna2_id"]).values, "_"
-            )
-            baseline_ant2_name = np.core.defchararray.add(
-                baseline_ant2_name,
-                ant_xds["station"].sel(antenna_id=xds["baseline_antenna2_id"]).values,
-            )
-
-            xds["baseline_antenna1_id"] = xr.DataArray(
-                baseline_ant1_name, dims="baseline_id"
-            )
-            xds["baseline_antenna2_id"] = xr.DataArray(
-                baseline_ant2_name, dims="baseline_id"
-            )
-            xds = xds.rename(
-                {
-                    "baseline_antenna1_id": "baseline_antenna1_name",
-                    "baseline_antenna2_id": "baseline_antenna2_name",
-                }
-            )
-        else:  # Single Dish
-            antenna_name = np.core.defchararray.add(
-                ant_xds["name"].sel(antenna_id=xds["antenna_id"]).values, "_"
-            )
-            antenna_name = np.core.defchararray.add(
-                antenna_name,
-                ant_xds["station"].sel(antenna_id=xds["antenna_id"]).values,
-            )
-            xds["antenna_id"] = xr.DataArray(antenna_name, dims="baseline_id")
-            xds = xds.rename({"antenna_id": "antenna_name"})
-    else:
-        if "baseline_antenna1_id" in xds:  # Interferometer
-            xds["baseline_antenna1_id"] = ant_xds["name"].sel(
-                antenna_id=xds["baseline_antenna1_id"]
-            )
-            xds["baseline_antenna2_id"] = ant_xds["name"].sel(
-                antenna_id=xds["baseline_antenna2_id"]
-            )
-            xds = xds.rename(
-                {
-                    "baseline_antenna1_id": "baseline_antenna1_name",
-                    "baseline_antenna2_id": "baseline_antenna2_name",
-                }
-            )
-        else:  # Single Dish
-            xds["antenna_id"] = ant_xds["name"].sel(antenna_id=xds["antenna_id"])
-            xds = xds.rename({"antenna_id": "antenna_name"})
+    if "baseline_antenna1_id" in xds:  # Interferometer
+        xds["baseline_antenna1_id"] = ant_xds["antenna_name"].sel(
+            antenna_id=xds["baseline_antenna1_id"]
+        )
+        xds["baseline_antenna2_id"] = ant_xds["antenna_name"].sel(
+            antenna_id=xds["baseline_antenna2_id"]
+        )
+        xds = xds.rename(
+            {
+                "baseline_antenna1_id": "baseline_antenna1_name",
+                "baseline_antenna2_id": "baseline_antenna2_name",
+            }
+        )
+    else:  # Single Dish
+        xds["antenna_id"] = ant_xds["antenna_name"].sel(antenna_id=xds["antenna_id"])
+        xds = xds.rename({"antenna_id": "antenna_name"})
 
     return xds