xradio 0.0.46__py3-none-any.whl → 0.0.48__py3-none-any.whl

xradio/measurement_set/_utils/_msv2/msv4_sub_xdss.py

@@ -127,9 +127,9 @@ def interpolate_to_time(
         xds = xds.interp(
             {time_name: interp_time.data}, method=method, assume_sorted=True
         )
-        # scan_number sneaks in as a coordinate of the main time axis, drop it
-        if "scan_number" in xds.coords:
-            xds = xds.drop_vars("scan_number")
+        # scan_name sneaks in as a coordinate of the main time axis, drop it
+        if "scan_name" in xds.coords:
+            xds = xds.drop_vars("scan_name")
         points_after = xds[time_name].size
         logger.debug(
             f"{message_prefix}: interpolating the time coordinate "
@@ -497,7 +497,7 @@ def prepare_generic_sys_cal_xds(generic_sys_cal_xds: xr.Dataset) -> xr.Dataset:
 def create_system_calibration_xds(
     in_file: str,
     main_xds_frequency: xr.DataArray,
-    ant_xds_name_ids: xr.DataArray,
+    ant_xds: xr.DataArray,
     sys_cal_interp_time: Union[xr.DataArray, None] = None,
 ):
     """
@@ -510,8 +510,8 @@ def create_system_calibration_xds(
     main_xds_frequency: xr.DataArray
         frequency array of the main xds (MSv4), containing among other things
         spectral_window_id and measures metadata
-    ant_xds_name_ids : xr.Dataset
-        antenna_name data array from antenna_xds, with name/id information
+    ant_xds : xr.Dataset
+        The antenna_xds that has information such as names, stations, etc., for coordinates
     sys_cal_interp_time: Union[xr.DataArray, None] = None,
         Time axis to interpolate the data vars to (usually main MSv4 time)
 
@@ -529,7 +529,7 @@ def create_system_calibration_xds(
             rename_ids=subt_rename_ids["SYSCAL"],
             taql_where=(
                 f" where (SPECTRAL_WINDOW_ID = {spectral_window_id})"
-                f" AND (ANTENNA_ID IN [{','.join(map(str, ant_xds_name_ids.antenna_id.values))}])"
+                f" AND (ANTENNA_ID IN [{','.join(map(str, ant_xds.antenna_id.values))}])"
             ),
         )
     except ValueError as _exc:
@@ -541,14 +541,14 @@ def create_system_calibration_xds(
 
     generic_sys_cal_xds = prepare_generic_sys_cal_xds(generic_sys_cal_xds)
 
-    mandatory_dimensions = ["antenna_name", "time_cal", "receptor_label"]
+    mandatory_dimensions = ["antenna_name", "time_system_cal", "receptor_label"]
     if "frequency" not in generic_sys_cal_xds.sizes:
         dims_all = mandatory_dimensions
     else:
-        dims_all = mandatory_dimensions + ["frequency_cal"]
+        dims_all = mandatory_dimensions + ["frequency_system_cal"]
 
     to_new_data_variables = {
-        "PHASE_DIFF": ["PHASE_DIFFERENCE", ["antenna_name", "time_cal"]],
+        "PHASE_DIFF": ["PHASE_DIFFERENCE", ["antenna_name", "time_system_cal"]],
         "TCAL": ["TCAL", dims_all],
         "TCAL_SPECTRUM": ["TCAL", dims_all],
         "TRX": ["TRX", dims_all],
@@ -564,27 +564,26 @@ def create_system_calibration_xds(
     }
 
     to_new_coords = {
-        "TIME": ["time_cal", ["time_cal"]],
+        "TIME": ["time_system_cal", ["time_system_cal"]],
         "receptor": ["receptor_label", ["receptor_label"]],
-        "frequency": ["frequency_cal", ["frequency_cal"]],
+        "frequency": ["frequency_system_cal", ["frequency_system_cal"]],
     }
 
     sys_cal_xds = xr.Dataset(attrs={"type": "system_calibration"})
-    coords = {
-        "antenna_name": ant_xds_name_ids.sel(
-            antenna_id=generic_sys_cal_xds["ANTENNA_ID"]
-        ).data,
-        "receptor_label": generic_sys_cal_xds.coords["receptor"].data,
+    ant_borrowed_coords = {
+        "antenna_name": ant_xds.coords["antenna_name"],
+        "receptor_label": ant_xds.coords["receptor_label"],
+        "polarization_type": ant_xds.coords["polarization_type"],
     }
-    sys_cal_xds = sys_cal_xds.assign_coords(coords)
+    sys_cal_xds = sys_cal_xds.assign_coords(ant_borrowed_coords)
     sys_cal_xds = convert_generic_xds_to_xradio_schema(
         generic_sys_cal_xds, sys_cal_xds, to_new_data_variables, to_new_coords
     )
 
     # Add frequency coord and its measures data, if present
-    if "frequency_cal" in dims_all:
+    if "frequency_system_cal" in dims_all:
         frequency_coord = {
-            "frequency_cal": generic_sys_cal_xds.coords["frequency"].data
+            "frequency_system_cal": generic_sys_cal_xds.coords["frequency"].data
         }
         sys_cal_xds = sys_cal_xds.assign_coords(frequency_coord)
         frequency_measure = {
@@ -592,10 +591,10 @@ def create_system_calibration_xds(
             "units": main_xds_frequency.attrs["units"],
             "observer": main_xds_frequency.attrs["observer"],
         }
-        sys_cal_xds.coords["frequency_cal"].attrs.update(frequency_measure)
+        sys_cal_xds.coords["frequency_system_cal"].attrs.update(frequency_measure)
 
     sys_cal_xds = rename_and_interpolate_to_time(
-        sys_cal_xds, "time_cal", sys_cal_interp_time, "system_calibration_xds"
+        sys_cal_xds, "time_system_cal", sys_cal_interp_time, "system_calibration_xds"
     )
 
     # correct expected types

xradio/measurement_set/load_processing_set.py

@@ -37,6 +37,7 @@ def load_processing_set(
         In memory representation of processing set (data is represented by Dask.arrays).
     """
     from xradio._utils.zarr.common import _open_dataset, _get_file_system_and_items
+    from xradio.measurement_set import MeasurementSetXds
 
     file_system, ms_store_list = _get_file_system_and_items(ps_store)
 
@@ -71,7 +72,7 @@ def load_processing_set(
                     "field_and_source_xds"
                 ] = field_and_source_xds_dict[data_group_name]
 
-        ps[ms_name] = xds
+        ps[ms_name] = MeasurementSetXds(xds)
 
     return ps
 

xradio/measurement_set/processing_set.py

@@ -110,7 +110,7 @@ class ProcessingSet(dict):
             "intents": [],
             "shape": [],
             "polarization": [],
-            "scan_number": [],
+            "scan_name": [],
             "spw_name": [],
             "field_name": [],
             "source_name": [],
@@ -129,9 +129,7 @@ class ProcessingSet(dict):
                 value.attrs["partition_info"]["spectral_window_name"]
             )
             summary_data["polarization"].append(value.polarization.values)
-            summary_data["scan_number"].append(
-                value.attrs["partition_info"]["scan_number"]
-            )
+            summary_data["scan_name"].append(value.attrs["partition_info"]["scan_name"])
             data_name = value.attrs["data_groups"][data_group]["correlated_data"]
 
             if "VISIBILITY" in data_name:
@@ -156,16 +154,27 @@ class ProcessingSet(dict):
             )
             summary_data["end_frequency"].append(to_list(value["frequency"].values)[-1])
 
-            if value[data_name].attrs["field_and_source_xds"].is_ephemeris:
+            if (
+                value[data_name].attrs["field_and_source_xds"].attrs["type"]
+                == "field_and_source_ephemeris"
+            ):
                 summary_data["field_coords"].append("Ephemeris")
+            # elif (
+            #     "time"
+            #     in value[data_name].attrs["field_and_source_xds"][center_name].coords
+            # ):
             elif (
-                "time"
-                in value[data_name].attrs["field_and_source_xds"][center_name].coords
+                value[data_name]
+                .attrs["field_and_source_xds"][center_name]["field_name"]
+                .size
+                > 1
             ):
                 summary_data["field_coords"].append("Multi-Phase-Center")
             else:
                 ra_dec_rad = (
-                    value[data_name].attrs["field_and_source_xds"][center_name].values
+                    value[data_name]
+                    .attrs["field_and_source_xds"][center_name]
+                    .values[0, :]
                 )
                 frame = (
                     value[data_name]
@@ -383,10 +392,7 @@ class ProcessingSet(dict):
 
     def get_combined_field_and_source_xds(self, data_group="base"):
         """
-        Combine the `field_and_source_xds` datasets from all Measurement Sets into a single dataset.
-
-        The combined `xarray.Dataset` will have a new dimension 'field_name', consolidating data from
-        each Measurement Set. Ephemeris data is handled separately.
+        Combine all non-ephemeris `field_and_source_xds` datasets from a Processing Set for a datagroup into a single dataset.
 
         Parameters
         ----------
@@ -395,22 +401,17 @@ class ProcessingSet(dict):
 
         Returns
         -------
-        tuple of xarray.Dataset
-            A tuple containing two `xarray.Dataset` objects:
-            - combined_field_and_source_xds: Combined dataset for standard fields.
-            - combined_ephemeris_field_and_source_xds: Combined dataset for ephemeris fields.
+        xarray.Dataset
+            combined_field_and_source_xds: Combined dataset for standard fields.
 
         Raises
         ------
         ValueError
             If the `field_and_source_xds` attribute is missing or improperly formatted in any Measurement Set.
         """
-        df = self.summary(data_group)
 
         combined_field_and_source_xds = xr.Dataset()
-        combined_ephemeris_field_and_source_xds = xr.Dataset()
         for ms_name, ms_xds in self.items():
-
             correlated_data_name = ms_xds.attrs["data_groups"][data_group][
                 "correlated_data"
             ]
@@ -421,75 +422,18 @@ class ProcessingSet(dict):
                 .copy(deep=True)
             )
 
-            if (
-                "line_name" in field_and_source_xds.coords
-            ):  # Not including line info since it is a function of spw.
-                field_and_source_xds = field_and_source_xds.drop_vars(
-                    ["LINE_REST_FREQUENCY", "LINE_SYSTEMIC_VELOCITY"], errors="ignore"
-                )
-                del field_and_source_xds["line_name"]
-                del field_and_source_xds["line_label"]
-
-            if "time" in field_and_source_xds.coords:
-                if "time" not in field_and_source_xds.field_name.dims:
-                    field_names = np.array(
-                        [field_and_source_xds.field_name.values.item()]
-                        * len(field_and_source_xds.time.values)
-                    )
-                    source_names = np.array(
-                        [field_and_source_xds.source_name.values.item()]
-                        * len(field_and_source_xds.time.values)
-                    )
-                    del field_and_source_xds["field_name"]
-                    del field_and_source_xds["source_name"]
-                    field_and_source_xds = field_and_source_xds.assign_coords(
-                        field_name=("time", field_names)
-                    )
-                    field_and_source_xds = field_and_source_xds.assign_coords(
-                        source_name=("time", source_names)
-                    )
-                field_and_source_xds = field_and_source_xds.swap_dims(
-                    {"time": "field_name"}
-                )
-                del field_and_source_xds["time"]
-            elif "time_ephemeris" in field_and_source_xds.coords:
-                if "time_ephemeris" not in field_and_source_xds.field_name.dims:
-                    field_names = np.array(
-                        [field_and_source_xds.field_name.values.item()]
-                        * len(field_and_source_xds.time_ephemeris.values)
-                    )
-                    source_names = np.array(
-                        [field_and_source_xds.source_name.values.item()]
-                        * len(field_and_source_xds.time_ephemeris.values)
-                    )
-                    del field_and_source_xds["field_name"]
-                    del field_and_source_xds["source_name"]
-                    field_and_source_xds = field_and_source_xds.assign_coords(
-                        field_name=("time_ephemeris", field_names)
-                    )
-                    field_and_source_xds = field_and_source_xds.assign_coords(
-                        source_name=("time_ephemeris", source_names)
-                    )
-                field_and_source_xds = field_and_source_xds.swap_dims(
-                    {"time_ephemeris": "field_name"}
-                )
-                del field_and_source_xds["time_ephemeris"]
-            else:
-                for dv_names in field_and_source_xds.data_vars:
-                    if "field_name" not in field_and_source_xds[dv_names].dims:
-                        field_and_source_xds[dv_names] = field_and_source_xds[
-                            dv_names
-                        ].expand_dims("field_name")
+            if not field_and_source_xds.attrs["type"] == "field_and_source_ephemeris":
 
-            if field_and_source_xds.is_ephemeris:
-                if len(combined_ephemeris_field_and_source_xds.data_vars) == 0:
-                    combined_ephemeris_field_and_source_xds = field_and_source_xds
-                else:
-                    combined_ephemeris_field_and_source_xds = xr.concat(
-                        [combined_ephemeris_field_and_source_xds, field_and_source_xds],
-                        dim="field_name",
+                if (
+                    "line_name" in field_and_source_xds.coords
+                ):  # Not including line info since it is a function of spw.
+                    field_and_source_xds = field_and_source_xds.drop_vars(
+                        ["LINE_REST_FREQUENCY", "LINE_SYSTEMIC_VELOCITY"],
+                        errors="ignore",
                     )
-            else:
+                    del field_and_source_xds["line_name"]
+                    del field_and_source_xds["line_label"]
+
                 if len(combined_field_and_source_xds.data_vars) == 0:
                     combined_field_and_source_xds = field_and_source_xds
                 else:
@@ -541,12 +485,86 @@ class ProcessingSet(dict):
                 combined_field_and_source_xds.field_name[min_index].values
             )
 
+        return combined_field_and_source_xds
+
+    def get_combined_field_and_source_xds_ephemeris(self, data_group="base"):
+        """
+        Combine all ephemeris `field_and_source_xds` datasets from a Processing Set for a datagroup into a single dataset.
+
+        Parameters
+        ----------
+        data_group : str, optional
+            The data group to process. Default is "base".
+
+        Returns
+        -------
+        xarray.Dataset
+            - combined_ephemeris_field_and_source_xds: Combined dataset for ephemeris fields.
+
+        Raises
+        ------
+        ValueError
+            If the `field_and_source_xds` attribute is missing or improperly formatted in any Measurement Set.
+        """
+
+        combined_ephemeris_field_and_source_xds = xr.Dataset()
+        for ms_name, ms_xds in self.items():
+
+            correlated_data_name = ms_xds.attrs["data_groups"][data_group][
+                "correlated_data"
+            ]
+
+            field_and_source_xds = (
+                ms_xds[correlated_data_name]
+                .attrs["field_and_source_xds"]
+                .copy(deep=True)
+            )
+
+            if field_and_source_xds.attrs["type"] == "field_and_source_ephemeris":
+
+                if (
+                    "line_name" in field_and_source_xds.coords
+                ):  # Not including line info since it is a function of spw.
+                    field_and_source_xds = field_and_source_xds.drop_vars(
+                        ["LINE_REST_FREQUENCY", "LINE_SYSTEMIC_VELOCITY"],
+                        errors="ignore",
+                    )
+                    del field_and_source_xds["line_name"]
+                    del field_and_source_xds["line_label"]
+
+                from xradio.measurement_set._utils._msv2.msv4_sub_xdss import (
+                    interpolate_to_time,
+                )
+
+                if "time_ephemeris" in field_and_source_xds:
+                    field_and_source_xds = interpolate_to_time(
+                        field_and_source_xds,
+                        field_and_source_xds.time,
+                        "field_and_source_xds",
+                        "time_ephemeris",
+                    )
+                    del field_and_source_xds["time_ephemeris"]
+                    field_and_source_xds = field_and_source_xds.rename(
+                        {"time_ephemeris": "time"}
+                    )
+
+                if "OBSERVER_POSITION" in field_and_source_xds:
+                    field_and_source_xds = field_and_source_xds.drop_vars(
+                        ["OBSERVER_POSITION"], errors="ignore"
+                    )
+
+                if len(combined_ephemeris_field_and_source_xds.data_vars) == 0:
+                    combined_ephemeris_field_and_source_xds = field_and_source_xds
+                else:
+
+                    combined_ephemeris_field_and_source_xds = xr.concat(
+                        [combined_ephemeris_field_and_source_xds, field_and_source_xds],
+                        dim="time",
+                    )
+
         if (len(combined_ephemeris_field_and_source_xds.data_vars) > 0) and (
             "FIELD_PHASE_CENTER" in combined_ephemeris_field_and_source_xds
         ):
-            combined_ephemeris_field_and_source_xds = (
-                combined_ephemeris_field_and_source_xds.drop_duplicates("field_name")
-            )
 
             from xradio._utils.coord_math import wrap_to_pi
 
@@ -556,7 +574,7 @@ class ProcessingSet(dict):
             )
             combined_ephemeris_field_and_source_xds["FIELD_OFFSET"] = xr.DataArray(
                 wrap_to_pi(offset.sel(sky_pos_label=["ra", "dec"])).values,
-                dims=["field_name", "sky_dir_label"],
+                dims=["time", "sky_dir_label"],
             )
             combined_ephemeris_field_and_source_xds["FIELD_OFFSET"].attrs = (
                 combined_ephemeris_field_and_source_xds["FIELD_PHASE_CENTER"].attrs
@@ -589,7 +607,7 @@ class ProcessingSet(dict):
                 combined_ephemeris_field_and_source_xds.field_name[min_index].values
             )
 
-        return combined_field_and_source_xds, combined_ephemeris_field_and_source_xds
+        return combined_ephemeris_field_and_source_xds
 
     def plot_phase_centers(self, label_all_fields=False, data_group="base"):
         """
@@ -615,8 +633,11 @@ class ProcessingSet(dict):
         ValueError
             If the combined datasets are empty or improperly formatted.
         """
-        combined_field_and_source_xds, combined_ephemeris_field_and_source_xds = (
-            self.get_combined_field_and_source_xds(data_group)
+        combined_field_and_source_xds = self.get_combined_field_and_source_xds(
+            data_group
+        )
+        combined_ephemeris_field_and_source_xds = (
+            self.get_combined_field_and_source_xds_ephemeris(data_group)
         )
         from matplotlib import pyplot as plt
 
@@ -669,6 +690,11 @@ class ProcessingSet(dict):
             center_field_name = combined_ephemeris_field_and_source_xds.attrs[
                 "center_field_name"
             ]
+
+            combined_ephemeris_field_and_source_xds = (
+                combined_ephemeris_field_and_source_xds.set_xindex("field_name")
+            )
+
             center_field = combined_ephemeris_field_and_source_xds.sel(
                 field_name=center_field_name
             )