xradio 0.0.40__py3-none-any.whl → 0.0.42__py3-none-any.whl

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

@@ -1,21 +1,28 @@
+import datetime
+import importlib
 import numcodecs
-import time
-from .._zarr.encoding import add_encoding
-from typing import Dict, Union
-import toolviper.utils.logger as logger
 import os
 import pathlib
+import time
+from typing import Dict, Union
 
 import numpy as np
 import xarray as xr
 
+import toolviper.utils.logger as logger
 from casacore import tables
-from xradio.vis._vis_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 xradio.vis._vis_utils._ms.create_antenna_xds import create_antenna_xds
-from xradio.vis._vis_utils._ms.create_field_and_source_xds import (
+from .msv4_info_dicts import create_info_dicts
+from xradio.measurement_set._utils._msv2.create_antenna_xds import (
+    create_antenna_xds,
+    create_gain_curve_xds,
+    create_phase_calibration_xds,
+)
+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
@@ -25,6 +32,7 @@ from .msv2_to_msv4_meta import (
     col_dims,
 )
 
+from .._zarr.encoding import add_encoding
 from .subtables import subt_rename_ids
 from ._tables.table_query import open_table_ro, open_query
 from ._tables.read import (
@@ -142,7 +150,7 @@ def mem_chunksize_to_dict_main(chunksize: float, xds: xr.Dataset) -> Dict[str, i
     polarization).
     """
 
-    sizeof_vis = itemsize_vis_spec(xds)
+    sizeof_vis = itemsize_spec(xds)
     size_all_pols = sizeof_vis * xds.sizes["polarization"]
     if size_all_pols / GiBYTES_TO_BYTES > chunksize:
         raise RuntimeError(
@@ -324,7 +332,7 @@ def find_baseline_or_antenna_var(xds: xr.Dataset) -> str:
     return baseline_or_antenna_name
 
 
-def itemsize_vis_spec(xds: xr.Dataset) -> int:
+def itemsize_spec(xds: xr.Dataset) -> int:
     """
     Size in bytes of one visibility (or spectrum) value.
     """
@@ -414,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)
@@ -474,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"] =
@@ -597,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
@@ -616,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
 
@@ -680,6 +701,28 @@ def create_taql_query(partition_info):
     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,
@@ -693,6 +736,7 @@ def convert_and_write_partition(
     pointing_interpolate: bool = False,
     ephemeris_interpolate: bool = False,
     phase_cal_interpolate: bool = False,
+    sys_cal_interpolate: bool = False,
     compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
     storage_backend="zarr",
     overwrite: bool = False,
@@ -705,7 +749,7 @@ def convert_and_write_partition(
         _description_
     out_file : str
         _description_
-    obs_mode : str
+    intents : str
         _description_
     ddi : int, optional
         _description_, by default 0
@@ -723,6 +767,10 @@ def convert_and_write_partition(
         _description_, by default None
     ephemeris_interpolate : bool, optional
         _description_, by default None
+    phase_cal_interpolate : bool, optional
+        _description_, by default None
+    sys_cal_interpolate : bool, optional
+        _description_, by default None
     compressor : numcodecs.abc.Codec, optional
         _description_, by default numcodecs.Zstd(level=2)
     storage_backend : str, optional
@@ -738,7 +786,7 @@ def convert_and_write_partition(
 
     taql_where = create_taql_query(partition_info)
     ddi = partition_info["DATA_DESC_ID"][0]
-    obs_mode = str(partition_info["OBS_MODE"][0])
+    intents = str(partition_info["OBS_MODE"][0])
 
     start = time.time()
     with open_table_ro(in_file) as mtable:
@@ -766,24 +814,32 @@ def convert_and_write_partition(
                 tb_tool.getcol("OBSERVATION_ID"), "OBSERVATION_ID"
             )
 
-            def get_observation_info(in_file, observation_id, obs_mode):
+            def get_observation_info(in_file, observation_id, intents):
                 generic_observation_xds = load_generic_table(
                     in_file,
                     "OBSERVATION",
                     taql_where=f" where (ROWID() IN [{str(observation_id)}])",
                 )
 
-                if obs_mode == "None":
-                    obs_mode = "obs_" + str(observation_id)
+                if intents == "None":
+                    intents = "obs_" + str(observation_id)
 
-                return generic_observation_xds["TELESCOPE_NAME"].values[0], obs_mode
+                return generic_observation_xds["TELESCOPE_NAME"].values[0], intents
 
-            telescope_name, obs_mode = get_observation_info(
-                in_file, observation_id, obs_mode
+            telescope_name, intents = get_observation_info(
+                in_file, observation_id, intents
             )
 
             start = time.time()
-            xds = xr.Dataset()
+            xds = xr.Dataset(
+                attrs={
+                    "creation_date": datetime.datetime.utcnow().isoformat(),
+                    "xradio_version": importlib.metadata.version("xradio"),
+                    "schema_version": "4.0.-9994",
+                    "type": "visibility",
+                }
+            )
+
             # interval = check_if_consistent(tb_tool.getcol("INTERVAL"), "INTERVAL")
             interval = tb_tool.getcol("INTERVAL")
 
@@ -826,7 +882,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 (
@@ -874,25 +929,56 @@ 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()
+            gain_curve_xds = create_gain_curve_xds(
+                in_file, xds.frequency.attrs["spectral_window_id"], ant_xds
+            )
+            logger.debug("Time gain_curve xds  " + str(time.time() - start))
+
+            start = time.time()
+            phase_calibration_xds = create_phase_calibration_xds(
+                in_file,
+                xds.frequency.attrs["spectral_window_id"],
+                ant_xds,
                 time_min_max,
                 phase_cal_interp_time,
             )
+            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)
+            # 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 = ant_xds.drop_vars(
-                "antenna_id"
-            )  # No longer needed after converting to name.
+            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")
 
-            logger.debug("Time ant xds  " + str(time.time() - start))
+            # Create system_calibration_xds
+            start = time.time()
+            if sys_cal_interpolate:
+                sys_cal_interp_time = xds.time.values
+            else:
+                sys_cal_interp_time = None
+            system_calibration_xds = create_system_calibration_xds(
+                in_file,
+                xds.frequency,
+                ant_xds_name_ids,
+                sys_cal_interp_time,
+            )
+            logger.debug("Time system_calibation " + str(time.time() - start))
 
             # Create weather_xds
             start = time.time()
-            weather_xds = create_weather_xds(in_file)
+            weather_xds = create_weather_xds(in_file, ant_xds_station_name_ids)
             logger.debug("Time weather " + str(time.time() - start))
 
             # Create pointing_xds
+            pointing_xds = xr.Dataset()
+            print("with_pointing", with_pointing)
             if with_pointing:
                 start = time.time()
                 if pointing_interpolate:
@@ -914,7 +1000,6 @@ def convert_and_write_partition(
                 )
 
             start = time.time()
-            xds.attrs["type"] = "visibility"
 
             # Time and frequency should always be increasing
             if len(xds.frequency) > 1 and xds.frequency[1] - xds.frequency[0] < 0:
@@ -943,7 +1028,7 @@ def convert_and_write_partition(
             # assert len(col_unique) == 1, col_name + " is not consistent."
             # return col_unique[0]
 
-            field_and_source_xds, source_id, num_lines = create_field_and_source_xds(
+            field_and_source_xds, source_id, _num_lines = create_field_and_source_xds(
                 in_file,
                 field_id,
                 xds.frequency.attrs["spectral_window_id"],
@@ -954,22 +1039,20 @@ 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)
 
-            if overwrite:
-                mode = "w"
-            else:
-                mode = "w-"
+            partition_info_misc_fields = {
+                "scan_id": scan_id,
+                "intents": intents,
+                "taql_where": taql_where,
+            }
+            info_dicts = create_info_dicts(
+                in_file, xds, field_and_source_xds, partition_info_misc_fields, tb_tool
+            )
+            xds.attrs.update(info_dicts)
 
+            # xds ready, prepare to write
+            start = time.time()
             main_chunksize = parse_chunksize(main_chunksize, "main", xds)
             add_encoding(xds, compressor=compressor, chunks=main_chunksize)
             logger.debug("Time add compressor and chunk " + str(time.time() - start))
@@ -979,52 +1062,58 @@ def convert_and_write_partition(
                 pathlib.Path(in_file).name.replace(".ms", "") + "_" + str(ms_v4_id),
             )
 
-            if "line_name" in field_and_source_xds.coords:
-                line_name = to_list(
-                    unique_1d(np.ravel(field_and_source_xds.line_name.values))
-                )
+            if overwrite:
+                mode = "w"
             else:
-                line_name = []
-
-            xds.attrs["partition_info"] = {
-                # "spectral_window_id": xds.frequency.attrs["spectral_window_id"],
-                "spectral_window_name": xds.frequency.attrs["spectral_window_name"],
-                # "field_id": to_list(unique_1d(field_id)),
-                "field_name": to_list(
-                    np.unique(field_and_source_xds.field_name.values)
-                ),
-                # "source_id": to_list(unique_1d(source_id)),
-                "line_name": line_name,
-                "scan_number": to_list(np.unique(scan_id)),
-                "source_name": to_list(
-                    np.unique(field_and_source_xds.source_name.values)
-                ),
-                "polarization_setup": to_list(xds.polarization.values),
-                "num_lines": num_lines,
-                "obs_mode": obs_mode.split(","),
-                "taql": taql_where,
-            }
+                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, "system_calibration_xds"),
+                        mode=mode,
+                    )
+
+                if gain_curve_xds:
+                    gain_curve_xds.to_zarr(
+                        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_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":
@@ -1056,7 +1145,11 @@ def antenna_ids_to_names(
             }
         )
     else:
-        xds["baseline_id"] = ant_xds["antenna_name"].sel(antenna_id=xds["baseline_id"])
+        # 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"]
+        )
         unwanted_coords_from_ant_xds = [
             "antenna_id",
             "antenna_name",
@@ -1080,7 +1173,7 @@ def add_data_groups(xds):
     xds.attrs["data_groups"] = {}
     if "VISIBILITY" in xds:
         xds.attrs["data_groups"]["base"] = {
-            "visibility": "VISIBILITY",
+            "correlated_data": "VISIBILITY",
             "flag": "FLAG",
             "weight": "WEIGHT",
             "uvw": "UVW",
@@ -1088,7 +1181,7 @@ def add_data_groups(xds):
 
     if "VISIBILITY_CORRECTED" in xds:
         xds.attrs["data_groups"]["corrected"] = {
-            "visibility": "VISIBILITY_CORRECTED",
+            "correlated_data": "VISIBILITY_CORRECTED",
             "flag": "FLAG",
             "weight": "WEIGHT",
             "uvw": "UVW",
@@ -1096,7 +1189,7 @@ def add_data_groups(xds):
 
     if "VISIBILITY_MODEL" in xds:
         xds.attrs["data_groups"]["model"] = {
-            "visibility": "VISIBILITY_MODEL",
+            "correlated_data": "VISIBILITY_MODEL",
             "flag": "FLAG",
             "weight": "WEIGHT",
             "uvw": "UVW",
@@ -1105,7 +1198,16 @@ def add_data_groups(xds):
     is_single_dish = False
     if "SPECTRUM" in xds:
         xds.attrs["data_groups"]["base"] = {
-            "spectrum": "SPECTRUM",
+            "correlated_data": "SPECTRUM",
+            "flag": "FLAG",
+            "weight": "WEIGHT",
+            "uvw": "UVW",
+        }
+        is_single_dish = True
+
+    if "SPECTRUM_MODEL" in xds:
+        xds.attrs["data_groups"]["model"] = {
+            "correlated_data": "SPECTRUM_MODEL",
             "flag": "FLAG",
             "weight": "WEIGHT",
             "uvw": "UVW",
@@ -1114,7 +1216,7 @@ def add_data_groups(xds):
 
     if "SPECTRUM_CORRECTED" in xds:
         xds.attrs["data_groups"]["corrected"] = {
-            "spectrum": "SPECTRUM_CORRECTED",
+            "correlated_data": "SPECTRUM_CORRECTED",
             "flag": "FLAG",
             "weight": "WEIGHT",
             "uvw": "UVW",