xradio 0.0.33__py3-none-any.whl → 0.0.36__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,
@@ -400,7 +405,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 +413,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 +452,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(
@@ -558,21 +564,16 @@ def create_data_variables(
             try:
                 start = time.time()
                 if col == "WEIGHT":
-                    xds[col_to_data_variable_names[col]] = xr.DataArray(
-                        np.tile(
-                            read_col_conversion(
-                                tb_tool,
-                                col,
-                                time_baseline_shape,
-                                tidxs,
-                                bidxs,
-                                use_table_iter,
-                            )[:, :, None, :],
-                            (1, 1, xds.sizes["frequency"], 1),
-                        ),
-                        dims=col_dims[col],
+                    xds = get_weight(
+                        xds,
+                        col,
+                        tb_tool,
+                        time_baseline_shape,
+                        tidxs,
+                        bidxs,
+                        use_table_iter,
+                        main_column_descriptions,
                     )
-
                 else:
                     xds[col_to_data_variable_names[col]] = xr.DataArray(
                         read_col_conversion(
@@ -585,20 +586,73 @@ def create_data_variables(
                         ),
                         dims=col_dims[col],
                     )
-                    logger.debug(
-                        "Time to read column "
-                        + str(col)
-                        + " : "
-                        + str(time.time() - start)
+
+                xds[col_to_data_variable_names[col]].attrs.update(
+                    create_attribute_metadata(col, main_column_descriptions)
+                )
+
+                logger.debug(
+                    "Time to read column " + str(col) + " : " + str(time.time() - start)
+                )
+            except Exception as e:
+                logger.debug("Could not load column", col)
+
+                if ("WEIGHT_SPECTRUM" == col) and (
+                    "WEIGHT" in col_names
+                ):  # Bogus WEIGHT_SPECTRUM column, need to use WEIGHT.
+                    xds = get_weight(
+                        xds,
+                        "WEIGHT",
+                        tb_tool,
+                        time_baseline_shape,
+                        tidxs,
+                        bidxs,
+                        use_table_iter,
+                        main_column_descriptions,
                     )
-            except:
-                # logger.debug("Could not load column",col)
-                # print("Could not load column", col)
-                continue
 
-            xds[col_to_data_variable_names[col]].attrs.update(
-                create_attribute_metadata(col, main_column_descriptions)
-            )
+
+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,
+    tb_tool,
+    time_baseline_shape,
+    tidxs,
+    bidxs,
+    use_table_iter,
+    main_column_descriptions,
+):
+    xds[col_to_data_variable_names[col]] = xr.DataArray(
+        np.tile(
+            read_col_conversion(
+                tb_tool,
+                col,
+                time_baseline_shape,
+                tidxs,
+                bidxs,
+                use_table_iter,
+            )[:, :, None, :],
+            (1, 1, xds.sizes["frequency"], 1),
+        ),
+        dims=col_dims[col],
+    )
+
+    xds[col_to_data_variable_names[col]].attrs.update(
+        create_attribute_metadata(col, main_column_descriptions)
+    )
+    return xds
 
 
 def create_taql_query(partition_info):
@@ -626,7 +680,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",
@@ -635,6 +689,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,
@@ -739,8 +794,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))
 
@@ -756,16 +822,30 @@ def convert_and_write_partition(
                 use_table_iter,
             )
 
+            # 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.
-                xds["WEIGHT"] = xr.DataArray(
-                    np.ones(xds.SPECTRUM.shape, dtype=np.float64),
-                    dims=xds.SPECTRUM.dims,
-                )
+                if is_single_dish:
+                    xds["WEIGHT"] = xr.DataArray(
+                        np.ones(xds.SPECTRUM.shape, dtype=np.float64),
+                        dims=xds.SPECTRUM.dims,
+                    )
+                else:
+                    xds["WEIGHT"] = xr.DataArray(
+                        np.ones(xds.VISIBILITY.shape, dtype=np.float64),
+                        dims=xds.VISIBILITY.dims,
+                    )
 
             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(
@@ -781,17 +861,26 @@ 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)
+            ant_xds = ant_xds.drop_vars(
+                "antenna_id"
+            )  # No longer needed after converting to name.
 
             logger.debug("Time ant xds  " + str(time.time() - start))
 
@@ -800,9 +889,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:
@@ -824,6 +911,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:
@@ -832,9 +920,6 @@ def convert_and_write_partition(
             if len(xds.time) > 1 and xds.time[1] - xds.time[0] < 0:
                 xds = xds.sel(time=slice(None, None, -1))
 
-            # Add data_groups and field_info
-            xds, is_single_dish = add_data_groups(xds)
-
             # Create field_and_source_xds (combines field, source and ephemeris data into one super dataset)
             start = time.time()
             if ephemeris_interpolate:
@@ -842,10 +927,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")
@@ -859,7 +940,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"],
@@ -893,11 +974,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"],
@@ -906,11 +996,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,
             }
 
@@ -926,7 +1019,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
                     )
@@ -945,69 +1038,26 @@ def convert_and_write_partition(
 
 
 def antenna_ids_to_names(xds, ant_xds):
+    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"})
-    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"})
+    if "baseline_antenna1_id" in xds:  # Interferometer
+        xds["baseline_antenna1_id"] = ant_xds["antenna_name"].sel(
+            antenna_id=xds["baseline_antenna1_id"]
+        )
+        xds["baseline_antenna2_id"] = ant_xds["antenna_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["antenna_name"].sel(antenna_id=xds["antenna_id"])
+        xds = xds.rename({"antenna_id": "antenna_name"})
 
     return xds