xradio 0.0.40__py3-none-any.whl → 0.0.42__py3-none-any.whl

xradio/{vis/_vis_utils/_ms → measurement_set/_utils/_msv2}/create_field_and_source_xds.py

@@ -6,20 +6,27 @@ import numpy as np
 import xarray as xr
 
 import toolviper.utils.logger as logger
-from xradio.vis._vis_utils._ms.msv4_sub_xdss import interpolate_to_time
-from xradio.vis._vis_utils._ms.subtables import subt_rename_ids
-from xradio.vis._vis_utils._ms._tables.read import (
+from xradio.measurement_set._utils._msv2.msv4_sub_xdss import interpolate_to_time
+from xradio.measurement_set._utils._msv2.subtables import subt_rename_ids
+from xradio.measurement_set._utils._msv2._tables.read import (
     convert_casacore_time_to_mjd,
     make_taql_where_between_min_max,
     load_generic_table,
 )
-from xradio._utils.common import cast_to_str, convert_to_si_units, add_position_offsets
+from xradio._utils.list_and_array import cast_to_str
+from xradio._utils.coord_math import (
+    convert_to_si_units,
+    add_position_offsets,
+    wrap_to_pi,
+)
+
 from xradio._utils.list_and_array import (
     check_if_consistent,
     unique_1d,
     to_np_array,
 )
 from xradio._utils.schema import (
+    casacore_to_msv4_measure_type,
     column_description_casacore_to_msv4_measure,
     convert_generic_xds_to_xradio_schema,
 )
@@ -160,7 +167,7 @@ def extract_ephemeris_info(
 
     # Get meta data.
     ephemeris_meta = ephemeris_xds.attrs["other"]["msv2"]["ctds_attrs"]
-    ephemris_column_description = ephemeris_xds.attrs["other"]["msv2"]["ctds_attrs"][
+    ephemeris_column_description = ephemeris_xds.attrs["other"]["msv2"]["ctds_attrs"][
         "column_descriptions"
     ]
 
@@ -169,33 +176,45 @@ def extract_ephemeris_info(
     ), "Only geocentric observer ephemeris are supported."
 
     if "posrefsys" in ephemeris_meta:
-        sky_coord_frame = ephemeris_meta["posrefsys"].replace("ICRF/", "")
+        # Note the phase center can be given as "J2000" or "J2000.0"
+        ref_frame = (
+            ephemeris_meta["posrefsys"]
+            .replace("ICRF/", "", 1)
+            .replace("J2000.0", "J2000", 1)
+        )
+        if ref_frame in casacore_to_msv4_measure_type["direction"].get("Ref_map", {}):
+            ref_frame = casacore_to_msv4_measure_type["direction"]["Ref_map"][ref_frame]
+        else:
+            logger.debug(
+                f"Unrecognized casacore direction reference frame found in posrefsys: {ref_frame}"
+            )
+        sky_coord_frame = ref_frame.lower()
     else:
-        sky_coord_frame = "ICRS"  # We will have to just assume this.
+        sky_coord_frame = "icrs"  # We will have to just assume this.
 
     # Find out witch keyword is used for units (UNIT/QuantumUnits)
-    if "UNIT" in ephemris_column_description["RA"]["keywords"]:
+    if "UNIT" in ephemeris_column_description["RA"]["keywords"]:
         unit_keyword = "UNIT"
     else:
         unit_keyword = "QuantumUnits"
 
     temp_xds = xr.Dataset()
 
-    # Add mandatory data: OBSERVATION_POSITION
-    observation_position = [
+    # Add mandatory data: OBSERVER_POSITION
+    observer_position = [
         ephemeris_meta["GeoLong"],
         ephemeris_meta["GeoLat"],
         ephemeris_meta["GeoDist"],
     ]
-    temp_xds["OBSERVATION_POSITION"] = xr.DataArray(
-        observation_position, dims=["ellipsoid_pos_label"]
+    temp_xds["OBSERVER_POSITION"] = xr.DataArray(
+        observer_position, dims=["ellipsoid_pos_label"]
     )
-    temp_xds["OBSERVATION_POSITION"].attrs.update(
+    temp_xds["OBSERVER_POSITION"].attrs.update(
         {
             "type": "location",
             "units": ["deg", "deg", "m"],
-            "data": observation_position,
-            "ellipsoid": "WGS84",
+            "data": observer_position,
+            "frame": "WGS84",
             "origin_object_name": "Earth",
             "coordinate_system": ephemeris_meta["obsloc"].lower(),
         }
@@ -210,13 +229,13 @@ def extract_ephemeris_info(
                 ephemeris_xds["Rho"].data,
             )
         ),
-        dims=["time_ephemeris_axis", "sky_pos_label"],
+        dims=["time_ephemeris", "sky_pos_label"],
     )
     # Have to use cast_to_str because the ephemeris table units are not consistently in a list or a string.
     sky_coord_units = [
-        cast_to_str(ephemris_column_description["RA"]["keywords"][unit_keyword]),
-        cast_to_str(ephemris_column_description["DEC"]["keywords"][unit_keyword]),
-        cast_to_str(ephemris_column_description["Rho"]["keywords"][unit_keyword]),
+        cast_to_str(ephemeris_column_description["RA"]["keywords"][unit_keyword]),
+        cast_to_str(ephemeris_column_description["DEC"]["keywords"][unit_keyword]),
+        cast_to_str(ephemeris_column_description["Rho"]["keywords"][unit_keyword]),
     ]
     temp_xds["SOURCE_LOCATION"].attrs.update(
         {"type": "sky_coord", "frame": sky_coord_frame, "units": sky_coord_units}
@@ -227,16 +246,17 @@ def extract_ephemeris_info(
     # Metadata has to be fixed manually. Alternatively, issues like
     # UNIT/QuantumUnits issue could be handled in convert_generic_xds_to_xradio_schema,
     # but for now preferring not to pollute that function.
+    time_ephemeris_dim = ["time_ephemeris"]
     to_new_data_variables = {
         # mandatory: SOURCE_RADIAL_VELOCITY
-        "RadVel": ["SOURCE_RADIAL_VELOCITY", ["time_ephemeris_axis"]],
+        "RadVel": ["SOURCE_RADIAL_VELOCITY", time_ephemeris_dim],
         # optional: data NORTH_POLE_POSITION_ANGLE and NORTH_POLE_ANGULAR_DISTANCE
-        "NP_ang": ["NORTH_POLE_POSITION_ANGLE", ["time_ephemeris_axis"]],
-        "NP_dist": ["NORTH_POLE_ANGULAR_DISTANCE", ["time_ephemeris_axis"]],
+        "NP_ang": ["NORTH_POLE_POSITION_ANGLE", time_ephemeris_dim],
+        "NP_dist": ["NORTH_POLE_ANGULAR_DISTANCE", time_ephemeris_dim],
         # optional: HELIOCENTRIC_RADIAL_VELOCITY
-        "rdot": ["HELIOCENTRIC_RADIAL_VELOCITY", ["time_ephemeris_axis"]],
+        "rdot": ["HELIOCENTRIC_RADIAL_VELOCITY", time_ephemeris_dim],
         # optional: OBSERVER_PHASE_ANGLE
-        "phang": ["OBSERVER_PHASE_ANGLE", ["time_ephemeris_axis"]],
+        "phang": ["OBSERVER_PHASE_ANGLE", time_ephemeris_dim],
     }
     convert_generic_xds_to_xradio_schema(
         ephemeris_xds, temp_xds, to_new_data_variables, {}
@@ -251,7 +271,7 @@ def extract_ephemeris_info(
                     "type": "quantity",
                     "units": [
                         cast_to_str(
-                            ephemris_column_description[generic_var_name]["keywords"][
+                            ephemeris_column_description[generic_var_name]["keywords"][
                                 unit_keyword
                             ]
                         )
@@ -259,8 +279,8 @@ def extract_ephemeris_info(
                 }
             )
 
-    # Add optional data: SUB_OBSERVER_POSITION and SUB_SOLAR_POSITION
-    if "DiskLong" in ephemris_column_description:
+    # Add optional data: SUB_OBSERVER_DIRECTION and SUB_SOLAR_POSITION
+    if "DiskLong" in ephemeris_column_description:
         key_lon = "DiskLong"
         key_lat = "DiskLat"
     else:
@@ -268,7 +288,7 @@ def extract_ephemeris_info(
         key_lat = "diskLat"
 
     if key_lon in ephemeris_xds.data_vars:
-        temp_xds["SUB_OBSERVER_POSITION"] = xr.DataArray(
+        temp_xds["SUB_OBSERVER_DIRECTION"] = xr.DataArray(
             np.column_stack(
                 (
                     ephemeris_xds[key_lon].data,
@@ -276,21 +296,21 @@ def extract_ephemeris_info(
                     np.zeros(ephemeris_xds[key_lon].shape),
                 )
             ),
-            dims=["time_ephemeris_axis", "ellipsoid_pos_label"],
+            dims=["time_ephemeris", "ellipsoid_pos_label"],
         )
 
-        temp_xds["SUB_OBSERVER_POSITION"].attrs.update(
+        temp_xds["SUB_OBSERVER_DIRECTION"].attrs.update(
             {
                 "type": "location",
-                "ellipsoid": "NA",
+                "frame": "Undefined",
                 "origin_object_name": ephemeris_meta["NAME"],
                 "coordinate_system": "planetodetic",
                 "units": [
                     cast_to_str(
-                        ephemris_column_description[key_lon]["keywords"][unit_keyword]
+                        ephemeris_column_description[key_lon]["keywords"][unit_keyword]
                     ),
                     cast_to_str(
-                        ephemris_column_description[key_lat]["keywords"][unit_keyword]
+                        ephemeris_column_description[key_lat]["keywords"][unit_keyword]
                     ),
                     "m",
                 ],
@@ -306,56 +326,56 @@ def extract_ephemeris_info(
                     ephemeris_xds["r"].data,
                 )
             ),
-            dims=["time_ephemeris_axis", "ellipsoid_pos_label"],
+            dims=["time_ephemeris", "ellipsoid_pos_label"],
         )
         temp_xds["SUB_SOLAR_POSITION"].attrs.update(
             {
                 "type": "location",
-                "ellipsoid": "NA",
+                "frame": "Undefined",
                 "origin_object_name": "Sun",
                 "coordinate_system": "planetodetic",
                 "units": [
                     cast_to_str(
-                        ephemris_column_description["SI_lon"]["keywords"][unit_keyword]
+                        ephemeris_column_description["SI_lon"]["keywords"][unit_keyword]
                     ),
                     cast_to_str(
-                        ephemris_column_description["SI_lat"]["keywords"][unit_keyword]
+                        ephemeris_column_description["SI_lat"]["keywords"][unit_keyword]
                     ),
                     cast_to_str(
-                        ephemris_column_description["r"]["keywords"][unit_keyword]
+                        ephemeris_column_description["r"]["keywords"][unit_keyword]
                     ),
                 ],
             }
         )
 
-    # We are using the "time_ephemeris_axis" label because it might not match the optional time axis of the source and field info. If ephemeris_interpolate=True then rename it to time.
+    # We are using the "time_ephemeris" label because it might not match the optional time axis of the source and field info. If ephemeris_interpolate=True then rename it to time.
     coords = {
         "ellipsoid_pos_label": ["lon", "lat", "dist"],
-        "time_ephemeris_axis": ephemeris_xds["time"].data,
+        "time_ephemeris": ephemeris_xds["time"].data,
         "sky_pos_label": ["ra", "dec", "dist"],
     }
     temp_xds = temp_xds.assign_coords(coords)
     time_coord_attrs = {
         "type": "time",
         "units": ["s"],
-        "scale": "UTC",
-        "format": "UNIX",
+        "scale": "utc",
+        "format": "unix",
     }
-    temp_xds["time_ephemeris_axis"].attrs.update(time_coord_attrs)
+    temp_xds["time_ephemeris"].attrs.update(time_coord_attrs)
 
     # Convert to si units and interpolate if ephemeris_interpolate=True:
     temp_xds = convert_to_si_units(temp_xds)
     temp_xds = interpolate_to_time(
-        temp_xds, interp_time, "field_and_source_xds", time_name="time_ephemeris_axis"
+        temp_xds, interp_time, "field_and_source_xds", time_name="time_ephemeris"
     )
 
-    # If we interpolate rename the time_ephemeris_axis axis to time.
+    # If we interpolate rename the time_ephemeris axis to time.
     if interp_time is not None:
-        time_coord = {"time": ("time_ephemeris_axis", interp_time.data)}
+        time_coord = {"time": ("time_ephemeris", interp_time.data)}
         temp_xds = temp_xds.assign_coords(time_coord)
         temp_xds.coords["time"].attrs.update(time_coord_attrs)
-        temp_xds = temp_xds.swap_dims({"time_ephemeris_axis": "time"}).drop_vars(
-            "time_ephemeris_axis"
+        temp_xds = temp_xds.swap_dims({"time_ephemeris": "time"}).drop_vars(
+            "time_ephemeris"
         )
 
     xds = xr.merge([xds, temp_xds])
@@ -373,21 +393,28 @@ def extract_ephemeris_info(
             interp_time is not None
         ), 'ephemeris_interpolate must be True if there is ephemeris data and multiple fields (this will occur if "FIELD_ID" is not in partition_scheme).'
 
+        field_phase_center = wrap_to_pi(
+            xds[center_dv].values + xds["SOURCE_LOCATION"][:, 0:2].values
+        )
+        field_phase_center = np.column_stack(
+            (field_phase_center, np.zeros(xds[center_dv].values.shape[0]))
+        )
+        field_phase_center[:, -1] = (
+            field_phase_center[:, -1] + xds["SOURCE_LOCATION"][:, -1].values
+        )
+
         xds[center_dv] = xr.DataArray(
-            add_position_offsets(
-                np.column_stack(
-                    (xds[center_dv].values, np.zeros(xds[center_dv].values.shape[0]))
-                ),
-                xds["SOURCE_LOCATION"].values,
-            ),
+            field_phase_center,
             dims=[xds["SOURCE_LOCATION"].dims[0], "sky_pos_label"],
         )
     else:
+        field_phase_center = (
+            np.append(xds[center_dv].values, 0) + xds["SOURCE_LOCATION"].values
+        )
+        field_phase_center[:, 0:2] = wrap_to_pi(field_phase_center[:, 0:2])
+
         xds[center_dv] = xr.DataArray(
-            add_position_offsets(
-                np.append(xds[center_dv].values, 0),
-                xds["SOURCE_LOCATION"].values,
-            ),
+            field_phase_center,
             dims=[xds["SOURCE_LOCATION"].dims[0], "sky_pos_label"],
         )
 
@@ -459,6 +486,65 @@ def make_line_dims_and_coords(
     return line_dims, line_coords
 
 
+def pad_missing_sources(
+    source_xds: xr.Dataset, unique_source_ids: np.array
+) -> xr.Dataset:
+    """
+    In some MSs there can be source IDs referenced from the field subtable which do not exist in
+    the source table: https://github.com/casangi/xradio/issues/266
+
+    This addresses the issue by padding/filling those IDs with "Unknown"/nan values. Produces a
+    source_xds that, in addition to the information loaded for the non-missing source IDs, has
+    padding for the IDs that are missing from the input MSv2 source table.
+    This function does not need to do anything when unique_source_ids is a single value
+    (partitioning by "FIELD_ID" or othwerwise single field/source)
+
+    Parameters:
+    ----------
+    xds: xr.Dataset
+        source dataset to fix/pad missing sources
+    unique_source_ids: np.array
+        IDs of the sources included in this partition
+
+    Returns:
+    -------
+    filled_source_xds : xr.Dataset
+        source dataset with padding in the originally missing sources
+    """
+
+    # Only fill gaps in multi-source xdss. If single source_id, no need to
+    if len(unique_source_ids) <= 1:
+        return source_xds
+
+    missing_source_ids = [
+        source_id
+        for source_id in unique_source_ids
+        if source_id not in source_xds.coords["SOURCE_ID"]
+    ]
+
+    # would like to use the new-ish xr.pad, but it creates issues with indices/coords and is
+    # also not free of overheads, as it for example changes all numeric types to float64
+    missing_source_xds = xr.full_like(source_xds.isel(SOURCE_ID=0), fill_value=np.nan)
+    pad_str = "Unknown"
+    pad_str_type = "<U9"
+    for var in missing_source_xds.data_vars:
+        if np.issubdtype(missing_source_xds.data_vars[var].dtype, np.str_):
+            # Avoid truncation to length of previously loaded strings
+            missing_source_xds[var] = missing_source_xds[var].astype(
+                np.dtype(pad_str_type)
+            )
+            missing_source_xds[var] = pad_str
+
+    concat_dim = "SOURCE_ID"
+    xdss_to_concat = [source_xds]
+    for missing_id in missing_source_ids:
+        missing_source_xds[concat_dim] = missing_id
+        xdss_to_concat.append(missing_source_xds)
+    filled_source_xds = xr.concat(xdss_to_concat, concat_dim).sortby(concat_dim)
+
+    return filled_source_xds
+
+
 def extract_source_info(
     xds: xr.Dataset,
     path: str,
@@ -541,6 +627,8 @@ def extract_source_info(
         "column_descriptions"
     ]
 
+    source_xds = pad_missing_sources(source_xds, unique_source_id)
+
     # 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.
     if len(source_id) == 1:

xradio/measurement_set/_utils/_msv2/msv4_info_dicts.py

@@ -0,0 +1,203 @@
+import numpy as np
+import xarray as xr
+
+from casacore import tables
+import toolviper.utils.logger as logger
+
+from .subtables import subt_rename_ids
+from ._tables.read import load_generic_table, convert_casacore_time
+from xradio._utils.list_and_array import check_if_consistent, unique_1d, to_list
+
+
+def create_info_dicts(
+    in_file: str,
+    xds: xr.Dataset,
+    field_and_source_xds: xr.Dataset,
+    partition_info_misc_fields: dict,
+    tb_tool: tables.table,
+) -> dict:
+    """
+    For an MSv4, produces several info dicts (partition_info, processor_info,
+    observation_info). The info dicts are returned in a dictionary that
+    contains them indexed by their corresponding keys, which can be used
+    directly to update the attrs dict of an MSv4.
+
+    Parameters:
+    -----------
+    in_file: str
+       path to the input MSv2
+    xds: xr.Dataset
+       main xds of the MSv4 being converted
+    field_and_source_xds: xr.Dataset
+       field_and_source_xds subdataset
+    partition_info_misc_fiels: dict
+       dict with several scalar fields for the partition_info dict that are
+       collected while processing the main MSv4 table. Expected: scan_id,
+       obs_mode, taql_where
+    tb_tool: tables.table
+       table (query) on the main table with an MSv4 query
+
+    Returns:
+    --------
+    info_dicts: dict
+       info dicts ready to be used to update the attrs of the MSv4
+    """
+
+    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 = []
+
+    info_dicts = {}
+    info_dicts["partition_info"] = {
+        # "spectral_window_id": xds.frequency.attrs["spectral_window_id"],
+        "spectral_window_name": xds.frequency.attrs["spectral_window_name"],
+        # "field_id": to_list(unique_1d(field_id)),
+        "field_name": to_list(np.unique(field_and_source_xds.field_name.values)),
+        "polarization_setup": to_list(xds.polarization.values),
+        "scan_number": to_list(np.unique(partition_info_misc_fields["scan_id"])),
+        "source_name": to_list(np.unique(field_and_source_xds.source_name.values)),
+        # "source_id": to_list(unique_1d(source_id)),
+        "intents": partition_info_misc_fields["intents"].split(","),
+        "taql": partition_info_misc_fields["taql_where"],
+        "line_name": line_name,
+    }
+
+    observation_id = check_if_consistent(
+        tb_tool.getcol("OBSERVATION_ID"), "OBSERVATION_ID"
+    )
+    info_dicts["observation_info"] = create_observation_info(in_file, observation_id)
+
+    processor_id = check_if_consistent(tb_tool.getcol("PROCESSOR_ID"), "PROCESSOR_ID")
+    info_dicts["processor_info"] = create_processor_info(in_file, processor_id)
+
+    return info_dicts
+
+
+def create_observation_info(in_file: str, observation_id: int):
+    """
+    Makes a dict with the observation info extracted from the PROCESSOR subtable.
+
+    Parameters
+    ----------
+    in_file: str
+       path to an input MSv2
+    processor_id: int
+        processor ID for one MSv4 dataset
+
+    Returns:
+    --------
+    observation_info: dict
+        observation description ready for the MSv4 observation_info attr
+    """
+
+    generic_observation_xds = load_generic_table(
+        in_file,
+        "OBSERVATION",
+        rename_ids=subt_rename_ids["OBSERVATION"],
+        taql_where=f" where ROWID() = {observation_id}",
+    )
+
+    observation_info = {
+        "observer": [generic_observation_xds["OBSERVER"].values[0]],
+        "release_date": str(
+            convert_casacore_time(generic_observation_xds["RELEASE_DATE"].values)[0]
+        ),
+    }
+    # could just assume lower:upper case but keeping explicit dict for now
+    mandatory_fields = {"project": "PROJECT"}
+    for field_msv4, row_msv2 in mandatory_fields.items():
+        observation_info[field_msv4] = generic_observation_xds[row_msv2].values[0]
+
+    exec_block_xds = None
+    try:
+        exec_block_xds = load_generic_table(in_file, "ASDM_EXECBLOCK")
+    except ValueError as exc:
+        logger.debug(
+            "Did not find the ASDM_EXECBLOCK subtable, not loading optional fields in observation_info"
+        )
+    if exec_block_xds:
+        exec_block_info = extract_exec_block_info(exec_block_xds)
+        observation_info.update(exec_block_info)
+
+    return observation_info
+
+
+def extract_exec_block_info(exec_block_xds: xr.Dataset) -> dict:
+    """
+    Get the (optional) fields of the observation_info that come from the
+    ASDM_EXECBLOCK subtable.
+
+    Note this does not parse strings like 'session_reference':
+    '<EntityRef entityId="uid://A001/X133d/X169f" partId="X00000000" entityTypeName="OUSStatus"'
+    We might want to simplify that to 'uid://A001/X133d/X169f', but keeping the
+    full string for now, as it has additional information such as the type.
+
+    Parameters
+    ----------
+    exec_block_xds: xr.Dataset
+        raw xds read from subtable ASDM_EXECBLOCK
+
+    Returns:
+    --------
+    exec_block_info: dict
+        Execution block description ready for the MSv4 observation_info dict
+    """
+
+    optional_fields = {
+        "execution_block_id": "execBlockId",
+        "execution_block_number": "execBlockNum",
+        "execution_block_UID": "execBlockUID",
+        "session_reference": "sessionReference",
+        "observing_script": "observingScript",
+        "observing_script_UID": "observingScriptUID",
+        "observing_log": "observingLog",
+    }
+
+    exec_block_info = {}
+    for field_msv4, row_msv2 in optional_fields.items():
+        if row_msv2 in exec_block_xds.data_vars:
+            msv2_value = exec_block_xds[row_msv2].values[0]
+            if isinstance(msv2_value, np.ndarray):
+                exec_block_info[field_msv4] = ",".join([log for log in msv2_value])
+            else:
+                exec_block_info[field_msv4] = msv2_value
+
+    return exec_block_info
+
+
+def create_processor_info(in_file: str, processor_id: int):
+    """
+    Makes a dict with the processor info extracted from the PROCESSOR subtable.
+
+    Parameters
+    ----------
+    in_file: str
+        path to an input MSv2
+    processor_id: int
+        processor ID for one MSv4 dataset
+
+    Returns:
+    --------
+    processor_info: dict
+        processor description ready for the MSv4 processor_info attr
+    """
+
+    generic_processor_xds = load_generic_table(
+        in_file,
+        "PROCESSOR",
+        rename_ids=subt_rename_ids["PROCESSOR"],
+        taql_where=f" where ROWID() = {processor_id}",
+    )
+
+    # Many telescopes (ASKAP, MeerKAT, SKA-Mid, VLBI, VLBA, ngEHT) seem to
+    # produce an empty PROCESSOR subtable
+    if len(generic_processor_xds.data_vars) <= 0:
+        processor_info = {"type": "", "sub_type": ""}
+    else:
+        processor_info = {
+            "type": generic_processor_xds["TYPE"].values[0],
+            "sub_type": generic_processor_xds["SUB_TYPE"].values[0],
+        }
+
+    return processor_info