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

xradio/vis/_vis_utils/_ms/msv4_sub_xdss.py

@@ -5,17 +5,10 @@ from typing import Tuple, Union
 import numpy as np
 import xarray as xr
 
-from .msv2_to_msv4_meta import column_description_casacore_to_msv4_measure
+from xradio._utils.schema import column_description_casacore_to_msv4_measure
 from .subtables import subt_rename_ids
 from ._tables.read import make_taql_where_between_min_max, load_generic_table
 
-from xradio._utils.list_and_array import (
-    check_if_consistent,
-    unique_1d,
-    to_list,
-    to_np_array,
-)
-
 
 def interpolate_to_time(
     xds: xr.Dataset,
@@ -63,219 +56,6 @@ def interpolate_to_time(
     return xds
 
 
-def create_ant_xds(
-    in_file: str,
-    spectral_window_id: int,
-    antenna_id: list,
-    feed_id: list,
-    telescope_name: str,
-):
-    """
-    Creates an Antenna Xarray Dataset from a MS v2 ANTENNA table.
-
-    Parameters
-    ----------
-    in_file : str
-        Input MS name.
-
-    Returns
-    -------
-    xr.Dataset
-        Antenna Xarray Dataset.
-    """
-    # generic_obs_xds = load_generic_table(
-    #     in_file,
-    #     "OBSERVATION",
-    #     taql_where=f" where (ROWID() IN [{','.join(map(str,unique_antenna_id))}])",  # order is not guaranteed
-    # )
-
-    # Dictionaries that define the conversion from MSv2 to MSv4:
-    to_new_data_variable_names = {
-        "POSITION": "ANTENNA_POSITION",
-        "OFFSET": "ANTENNA_FEED_OFFSET",
-        "DISH_DIAMETER": "ANTENNA_DISH_DIAMETER",
-    }
-    data_variable_dims = {
-        "POSITION": ["antenna_id", "cartesian_pos_label"],
-        "OFFSET": ["antenna_id", "cartesian_pos_label"],
-        "DISH_DIAMETER": ["antenna_id"],
-    }
-    to_new_coord_names = {
-        "NAME": "name",
-        "STATION": "station",
-        "MOUNT": "mount",
-        "PHASED_ARRAY_ID": "phased_array_id",
-    }
-
-    # coord_dims = {
-    #     "name": ["antenna_id"],
-    #     "station": ["antenna_id"],
-    #     "mount": ["antenna_id"],
-    #     "phased_array_id": ["antenna_id"],
-    #     "POSITION": "POSITION",
-    #     "OFFSET": "FEED_OFFSET",
-    #     "DISH_DIAMETER": "DISH_DIAMETER",
-    # }
-    # data_variable_dims = {
-    #     "POSITION": ["antenna_id", "xyz_label"],
-    #     "OFFSET": ["antenna_id", "xyz_label"],
-    #     "DISH_DIAMETER": ["antenna_id"],
-    # }
-    # to_new_coord_names = {
-    #     "NAME": "name",
-    #     "STATION": "station",
-    #     "TYPE": "type",
-    #     "MOUNT": "mount",
-    #     "PHASED_ARRAY_ID": "phased_array_id",
-    # }
-    coord_dims = {
-        "NAME": ["antenna_id"],
-        "STATION": ["antenna_id"],
-        "TYPE": ["antenna_id"],
-        "MOUNT": ["antenna_id"],
-        "PHASED_ARRAY_ID": ["antenna_id"],
-    }
-
-    # Read ANTENNA table into a Xarray Dataset.
-    unique_antenna_id = unique_1d(
-        antenna_id
-    )  # Also ensures that it is sorted otherwise TaQL will give wrong results.
-
-    generic_ant_xds = load_generic_table(
-        in_file,
-        "ANTENNA",
-        rename_ids=subt_rename_ids["ANTENNA"],
-        taql_where=f" where (ROWID() IN [{','.join(map(str,unique_antenna_id))}])",  # order is not guaranteed
-    )
-    generic_ant_xds = generic_ant_xds.assign_coords({"antenna_id": unique_antenna_id})
-    generic_ant_xds = generic_ant_xds.sel(
-        antenna_id=antenna_id, drop=False
-    )  # Make sure the antenna_id order is correct.
-
-    ant_column_description = generic_ant_xds.attrs["other"]["msv2"]["ctds_attrs"][
-        "column_descriptions"
-    ]
-
-    # ['OFFSET', 'POSITION', 'DISH_DIAMETER', 'FLAG_ROW', 'MOUNT', 'NAME', 'STATION']
-    ant_xds = xr.Dataset()
-    ant_xds = ant_xds.assign_coords(
-        {"antenna_id": antenna_id, "cartesian_pos_label": ["x", "y", "z"]}
-    )
-
-    coords = {}
-    for key in generic_ant_xds:
-        msv4_measure = column_description_casacore_to_msv4_measure(
-            ant_column_description[key.upper()]
-        )
-        if key in to_new_data_variable_names:
-            ant_xds[to_new_data_variable_names[key]] = xr.DataArray(
-                generic_ant_xds[key].data, dims=data_variable_dims[key]
-            )
-
-            if msv4_measure:
-                ant_xds[to_new_data_variable_names[key]].attrs.update(msv4_measure)
-
-            if key in ["DISH_DIAMETER"]:
-                ant_xds[to_new_data_variable_names[key]].attrs.update(
-                    {"units": ["m"], "type": "quantity"}
-                )
-
-        if key in to_new_coord_names:
-            coords[to_new_coord_names[key]] = (
-                coord_dims[key],
-                generic_ant_xds[key].data,
-            )
-
-    ant_xds["ANTENNA_FEED_OFFSET"].attrs["type"] = "earth_location_offset"
-    ant_xds["ANTENNA_FEED_OFFSET"].attrs["coordinate_system"] = "geocentric"
-    ant_xds["ANTENNA_POSITION"].attrs["coordinate_system"] = "geocentric"
-
-    # Extract feed information
-    generic_feed_xds = load_generic_table(
-        in_file,
-        "FEED",
-        rename_ids=subt_rename_ids["FEED"],
-        taql_where=f" where (ANTENNA_ID IN [{','.join(map(str, unique_antenna_id))}]) AND (FEED_ID IN [{','.join(map(str, feed_id))}])",
-    )  # Some Lofar and MeerKAT data have the spw column set to -1 so we can't use '(SPECTRAL_WINDOW_ID = {spectral_window_id})'
-
-    if "SPECTRAL_WINDOW_ID" in generic_feed_xds and not all(
-        generic_feed_xds.SPECTRAL_WINDOW_ID == -1
-    ):
-        generic_feed_xds = generic_feed_xds.where(
-            generic_feed_xds.SPECTRAL_WINDOW_ID == spectral_window_id, drop=True
-        )
-    if "row" in generic_feed_xds and len(generic_feed_xds.row) > 0:
-        # Some times the feed table is empty (this is the case with ALMA spw WVR#NOMINAL).
-        assert len(generic_feed_xds.ANTENNA_ID) == len(
-            ant_xds.antenna_id
-        ), "Can only process feed table with a single time entry for an antenna and spectral_window_id."
-        generic_feed_xds = generic_feed_xds.set_xindex(
-            "ANTENNA_ID"
-        )  # Allows for non-dimension coordinate selection.
-        generic_feed_xds = generic_feed_xds.sel(
-            ANTENNA_ID=ant_xds.antenna_id
-        )  # Make sure the antenna_id is in the same order as the xds.
-
-        num_receptors = np.ravel(generic_feed_xds.NUM_RECEPTORS)
-        num_receptors = unique_1d(num_receptors[~np.isnan(num_receptors)])
-
-        assert (
-            len(num_receptors) == 1
-        ), "The number of receptors must be constant in feed table."
-
-        feed_column_description = generic_feed_xds.attrs["other"]["msv2"]["ctds_attrs"][
-            "column_descriptions"
-        ]
-
-        to_new_data_variable_names = {
-            "BEAM_OFFSET": "BEAM_OFFSET",
-            "RECEPTOR_ANGLE": "RECEPTOR_ANGLE",
-            "POLARIZATION_TYPE": "POLARIZATION_TYPE",
-            # "pol_response": "POLARIZATION_RESPONSE", ?repeated dim creates problems.
-            "FOCUS_LENGTH": "FOCUS_LENGTH",  # optional
-            # "position": "ANTENNA_FEED_OFFSET" #Will be added to the existing position in ant_xds
-        }
-
-        data_variable_dims = {
-            "BEAM_OFFSET": ["antenna_id", "receptor_name", "sky_dir_label"],
-            "RECEPTOR_ANGLE": ["antenna_id", "receptor_name"],
-            "POLARIZATION_TYPE": ["antenna_id", "receptor_name"],
-            # "pol_response": ["antenna_id", "receptor_name", "receptor_name_"],
-            "FOCUS_LENGTH": ["antenna_id"],
-            # "position": ["antenna_id", "cartesian_pos_label"],
-        }
-
-        for key in generic_feed_xds:
-            msv4_measure = column_description_casacore_to_msv4_measure(
-                feed_column_description[key.upper()]
-            )
-            if key in to_new_data_variable_names:
-                ant_xds[to_new_data_variable_names[key]] = xr.DataArray(
-                    generic_feed_xds[key].data, dims=data_variable_dims[key]
-                )
-
-                if msv4_measure:
-                    ant_xds[to_new_data_variable_names[key]].attrs.update(msv4_measure)
-
-            if key in to_new_coord_names:
-                coords[to_new_coord_names[key]] = (
-                    coord_dims[key],
-                    generic_feed_xds[key].data,
-                )
-
-        ant_xds["ANTENNA_FEED_OFFSET"] = (
-            ant_xds["ANTENNA_FEED_OFFSET"] + generic_ant_xds["POSITION"].data
-        )
-
-        coords["receptor_name"] = np.arange(ant_xds.sizes["receptor_name"]).astype(str)
-
-    ant_xds = ant_xds.assign_coords(coords)
-
-    ant_xds.attrs["overall_telescope_name"] = telescope_name
-
-    return ant_xds
-
-
 def create_weather_xds(in_file: str):
     """
     Creates a Weather Xarray Dataset from a MS v2 WEATHER table.
@@ -341,11 +121,13 @@ def create_weather_xds(in_file: str):
     # ['ANTENNA_ID', 'TIME', 'INTERVAL', 'H2O', 'IONOS_ELECTRON',
     #  'PRESSURE', 'REL_HUMIDITY', 'TEMPERATURE', 'DEW_POINT',
     #  'WIND_DIRECTION', 'WIND_SPEED']
-    weather_xds = xr.Dataset()
-
+    weather_xds = xr.Dataset(attrs={"type": "weather"})
+    time_attrs = column_description_casacore_to_msv4_measure(
+        weather_column_description["TIME"]
+    )
     coords = {
         "station_id": generic_weather_xds["STATION_ID"].data,
-        "time": generic_weather_xds["TIME"].data,
+        "time": ("time", generic_weather_xds["TIME"].data, time_attrs),
     }
     for key in generic_weather_xds:
         msv4_measure = column_description_casacore_to_msv4_measure(
@@ -401,6 +183,7 @@ def create_weather_xds(in_file: str):
 
 def create_pointing_xds(
     in_file: str,
+    ant_xds_name_ids: xr.DataArray,
     time_min_max: Union[Tuple[np.float64, np.float64], None],
     interp_time: Union[xr.DataArray, None] = None,
 ) -> xr.Dataset:
@@ -414,6 +197,8 @@ def create_pointing_xds(
     ----------
     in_file : str
         Input MS name.
+    ant_xds_name_ids : xr.Dataset
+        antenna_name data array from antenna_xds, with name/id information
     time_min_max : tuple
         min / max times values to constrain loading (from the TIME column)
     interp_time : Union[xr.DataArray, None] (Default value = None)
@@ -440,19 +225,19 @@ def create_pointing_xds(
         # "on_source": "ON_SOURCE",   # removed
         "OVER_THE_TOP": "OVER_THE_TOP",
     }
-    time_ant_dims = ["time", "antenna_id"]
-    time_ant_dir_dims = time_ant_dims + ["direction"]
+    time_ant_dims = ["time", "antenna_name"]
+    time_ant_dir_dims = time_ant_dims + ["sky_dir_label"]
     data_variable_dims = {
-        # "name": ["time", "antenna_id"],   # removed
-        # "time_origin": ["time", "antenna_id"],   # removed?
+        # "name": ["time", "antenna_name"],   # removed
+        # "time_origin": ["time", "antenna_name"],   # removed?
         "DIRECTION": time_ant_dir_dims,
         "ENCODER": time_ant_dir_dims,
         "TARGET": time_ant_dir_dims,
         "POINTING_OFFSET": time_ant_dir_dims,
         "SOURCE_OFFSET": time_ant_dir_dims,
-        # "pointing_model_id": ["time", "antenna_id"],   # removed
-        # "tracking": ["time", "antenna_id"],   # => attribute
-        # "on_source": ["time", "antenna_id"],  # removed
+        # "pointing_model_id": ["time", "antenna_name"],   # removed
+        # "tracking": ["time", "antenna_name"],   # => attribute
+        # "on_source": ["time", "antenna_name"],  # removed
         "OVER_THE_TOP": time_ant_dims,
     }
     # Unused here
@@ -483,7 +268,7 @@ def create_pointing_xds(
         "ctds_attrs"
     ]["column_descriptions"]
 
-    pointing_xds = xr.Dataset()
+    pointing_xds = xr.Dataset(attrs={"type": "pointing"})
     for key in generic_pointing_xds:
         if key in to_new_data_variable_names:
             data_var_name = to_new_data_variable_names[key]
@@ -532,8 +317,10 @@ def create_pointing_xds(
 
     coords = {
         "time": generic_pointing_xds["TIME"].values,
-        "antenna_id": np.arange(generic_pointing_xds.sizes["ANTENNA_ID"]),
-        "direction": ["ra", "dec"],
+        "antenna_name": ant_xds_name_ids.sel(
+            antenna_id=generic_pointing_xds["ANTENNA_ID"]
+        ).data,
+        "sky_dir_label": ["ra", "dec"],
     }
     pointing_xds = pointing_xds.assign_coords(coords)
 
@@ -550,7 +337,7 @@ def create_pointing_xds(
     if move_target_as_attr:
         target = generic_pointing_xds.data_vars["TARGET"]
         pointing_xds.attrs["target"] = {
-            "dims": ["direction"],
+            "dims": ["sky_dir_label"],
             "data": target.values[0, 0].tolist(),
             "attrs": column_description_casacore_to_msv4_measure(
                 pointing_column_descriptions["TARGET"]

xradio/vis/_vis_utils/_utils/xds_helper.py

@@ -171,11 +171,13 @@ def expand_xds(xds: xr.Dataset) -> xr.Dataset:
     assert "baseline" not in xds.coords
 
     txds = xds.copy()
+
     unique_baselines, baselines = np.unique(
         [txds.baseline_ant1_id.values, txds.baseline_ant2_id.values],
         axis=1,
         return_inverse=True,
     )
+
     txds["baseline"] = xr.DataArray(baselines.astype("int32"), dims=["row"])
 
     try:
@@ -184,7 +186,6 @@ def expand_xds(xds: xr.Dataset) -> xr.Dataset:
             .unstack("row")
             .transpose("time", "baseline", ...)
         )
-
         # unstack changes type to float when it needs to introduce NaNs, so
         # we need to reset to the proper type. Avoid if possible, as the
         # astype are costly
@@ -224,8 +225,15 @@ def flatten_xds(xds: xr.Dataset) -> xr.Dataset:
         txds = xds.stack({"row": ("time", "baseline")}).transpose("row", ...)
         # compute for issue https://github.com/hainegroup/oceanspy/issues/332
         # drop=True silently does compute (or at least used to)
+
+        # Skip this step for now since on Mac nan_int=0. See issue https://github.com/casangi/xradio/issues/219
+        # txds = txds.where(
+        #     ((txds.STATE_ID != nan_int) & (txds.FIELD_ID != nan_int)).compute(),
+        #     drop=True,
+        # )  # .unify_chunks()
+
         txds = txds.where(
-            ((txds.STATE_ID != nan_int) & (txds.FIELD_ID != nan_int)).compute(),
+            ~np.isnan(txds["EXPOSURE"]).compute(),
             drop=True,
         )  # .unify_chunks()
 

xradio/vis/convert_msv2_to_processing_set.py

@@ -17,9 +17,10 @@ def convert_msv2_to_processing_set(
     pointing_chunksize: Union[Dict, float, None] = None,
     pointing_interpolate: bool = False,
     ephemeris_interpolate: bool = False,
+    phase_cal_interpolate: bool = False,
     use_table_iter: bool = False,
     compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
-    storage_backend="zarr",
+    storage_backend: str = "zarr",
     parallel: bool = False,
     overwrite: bool = False,
 ):
@@ -78,6 +79,8 @@ def convert_msv2_to_processing_set(
             + str(partition_info["SCAN_NUMBER"])
         )
 
+        # prepend '0' to ms_v4_id as needed
+        ms_v4_id = f"{ms_v4_id:0>{len(str(len(partitions) - 1))}}"
         if parallel:
             delayed_list.append(
                 dask.delayed(convert_and_write_partition)(
@@ -92,6 +95,7 @@ def convert_msv2_to_processing_set(
                     pointing_chunksize=pointing_chunksize,
                     pointing_interpolate=pointing_interpolate,
                     ephemeris_interpolate=ephemeris_interpolate,
+                    phase_cal_interpolate=phase_cal_interpolate,
                     compressor=compressor,
                     overwrite=overwrite,
                 )
@@ -109,6 +113,7 @@ def convert_msv2_to_processing_set(
                 pointing_chunksize=pointing_chunksize,
                 pointing_interpolate=pointing_interpolate,
                 ephemeris_interpolate=ephemeris_interpolate,
+                phase_cal_interpolate=phase_cal_interpolate,
                 compressor=compressor,
                 overwrite=overwrite,
             )

xradio/vis/load_processing_set.py

@@ -36,9 +36,9 @@ def load_processing_set(
     processing_set
         In memory representation of processing set (data is represented by Dask.arrays).
     """
-    from xradio._utils.zarr.common import _open_dataset, _get_ms_stores_and_file_system
+    from xradio._utils.zarr.common import _open_dataset, _get_file_system_and_items
 
-    file_system, ms_store_list = _get_ms_stores_and_file_system(ps_store)
+    file_system, ms_store_list = _get_file_system_and_items(ps_store)
 
     ps = processing_set()
     for ms_name, ms_xds_isel in sel_parms.items():

xradio/vis/read_processing_set.py

@@ -2,7 +2,7 @@ import os
 
 from ._processing_set import processing_set
 import graphviper.utils.logger as logger
-from xradio._utils.zarr.common import _open_dataset, _get_ms_stores_and_file_system
+from xradio._utils.zarr.common import _open_dataset, _get_file_system_and_items
 import s3fs
 
 
@@ -25,7 +25,7 @@ def read_processing_set(
     processing_set
         Lazy representation of processing set (data is represented by Dask.arrays).
     """
-    file_system, ms_store_list = _get_ms_stores_and_file_system(ps_store)
+    file_system, ms_store_list = _get_file_system_and_items(ps_store)
 
     ps = processing_set()
     data_group = "base"
@@ -38,7 +38,7 @@ def read_processing_set(
         data_groups = xds.attrs["data_groups"]
 
         if (obs_modes is None) or (
-            xds.attrs["partition_info"]["obs_mode"] in obs_modes
+            bool(set(xds.attrs["partition_info"]["obs_mode"]).intersection(obs_modes))
         ):
             sub_xds_dict, field_and_source_xds_dict = _read_sub_xds(
                 ms_store, file_system=file_system, data_groups=data_groups
@@ -94,6 +94,9 @@ def _read_sub_xds(ms_store, file_system, data_groups, load=False):
         xds = _open_dataset(
             os.path.join(ms_store, n), load=load, file_system=file_system
         )
+        # Skip empty tables
+        if not xds.coords and not xds.data_vars:
+            continue
         if n in field_dict.keys():
             field_and_source_xds_dict[field_dict[n]] = xds
         else: