xradio 0.0.56__py3-none-any.whl → 0.0.59__py3-none-any.whl

xradio/measurement_set/_utils/_msv2/conversion.py

@@ -1,11 +1,13 @@
+from collections import deque
 import datetime
 import importlib
 import numcodecs
 import os
 import pathlib
 import time
-from typing import Dict, Union
+from typing import Callable, Dict, Union
 
+import dask.array as da
 import numpy as np
 import xarray as xr
 import traceback
@@ -52,7 +54,9 @@ from ._tables.read import (
 )
 from ._tables.read_main_table import get_baselines, get_baseline_indices, get_utimes_tol
 from .._utils.stokes_types import stokes_types
-from xradio._utils.list_and_array import check_if_consistent, unique_1d, to_list
+
+from xradio._utils.list_and_array import check_if_consistent, unique_1d
+from xradio._utils.dict_helpers import make_spectral_coord_reference_dict, make_quantity
 
 
 def parse_chunksize(
@@ -218,7 +222,6 @@ def mem_chunksize_to_dict_main_balanced(
         dictionary of chunk sizes (as dim->size)
     """
 
-    dim_names = [name for name in xds_dim_sizes.keys()]
     dim_sizes = [size for size in xds_dim_sizes.values()]
     # Fix fourth dim (polarization) to all (not free to auto-calculate)
     free_dims_mask = np.array([True, True, True, False])
@@ -429,8 +432,49 @@ 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, scan_id
-):
+    xds: xr.Dataset,
+    in_file: str,
+    ddi: int,
+    utime: np.ndarray,
+    interval: np.ndarray,
+    baseline_ant1_id: np.ndarray,
+    baseline_ant2_id: np.ndarray,
+    scan_id: np.ndarray,
+) -> tuple[xr.Dataset, int]:
+    """
+    Creates coordinates of a VisibilityXds/SpectrumXds and assigns them to the input
+    correlated dataset.
+
+    Parameters
+    ----------
+    xds :
+        dataset to add the coords to
+    in_file :
+        path to input MSv2
+    ddi :
+        DDI index (row) for this MSv4
+    utime :
+        unique times, for the time coordinate
+    interval :
+        interval col values from the MSv2, for the integration_time attribute
+        of the time coord
+    baseline_ant1_id :
+        ANTENNA1 ids to be used as coord
+    baseline_ant2_id :
+        ANTENNA2 ids to be used as coord
+    scan_id :
+        SCAN_ID values from MSv2, for the scan_name coord
+
+    Returns
+    -------
+    tuple[xr.Dataset, int]
+        A tuple of:
+        - The input dataset with coordinates added and populated with all MSv4 schema
+          attributes.
+        - The MSv2 spectral_window_id of this DDI/MSv4, which is no longer added to
+          the frequency coord but is required to create other secondary xdss (antenna,
+          gain_curve, phase_calibration, system_calibration, field_and_source).
+    """
     coords = {
         "time": utime,
         "baseline_antenna1_id": ("baseline_id", baseline_ant1_id),
@@ -484,16 +528,16 @@ def create_coordinates(
         spw_name = spw_name + "_" + str(spectral_window_id)
 
     xds.frequency.attrs["spectral_window_name"] = spw_name
+    xds.frequency.attrs["spectral_window_intent"] = "UNSPECIFIED"
     msv4_measure = column_description_casacore_to_msv4_measure(
         freq_column_description["REF_FREQUENCY"],
         ref_code=spectral_window_xds["MEAS_FREQ_REF"].data,
     )
-    xds.frequency.attrs["reference_frequency"] = {
-        "dims": [],
-        "data": float(spectral_window_xds.REF_FREQUENCY.values),
-        "attrs": msv4_measure,
-    }
-    xds.frequency.attrs["spectral_window_id"] = spectral_window_id
+    xds.frequency.attrs["reference_frequency"] = make_spectral_coord_reference_dict(
+        float(spectral_window_xds.REF_FREQUENCY.values),
+        msv4_measure["units"],
+        msv4_measure["observer"],
+    )
 
     # Add if doppler table is present
     # xds.frequency.attrs["doppler_velocity"] =
@@ -512,14 +556,9 @@ def create_coordinates(
         freq_column_description["CHAN_WIDTH"],
         ref_code=spectral_window_xds["MEAS_FREQ_REF"].data,
     )
-    if not msv4_measure:
-        msv4_measure["type"] = "quantity"
-        msv4_measure["units"] = ["Hz"]
-    xds.frequency.attrs["channel_width"] = {
-        "dims": [],
-        "data": np.abs(unique_chan_width[0]),
-        "attrs": msv4_measure,
-    }
+    xds.frequency.attrs["channel_width"] = make_quantity(
+        np.abs(unique_chan_width[0]), msv4_measure["units"] if msv4_measure else "Hz"
+    )
 
     ###### Create Time Coordinate ######
     main_table_attrs = extract_table_attributes(in_file)
@@ -532,16 +571,11 @@ def create_coordinates(
     msv4_measure = column_description_casacore_to_msv4_measure(
         main_column_descriptions["INTERVAL"]
     )
-    if not msv4_measure:
-        msv4_measure["type"] = "quantity"
-        msv4_measure["units"] = ["s"]
-    xds.time.attrs["integration_time"] = {
-        "dims": [],
-        "data": interval,
-        "attrs": msv4_measure,
-    }
-    xds.time.attrs["effective_integration_time"] = "EFFECTIVE_INTEGRATION_TIME"
-    return xds
+    xds.time.attrs["integration_time"] = make_quantity(
+        interval, msv4_measure["units"] if msv4_measure else "s"
+    )
+
+    return xds, spectral_window_id
 
 
 def find_min_max_times(tb_tool: tables.table, taql_where: str) -> tuple:
@@ -584,91 +618,115 @@ def create_data_variables(
     parallel_mode,
     main_chunksize,
 ):
-
-    # Get time chunks
-    time_chunksize = None
-    if parallel_mode == "time":
-        try:
-            time_chunksize = main_chunksize["time"]
-        except KeyError:
-            # If time isn't chunked then `read_col_conversion_dask` is slower than `read_col_conversion_numpy`
-            logger.warning(
-                "'time' isn't specified in `main_chunksize`. Defaulting to `parallel_mode = 'none'`."
-            )
-            parallel_mode = "none"
-
-    # Set read_col_conversion from value of `parallel_mode` argument
-    # TODO: To make this compatible with multi-node conversion, `read_col_conversion_dask` and TableManager must be pickled.
-    # Casacore will make this difficult
-    global read_col_conversion
-    if parallel_mode == "time":
-        read_col_conversion = read_col_conversion_dask
-    else:
-        read_col_conversion = read_col_conversion_numpy
+    time_chunksize = main_chunksize.get("time", None) if main_chunksize else None
+    if parallel_mode == "time" and time_chunksize is None:
+        logger.warning(
+            "'time' isn't specified in `main_chunksize`. Defaulting to `parallel_mode = 'none'`."
+        )
+        parallel_mode = "none"
 
     # Create Data Variables
     with table_manager.get_table() as tb_tool:
         col_names = tb_tool.colnames()
 
+    target_cols = set(col_names) & set(col_to_data_variable_names.keys())
+    if target_cols.issuperset({"WEIGHT", "WEIGHT_SPECTRUM"}):
+        target_cols.remove("WEIGHT")
+
     main_table_attrs = extract_table_attributes(in_file)
     main_column_descriptions = main_table_attrs["column_descriptions"]
-    for col in col_names:
-        if col in col_to_data_variable_names:
-            if (col == "WEIGHT") and ("WEIGHT_SPECTRUM" in col_names):
-                continue
-            try:
-                start = time.time()
-                if col == "WEIGHT":
-                    xds = get_weight(
-                        xds,
-                        col,
-                        table_manager,
-                        time_baseline_shape,
-                        tidxs,
-                        bidxs,
-                        use_table_iter,
-                        main_column_descriptions,
-                        time_chunksize,
-                    )
-                else:
-                    xds[col_to_data_variable_names[col]] = xr.DataArray(
-                        read_col_conversion(
-                            table_manager,
-                            col,
-                            time_baseline_shape,
-                            tidxs,
-                            bidxs,
-                            use_table_iter,
-                            time_chunksize,
-                        ),
-                        dims=col_dims[col],
-                    )
-
-                xds[col_to_data_variable_names[col]].attrs.update(
-                    create_attribute_metadata(col, main_column_descriptions)
+
+    # Use a double-ended queue in case WEIGHT_SPECTRUM conversion fails, and
+    # we need to add WEIGHT to list of columns to convert during iteration
+    target_cols = deque(target_cols)
+
+    while target_cols:
+        col = target_cols.popleft()
+        datavar_name = col_to_data_variable_names[col]
+        read_col_conversion = get_read_col_conversion_function(col, parallel_mode)
+
+        try:
+            start = time.time()
+            col_data = read_col_conversion(
+                table_manager,
+                col,
+                time_baseline_shape,
+                tidxs,
+                bidxs,
+                use_table_iter,
+                time_chunksize,
+            )
+
+            if col == "TIME_CENTROID":
+                col_data = convert_casacore_time(col_data, False)
+
+            elif col == "WEIGHT":
+                col_data = repeat_weight_array(
+                    col_data, parallel_mode, xds.sizes, main_chunksize
                 )
 
+            xds[datavar_name] = xr.DataArray(
+                col_data,
+                dims=col_dims[col],
+                attrs=create_attribute_metadata(col, main_column_descriptions),
+            )
+            logger.debug(f"Time to read column {col} : {time.time() - start}")
+
+        except Exception as exc:
+            logger.debug(f"Could not load column {col}, exception: {exc}")
+            logger.debug(traceback.format_exc())
+
+            if col == "WEIGHT_SPECTRUM" and "WEIGHT" in col_names:
                 logger.debug(
-                    "Time to read column " + str(col) + " : " + str(time.time() - start)
+                    "Failed to convert WEIGHT_SPECTRUM column: "
+                    "will attempt to use WEIGHT instead"
                 )
-            except Exception as exc:
-                logger.debug(f"Could not load column {col}, exception: {exc}")
-                logger.debug(traceback.format_exc())
-
-                if ("WEIGHT_SPECTRUM" == col) and (
-                    "WEIGHT" in col_names
-                ):  # Bogus WEIGHT_SPECTRUM column, need to use WEIGHT.
-                    xds = get_weight(
-                        xds,
-                        "WEIGHT",
-                        table_manager,
-                        time_baseline_shape,
-                        tidxs,
-                        bidxs,
-                        use_table_iter,
-                        main_column_descriptions,
-                        time_chunksize,
-                    )
+                target_cols.append("WEIGHT")
+
+
+def get_read_col_conversion_function(col_name: str, parallel_mode: str) -> Callable:
+    """
+    Returns the appropriate read_col_conversion function: use the dask version
+    for large columns and parallel_mode="time", or the numpy version otherwise.
+    """
+    large_columns = {
+        "DATA",
+        "CORRECTED_DATA",
+        "MODEL_DATA",
+        "WEIGHT_SPECTRUM",
+        "WEIGHT",
+        "FLAG",
+    }
+    return (
+        read_col_conversion_dask
+        if parallel_mode == "time" and col_name in large_columns
+        else read_col_conversion_numpy
+    )
+
+
+def repeat_weight_array(
+    weight_arr,
+    parallel_mode: str,
+    main_sizes: dict[str, int],
+    main_chunksize: dict[str, int],
+):
+    """
+    Repeat the weights read from the WEIGHT column along the frequency dimension.
+    Returns a dask array if parallel_mode="time", or a numpy array otherwise.
+    """
+    reshaped_arr = weight_arr[:, :, None, :]
+    repeats = (1, 1, main_sizes["frequency"], 1)
+
+    if parallel_mode == "time":
+        result = da.tile(reshaped_arr, repeats)
+        # da.tile() adds each repeat as a separate chunk, so rechunking is necessary
+        chunksizes = tuple(
+            main_chunksize.get(dim, main_sizes[dim])
+            for dim in ("time", "baseline_id", "frequency", "polarization")
+        )
+        return result.rechunk(chunksizes)
+
+    return np.tile(reshaped_arr, repeats)
 
 
 def add_missing_data_var_attrs(xds):
@@ -681,7 +739,7 @@ def add_missing_data_var_attrs(xds):
     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"] = [""]
+            xds.data_vars[var_name].attrs["units"] = ""
 
     vis_var_names = ["VISIBILITY_MODEL"]
     for var_name in vis_var_names:
@@ -692,44 +750,11 @@ def add_missing_data_var_attrs(xds):
                     "VISIBILITY"
                 ].attrs["units"]
             else:
-                xds.data_vars[var_name].attrs["units"] = [""]
+                xds.data_vars[var_name].attrs["units"] = ""
 
     return xds
 
 
-def get_weight(
-    xds,
-    col,
-    table_manager,
-    time_baseline_shape,
-    tidxs,
-    bidxs,
-    use_table_iter,
-    main_column_descriptions,
-    time_chunksize,
-):
-    xds[col_to_data_variable_names[col]] = xr.DataArray(
-        np.tile(
-            read_col_conversion(
-                table_manager,
-                col,
-                time_baseline_shape,
-                tidxs,
-                bidxs,
-                use_table_iter,
-                time_chunksize,
-            )[:, :, 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_where(partition_info: dict):
     main_par_table_cols = [
         "DATA_DESC_ID",
@@ -771,9 +796,9 @@ def fix_uvw_frame(
     """
     if xds.UVW.attrs["frame"] == "ITRF":
         if is_single_dish:
-            center_var = "FIELD_REFERENCE_CENTER"
+            center_var = "FIELD_REFERENCE_CENTER_DIRECTION"
         else:
-            center_var = "FIELD_PHASE_CENTER"
+            center_var = "FIELD_PHASE_CENTER_DIRECTION"
 
         xds.UVW.attrs["frame"] = field_and_source_xds[center_var].attrs["frame"]
 
@@ -800,11 +825,9 @@ def estimate_memory_for_partition(in_file: str, partition: dict) -> float:
                 if "shape" in col_descr and isinstance(col_descr["shape"], np.ndarray):
                     # example: "shape": array([15,  4]) => gives pols x channels
                     cells_in_row = col_descr["shape"].prod()
-                    npols = col_descr["shape"][-1]
                 else:
                     first_row = np.array(tb_tool.col(data_col)[0])
                     cells_in_row = np.prod(first_row.shape)
-                    npols = first_row.shape[-1]
 
                 if col_descr["valueType"] == "complex":
                     # Assume. Otherwise, read first column and get the itemsize:
@@ -911,7 +934,6 @@ def estimate_memory_for_partition(in_file: str, partition: dict) -> float:
     taql_partition = create_taql_query_where(partition)
     taql_main = f"select * from $mtable {taql_partition}"
     with open_table_ro(in_file) as mtable:
-        col_names = mtable.colnames()
         with open_query(mtable, taql_main) as tb_tool:
             # Do not feel tempted to rely on nrows. nrows tends to underestimate memory when baselines are missing.
             # For some EVN datasets that can easily underestimate by a 50%
@@ -940,6 +962,7 @@ def estimate_memory_for_partition(in_file: str, partition: dict) -> float:
         + calculate_term_other_msv2_indices(msv2_nrows)
         + calculate_term_sub_xds(estimate_main_xds)
         + calculate_term_to_zarr(estimate_main_xds)
+        + calculate_term_attrs(estimate_main_xds)
     )
     estimate /= GiBYTES_TO_BYTES
 
@@ -1104,7 +1127,7 @@ def convert_and_write_partition(
         scan_id[tidxs, bidxs] = tb_tool.getcol("SCAN_NUMBER")
         scan_id = np.max(scan_id, axis=1)
 
-        xds = create_coordinates(
+        xds, spectral_window_id = create_coordinates(
             xds,
             in_file,
             ddi,
@@ -1172,7 +1195,7 @@ def convert_and_write_partition(
 
         ant_xds = create_antenna_xds(
             in_file,
-            xds.frequency.attrs["spectral_window_id"],
+            spectral_window_id,
             antenna_id,
             feed_id,
             telescope_name,
@@ -1181,9 +1204,7 @@ def convert_and_write_partition(
         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
-        )
+        gain_curve_xds = create_gain_curve_xds(in_file, spectral_window_id, ant_xds)
         logger.debug("Time gain_curve xds  " + str(time.time() - start))
 
         start = time.time()
@@ -1193,7 +1214,7 @@ def convert_and_write_partition(
             phase_cal_interp_time = None
         phase_calibration_xds = create_phase_calibration_xds(
             in_file,
-            xds.frequency.attrs["spectral_window_id"],
+            spectral_window_id,
             ant_xds,
             time_min_max,
             phase_cal_interp_time,
@@ -1208,6 +1229,7 @@ def convert_and_write_partition(
             sys_cal_interp_time = None
         system_calibration_xds = create_system_calibration_xds(
             in_file,
+            spectral_window_id,
             xds.frequency,
             ant_xds,
             sys_cal_interp_time,
@@ -1269,10 +1291,6 @@ def convert_and_write_partition(
 
         # Create field_and_source_xds (combines field, source and ephemeris data into one super dataset)
         start = time.time()
-        if ephemeris_interpolate:
-            ephemeris_interp_time = xds.time.values
-        else:
-            ephemeris_interp_time = None
 
         # if "FIELD_ID" not in partition_scheme:
         #     field_id = np.full(time_baseline_shape, -42, dtype=int)
@@ -1298,7 +1316,7 @@ def convert_and_write_partition(
             create_field_and_source_xds(
                 in_file,
                 field_id,
-                xds.frequency.attrs["spectral_window_id"],
+                spectral_window_id,
                 field_times,
                 is_single_dish,
                 time_min_max,
@@ -1330,7 +1348,7 @@ def convert_and_write_partition(
         add_encoding(xds, compressor=compressor, chunks=main_chunksize)
         logger.debug("Time add compressor and chunk " + str(time.time() - start))
 
-        file_name = os.path.join(
+        os.path.join(
             out_file,
             pathlib.Path(in_file).name.replace(".ms", "") + "_" + str(ms_v4_id),
         )
@@ -1377,7 +1395,7 @@ def convert_and_write_partition(
             ms_xdt["/phased_array_xds"] = phased_array_xds
 
         if storage_backend == "zarr":
-            ms_xdt.to_zarr(store=os.path.join(out_file, ms_v4_name))
+            ms_xdt.to_zarr(store=os.path.join(out_file, ms_v4_name), mode=mode)
         elif storage_backend == "netcdf":
             # xds.to_netcdf(path=file_name+"/MAIN", mode=mode) #Does not work
             raise

xradio/measurement_set/_utils/_msv2/create_antenna_xds.py

@@ -1,5 +1,3 @@
-import toolviper.utils.logger as logger
-import time
 from typing import Tuple, Union
 
 import numpy as np
@@ -10,7 +8,6 @@ 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,
     make_taql_where_between_min_max,
     table_exists,
 )
@@ -20,17 +17,15 @@ from xradio.measurement_set._utils._msv2.msv4_sub_xdss import (
 )
 
 from xradio._utils.list_and_array import (
-    check_if_consistent,
     unique_1d,
-    to_list,
-    to_np_array,
 )
+from xradio._utils.dict_helpers import make_quantity_attrs
 
 
 def create_antenna_xds(
     in_file: str,
     spectral_window_id: int,
-    antenna_id: list,
+    antenna_id: np.ndarray,
     feed_id: list,
     telescope_name: str,
     partition_polarization: xr.DataArray,
@@ -44,8 +39,8 @@ def create_antenna_xds(
         Path to the input MSv2.
     spectral_window_id : int
         Spectral window ID.
-    antenna_id : list
-        List of antenna IDs.
+    antenna_id : np.ndarray
+        Antenna IDs.
     feed_id : list
         List of feed IDs.
     telescope_name : str
@@ -82,7 +77,7 @@ def create_antenna_xds(
 
 
 def extract_antenna_info(
-    ant_xds: xr.Dataset, in_file: str, antenna_id: list, telescope_name: str
+    ant_xds: xr.Dataset, in_file: str, antenna_id: np.ndarray, telescope_name: str
 ) -> xr.Dataset:
     """Reformats MSv2 Antenna table content to MSv4 schema.
 
@@ -92,8 +87,8 @@ def extract_antenna_info(
         The dataset that will be updated with antenna information.
     in_file : str
         Path to the input MSv2.
-    antenna_id : list
-        A list of antenna IDs to extract information for.
+    antenna_id : np.array
+        Antenna IDs to extract information for.
     telescope_name : str
         The name of the telescope.
 
@@ -138,7 +133,7 @@ def extract_antenna_info(
         generic_ant_xds, ant_xds, to_new_data_variables, to_new_coords
     )
 
-    ant_xds["ANTENNA_DISH_DIAMETER"].attrs.update({"units": ["m"], "type": "quantity"})
+    ant_xds["ANTENNA_DISH_DIAMETER"].attrs.update(make_quantity_attrs(["m"]))
 
     ant_xds["ANTENNA_POSITION"].attrs["coordinate_system"] = "geocentric"
     ant_xds["ANTENNA_POSITION"].attrs["origin_object_name"] = "earth"
@@ -507,9 +502,7 @@ def create_phase_calibration_xds(
     phase_cal_xds = phase_cal_xds.assign_coords(ant_borrowed_coords | tone_label_coord)
 
     # Adjust expected types
-    phase_cal_xds["time_phase_cal"] = (
-        phase_cal_xds.time_phase_cal.astype("float64").astype("float64") / 10**9
-    )
+    phase_cal_xds["time_phase_cal"] = phase_cal_xds.time_phase_cal
 
     phase_cal_xds = rename_and_interpolate_to_time(
         phase_cal_xds, "time_phase_cal", phase_cal_interp_time, "phase_cal_xds"