xradio 0.0.34__py3-none-any.whl → 0.0.37__py3-none-any.whl

xradio/vis/_vis_utils/_ms/conversion.py

@@ -4,17 +4,22 @@ from .._zarr.encoding import add_encoding
 from typing import Dict, Union
 import graphviper.utils.logger as logger
 import os
+import pathlib
 
 import numpy as np
 import xarray as xr
 
 from casacore import tables
-from .msv4_sub_xdss import create_ant_xds, create_pointing_xds, create_weather_xds
+from xradio.vis._vis_utils._ms.msv4_sub_xdss import (
+    create_pointing_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 (
     create_field_and_source_xds,
 )
+from xradio._utils.schema import column_description_casacore_to_msv4_measure
 from .msv2_to_msv4_meta import (
-    column_description_casacore_to_msv4_measure,
     create_attribute_metadata,
     col_to_data_variable_names,
     col_dims,
@@ -75,7 +80,7 @@ def check_chunksize(chunksize: dict, xds_type: str) -> None:
         allowed_dims = [
             "time",
             "baseline_id",
-            "antenna_id",
+            "antenna_name",
             "frequency",
             "polarization",
         ]
@@ -133,7 +138,7 @@ def mem_chunksize_to_dict_main(chunksize: float, xds: xr.Dataset) -> Dict[str, i
     It presently relies on the logic of mem_chunksize_to_dict_main_balanced() to find a
     balanced list of dimension sizes for the chunks
 
-    Assumes these relevant dims: (time, antenna_id/baseline_id, frequency,
+    Assumes these relevant dims: (time, antenna_name/baseline_id, frequency,
     polarization).
     """
 
@@ -144,11 +149,11 @@ def mem_chunksize_to_dict_main(chunksize: float, xds: xr.Dataset) -> Dict[str, i
             "Cannot calculate chunk sizes when memory bound ({chunksize}) does not even allow all polarizations in one chunk"
         )
 
-    baseline_or_antenna_id = find_baseline_or_antenna_var(xds)
-    total_size = calc_used_gb(xds.sizes, baseline_or_antenna_id, sizeof_vis)
+    baseline_or_antenna_name = find_baseline_or_antenna_var(xds)
+    total_size = calc_used_gb(xds.sizes, baseline_or_antenna_name, sizeof_vis)
 
     ratio = chunksize / total_size
-    chunked_dims = ["time", baseline_or_antenna_id, "frequency", "polarization"]
+    chunked_dims = ["time", baseline_or_antenna_name, "frequency", "polarization"]
     if ratio >= 1:
         result = {dim: xds.sizes[dim] for dim in chunked_dims}
         logger.debug(
@@ -157,14 +162,17 @@ def mem_chunksize_to_dict_main(chunksize: float, xds: xr.Dataset) -> Dict[str, i
     else:
         xds_dim_sizes = {k: xds.sizes[k] for k in chunked_dims}
         result = mem_chunksize_to_dict_main_balanced(
-            chunksize, xds_dim_sizes, baseline_or_antenna_id, sizeof_vis
+            chunksize, xds_dim_sizes, baseline_or_antenna_name, sizeof_vis
         )
 
     return result
 
 
 def mem_chunksize_to_dict_main_balanced(
-    chunksize: float, xds_dim_sizes: dict, baseline_or_antenna_id: str, sizeof_vis: int
+    chunksize: float,
+    xds_dim_sizes: dict,
+    baseline_or_antenna_name: str,
+    sizeof_vis: int,
 ) -> Dict[str, int]:
     """
     Assumes the ratio is <1 and all pols can fit in memory (from
@@ -231,7 +239,7 @@ def mem_chunksize_to_dict_main_balanced(
             dim_chunksizes[idx] += int_delta
         used = np.prod(dim_chunksizes) * sizeof_vis / GiBYTES_TO_BYTES
 
-    chunked_dim_names = ["time", baseline_or_antenna_id, "frequency", "polarization"]
+    chunked_dim_names = ["time", baseline_or_antenna_name, "frequency", "polarization"]
     dim_chunksizes_int = [int(v) for v in dim_chunksizes]
     result = dict(zip(chunked_dim_names, dim_chunksizes_int))
 
@@ -309,11 +317,11 @@ def mem_chunksize_to_dict_pointing(chunksize: float, xds: xr.Dataset) -> Dict[st
 
 def find_baseline_or_antenna_var(xds: xr.Dataset) -> str:
     if "baseline_id" in xds.coords:
-        baseline_or_antenna_id = "baseline_id"
-    elif "antenna_id" in xds.coords:
-        baseline_or_antenna_id = "antenna_id"
+        baseline_or_antenna_name = "baseline_id"
+    elif "antenna_name" in xds.coords:
+        baseline_or_antenna_name = "antenna_name"
 
-    return baseline_or_antenna_id
+    return baseline_or_antenna_name
 
 
 def itemsize_vis_spec(xds: xr.Dataset) -> int:
@@ -347,11 +355,11 @@ def itemsize_pointing_spec(xds: xr.Dataset) -> int:
 
 
 def calc_used_gb(
-    chunksizes: dict, baseline_or_antenna_id: str, sizeof_vis: int
+    chunksizes: dict, baseline_or_antenna_name: str, sizeof_vis: int
 ) -> float:
     return (
         chunksizes["time"]
-        * chunksizes[baseline_or_antenna_id]
+        * chunksizes[baseline_or_antenna_name]
         * chunksizes["frequency"]
         * chunksizes["polarization"]
         * sizeof_vis
@@ -400,7 +408,7 @@ def calc_indx_for_row_split(tb_tool, taql_where):
 
 
 def create_coordinates(
-    xds, in_file, ddi, utime, interval, baseline_ant1_id, baseline_ant2_id
+    xds, in_file, ddi, utime, interval, baseline_ant1_id, baseline_ant2_id, scan_id
 ):
     coords = {
         "time": utime,
@@ -408,6 +416,7 @@ def create_coordinates(
         "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),
     }
 
     ddi_xds = load_generic_table(in_file, "DATA_DESCRIPTION").sel(row=ddi)
@@ -446,12 +455,12 @@ def create_coordinates(
     )
     xds.frequency.attrs.update(msv4_measure)
 
-    if (spectral_window_xds.NAME.values.item() is None) or (
-        spectral_window_xds.NAME.values.item() == "none"
-    ):
+    spw_name = spectral_window_xds.NAME.values.item()
+    if (spw_name is None) or (spw_name == "none") or (spw_name == ""):
         spw_name = "spw_" + str(spectral_window_id)
     else:
-        spw_name = spectral_window_xds.NAME.values.item()
+        # spw_name = spectral_window_xds.NAME.values.item()
+        spw_name = spw_name + "_" + str(spectral_window_id)
 
     xds.frequency.attrs["spectral_window_name"] = spw_name
     msv4_measure = column_description_casacore_to_msv4_measure(
@@ -588,7 +597,7 @@ def create_data_variables(
                 logger.debug(
                     "Time to read column " + str(col) + " : " + str(time.time() - start)
                 )
-            except:
+            except Exception as e:
                 logger.debug("Could not load column", col)
 
                 if ("WEIGHT_SPECTRUM" == col) and (
@@ -606,6 +615,18 @@ 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"""
+    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"]
+
+    return xds
+
+
 def get_weight(
     xds,
     col,
@@ -662,7 +683,7 @@ def create_taql_query(partition_info):
 def convert_and_write_partition(
     in_file: str,
     out_file: str,
-    ms_v4_id: int,
+    ms_v4_id: Union[int, str],
     partition_info: Dict,
     use_table_iter: bool,
     partition_scheme: str = "ddi_intent_field",
@@ -671,6 +692,7 @@ def convert_and_write_partition(
     pointing_chunksize: Union[Dict, float, None] = None,
     pointing_interpolate: bool = False,
     ephemeris_interpolate: bool = False,
+    phase_cal_interpolate: bool = False,
     compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
     storage_backend="zarr",
     overwrite: bool = False,
@@ -715,7 +737,6 @@ def convert_and_write_partition(
     """
 
     taql_where = create_taql_query(partition_info)
-    # print("taql_where", taql_where)
     ddi = partition_info["DATA_DESC_ID"][0]
     obs_mode = str(partition_info["OBS_MODE"][0])
 
@@ -775,8 +796,19 @@ def convert_and_write_partition(
             else:
                 interval = interval_unique[0]
 
+            scan_id = np.full(time_baseline_shape, -42, dtype=int)
+            scan_id[tidxs, bidxs] = tb_tool.getcol("SCAN_NUMBER")
+            scan_id = np.max(scan_id, axis=1)
+
             xds = create_coordinates(
-                xds, in_file, ddi, utime, interval, baseline_ant1_id, baseline_ant2_id
+                xds,
+                in_file,
+                ddi,
+                utime,
+                interval,
+                baseline_ant1_id,
+                baseline_ant2_id,
+                scan_id,
             )
             logger.debug("Time create coordinates " + str(time.time() - start))
 
@@ -792,9 +824,11 @@ def convert_and_write_partition(
                 use_table_iter,
             )
 
-            # Add data_groups and field_info
+            # Add data_groups
             xds, is_single_dish = add_data_groups(xds)
 
+            xds = add_missing_data_var_attrs(xds)
+
             if (
                 "WEIGHT" not in xds.data_vars
             ):  # Some single dish datasets don't have WEIGHT.
@@ -811,6 +845,9 @@ def convert_and_write_partition(
 
             logger.debug("Time create data variables " + str(time.time() - start))
 
+            # To constrain the time range to load (in pointing, ephemerides, phase_cal data_vars)
+            time_min_max = find_min_max_times(tb_tool, taql_where)
+
             # Create ant_xds
             start = time.time()
             feed_id = unique_1d(
@@ -826,17 +863,27 @@ def convert_and_write_partition(
                     [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_ant_xds(
+            ant_xds = create_antenna_xds(
                 in_file,
                 xds.frequency.attrs["spectral_window_id"],
                 antenna_id,
                 feed_id,
                 telescope_name,
+                time_min_max,
+                phase_cal_interp_time,
             )
 
             # Change antenna_ids to antenna_names
-            xds = antenna_ids_to_names(xds, ant_xds)
+            xds = antenna_ids_to_names(xds, ant_xds, is_single_dish)
+            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.
 
             logger.debug("Time ant xds  " + str(time.time() - start))
 
@@ -845,9 +892,7 @@ def convert_and_write_partition(
             weather_xds = create_weather_xds(in_file)
             logger.debug("Time weather " + str(time.time() - start))
 
-            # To constrain the time range to load (in pointing, ephemerides data_vars)
-            time_min_max = find_min_max_times(tb_tool, taql_where)
-
+            # Create pointing_xds
             if with_pointing:
                 start = time.time()
                 if pointing_interpolate:
@@ -855,7 +900,7 @@ def convert_and_write_partition(
                 else:
                     pointing_interp_time = None
                 pointing_xds = create_pointing_xds(
-                    in_file, time_min_max, pointing_interp_time
+                    in_file, ant_xds_name_ids, time_min_max, pointing_interp_time
                 )
                 pointing_chunksize = parse_chunksize(
                     pointing_chunksize, "pointing", pointing_xds
@@ -869,6 +914,7 @@ 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:
@@ -884,10 +930,6 @@ def convert_and_write_partition(
             else:
                 ephemeris_interp_time = None
 
-            scan_id = np.full(time_baseline_shape, -42, dtype=int)
-            scan_id[tidxs, bidxs] = tb_tool.getcol("SCAN_NUMBER")
-            scan_id = np.max(scan_id, axis=1)
-
             if "FIELD_ID" not in partition_scheme:
                 field_id = np.full(time_baseline_shape, -42, dtype=int)
                 field_id[tidxs, bidxs] = tb_tool.getcol("FIELD_ID")
@@ -901,7 +943,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 = 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"],
@@ -914,7 +956,6 @@ def convert_and_write_partition(
 
             # 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.
-            # print('##################',field_and_source_xds)
             if is_single_dish:
                 xds.UVW.attrs["frame"] = field_and_source_xds[
                     "FIELD_REFERENCE_CENTER"
@@ -935,11 +976,20 @@ def convert_and_write_partition(
 
             file_name = os.path.join(
                 out_file,
-                out_file.replace(".vis.zarr", "").replace(".zarr", "").split("/")[-1]
+                pathlib.Path(out_file)
+                .name.replace(".vis.zarr", "")
+                .replace(".zarr", "")
                 + "_"
                 + 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))
+                )
+            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"],
@@ -948,11 +998,14 @@ def convert_and_write_partition(
                     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),
-                "obs_mode": obs_mode,
+                "num_lines": num_lines,
+                "obs_mode": obs_mode.split(","),
                 "taql": taql_where,
             }
 
@@ -968,7 +1021,7 @@ def convert_and_write_partition(
                         mode=mode,
                     )
 
-                if with_pointing:
+                if with_pointing and len(pointing_xds.data_vars) > 1:
                     pointing_xds.to_zarr(
                         store=os.path.join(file_name, "POINTING"), mode=mode
                     )
@@ -986,70 +1039,43 @@ def convert_and_write_partition(
     # logger.info("Saved ms_v4 " + file_name + " in " + str(time.time() - start_with) + "s")
 
 
-def antenna_ids_to_names(xds, ant_xds):
+def antenna_ids_to_names(
+    xds: xr.Dataset, ant_xds: xr.Dataset, is_single_dish: bool
+) -> xr.Dataset:
+    ant_xds = ant_xds.set_xindex(
+        "antenna_id"
+    )  # Allows for non-dimension coordinate selection.
 
-    if ant_xds.attrs["overall_telescope_name"] in ["ALMA", "VLA", "NOEMA", "EVLA"]:
-        moving_antennas = True
-    else:
-        moving_antennas = False
-
-    if moving_antennas:
-        if "baseline_antenna1_id" in xds:  # Interferometer
-
-            baseline_ant1_name = np.core.defchararray.add(
-                ant_xds["name"].sel(antenna_id=xds["baseline_antenna1_id"]).values, "_"
-            )
-            baseline_ant1_name = np.core.defchararray.add(
-                baseline_ant1_name,
-                ant_xds["station"].sel(antenna_id=xds["baseline_antenna1_id"]).values,
-            )
-            baseline_ant2_name = np.core.defchararray.add(
-                ant_xds["name"].sel(antenna_id=xds["baseline_antenna2_id"]).values, "_"
-            )
-            baseline_ant2_name = np.core.defchararray.add(
-                baseline_ant2_name,
-                ant_xds["station"].sel(antenna_id=xds["baseline_antenna2_id"]).values,
-            )
-
-            xds["baseline_antenna1_id"] = xr.DataArray(
-                baseline_ant1_name, dims="baseline_id"
-            )
-            xds["baseline_antenna2_id"] = xr.DataArray(
-                baseline_ant2_name, dims="baseline_id"
-            )
-            xds = xds.rename(
-                {
-                    "baseline_antenna1_id": "baseline_antenna1_name",
-                    "baseline_antenna2_id": "baseline_antenna2_name",
-                }
-            )
-        else:  # Single Dish
-            antenna_name = np.core.defchararray.add(
-                ant_xds["name"].sel(antenna_id=xds["antenna_id"]).values, "_"
-            )
-            antenna_name = np.core.defchararray.add(
-                antenna_name,
-                ant_xds["station"].sel(antenna_id=xds["antenna_id"]).values,
-            )
-            xds["antenna_id"] = xr.DataArray(antenna_name, dims="baseline_id")
-            xds = xds.rename({"antenna_id": "antenna_name"})
+    if not is_single_dish:  # Interferometer
+        xds["baseline_antenna1_id"].data = ant_xds["antenna_name"].sel(
+            antenna_id=xds["baseline_antenna1_id"].data
+        )
+        xds["baseline_antenna2_id"].data = ant_xds["antenna_name"].sel(
+            antenna_id=xds["baseline_antenna2_id"].data
+        )
+        xds = xds.rename(
+            {
+                "baseline_antenna1_id": "baseline_antenna1_name",
+                "baseline_antenna2_id": "baseline_antenna2_name",
+            }
+        )
     else:
-        if "baseline_antenna1_id" in xds:  # Interferometer
-            xds["baseline_antenna1_id"] = ant_xds["name"].sel(
-                antenna_id=xds["baseline_antenna1_id"]
-            )
-            xds["baseline_antenna2_id"] = ant_xds["name"].sel(
-                antenna_id=xds["baseline_antenna2_id"]
-            )
-            xds = xds.rename(
-                {
-                    "baseline_antenna1_id": "baseline_antenna1_name",
-                    "baseline_antenna2_id": "baseline_antenna2_name",
-                }
-            )
-        else:  # Single Dish
-            xds["antenna_id"] = ant_xds["name"].sel(antenna_id=xds["antenna_id"])
-            xds = xds.rename({"antenna_id": "antenna_name"})
+        xds["baseline_id"] = ant_xds["antenna_name"].sel(antenna_id=xds["baseline_id"])
+        unwanted_coords_from_ant_xds = [
+            "antenna_id",
+            "antenna_name",
+            "mount",
+            "station",
+        ]
+        for unwanted_coord in unwanted_coords_from_ant_xds:
+            xds = xds.drop_vars(unwanted_coord)
+        xds = xds.rename({"baseline_id": "antenna_name"})
+
+        # drop more vars that seem unwanted in main_sd_xds, but there shouuld be a better way
+        # of not creating them in the first place
+        unwanted_coords_sd = ["baseline_antenna1_id", "baseline_antenna2_id"]
+        for unwanted_coord in unwanted_coords_sd:
+            xds = xds.drop_vars(unwanted_coord)
 
     return xds