xradio 0.0.34__py3-none-any.whl → 0.0.37__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 (
@@ -32,7 +30,7 @@ def create_field_and_source_xds(
     is_single_dish: bool,
     time_min_max: Tuple[np.float64, np.float64],
     ephemeris_interp_time: Union[xr.DataArray, None] = None,
-):
+) -> tuple[xr.Dataset, int]:
     """
     Create a field and source xarray dataset (xds) from the given input file, field ID, and spectral window ID.
     Data is extracted from the FIELD and SOURCE tables and if there is ephemeris data, it is also extracted.
@@ -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]
@@ -55,14 +57,16 @@ def create_field_and_source_xds(
     -------
     field_and_source_xds : xr.Dataset
         The xarray dataset containing the field and source information.
+    num_lines : int
+        Sum of num_lines for all unique sources.
     """
 
     start_time = time.time()
 
-    field_and_source_xds = xr.Dataset()
+    field_and_source_xds = xr.Dataset(attrs={"type": "field_and_source"})
 
     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 +81,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 +99,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 +123,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 +176,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 +194,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 +202,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 +244,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 +259,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 +289,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 +319,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 +344,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 +360,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 +374,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,25 +417,27 @@ 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
 
 
-def extract_source_info(xds, path, source_id, spectral_window_id):
+def extract_source_info(
+    xds: xr.Dataset, path: str, source_id: int, spectral_window_id: int
+) -> tuple[xr.Dataset, int]:
     """
     Extracts source information from the given path and adds it to the xarray dataset.
 
@@ -442,6 +456,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 +471,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 +497,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 +553,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 +584,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 +628,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 +775,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(