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

xradio/vis/_vis_utils/_ms/{_tables/create_field_and_source_xds.py → create_field_and_source_xds.py}

@@ -13,14 +13,12 @@ from xradio.vis._vis_utils._ms.subtables import subt_rename_ids
 from xradio.vis._vis_utils._ms._tables.read import (
     convert_casacore_time_to_mjd,
     make_taql_where_between_min_max,
-    read_generic_table,
+    load_generic_table,
 )
-from xradio.vis._vis_utils._ms._tables.table_query import open_table_ro
 import graphviper.utils.logger as logger
 from xradio._utils.list_and_array import (
     check_if_consistent,
     unique_1d,
-    to_list,
     to_np_array,
 )
 from xradio._utils.common import cast_to_str, convert_to_si_units, add_position_offsets
@@ -87,6 +85,16 @@ def create_field_and_source_xds(
         f"create_field_and_source_xds() execution time {time.time() - start_time:0.2f} s"
     )
 
+    # Check if we can drop time axis. The phase centers are repeated.
+    if field_times is not None:
+        if is_single_dish:
+            center_dv = "FIELD_REFERENCE_CENTER"
+        else:
+            center_dv = "FIELD_PHASE_CENTER"
+
+        if np.unique(field_and_source_xds[center_dv], axis=0).shape[0] == 1:
+            field_and_source_xds = field_and_source_xds.isel(time=0).drop_vars("time")
+
     return field_and_source_xds, source_id
 
 
@@ -131,7 +139,7 @@ def extract_ephemeris_info(
     taql_time_range = make_taql_where_between_min_max(
         min_max_mjd, path, table_name, "MJD"
     )
-    ephemeris_xds = read_generic_table(
+    ephemeris_xds = load_generic_table(
         path, table_name, timecols=["MJD"], taql_where=taql_time_range
     )
 
@@ -141,6 +149,7 @@ def extract_ephemeris_info(
     ephemeris_xds = ephemeris_xds.isel(
         ephemeris_id=0
     )  # Collapse the ephemeris_id dimension.
+    # Data varaibles  ['time', 'RA', 'DEC', 'Rho', 'RadVel', 'NP_ang', 'NP_dist', 'DiskLong', 'DiskLat', 'Sl_lon', 'Sl_lat', 'r', 'rdot', 'phang']
 
     # Get meta data.
     ephemeris_meta = ephemeris_xds.attrs["other"]["msv2"]["ctds_attrs"]
@@ -177,9 +186,9 @@ def extract_ephemeris_info(
     temp_xds["SOURCE_POSITION"] = xr.DataArray(
         np.column_stack(
             (
-                ephemeris_xds["ra"].data,
-                ephemeris_xds["dec"].data,
-                ephemeris_xds["rho"].data,
+                ephemeris_xds["RA"].data,
+                ephemeris_xds["DEC"].data,
+                ephemeris_xds["Rho"].data,
             )
         ),
         dims=["ephem_time", "sky_pos_label"],
@@ -196,7 +205,7 @@ def extract_ephemeris_info(
 
     # Add mandatory data: SOURCE_RADIAL_VELOCITY
     temp_xds["SOURCE_RADIAL_VELOCITY"] = xr.DataArray(
-        ephemeris_xds["radvel"].data, dims=["ephem_time"]
+        ephemeris_xds["RadVel"].data, dims=["ephem_time"]
     )
     temp_xds["SOURCE_RADIAL_VELOCITY"].attrs.update(
         {
@@ -230,9 +239,9 @@ def extract_ephemeris_info(
     )  # I think the units are ['deg','deg','m'] and 'WGS84'.
 
     # Add optional data NORTH_POLE_POSITION_ANGLE and NORTH_POLE_ANGULAR_DISTANCE
-    if "np_ang" in ephemeris_xds.data_vars:
+    if "NP_ang" in ephemeris_xds.data_vars:
         temp_xds["NORTH_POLE_POSITION_ANGLE"] = xr.DataArray(
-            ephemeris_xds["np_ang"].data, dims=["ephem_time"]
+            ephemeris_xds["NP_ang"].data, dims=["ephem_time"]
         )
         temp_xds["NORTH_POLE_POSITION_ANGLE"].attrs.update(
             {
@@ -245,9 +254,9 @@ def extract_ephemeris_info(
             }
         )
 
-    if "np_dist" in ephemeris_xds.data_vars:
+    if "NP_dist" in ephemeris_xds.data_vars:
         temp_xds["NORTH_POLE_ANGULAR_DISTANCE"] = xr.DataArray(
-            ephemeris_xds["np_dist"].data, dims=["ephem_time"]
+            ephemeris_xds["NP_dist"].data, dims=["ephem_time"]
         )
         temp_xds["NORTH_POLE_ANGULAR_DISTANCE"].attrs.update(
             {
@@ -261,25 +270,25 @@ def extract_ephemeris_info(
         )
 
     # Add optional data: SUB_OBSERVER_POSITION and SUB_SOLAR_POSITION
-    if "disklong" in ephemeris_xds.data_vars:
+    if "DiskLong" in ephemris_column_description:
+        key_lon = "DiskLong"
+        key_lat = "DiskLat"
+    else:
+        key_lon = "diskLong"
+        key_lat = "diskLat"
+
+    if key_lon in ephemeris_xds.data_vars:
         temp_xds["SUB_OBSERVER_POSITION"] = xr.DataArray(
             np.column_stack(
                 (
-                    ephemeris_xds["disklong"].data,
-                    ephemeris_xds["disklat"].data,
-                    np.zeros(ephemeris_xds["disklong"].shape),
+                    ephemeris_xds[key_lon].data,
+                    ephemeris_xds[key_lat].data,
+                    np.zeros(ephemeris_xds[key_lon].shape),
                 )
             ),
             dims=["ephem_time", "ellipsoid_pos_label"],
         )
 
-        if "DiskLong" in ephemris_column_description:
-            units_key_lon = "DiskLong"
-            units_key_lat = "DiskLat"
-        else:
-            units_key_lon = "diskLong"
-            units_key_lat = "diskLat"
-
         temp_xds["SUB_OBSERVER_POSITION"].attrs.update(
             {
                 "type": "location",
@@ -288,26 +297,22 @@ def extract_ephemeris_info(
                 "coordinate_system": "planetodetic",
                 "units": [
                     cast_to_str(
-                        ephemris_column_description[units_key_lon]["keywords"][
-                            unit_keyword
-                        ]
+                        ephemris_column_description[key_lon]["keywords"][unit_keyword]
                     ),
                     cast_to_str(
-                        ephemris_column_description[units_key_lat]["keywords"][
-                            unit_keyword
-                        ]
+                        ephemris_column_description[key_lat]["keywords"][unit_keyword]
                     ),
                     "m",
                 ],
             }
         )
 
-    if "si_lon" in ephemeris_xds.data_vars:
+    if "SI_lon" in ephemeris_xds.data_vars:
         temp_xds["SUB_SOLAR_POSITION"] = xr.DataArray(
             np.column_stack(
                 (
-                    ephemeris_xds["si_lon"].data,
-                    ephemeris_xds["si_lat"].data,
+                    ephemeris_xds["SI_lon"].data,
+                    ephemeris_xds["SI_lat"].data,
                     ephemeris_xds["r"].data,
                 )
             ),
@@ -413,7 +418,7 @@ def extract_ephemeris_info(
             dims=[xds["SOURCE_POSITION"].dims[0], "sky_pos_label"],
         )
 
-    xds["FIELD_PHASE_CENTER"].attrs.update(xds["SOURCE_POSITION"].attrs)
+    xds[center_dv].attrs.update(xds["SOURCE_POSITION"].attrs)
 
     return xds
 
@@ -438,7 +443,7 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
     xds : xr.Dataset
         The xarray dataset with the added source information.
     """
-
+    coords = {}
     is_ephemeris = xds.attrs[
         "is_ephemeris"
     ]  # If ephemeris data is present we ignore the SOURCE_DIRECTION in the source table.
@@ -452,12 +457,10 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
         )  # Need to add this for ps.summary() to work.
         return xds
 
-    from xradio._utils.list_and_array import check_if_consistent, unique_1d
-
     unique_source_id = unique_1d(source_id)
     taql_where = f"where (SOURCE_ID IN [{','.join(map(str, unique_source_id))}]) AND (SPECTRAL_WINDOW_ID = {spectral_window_id})"
 
-    source_xds = read_generic_table(
+    source_xds = load_generic_table(
         path,
         "SOURCE",
         ignore=["SOURCE_MODEL"],  # Trying to read SOURCE_MODEL causes an error.
@@ -466,7 +469,7 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
 
     if len(source_xds.data_vars) == 0:  # The source xds is empty.
         logger.warning(
-            f"SOURCE table empty for source_id {source_id} and spectral_window_id {spectral_window_id}."
+            f"SOURCE table empty for (unique) source_id {unique_source_id} and spectral_window_id {spectral_window_id}."
         )
         xds = xds.assign_coords(
             {"source_name": "Unknown"}
@@ -474,36 +477,35 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
         return xds
 
     assert (
-        len(source_xds.spectral_window_id) == 1
+        len(source_xds.SPECTRAL_WINDOW_ID) == 1
     ), "Can only process source table with a single spectral_window_id for a given MSv4 partition."
 
     # This source table time is not the same as the time in the field_and_source_xds that is derived from the main MSv4 time axis.
     # The source_id maps to the time axis in the field_and_source_xds. That is why "if len(source_id) == 1" is used to check if there should be a time axis.
-    assert (
-        len(source_xds.time) == 1
+    assert len(source_xds.TIME) <= len(
+        unique_source_id
     ), "Can only process source table with a single time entry for a source_id and spectral_window_id."
 
-    source_xds = source_xds.isel(time=0, spectral_window_id=0, drop=True)
+    source_xds = source_xds.isel(TIME=0, SPECTRAL_WINDOW_ID=0, drop=True)
     source_column_description = source_xds.attrs["other"]["msv2"]["ctds_attrs"][
         "column_descriptions"
     ]
 
     # Get source name (the time axis is optional and will probably be required if the partition scheme does not include 'FIELD_ID' or 'SOURCE_ID'.).
     # Note again that this optional time axis has nothing to do with the original time axis in the source table that we drop.
-    coords = {}
     if len(source_id) == 1:
-        source_xds = source_xds.sel(source_id=source_id[0])
+        source_xds = source_xds.sel(SOURCE_ID=source_id[0])
         coords["source_name"] = (
-            source_xds["name"].values.item() + "_" + str(source_id[0])
+            source_xds["NAME"].values.item() + "_" + str(source_id[0])
         )
         direction_dims = ["sky_dir_label"]
         # coords["source_id"] = source_id[0]
     else:
-        source_xds = source_xds.sel(source_id=source_id)
+        source_xds = source_xds.sel(SOURCE_ID=source_id)
         coords["source_name"] = (
             "time",
             np.char.add(
-                source_xds["name"].data, np.char.add("_", source_id.astype(str))
+                source_xds["NAME"].data, np.char.add("_", source_id.astype(str))
             ),
         )
         direction_dims = ["time", "sky_dir_label"]
@@ -511,41 +513,67 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
 
     # If ephemeris data is present we ignore the SOURCE_DIRECTION.
     if not is_ephemeris:
+        direction_msv2_col = "DIRECTION"
         msv4_measure = column_description_casacore_to_msv4_measure(
-            source_column_description["DIRECTION"]
-        )
-        xds["SOURCE_DIRECTION"] = xr.DataArray(
-            source_xds["direction"].data, dims=direction_dims
+            source_column_description[direction_msv2_col]
         )
+
+        msv2_direction_dims = source_xds[direction_msv2_col].dims
+        if (
+            len(msv2_direction_dims) == len(direction_dims) + 1
+            and "dim_1" in msv2_direction_dims
+            and "dim_2" in msv2_direction_dims
+        ):
+            # CASA simulator produces transposed direction values, adding an
+            # unexpected dimension. Drop it (https://github.com/casangi/xradio/issues/#196)
+            direction_var = source_xds[direction_msv2_col].isel(dim_1=0, drop=True)
+        else:
+            direction_var = source_xds[direction_msv2_col]
+
+        xds["SOURCE_DIRECTION"] = xr.DataArray(direction_var.data, dims=direction_dims)
         xds["SOURCE_DIRECTION"].attrs.update(msv4_measure)
 
     # Do we have line data:
-    if source_xds["num_lines"].data.ndim == 0:
-        num_lines = np.array([source_xds["num_lines"].data.item()])
+    if source_xds["NUM_LINES"].data.ndim == 0:
+        num_lines = np.array([source_xds["NUM_LINES"].data.item()])
     else:
-        num_lines = source_xds["num_lines"].data
+        num_lines = source_xds["NUM_LINES"].data
 
     if any(num_lines > 0):
 
+        # Transition is an optional column and occasionally not populated
+        if "TRANSITION" in source_xds.data_vars:
+            transition_var_data = source_xds["TRANSITION"].data
+        else:
+            transition_var_data = np.zeros(source_xds["DIRECTION"].shape, dtype="str")
+
+        # if TRANSITION is left empty (or otherwise incomplete), and num_lines > 1,
+        # the data_vars expect a "num_lines" size in the last dimension
+        vars_shape = transition_var_data.shape[:-1] + (np.max(num_lines),)
+        if transition_var_data.shape == vars_shape:
+            coords_lines_data = transition_var_data
+        else:
+            coords_lines_data = np.broadcast_to(
+                transition_var_data, max(transition_var_data.shape, vars_shape)
+            )
+
         if len(source_id) == 1:
-            coords_lines = {"line_name": source_xds["transition"].data}
+            coords_lines = {"line_name": coords_lines_data}
             xds = xds.assign_coords(coords_lines)
             line_dims = ["line_label"]
         else:
-            coords_lines = {
-                "line_name": (("time", "line_label"), source_xds["transition"].data)
-            }
+            coords_lines = {"line_name": (("time", "line_label"), coords_lines_data)}
             xds = xds.assign_coords(coords_lines)
             line_dims = ["time", "line_label"]
 
         optional_data_variables = {
-            "rest_frequency": "LINE_REST_FREQUENCY",
-            "sysvel": "LINE_SYSTEMIC_VELOCITY",
+            "REST_FREQUENCY": "LINE_REST_FREQUENCY",
+            "SYSVEL": "LINE_SYSTEMIC_VELOCITY",
         }
         for generic_name, msv4_name in optional_data_variables.items():
             if generic_name in source_xds:
                 msv4_measure = column_description_casacore_to_msv4_measure(
-                    source_column_description[generic_name.upper()]
+                    source_column_description[generic_name]
                 )
 
                 xds[msv4_name] = xr.DataArray(
@@ -555,7 +583,7 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
 
     # Need to add doppler info if present. Add check.
     try:
-        doppler_xds = read_generic_table(
+        doppler_xds = load_generic_table(
             path,
             "DOPPLER",
         )
@@ -593,13 +621,13 @@ def create_field_info_and_check_ephemeris(
     ephemeris_table_name : str
         The name of the ephemeris table.
     """
+    coords = {}
 
-    # Federico do know how to do this taql query?
-    unqiue_field_id = unique_1d(
+    unique_field_id = unique_1d(
         field_id
-    )  # field_ids can be repeated so that the time mapping is correct if there are multiple fields. The read_generic_table required unique field_ids.
-    taql_where = f"where (ROWID() IN [{','.join(map(str, unqiue_field_id))}])"
-    field_xds = read_generic_table(
+    )  # field_ids can be repeated so that the time mapping is correct if there are multiple fields. The load_generic_table required unique field_ids.
+    taql_where = f"where (ROWID() IN [{','.join(map(str, unique_field_id))}])"
+    field_xds = load_generic_table(
         in_file,
         "FIELD",
         rename_ids=subt_rename_ids["FIELD"],
@@ -611,9 +639,11 @@ def create_field_info_and_check_ephemeris(
     ), "Polynomial field positions not supported. Please open an issue on https://github.com/casangi/xradio/issues so that we can add support for this."
     field_xds = field_xds.isel(poly_id=0, drop=True)
     # field_xds = field_xds.assign_coords({'field_id':field_xds['field_id'].data})
-    field_xds = field_xds.assign_coords({"field_id": unqiue_field_id})
-    field_xds = field_xds.sel(field_id=field_id, drop=False)
-    source_id = to_np_array(field_xds.source_id.values)
+    field_xds = field_xds.assign_coords({"field_id": unique_field_id})
+    field_xds = field_xds.sel(
+        field_id=field_id, drop=False
+    )  # Make sure field_id match up with time axis (duplicate fields are allowed).
+    source_id = to_np_array(field_xds.SOURCE_ID.values)
 
     ephemeris_table_name = None
     ephemeris_path = None
@@ -623,10 +653,9 @@ def create_field_info_and_check_ephemeris(
     )
 
     # Need to check if ephemeris_id is present and if ephemeris table is present.
-    if "ephemeris_id" in field_xds:
-        ephemeris_id = check_if_consistent(
-            field_xds.ephemeris_id, "ephemeris_id"
-        )  # int(field_xds["ephemeris_id"].data)
+    if "EPHEMERIS_ID" in field_xds:
+        # Note: this assumes partition_scheme includes "FIELD_ID"
+        ephemeris_id = check_if_consistent(field_xds.EPHEMERIS_ID, "EPHEMERIS_ID")
 
         if ephemeris_id > -1:
             files = os.listdir(os.path.join(in_file, "FIELD"))
@@ -654,48 +683,48 @@ def create_field_info_and_check_ephemeris(
 
     if is_single_dish:
         field_data_variables = {
-            "reference_dir": "FIELD_REFERENCE_CENTER",
+            "REFERENCE_DIR": "FIELD_REFERENCE_CENTER",
         }
     else:
         field_data_variables = {
-            # "delay_dir": "FIELD_DELAY_CENTER",
-            "phase_dir": "FIELD_PHASE_CENTER",
-            # "reference_dir": "FIELD_REFERENCE_CENTER",
+            # "DELAY_DIR": "FIELD_DELAY_CENTER",
+            "PHASE_DIR": "FIELD_PHASE_CENTER",
+            # "REFERENCE_DIR": "FIELD_REFERENCE_CENTER",
         }
 
     field_measures_type = "sky_coord"
 
-    coords = {}
     coords["sky_dir_label"] = ["ra", "dec"]
     field_column_description = field_xds.attrs["other"]["msv2"]["ctds_attrs"][
         "column_descriptions"
     ]
 
-    coords = {}
     # field_times is the same as the time axis in the main MSv4 dataset and is used if more than one field is present.
     if field_times is not None:
         coords["time"] = field_times
         dims = ["time", "sky_dir_label"]
         coords["field_name"] = (
             "time",
-            np.char.add(field_xds["name"].data, np.char.add("_", field_id.astype(str))),
+            np.char.add(field_xds["NAME"].data, np.char.add("_", field_id.astype(str))),
         )
         # coords["field_id"] = ("time", field_id)
     else:
-        coords["field_name"] = field_xds["name"].values.item() + "_" + str(field_id)
+        coords["field_name"] = field_xds["NAME"].values.item() + "_" + str(field_id)
         # coords["field_id"] = field_id
         dims = ["sky_dir_label"]
 
     for generic_name, msv4_name in field_data_variables.items():
 
-        if field_xds.get("delaydir_ref") is None:
+        delay_dir_ref_col = "DelayDir_Ref"
+        if field_xds.get(delay_dir_ref_col) is None:
             delaydir_ref = None
         else:
             delaydir_ref = check_if_consistent(
-                field_xds.get("delaydir_ref"), "delaydir_ref"
+                field_xds.get(delay_dir_ref_col), delay_dir_ref_col
             )
+
         msv4_measure = column_description_casacore_to_msv4_measure(
-            field_column_description[generic_name.upper()], ref_code=delaydir_ref
+            field_column_description[generic_name], ref_code=delaydir_ref
         )
 
         field_and_source_xds[msv4_name] = xr.DataArray.from_dict(

xradio/vis/_vis_utils/_ms/descr.py

@@ -5,7 +5,7 @@ import numpy as np
 import pandas as pd
 import xarray as xr
 
-from ._tables.read import read_generic_table, read_flat_col_chunk
+from ._tables.read import load_generic_table, read_flat_col_chunk
 from ._tables.table_query import open_query, open_table_ro
 from xradio._utils.list_and_array import unique_1d
 
@@ -45,7 +45,7 @@ def describe_ms(
     ]:
         raise ValueError("invalid mode, must be summary, flat or expanded")
 
-    ddi_xds = read_generic_table(infile, "DATA_DESCRIPTION")
+    ddi_xds = load_generic_table(infile, "DATA_DESCRIPTION")
     ddis = list(ddi_xds.row.values) if rowmap is None else list(rowmap.keys())
     summary: Union[pd.DataFrame, Dict] = []
     if mode == "summary":
@@ -100,8 +100,8 @@ def populate_ms_descr(
     -------
     pd.DataFrame
     """
-    spw_ids = ddi_xds.spectral_window_id.values
-    pol_ids = ddi_xds.polarization_id.values
+    spw_ids = ddi_xds.SPECTRAL_WINDOW_ID.values
+    pol_ids = ddi_xds.POLARIZATION_ID.values
     sdf = {
         "ddi": ddi,
         "spw_id": spw_ids[ddi],
@@ -110,8 +110,8 @@ def populate_ms_descr(
     }
 
     # figure out characteristics of main table from select subtables (must all be present)
-    spw_xds = read_generic_table(infile, "SPECTRAL_WINDOW")
-    pol_xds = read_generic_table(infile, "POLARIZATION")
+    spw_xds = load_generic_table(infile, "SPECTRAL_WINDOW")
+    pol_xds = load_generic_table(infile, "POLARIZATION")
 
     if mode in ["expanded", "summary"]:
         times = (
@@ -134,8 +134,8 @@ def populate_ms_descr(
             }
         )
 
-    chans = spw_xds.num_chan.values
-    pols = pol_xds.num_corr.values
+    chans = spw_xds.NUM_CHAN.values
+    pols = pol_xds.NUM_CORR.values
     sdf.update(
         {
             "chans": (