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

xradio/vis/_vis_utils/_ms/create_field_and_source_xds.py

@@ -5,9 +5,7 @@ from typing import Tuple, Union
 import numpy as np
 import xarray as xr
 
-from xradio.vis._vis_utils._ms.msv2_to_msv4_meta import (
-    column_description_casacore_to_msv4_measure,
-)
+from xradio._utils.schema import column_description_casacore_to_msv4_measure
 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 (
@@ -46,6 +44,10 @@ def create_field_and_source_xds(
         The field ids to select.
     spectral_window_id : int
         The ID of the spectral window.
+    field_times: list
+        Time data for field. It is the same as the time axis in the main MSv4 dataset and is used if more than one field is present.
+    is_single_dish: bool
+        whether the main xds has single-dish (SPECTRUM) data
     time_min_max : Tuple[np.float64, np.float46]
         Min / max times to constrain loading (usually to the time range relevant to an MSv4)
     ephemeris_interp_time : Union[xr.DataArray, None]
@@ -62,7 +64,7 @@ def create_field_and_source_xds(
     field_and_source_xds = xr.Dataset()
 
     field_and_source_xds, ephemeris_path, ephemeris_table_name, source_id = (
-        create_field_info_and_check_ephemeris(
+        extract_field_info_and_check_ephemeris(
             field_and_source_xds, in_file, field_id, field_times, is_single_dish
         )
     )
@@ -77,7 +79,7 @@ def create_field_and_source_xds(
             ephemeris_interp_time,
         )
 
-    field_and_source_xds = extract_source_info(
+    field_and_source_xds, num_lines = extract_source_info(
         field_and_source_xds, in_file, source_id, spectral_window_id
     )
 
@@ -95,7 +97,7 @@ def create_field_and_source_xds(
         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
+    return field_and_source_xds, source_id, num_lines
 
 
 def extract_ephemeris_info(
@@ -119,7 +121,7 @@ def extract_ephemeris_info(
         The name of the ephemeris table.
     time_min_max : Tuple[np.float46, np.float64]
         Min / max times to constrain loading (usually to the time range relevant to an MSv4)
-    ephemeris_interp_time : Union[xr.DataArray, None]
+    interp_time : Union[xr.DataArray, None]
         Time axis to interpolate the data vars to (usually main MSv4 time)
 
     Returns:
@@ -172,18 +174,17 @@ def extract_ephemeris_info(
     else:
         unit_keyword = "QuantumUnits"
 
+    # 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.
     coords = {
         "ellipsoid_pos_label": ["lon", "lat", "dist"],
-        "ephem_time": ephemeris_xds[
-            "time"
-        ].data,  # We are using the "ephem_time" 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.
+        "time_ephemeris_axis": ephemeris_xds["time"].data,
         "sky_pos_label": ["ra", "dec", "dist"],
     }
 
     temp_xds = xr.Dataset()
 
-    # Add mandatory data: SOURCE_POSITION
-    temp_xds["SOURCE_POSITION"] = xr.DataArray(
+    # Add mandatory data: SOURCE_LOCATION (POSITION / sky_pos_label)
+    temp_xds["SOURCE_LOCATION"] = xr.DataArray(
         np.column_stack(
             (
                 ephemeris_xds["RA"].data,
@@ -191,7 +192,7 @@ def extract_ephemeris_info(
                 ephemeris_xds["Rho"].data,
             )
         ),
-        dims=["ephem_time", "sky_pos_label"],
+        dims=["time_ephemeris_axis", "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 = [
@@ -199,13 +200,13 @@ def extract_ephemeris_info(
         cast_to_str(ephemris_column_description["DEC"]["keywords"][unit_keyword]),
         cast_to_str(ephemris_column_description["Rho"]["keywords"][unit_keyword]),
     ]
-    temp_xds["SOURCE_POSITION"].attrs.update(
+    temp_xds["SOURCE_LOCATION"].attrs.update(
         {"type": "sky_coord", "frame": sky_coord_frame, "units": sky_coord_units}
     )
 
     # 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=["time_ephemeris_axis"]
     )
     temp_xds["SOURCE_RADIAL_VELOCITY"].attrs.update(
         {
@@ -241,7 +242,7 @@ def extract_ephemeris_info(
     # Add optional data NORTH_POLE_POSITION_ANGLE and NORTH_POLE_ANGULAR_DISTANCE
     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=["time_ephemeris_axis"]
         )
         temp_xds["NORTH_POLE_POSITION_ANGLE"].attrs.update(
             {
@@ -256,7 +257,7 @@ def extract_ephemeris_info(
 
     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=["time_ephemeris_axis"]
         )
         temp_xds["NORTH_POLE_ANGULAR_DISTANCE"].attrs.update(
             {
@@ -286,7 +287,7 @@ def extract_ephemeris_info(
                     np.zeros(ephemeris_xds[key_lon].shape),
                 )
             ),
-            dims=["ephem_time", "ellipsoid_pos_label"],
+            dims=["time_ephemeris_axis", "ellipsoid_pos_label"],
         )
 
         temp_xds["SUB_OBSERVER_POSITION"].attrs.update(
@@ -316,7 +317,7 @@ def extract_ephemeris_info(
                     ephemeris_xds["r"].data,
                 )
             ),
-            dims=["ephem_time", "ellipsoid_pos_label"],
+            dims=["time_ephemeris_axis", "ellipsoid_pos_label"],
         )
         temp_xds["SUB_SOLAR_POSITION"].attrs.update(
             {
@@ -341,7 +342,7 @@ def extract_ephemeris_info(
     # Add optional data: HELIOCENTRIC_RADIAL_VELOCITY
     if "rdot" in ephemeris_xds.data_vars:
         temp_xds["HELIOCENTRIC_RADIAL_VELOCITY"] = xr.DataArray(
-            ephemeris_xds["rdot"].data, dims=["ephem_time"]
+            ephemeris_xds["rdot"].data, dims=["time_ephemeris_axis"]
         )
         temp_xds["HELIOCENTRIC_RADIAL_VELOCITY"].attrs.update(
             {
@@ -357,7 +358,7 @@ def extract_ephemeris_info(
     # Add optional data: OBSERVER_PHASE_ANGLE
     if "phang" in ephemeris_xds.data_vars:
         temp_xds["OBSERVER_PHASE_ANGLE"] = xr.DataArray(
-            ephemeris_xds["phang"].data, dims=["ephem_time"]
+            ephemeris_xds["phang"].data, dims=["time_ephemeris_axis"]
         )
         temp_xds["OBSERVER_PHASE_ANGLE"].attrs.update(
             {
@@ -371,24 +372,33 @@ def extract_ephemeris_info(
         )
 
     temp_xds = temp_xds.assign_coords(coords)
-    temp_xds["ephem_time"].attrs.update(
-        {"type": "time", "units": ["s"], "scale": "UTC", "format": "UNIX"}
-    )
+    time_coord_attrs = {
+        "type": "time",
+        "units": ["s"],
+        "scale": "UTC",
+        "format": "UNIX",
+    }
+    temp_xds["time_ephemeris_axis"].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="ephem_time"
+        temp_xds, interp_time, "field_and_source_xds", time_name="time_ephemeris_axis"
     )
 
-    # If we interpolate rename the ephem_time axis to time.
+    # If we interpolate rename the time_ephemeris_axis axis to time.
     if interp_time is not None:
-        temp_xds = temp_xds.swap_dims({"ephem_time": "time"}).drop_vars("ephem_time")
+        time_coord = {"time": ("time_ephemeris_axis", 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"
+        )
 
     xds = xr.merge([xds, temp_xds])
 
-    # Add the SOURCE_POSITION to the FIELD_PHASE_CENTER or FIELD_REFERENCE_CENTER. Ephemeris obs: When loaded from the MSv2 field table the FIELD_REFERENCE_CENTER or FIELD_PHASE_CENTER only contain an offset from the SOURCE_POSITION.
-    # We also need to add a distance dimension to the FIELD_PHASE_CENTER or FIELD_REFERENCE_CENTER to match the SOURCE_POSITION.
+    # Add the SOURCE_LOCATION to the FIELD_PHASE_CENTER or FIELD_REFERENCE_CENTER. Ephemeris obs: When loaded from the MSv2 field table the FIELD_REFERENCE_CENTER or FIELD_PHASE_CENTER only contain an offset from the SOURCE_LOCATION.
+    # We also need to add a distance dimension to the FIELD_PHASE_CENTER or FIELD_REFERENCE_CENTER to match the SOURCE_LOCATION.
     # FIELD_PHASE_CENTER is used for interferometer data and FIELD_REFERENCE_CENTER is used for single dish data.
     if is_single_dish:
         center_dv = "FIELD_REFERENCE_CENTER"
@@ -405,20 +415,20 @@ def extract_ephemeris_info(
                 np.column_stack(
                     (xds[center_dv].values, np.zeros(xds[center_dv].values.shape[0]))
                 ),
-                xds["SOURCE_POSITION"].values,
+                xds["SOURCE_LOCATION"].values,
             ),
-            dims=[xds["SOURCE_POSITION"].dims[0], "sky_pos_label"],
+            dims=[xds["SOURCE_LOCATION"].dims[0], "sky_pos_label"],
         )
     else:
         xds[center_dv] = xr.DataArray(
             add_position_offsets(
                 np.append(xds[center_dv].values, 0),
-                xds["SOURCE_POSITION"].values,
+                xds["SOURCE_LOCATION"].values,
             ),
-            dims=[xds["SOURCE_POSITION"].dims[0], "sky_pos_label"],
+            dims=[xds["SOURCE_LOCATION"].dims[0], "sky_pos_label"],
         )
 
-    xds[center_dv].attrs.update(xds["SOURCE_POSITION"].attrs)
+    xds[center_dv].attrs.update(xds["SOURCE_LOCATION"].attrs)
 
     return xds
 
@@ -442,6 +452,8 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
     -------
     xds : xr.Dataset
         The xarray dataset with the added source information.
+    num_lines : int
+        Sum of num_lines for all unique sources extracted.
     """
     coords = {}
     is_ephemeris = xds.attrs[
@@ -455,7 +467,14 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
         xds = xds.assign_coords(
             {"source_name": "Unknown"}
         )  # Need to add this for ps.summary() to work.
-        return xds
+        return xds, 0
+
+    if not os.path.isdir(os.path.join(path, "SOURCE")):
+        logger.warning(
+            f"Could not find SOURCE table for source_id {source_id}. Source information will not be included in the field_and_source_xds."
+        )
+        xds = xds.assign_coords({"source_name": "Unknown"})
+        return xds, 0
 
     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})"
@@ -474,7 +493,7 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
         xds = xds.assign_coords(
             {"source_name": "Unknown"}
         )  # Need to add this for ps.summary() to work.
-        return xds
+        return xds, 0
 
     assert (
         len(source_xds.SPECTRAL_WINDOW_ID) == 1
@@ -530,8 +549,12 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
         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)
+        # SOURCE_LOCATION (DIRECTION / sky_dir_label)
+        xds["SOURCE_LOCATION"] = xr.DataArray(direction_var.data, dims=direction_dims)
+        location_msv4_measure = column_description_casacore_to_msv4_measure(
+            source_column_description[direction_msv2_col]
+        )
+        xds["SOURCE_LOCATION"].attrs.update(location_msv4_measure)
 
     # Do we have line data:
     if source_xds["NUM_LINES"].data.ndim == 0:
@@ -557,12 +580,19 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
                 transition_var_data, max(transition_var_data.shape, vars_shape)
             )
 
+        line_label_data = np.arange(coords_lines_data.shape[-1]).astype(str)
         if len(source_id) == 1:
-            coords_lines = {"line_name": coords_lines_data}
+            coords_lines = {
+                "line_name": ("line_label", coords_lines_data),
+                "line_label": line_label_data,
+            }
             xds = xds.assign_coords(coords_lines)
             line_dims = ["line_label"]
         else:
-            coords_lines = {"line_name": (("time", "line_label"), coords_lines_data)}
+            coords_lines = {
+                "line_name": (("time", "line_label"), coords_lines_data),
+                "line_label": line_label_data,
+            }
             xds = xds.assign_coords(coords_lines)
             line_dims = ["time", "line_label"]
 
@@ -594,10 +624,13 @@ def extract_source_info(xds, path, source_id, spectral_window_id):
         pass
 
     xds = xds.assign_coords(coords)
-    return xds
 
+    _, unique_source_ids_indices = np.unique(source_xds.SOURCE_ID, return_index=True)
 
-def create_field_info_and_check_ephemeris(
+    return xds, np.sum(num_lines[unique_source_ids_indices])
+
+
+def extract_field_info_and_check_ephemeris(
     field_and_source_xds, in_file, field_id, field_times, is_single_dish
 ):
     """
@@ -738,4 +771,13 @@ def create_field_info_and_check_ephemeris(
         field_and_source_xds[msv4_name].attrs["type"] = field_measures_type
 
     field_and_source_xds = field_and_source_xds.assign_coords(coords)
+    if "time" in field_and_source_xds:
+        time_column_description = field_xds.attrs["other"]["msv2"]["ctds_attrs"][
+            "column_descriptions"
+        ]["TIME"]
+        time_msv4_measure = column_description_casacore_to_msv4_measure(
+            time_column_description
+        )
+        field_and_source_xds.coords["time"].attrs.update(time_msv4_measure)
+
     return field_and_source_xds, ephemeris_path, ephemeris_table_name, source_id

xradio/vis/_vis_utils/_ms/msv2_to_msv4_meta.py

@@ -1,4 +1,5 @@
 import graphviper.utils.logger as logger
+from xradio._utils.schema import column_description_casacore_to_msv4_measure
 
 col_to_data_variable_names = {
     "FLOAT_DATA": "SPECTRUM",
@@ -30,111 +31,6 @@ col_to_coord_names = {
     "ANTENNA2": "baseline_ant2_id",
 }
 
-# Map casacore measures to astropy
-casacore_to_msv4_measure_type = {
-    "quanta": {
-        "type": "quantity",
-    },
-    "direction": {"type": "sky_coord", "Ref": "frame", "Ref_map": {"J2000": "fk5"}},
-    "epoch": {"type": "time", "Ref": "scale", "Ref_map": {"UTC": "utc"}},
-    "frequency": {
-        "type": "spectral_coord",
-        "Ref": "frame",
-        "Ref_map": {
-            "REST": "REST",
-            "LSRK": "LSRK",
-            "LSRD": "LSRD",
-            "BARY": "BARY",
-            "GEO": "GEO",
-            "TOPO": "TOPO",
-            "GALACTO": "GALACTO",
-            "LGROUP": "LGROUP",
-            "CMB": "CMB",
-            "Undefined": "Undefined",
-        },
-    },
-    "position": {
-        "type": "earth_location",
-        "Ref": "ellipsoid",
-        "Ref_map": {"ITRF": "GRS80"},
-    },
-    "uvw": {"type": "uvw", "Ref": "frame", "Ref_map": {"ITRF": "GRS80"}},
-    "radialvelocity": {"type": "quantity"},
-}
-
-casa_frequency_frames = [
-    "REST",
-    "LSRK",
-    "LSRD",
-    "BARY",
-    "GEO",
-    "TOPO",
-    "GALACTO",
-    "LGROUP",
-    "CMB",
-    "Undefined",
-]
-
-casa_frequency_frames_codes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 64]
-
-
-def column_description_casacore_to_msv4_measure(
-    casacore_column_description, ref_code=None, time_format="UNIX"
-):
-    import numpy as np
-
-    msv4_measure = {}
-    if "MEASINFO" in casacore_column_description["keywords"]:
-        measinfo = casacore_column_description["keywords"]["MEASINFO"]
-
-        # Get conversion information
-        msv4_measure_conversion = casacore_to_msv4_measure_type[measinfo["type"]]
-
-        # Convert type, copy unit
-        msv4_measure["type"] = msv4_measure_conversion["type"]
-        msv4_measure["units"] = list(
-            casacore_column_description["keywords"]["QuantumUnits"]
-        )
-
-        # Reference frame to convert?
-        if "Ref" in msv4_measure_conversion:
-            # Find reference frame
-            if "TabRefCodes" in measinfo:
-                ref_index = np.where(measinfo["TabRefCodes"] == ref_code)[0][0]
-                casa_ref = measinfo["TabRefTypes"][ref_index]
-            elif "Ref" in measinfo:
-                casa_ref = measinfo["Ref"]
-            elif measinfo["type"] == "frequency":
-                # Some MSv2 don't have the "TabRefCodes".
-                ref_index = np.where(casa_frequency_frames_codes == ref_code)[0][0]
-                casa_ref = casa_frequency_frames[ref_index]
-            else:
-                logger.debug(
-                    f"Could not determine {measinfo['type']} measure "
-                    "reference frame!"
-                )
-
-            # Convert into MSv4 representation of reference frame, warn if unknown
-            if casa_ref in msv4_measure_conversion.get("Ref_map", {}):
-                casa_ref = msv4_measure_conversion["Ref_map"][casa_ref]
-            else:
-                logger.debug(
-                    f"Unknown reference frame for {measinfo['type']} "
-                    f"measure, using verbatim: {casa_ref}"
-                )
-
-            msv4_measure[msv4_measure_conversion["Ref"]] = casa_ref
-
-        if msv4_measure["type"] == "time":
-            msv4_measure["format"] = time_format
-    elif "QuantumUnits" in casacore_column_description["keywords"]:
-        msv4_measure = {
-            "type": "quantity",
-            "units": list(casacore_column_description["keywords"]["QuantumUnits"]),
-        }
-
-    return msv4_measure
-
 
 def create_attribute_metadata(col, main_column_descriptions):
     attrs_metadata = column_description_casacore_to_msv4_measure(