xradio 0.0.41__py3-none-any.whl → 0.0.43__py3-none-any.whl

xradio/{correlated_data/_utils/_ms → measurement_set/_utils/_msv2}/conversion.py

@@ -11,18 +11,18 @@ import xarray as xr
 
 import toolviper.utils.logger as logger
 from casacore import tables
-from xradio.correlated_data._utils._ms.msv4_sub_xdss import (
+from xradio.measurement_set._utils._msv2.msv4_sub_xdss import (
     create_pointing_xds,
     create_system_calibration_xds,
     create_weather_xds,
 )
 from .msv4_info_dicts import create_info_dicts
-from xradio.correlated_data._utils._ms.create_antenna_xds import (
+from xradio.measurement_set._utils._msv2.create_antenna_xds import (
     create_antenna_xds,
     create_gain_curve_xds,
     create_phase_calibration_xds,
 )
-from xradio.correlated_data._utils._ms.create_field_and_source_xds import (
+from xradio.measurement_set._utils._msv2.create_field_and_source_xds import (
     create_field_and_source_xds,
 )
 from xradio._utils.schema import column_description_casacore_to_msv4_measure
@@ -422,9 +422,9 @@ def create_coordinates(
         "time": utime,
         "baseline_antenna1_id": ("baseline_id", baseline_ant1_id),
         "baseline_antenna2_id": ("baseline_id", baseline_ant2_id),
-        "uvw_label": ["u", "v", "w"],
         "baseline_id": np.arange(len(baseline_ant1_id)),
         "scan_number": ("time", scan_id),
+        "uvw_label": ["u", "v", "w"],
     }
 
     ddi_xds = load_generic_table(in_file, "DATA_DESCRIPTION").sel(row=ddi)
@@ -482,7 +482,6 @@ def create_coordinates(
     }
     xds.frequency.attrs["spectral_window_id"] = spectral_window_id
 
-    # xds.frequency.attrs["effective_channel_width"] = "EFFECTIVE_CHANNEL_WIDTH"
     # Add if doppler table is present
     # xds.frequency.attrs["doppler_velocity"] =
     # xds.frequency.attrs["doppler_type"] =
@@ -605,8 +604,8 @@ def create_data_variables(
                 logger.debug(
                     "Time to read column " + str(col) + " : " + str(time.time() - start)
                 )
-            except Exception as e:
-                logger.debug("Could not load column", col)
+            except Exception as exc:
+                logger.debug(f"Could not load column {col}, exception: {exc}")
 
                 if ("WEIGHT_SPECTRUM" == col) and (
                     "WEIGHT" in col_names
@@ -624,13 +623,27 @@ def create_data_variables(
 
 
 def add_missing_data_var_attrs(xds):
-    """Adds in attributes expected metadata that cannot be found
-    in the input MSv2. For now specifically for missing
-    single-dish/SPECTRUM metadata"""
+    """
+    Adds in the xds attributes expected metadata that cannot be found in the input MSv2.
+    For now:
+    - missing single-dish/SPECTRUM metadata
+    - missing interferometry/VISIBILITY_MODEL metadata
+    """
     data_var_names = ["SPECTRUM", "SPECTRUM_CORRECTED"]
     for var_name in data_var_names:
         if var_name in xds.data_vars:
-            xds.data_vars[var_name].attrs["units"] = ["Jy"]
+            xds.data_vars[var_name].attrs["units"] = [""]
+
+    vis_var_names = ["VISIBILITY_MODEL"]
+    for var_name in vis_var_names:
+        if var_name in xds.data_vars and "units" not in xds.data_vars[var_name].attrs:
+            # Assume MODEL uses the same units
+            if "VISIBILITY" in xds.data_vars:
+                xds.data_vars[var_name].attrs["units"] = xds.data_vars[
+                    "VISIBILITY"
+                ].attrs["units"]
+            else:
+                xds.data_vars[var_name].attrs["units"] = [""]
 
     return xds
 
@@ -674,6 +687,7 @@ def create_taql_query(partition_info):
         "FIELD_ID",
         "SCAN_NUMBER",
         "STATE_ID",
+        "ANTENNA1",
     ]
 
     taql_where = "WHERE "
@@ -683,11 +697,38 @@ def create_taql_query(partition_info):
                 taql_where
                 + f"({col_name} IN [{','.join(map(str, partition_info[col_name]))}]) AND"
             )
+            if col_name == "ANTENNA1":
+                taql_where = (
+                    taql_where
+                    + f"(ANTENNA2 IN [{','.join(map(str, partition_info[col_name]))}]) AND"
+                )
     taql_where = taql_where[:-3]
 
     return taql_where
 
 
+def fix_uvw_frame(
+    xds: xr.Dataset, field_and_source_xds: xr.Dataset, is_single_dish: bool
+) -> xr.Dataset:
+    """
+    Fix UVW frame
+
+    From CASA fixvis docs: clean and the im tool ignore the reference frame claimed by the UVW column (it is often
+    mislabelled as ITRF when it is really FK5 (J2000)) and instead assume the (u, v, w)s are in the same frame as the phase
+    tracking center. calcuvw does not yet force the UVW column and field centers to use the same reference frame!
+    Blank = use the phase tracking frame of vis.
+    """
+    if xds.UVW.attrs["frame"] == "ITRF":
+        if is_single_dish:
+            center_var = "FIELD_REFERENCE_CENTER"
+        else:
+            center_var = "FIELD_PHASE_CENTER"
+
+        xds.UVW.attrs["frame"] = field_and_source_xds[center_var].attrs["frame"]
+
+    return xds
+
+
 def convert_and_write_partition(
     in_file: str,
     out_file: str,
@@ -798,9 +839,9 @@ def convert_and_write_partition(
             start = time.time()
             xds = xr.Dataset(
                 attrs={
-                    "creation_date": datetime.datetime.now().isoformat(),
+                    "creation_date": datetime.datetime.utcnow().isoformat(),
                     "xradio_version": importlib.metadata.version("xradio"),
-                    "schema_version": "4.0.-9999",
+                    "schema_version": "4.0.-9994",
                     "type": "visibility",
                 }
             )
@@ -847,7 +888,6 @@ def convert_and_write_partition(
 
             # Add data_groups
             xds, is_single_dish = add_data_groups(xds)
-
             xds = add_missing_data_var_attrs(xds)
 
             if (
@@ -895,8 +935,8 @@ def convert_and_write_partition(
                 antenna_id,
                 feed_id,
                 telescope_name,
+                xds.polarization,
             )
-
             logger.debug("Time antenna xds  " + str(time.time() - start))
 
             start = time.time()
@@ -916,7 +956,10 @@ def convert_and_write_partition(
             logger.debug("Time phase_calibration xds  " + str(time.time() - start))
 
             # Change antenna_ids to antenna_names
-            xds = antenna_ids_to_names(xds, ant_xds, is_single_dish)
+            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")
@@ -943,6 +986,7 @@ def convert_and_write_partition(
             logger.debug("Time weather " + str(time.time() - start))
 
             # Create pointing_xds
+            pointing_xds = xr.Dataset()
             if with_pointing:
                 start = time.time()
                 if pointing_interpolate:
@@ -1003,16 +1047,7 @@ def convert_and_write_partition(
             )
             logger.debug("Time field_and_source_xds " + str(time.time() - start))
 
-            # Fix UVW frame
-            # From CASA fixvis docs: clean and the im tool ignore the reference frame claimed by the UVW column (it is often mislabelled as ITRF when it is really FK5 (J2000)) and instead assume the (u, v, w)s are in the same frame as the phase tracking center. calcuvw does not yet force the UVW column and field centers to use the same reference frame! Blank = use the phase tracking frame of vis.
-            if is_single_dish:
-                xds.UVW.attrs["frame"] = field_and_source_xds[
-                    "FIELD_REFERENCE_CENTER"
-                ].attrs["frame"]
-            else:
-                xds.UVW.attrs["frame"] = field_and_source_xds[
-                    "FIELD_PHASE_CENTER"
-                ].attrs["frame"]
+            xds = fix_uvw_frame(xds, field_and_source_xds, is_single_dish)
 
             partition_info_misc_fields = {
                 "scan_id": scan_id,
@@ -1040,41 +1075,53 @@ def convert_and_write_partition(
             else:
                 mode = "w-"
 
+            if is_single_dish:
+                xds.attrs["type"] = "spectrum"
+                xds = xds.drop_vars(["UVW"])
+                del xds["uvw_label"]
+            else:
+                if any("WVR" in s for s in intents):
+                    xds.attrs["type"] = "wvr"
+                else:
+                    xds.attrs["type"] = "visibility"
+
             start = time.time()
             if storage_backend == "zarr":
-                xds.to_zarr(store=os.path.join(file_name, "MAIN"), mode=mode)
-                ant_xds.to_zarr(store=os.path.join(file_name, "ANTENNA"), mode=mode)
+                xds.to_zarr(store=os.path.join(file_name, "correlated_xds"), mode=mode)
+                ant_xds.to_zarr(store=os.path.join(file_name, "antenna_xds"), mode=mode)
                 for group_name in xds.attrs["data_groups"]:
                     field_and_source_xds.to_zarr(
                         store=os.path.join(
-                            file_name, f"FIELD_AND_SOURCE_{group_name.upper()}"
+                            file_name, f"field_and_source_xds_{group_name}"
                         ),
                         mode=mode,
                     )
 
-                if with_pointing and len(pointing_xds.data_vars) > 1:
+                if with_pointing and len(pointing_xds.data_vars) > 0:
                     pointing_xds.to_zarr(
-                        store=os.path.join(file_name, "POINTING"), mode=mode
+                        store=os.path.join(file_name, "pointing_xds"), mode=mode
                     )
 
                 if system_calibration_xds:
                     system_calibration_xds.to_zarr(
-                        store=os.path.join(file_name, "SYSCAL"), mode=mode
+                        store=os.path.join(file_name, "system_calibration_xds"),
+                        mode=mode,
                     )
 
                 if gain_curve_xds:
                     gain_curve_xds.to_zarr(
-                        store=os.path.join(file_name, "GAIN_CURVE"), mode=mode
+                        store=os.path.join(file_name, "gain_curve_xds"), mode=mode
                     )
 
                 if phase_calibration_xds:
                     phase_calibration_xds.to_zarr(
-                        store=os.path.join(file_name, "PHASE_CAL"), mode=mode
+                        store=os.path.join(file_name, "phase_calibration_xds"),
+                        mode=mode,
                     )
 
                 if weather_xds:
                     weather_xds.to_zarr(
-                        store=os.path.join(file_name, "WEATHER"), mode=mode
+                        store=os.path.join(file_name, "weather_xds"), mode=mode
                     )
 
             elif storage_backend == "netcdf":
@@ -1086,8 +1133,31 @@ def convert_and_write_partition(
 
 
 def antenna_ids_to_names(
-    xds: xr.Dataset, ant_xds: xr.Dataset, is_single_dish: bool
+    xds: xr.Dataset,
+    ant_xds: xr.Dataset,
+    is_single_dish: bool,
+    with_antenna_partitioning,
 ) -> xr.Dataset:
+    """
+    Turns the antenna_ids that we get from MSv2 into MSv4 antenna_name
+
+    Parameters
+    ----------
+    xds: xr.Dataset
+        A main xds (MSv4)
+    ant_xds: xr.Dataset
+        The antenna_xds for this MSv4
+    is_single_dish: bool
+        Whether a single-dish ("spectrum" data) dataset
+    with_antenna_partitioning: bool
+        Whether the MSv4 partitions include the antenna axis => only
+        one antenna (and implicitly one 'baseline' - auto-correlation)
+
+    Returns
+    ----------
+    xr.Dataset
+        The main xds with antenna_id replaced with antenna_name
+    """
     ant_xds = ant_xds.set_xindex(
         "antenna_id"
     )  # Allows for non-dimension coordinate selection.
@@ -1106,7 +1176,15 @@ def antenna_ids_to_names(
             }
         )
     else:
-        xds["baseline_id"] = ant_xds["antenna_name"].sel(antenna_id=xds["baseline_id"])
+        if not with_antenna_partitioning:
+            # baseline_antenna1_id will be removed soon below, but it is useful here to know the actual antenna_ids,
+            # as opposed to the baseline_ids which can mismatch when data is missing for some antennas
+            xds["baseline_id"] = ant_xds["antenna_name"].sel(
+                antenna_id=xds["baseline_antenna1_id"]
+            )
+        else:
+            xds["baseline_id"] = ant_xds["antenna_name"]
+
         unwanted_coords_from_ant_xds = [
             "antenna_id",
             "antenna_name",

xradio/{correlated_data/_utils/_ms → measurement_set/_utils/_msv2}/create_antenna_xds.py

@@ -6,8 +6,8 @@ import numpy as np
 import xarray as xr
 import os
 
-from xradio.correlated_data._utils._ms.subtables import subt_rename_ids
-from xradio.correlated_data._utils._ms._tables.read import (
+from xradio.measurement_set._utils._msv2.subtables import subt_rename_ids
+from xradio.measurement_set._utils._msv2._tables.read import (
     load_generic_table,
     convert_casacore_time,
     convert_casacore_time_to_mjd,
@@ -15,7 +15,7 @@ from xradio.correlated_data._utils._ms._tables.read import (
     table_exists,
 )
 from xradio._utils.schema import convert_generic_xds_to_xradio_schema
-from xradio.correlated_data._utils._ms.msv4_sub_xdss import interpolate_to_time
+from xradio.measurement_set._utils._msv2.msv4_sub_xdss import interpolate_to_time
 
 from xradio._utils.list_and_array import (
     check_if_consistent,
@@ -31,6 +31,7 @@ def create_antenna_xds(
     antenna_id: list,
     feed_id: list,
     telescope_name: str,
+    partition_polarization: xr.DataArray,
 ) -> xr.Dataset:
     """
     Create an Xarray Dataset containing antenna information.
@@ -47,6 +48,8 @@ def create_antenna_xds(
         List of feed IDs.
     telescope_name : str
         Name of the telescope.
+    partition_polarization: xr.DataArray
+        Polarization labels of this partition, needed if that info is not present in FEED
 
     Returns
     ----------
@@ -59,6 +62,18 @@ def create_antenna_xds(
     ant_xds = extract_feed_info(
         ant_xds, in_file, antenna_id, feed_id, spectral_window_id
     )
+    # Needed for special SPWs such as ALMA WVR or CHANNEL_AVERAGE data (have no feed info)
+    if "polarization_type" not in ant_xds:
+        pols_chars = list(partition_polarization.values[0])
+        pols_labels = [f"pol_{idx}" for idx in np.arange(0, len(pols_chars))]
+        ant_xds = ant_xds.assign_coords(receptor_label=pols_labels)
+        pol_type_values = [pols_chars] * len(ant_xds.antenna_name)
+        ant_xds = ant_xds.assign_coords(
+            polarization_type=(
+                ["antenna_name", "receptor_label"],
+                pol_type_values,
+            )
+        )
 
     ant_xds.attrs["overall_telescope_name"] = telescope_name
     return ant_xds
@@ -87,7 +102,6 @@ def extract_antenna_info(
     """
     to_new_data_variables = {
         "POSITION": ["ANTENNA_POSITION", ["antenna_name", "cartesian_pos_label"]],
-        "OFFSET": ["ANTENNA_FEED_OFFSET", ["antenna_name", "cartesian_pos_label"]],
         "DISH_DIAMETER": ["ANTENNA_DISH_DIAMETER", ["antenna_name"]],
     }
 
@@ -95,7 +109,7 @@ def extract_antenna_info(
         "NAME": ["antenna_name", ["antenna_name"]],
         "STATION": ["station", ["antenna_name"]],
         "MOUNT": ["mount", ["antenna_name"]],
-        "PHASED_ARRAY_ID": ["phased_array_id", ["antenna_name"]],
+        # "PHASED_ARRAY_ID": ["phased_array_id", ["antenna_name"]],
         "antenna_id": ["antenna_id", ["antenna_name"]],
     }
 
@@ -124,9 +138,8 @@ def extract_antenna_info(
 
     ant_xds["ANTENNA_DISH_DIAMETER"].attrs.update({"units": ["m"], "type": "quantity"})
 
-    ant_xds["ANTENNA_FEED_OFFSET"].attrs["type"] = "earth_location_offset"
-    ant_xds["ANTENNA_FEED_OFFSET"].attrs["coordinate_system"] = "geocentric"
     ant_xds["ANTENNA_POSITION"].attrs["coordinate_system"] = "geocentric"
+    ant_xds["ANTENNA_POSITION"].attrs["origin_object_name"] = "earth"
 
     if telescope_name in ["ALMA", "VLA", "NOEMA", "EVLA"]:
         # antenna_name = ant_xds["antenna_name"].values + "_" + ant_xds["station"].values
@@ -202,12 +215,17 @@ def extract_feed_info(
         taql_where=f" where (ANTENNA_ID IN [{','.join(map(str, ant_xds.antenna_id.values))}]) AND (FEED_ID IN [{','.join(map(str, feed_id))}])",
     )  # Some Lofar and MeerKAT data have the spw column set to -1 so we can't use '(SPECTRAL_WINDOW_ID = {spectral_window_id})'
 
+    if not generic_feed_xds:
+        # Some MSv2 have a FEED table that does not cover all antenna_id (and feed_id)
+        return ant_xds
+
     feed_spw = np.unique(generic_feed_xds.SPECTRAL_WINDOW_ID)
     if len(feed_spw) == 1 and feed_spw[0] == -1:
         generic_feed_xds = generic_feed_xds.isel(SPECTRAL_WINDOW_ID=0, drop=True)
     else:
         if spectral_window_id not in feed_spw:
-            return ant_xds  # For some spw the feed table is empty (this is the case with ALMA spw WVR#NOMINAL).
+            # For some spw the feed table is empty (this is the case with ALMA spw WVR#NOMINAL).
+            return ant_xds
         else:
             generic_feed_xds = generic_feed_xds.sel(
                 SPECTRAL_WINDOW_ID=spectral_window_id, drop=True
@@ -228,14 +246,14 @@ def extract_feed_info(
     ), "The number of receptors must be constant in feed table."
 
     to_new_data_variables = {
-        "BEAM_OFFSET": [
-            "BEAM_OFFSET",
-            ["antenna_name", "receptor_label", "sky_dir_label"],
+        "RECEPTOR_ANGLE": [
+            "ANTENNA_RECEPTOR_ANGLE",
+            ["antenna_name", "receptor_label"],
         ],
-        "RECEPTOR_ANGLE": ["RECEPTOR_ANGLE", ["antenna_name", "receptor_label"]],
-        # "pol_response": ["POLARIZATION_RESPONSE", ["antenna_name", "receptor_label", "receptor_name_"]] #repeated dim creates problems.
-        "FOCUS_LENGTH": ["FOCUS_LENGTH", ["antenna_name"]],  # optional
-        # "position": ["ANTENNA_FEED_OFFSET",["antenna_name", "cartesian_pos_label"]] #Will be added to the existing position in ant_xds
+        "FOCUS_LENGTH": [
+            "ANTENNA_FOCUS_LENGTH",
+            ["antenna_name"],
+        ],  # optional
     }
 
     to_new_coords = {
@@ -249,29 +267,31 @@ def extract_feed_info(
         to_new_coords=to_new_coords,
     )
 
-    # print('ant_xds["ANTENNA_FEED_OFFSET"]',ant_xds["ANTENNA_FEED_OFFSET"].data)
-    # print('generic_feed_xds["POSITION"].data',generic_feed_xds["POSITION"].data)
-    feed_offset_attrs = ant_xds["ANTENNA_FEED_OFFSET"].attrs
-    ant_xds["ANTENNA_FEED_OFFSET"] = (
-        ant_xds["ANTENNA_FEED_OFFSET"] + generic_feed_xds["POSITION"].data
-    )
-    # recover attrs after arithmetic operation
-    ant_xds["ANTENNA_FEED_OFFSET"].attrs.update(feed_offset_attrs)
-
-    coords = {}
     # coords["receptor_label"] = "pol_" + np.arange(ant_xds.sizes["receptor_label"]).astype(str) #Works on laptop but fails in github test runner.
-    coords["receptor_label"] = np.array(
-        list(
-            map(
-                lambda x, y: x + "_" + y,
-                ["pol"] * ant_xds.sizes["receptor_label"],
-                np.arange(ant_xds.sizes["receptor_label"]).astype(str),
-            )
+    coords = {
+        "receptor_label": np.array(
+            list(
+                map(
+                    lambda x, y: x + "_" + y,
+                    ["pol"] * ant_xds.sizes["receptor_label"],
+                    np.arange(ant_xds.sizes["receptor_label"]).astype(str),
+                )
+            ),
+            dtype=str,
         )
-    )
+    }
 
-    coords["sky_dir_label"] = ["ra", "dec"]
     ant_xds = ant_xds.assign_coords(coords)
+
+    # Correct to expected types. Some ALMA-SD (at least) leave receptor_label, polarization_type columns
+    # in the MS empty, causing a type mismatch
+    if (
+        "polarization_type" in ant_xds.coords
+        and ant_xds.coords["polarization_type"].dtype != str
+    ):
+        ant_xds.coords["polarization_type"] = ant_xds.coords[
+            "polarization_type"
+        ].astype(str)
     return ant_xds
 
 
@@ -370,6 +390,11 @@ def create_gain_curve_xds(
         }
     )
 
+    # correct expected types (for example "GAIN_CURVE" can be float32)
+    for data_var in gain_curve_xds:
+        if gain_curve_xds.data_vars[data_var].dtype != np.float64:
+            gain_curve_xds[data_var] = gain_curve_xds[data_var].astype(np.float64)
+
     return gain_curve_xds
 
 
@@ -445,7 +470,7 @@ def create_phase_calibration_xds(
         "TIME": ["time_phase_cal", ["time_phase_cal"]],
     }
 
-    phase_cal_xds = xr.Dataset(attrs={"type": "phase_cal"})
+    phase_cal_xds = xr.Dataset(attrs={"type": "phase_calibration"})
     phase_cal_xds = convert_generic_xds_to_xradio_schema(
         generic_phase_cal_xds, phase_cal_xds, to_new_data_variables, to_new_coords
     )
@@ -494,8 +519,8 @@ def create_phase_calibration_xds(
     time_coord_attrs = {
         "type": "time",
         "units": ["s"],
-        "scale": "UTC",
-        "format": "UNIX",
+        "scale": "utc",
+        "format": "unix",
     }
 
     # If we interpolate rename the time_phase_cal axis to time.