Merge remote-tracking branch 'origin/333_visibilityxds_scan_number' i…

…nto review_schema_changes

src/xradio/measurement_set/_utils/_msv2/conversion.py

@@ -423,7 +423,7 @@ def create_coordinates(
         "baseline_antenna1_id": ("baseline_id", baseline_ant1_id),
         "baseline_antenna2_id": ("baseline_id", baseline_ant2_id),
         "baseline_id": np.arange(len(baseline_ant1_id)),
-        "scan_number": ("time", scan_id),
+        "scan_name": ("time", scan_id.astype(str)),
         "uvw_label": ["u", "v", "w"],
     }
 
@@ -1029,7 +1029,7 @@ def get_observation_info(in_file, observation_id, intents):
                         datetime.timezone.utc
                     ).isoformat(),
                     "xradio_version": importlib.metadata.version("xradio"),
-                    "schema_version": "4.0.-9991",
+                    "schema_version": "4.0.-9989",
                     "type": "visibility",
                 }
             )
@@ -1112,10 +1112,6 @@ def get_observation_info(in_file, observation_id, intents):
                     [xds["baseline_antenna1_id"].data, xds["baseline_antenna2_id"].data]
                 )
             )
-            if phase_cal_interpolate:
-                phase_cal_interp_time = xds.time.values
-            else:
-                phase_cal_interp_time = None
 
             ant_xds = create_antenna_xds(
                 in_file,
@@ -1134,6 +1130,10 @@ def get_observation_info(in_file, observation_id, intents):
             logger.debug("Time gain_curve xds  " + str(time.time() - start))
 
             start = time.time()
+            if phase_cal_interpolate:
+                phase_cal_interp_time = xds.time.values
+            else:
+                phase_cal_interp_time = None
             phase_calibration_xds = create_phase_calibration_xds(
                 in_file,
                 xds.frequency.attrs["spectral_window_id"],
@@ -1143,17 +1143,6 @@ def get_observation_info(in_file, observation_id, intents):
             )
             logger.debug("Time phase_calibration xds  " + str(time.time() - start))
 
-            # Change antenna_ids to antenna_names
-            with_antenna_partitioning = "ANTENNA1" in partition_info
-            xds = antenna_ids_to_names(
-                xds, ant_xds, is_single_dish, with_antenna_partitioning
-            )
-            # but before, keep the name-id arrays, we need them for the pointing and weather xds
-            ant_xds_name_ids = ant_xds["antenna_name"].set_xindex("antenna_id")
-            ant_xds_station_name_ids = ant_xds["station"].set_xindex("antenna_id")
-            # No longer needed after converting to name.
-            ant_xds = ant_xds.drop_vars("antenna_id")
-
             # Create system_calibration_xds
             start = time.time()
             if sys_cal_interpolate:
@@ -1163,11 +1152,22 @@ def get_observation_info(in_file, observation_id, intents):
             system_calibration_xds = create_system_calibration_xds(
                 in_file,
                 xds.frequency,
-                ant_xds_name_ids,
+                ant_xds,
                 sys_cal_interp_time,
             )
             logger.debug("Time system_calibation " + str(time.time() - start))
 
+            # Change antenna_ids to antenna_names
+            with_antenna_partitioning = "ANTENNA1" in partition_info
+            xds = antenna_ids_to_names(
+                xds, ant_xds, is_single_dish, with_antenna_partitioning
+            )
+            # but before, keep the name-id arrays, we need them for the pointing and weather xds
+            ant_xds_name_ids = ant_xds["antenna_name"].set_xindex("antenna_id")
+            ant_xds_station_name_ids = ant_xds["station"].set_xindex("antenna_id")
+            # No longer needed after converting to name.
+            ant_xds = ant_xds.drop_vars("antenna_id")
+
             # Create weather_xds
             start = time.time()
             weather_xds = create_weather_xds(in_file, ant_xds_station_name_ids)
@@ -1238,7 +1238,7 @@ def get_observation_info(in_file, observation_id, intents):
             xds = fix_uvw_frame(xds, field_and_source_xds, is_single_dish)
 
             partition_info_misc_fields = {
-                "scan_id": scan_id,
+                "scan_name": xds.coords["scan_name"].data,
                 "intents": intents,
                 "taql_where": taql_where,
             }

src/xradio/measurement_set/_utils/_msv2/msv4_info_dicts.py

@@ -55,7 +55,7 @@ def create_info_dicts(
         # "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"])),
+        "scan_name": to_list(np.unique(partition_info_misc_fields["scan_name"])),
         "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(","),

src/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:
@@ -570,13 +570,12 @@ def create_system_calibration_xds(
     }
 
     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
     )

src/xradio/measurement_set/processing_set.py

@@ -110,7 +110,7 @@ def _summary(self, data_group="base"):
             "intents": [],
             "shape": [],
             "polarization": [],
-            "scan_number": [],
+            "scan_name": [],
             "spw_name": [],
             "field_name": [],
             "source_name": [],
@@ -129,9 +129,7 @@ def _summary(self, data_group="base"):
                 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:

src/xradio/measurement_set/schema.py

@@ -1207,8 +1207,8 @@ class PartitionInfoDict:
     """ List of all field names """
     polarization_setup: list[str]
     """ List of polrization bases. """
-    scan_number: list[int]
-    """ List of scan numbers. """
+    scan_name: list[str]
+    """ List of scan names. """
     source_name: list[str]
     """ List of source names. """
     # source_id: mising / remove for good?
@@ -1368,7 +1368,7 @@ class GainCurveXds:
     """ Useful when data is combined from mutiple arrays for example ACA + ALMA. """
     receptor_label: Coord[ReceptorLabel, str]
     """ Names of receptors """
-    polarization_type: Optional[Coord[tuple[AntennaName, ReceptorLabel], str]]
+    polarization_type: Coord[tuple[AntennaName, ReceptorLabel], str]
     """ Polarization type to which each receptor responds (e.g. ”R”,”L”,”X” or ”Y”).
     This is the receptor polarization type as recorded in the final correlated data (e.g. ”RR”); i.e.
     as measured after all polarization combiners. ['X','Y'], ['R','L'] """
@@ -1425,7 +1425,7 @@ class PhaseCalibrationXds:
     """ Useful when data is combined from mutiple arrays for example ACA + ALMA. """
     receptor_label: Coord[ReceptorLabel, str]
     """ Names of receptors """
-    polarization_type: Optional[Coord[tuple[AntennaName, ReceptorLabel], str]]
+    polarization_type: Coord[tuple[AntennaName, ReceptorLabel], str]
     """ Polarization type to which each receptor responds (e.g. ”R”,”L”,”X” or ”Y”).
     This is the receptor polarization type as recorded in the final correlated data (e.g. ”RR”); i.e.
     as measured after all polarization combiners. ['X','Y'], ['R','L'] """
@@ -1705,7 +1705,12 @@ class SystemCalibrationXds:
     # Coordinates
     antenna_name: Coordof[AntennaNameArray]
     """ Antenna identifier """
-    receptor_label: Coord[ReceptorLabel, numpy.int64]
+    receptor_label: Coord[ReceptorLabel, str]
+    """ Names of receptors """
+    polarization_type: Coord[tuple[AntennaName, ReceptorLabel], str]
+    """ Polarization type to which each receptor responds (e.g. ”R”,”L”,”X” or ”Y”).
+    This is the receptor polarization type as recorded in the final correlated data (e.g. ”RR”); i.e.
+    as measured after all polarization combiners. ['X','Y'], ['R','L'] """
     """  """
     time: Optional[Coordof[TimeInterpolatedCoordArray]] = None
     """ Midpoint of time for which this set of parameters is accurate. Labeled 'time' when interpolating to main time axis """
@@ -1882,8 +1887,8 @@ class VisibilityXds:
     """
     uvw_label: Optional[Coordof[UvwLabelArray]] = None
     """ u,v,w """
-    scan_number: Optional[Coord[Time, Union[numpy.int64, numpy.int32]]] = None
-    """Arbitary scan number to identify data taken in the same logical scan."""
+    scan_name: Optional[Coord[Time, str]] = None
+    """Arbitary scan name to identify data taken in the same logical scan."""
 
     # --- Optional data variables / arrays ---
 
@@ -1976,8 +1981,8 @@ class SpectrumXds:
     """
     If the polarizations are not constant over baseline
     """
-    scan_number: Optional[Coord[Time, Union[numpy.int64, numpy.int32]]] = None
-    """Arbitary scan number to identify data taken in the same logical scan."""
+    scan_name: Optional[Coord[Time, str]] = None
+    """Arbitary scan name to identify data taken in the same logical scan."""
 
     # SPECTRUM_CORRECTED: Optional[Dataof[SpectrumArray]] = None