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

xradio/measurement_set/_utils/_msv2/msv4_sub_xdss.py

@@ -0,0 +1,550 @@
+import toolviper.utils.logger as logger
+import os
+import time
+from typing import Tuple, Union
+
+import numpy as np
+import xarray as xr
+
+from xradio._utils.coord_math import convert_to_si_units
+from xradio._utils.schema import (
+    column_description_casacore_to_msv4_measure,
+    convert_generic_xds_to_xradio_schema,
+)
+from .subtables import subt_rename_ids
+from ._tables.read import (
+    load_generic_table,
+    make_taql_where_between_min_max,
+    table_exists,
+    table_has_column,
+)
+
+
+def interpolate_to_time(
+    xds: xr.Dataset,
+    interp_time: Union[xr.DataArray, None],
+    message_prefix: str,
+    time_name: str = "time",
+) -> xr.Dataset:
+    """
+    Interpolate the time coordinate of the input xarray dataset to the
+    a data array. This can be used for example to interpolate a pointing_xds
+    to the time coord of the (main) MSv4, or similarly the ephemeris
+    data variables of a field_and_source_xds.
+
+    Uses interpolation method "linear", unless the source number of points is
+    1 in which case "nearest" is used, to avoid divide-by-zero issues.
+
+    Parameters:
+    ----------
+    xds : xr.Dataset
+        Xarray dataset to interpolate (presumably a pointing_xds or an xds of
+        ephemeris variables)
+    interp_time : Union[xr.DataArray, None]
+        Time axis to interpolate the dataset to (usually main MSv4 time)
+    message_prefix: str
+        A prefix for info/debug/etc. messages
+
+    Returns:
+    -------
+    interpolated_xds : xr.Dataset
+        xarray dataset with time axis interpolated to interp_time.
+    """
+    if interp_time is not None:
+        points_before = xds[time_name].size
+        if points_before > 1:
+            method = "linear"
+        else:
+            method = "nearest"
+        xds = xds.interp({time_name: interp_time}, method=method, assume_sorted=True)
+        # scan_number sneaks in as a coordinate of the main time axis, drop it
+        if "scan_number" in xds.coords:
+            xds = xds.drop_vars("scan_number")
+        points_after = xds[time_name].size
+        logger.debug(
+            f"{message_prefix}: interpolating the time coordinate "
+            f"from {points_before} to {points_after} points"
+        )
+
+    return xds
+
+
+def make_taql_where_weather(
+    in_file: str, ant_xds_station_name_ids: xr.DataArray
+) -> str:
+    """
+    The use of taql_where with WEATHER is complicated because of ALMA and its
+    (ab)use of NS_WX_STATION_ID vs. ANTENNA_ID (all -1) (see read.py). We
+    cannot simply use a 'WHERE ANTENNA_ID=...'. This function produces a TaQL
+    where string that uses ANTENNA_ID/NS_WX_STATION_ID depending on waht is
+    found in the WEATHER subtable.
+
+    Parameters
+    ----------
+    in_file : str
+        Input MS name.
+    ant_xds_station_name_ids : xr.DataArray
+        station name data array from antenna_xds, with name/id information
+
+    Returns
+    -------
+    taql_where: str
+        WHERE substring safe to use in taql_where when loading WEATHER subtable
+    """
+    weather_path = os.path.join(in_file, "WEATHER")
+    if table_exists(weather_path) and table_has_column(
+        weather_path, "NS_WX_STATION_ID"
+    ):
+        taql_where = f"WHERE (NS_WX_STATION_ID IN [{','.join(map(str, ant_xds_station_name_ids.antenna_id.values))}])"
+    else:
+        taql_where = f"WHERE (ANTENNA_ID IN [{','.join(map(str, ant_xds_station_name_ids.antenna_id.values))}])"
+
+    return taql_where
+
+
+def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
+    """
+    Creates a Weather Xarray Dataset from a MS v2 WEATHER table.
+
+    Parameters
+    ----------
+    in_file : str
+        Input MS name.
+    ant_xds_station_name_ids : xr.DataArray
+        station name data array from antenna_xds, with name/id information
+
+    Returns
+    -------
+    xr.Dataset
+        Weather Xarray Dataset.
+    """
+
+    try:
+        taql_where = make_taql_where_weather(in_file, ant_xds_station_name_ids)
+        generic_weather_xds = load_generic_table(
+            in_file,
+            "WEATHER",
+            rename_ids=subt_rename_ids["WEATHER"],
+            taql_where=taql_where,
+        )
+    except ValueError as _exc:
+        return None
+
+    if not generic_weather_xds.data_vars:
+        # for example when the weather subtable only has info for antennas/stations
+        # not present in the MSv4 (no overlap between antennas loaded in ant_xds and weather)
+        return None
+
+    weather_xds = xr.Dataset(attrs={"type": "weather"})
+    stations_present = ant_xds_station_name_ids.sel(
+        antenna_id=generic_weather_xds["ANTENNA_ID"]
+    )
+    coords = {
+        "station_name": stations_present.data,
+        "antenna_name": stations_present.coords["antenna_name"].data,
+    }
+    weather_xds = weather_xds.assign_coords(coords)
+
+    dims_station_time = ["station_name", "time_weather"]
+    to_new_data_variables = {
+        "H20": ["H2O", dims_station_time],
+        "IONOS_ELECTRON": ["IONOS_ELECTRON", dims_station_time],
+        "PRESSURE": ["PRESSURE", dims_station_time],
+        "REL_HUMIDITY": ["REL_HUMIDITY", dims_station_time],
+        "TEMPERATURE": ["TEMPERATURE", dims_station_time],
+        "DEW_POINT": ["DEW_POINT", dims_station_time],
+        "WIND_DIRECTION": ["WIND_DIRECTION", dims_station_time],
+        "WIND_SPEED": ["WIND_SPEED", dims_station_time],
+    }
+
+    to_new_coords = {
+        "TIME": ["time_weather", ["time_weather"]],
+    }
+
+    weather_xds = convert_generic_xds_to_xradio_schema(
+        generic_weather_xds, weather_xds, to_new_data_variables, to_new_coords
+    )
+
+    # TODO: option to interpolate to main time
+
+    # PRESSURE: hPa in MSv2 specs and some MSs => Pa
+    weather_xds = convert_to_si_units(weather_xds)
+
+    # correct expected types (for example "IONOS_ELECTRON", "PRESSURE" can be float32)
+    for data_var in weather_xds:
+        if weather_xds.data_vars[data_var].dtype != np.float64:
+            weather_xds[data_var] = weather_xds[data_var].astype(np.float64)
+
+    return weather_xds
+
+
+def correct_generic_pointing_xds(
+    generic_pointing_xds: xr.Dataset, to_new_data_variables: dict[str, list]
+) -> xr.Dataset:
+    """
+    Takes a (generic) pointing_xds as read from a POINTING subtable of an MSv2
+    and tries to correct several deviations from the MSv2 specs seen in
+    common test data.
+    The problems fixed here include wrong dimensions:
+    - for example transposed dimensions with respect to the MSv2 specs (output
+    from CASA simulator),
+    - missing/additional unexpected dimensions when some of the columns are
+    empty (in the sense of "empty casacore cells").
+
+    This function modifies the data arrays of the data vars affected by such
+    issues.
+
+    Parameters
+    ----------
+    generic_pointing_xds: xr.Dataset
+        The generic pointing dataset (loaded from MSv2) to be fixed
+    to_new_data_variables: dict
+        The dict used for convert_generic_xds_to_xradio_schema, which gives all
+        the data variables relevant for the final MSv4 dataset.
+
+    Returns:
+    --------
+    xr.Dataset
+        Corrected dataset with dimensions conforming to MSv2 specs.
+    """
+
+    correct_pointing_xds = generic_pointing_xds.copy()
+
+    for data_var_name in generic_pointing_xds:
+        if data_var_name in to_new_data_variables:
+            # Corrects dim sizes of "empty cell" variables, such as empty DIRECTION, TARGET, etc.
+            if (
+                "dim_2" in generic_pointing_xds.sizes
+                and generic_pointing_xds.sizes["dim_2"] == 0
+            ):
+                # When all direction variables are "empty"
+                data_var_data = xr.DataArray(
+                    [[[[np.nan, np.nan]]]],
+                    dims=generic_pointing_xds.dims,
+                ).isel(n_polynomial=0, drop=True)
+                correct_pointing_xds[data_var_name].data = data_var_data
+
+            elif (
+                "dir" in generic_pointing_xds.sizes
+                and generic_pointing_xds.sizes["dir"] == 0
+            ):
+                # When some direction variables are "empty" but some are populated properly
+                if "dim_2" in generic_pointing_xds[key].sizes:
+                    data_var_data = xr.DataArray(
+                        generic_pointing_xds[key].values,
+                        dims=generic_pointing_xds[key].dims,
+                    )
+                else:
+                    shape = tuple(
+                        generic_pointing_xds.sizes[dim]
+                        for dim in ["TIME", "ANTENNA_ID"]
+                    ) + (2,)
+                    data_var_data = xr.DataArray(
+                        np.full(shape, np.nan),
+                        dims=generic_pointing_xds[key].dims,
+                    )
+                correct_pointing_xds[data_var_name].data = data_var_data
+
+    return correct_pointing_xds
+
+
+def create_pointing_xds(
+    in_file: str,
+    ant_xds_name_ids: xr.DataArray,
+    time_min_max: Union[Tuple[np.float64, np.float64], None],
+    interp_time: Union[xr.DataArray, None] = None,
+) -> xr.Dataset:
+    """
+    Creates a Pointing Xarray Dataset from an MS v2 POINTING (sub)table.
+
+    WIP: details of a few direction variables (and possibly moving some to attributes) to be
+    settled (see MSv4 spreadsheet).
+
+    Parameters
+    ----------
+    in_file : str
+        Input MS name.
+    ant_xds_name_ids : xr.DataArray
+        antenna_name data array from antenna_xds, with name/id information
+    time_min_max : tuple
+        min / max times values to constrain loading (from the TIME column)
+    interp_time : Union[xr.DataArray, None] (Default value = None)
+        interpolate time to this (presumably main dataset time)
+
+    Returns
+    -------
+    xr.Dataset
+         Pointing Xarray dataset
+    """
+    start = time.time()
+
+    taql_time_range = make_taql_where_between_min_max(
+        time_min_max, in_file, "POINTING", "TIME"
+    )
+
+    if taql_time_range is None:
+        taql_where = f"WHERE (ANTENNA_ID IN [{','.join(map(str, ant_xds_name_ids.antenna_id.values))}])"
+    else:
+        taql_where = (
+            taql_time_range
+            + f" AND (ANTENNA_ID IN [{','.join(map(str, ant_xds_name_ids.antenna_id.values))}])"
+        )
+    # Read POINTING table into a Xarray Dataset.
+    generic_pointing_xds = load_generic_table(
+        in_file,
+        "POINTING",
+        rename_ids=subt_rename_ids["POINTING"],
+        taql_where=taql_where,
+    )
+
+    if not generic_pointing_xds.data_vars:
+        # apparently empty MS/POINTING table => produce empty xds
+        return xr.Dataset()
+
+    # Checking a simple way of using only the one single coefficient of the polynomials
+    if "n_polynomial" in generic_pointing_xds.sizes:
+        size = generic_pointing_xds.sizes["n_polynomial"]
+        if size == 1:
+            generic_pointing_xds = generic_pointing_xds.sel({"n_polynomial": 0})
+        elif size == 0:
+            generic_pointing_xds = generic_pointing_xds.drop_dims("n_polynomial")
+
+    time_ant_dims = ["time", "antenna_name"]
+    time_ant_dir_dims = time_ant_dims + ["local_sky_dir_label"]
+    to_new_data_variables = {
+        "DIRECTION": ["POINTING_BEAM", time_ant_dir_dims],
+        "ENCODER": ["POINTING_DISH_MEASURED", time_ant_dir_dims],
+        "OVER_THE_TOP": ["POINTING_OVER_THE_TOP", time_ant_dims],
+    }
+
+    to_new_coords = {
+        "TIME": ["time", ["time"]],
+        "dim_2": ["local_sky_dir_label", ["local_sky_dir_label"]],
+    }
+
+    generic_pointing_xds = correct_generic_pointing_xds(
+        generic_pointing_xds, to_new_data_variables
+    )
+    pointing_xds = xr.Dataset(attrs={"type": "pointing"})
+    coords = {
+        "antenna_name": ant_xds_name_ids.sel(
+            antenna_id=generic_pointing_xds["ANTENNA_ID"]
+        ).data,
+        "local_sky_dir_label": ["az", "alt"],
+    }
+    pointing_xds = pointing_xds.assign_coords(coords)
+    pointing_xds = convert_generic_xds_to_xradio_schema(
+        generic_pointing_xds, pointing_xds, to_new_data_variables, to_new_coords
+    )
+
+    pointing_xds = interpolate_to_time(pointing_xds, interp_time, "pointing_xds")
+
+    logger.debug(f"create_pointing_xds() execution time {time.time() - start:0.2f} s")
+
+    return pointing_xds
+
+
+def prepare_generic_sys_cal_xds(generic_sys_cal_xds: xr.Dataset) -> xr.Dataset:
+    """
+    A generic_sys_cal_xds loaded with load_generic_table() cannot be easily
+    used in convert_generic_xds_to_xradio_schema() to produce an MSv4
+    sys_cal_xds dataset, as their structure differs in dimensions and order
+    of dimensions.
+    This function performs various prepareation steps, such as:
+    - filter out dimensions not neeed for an individual MSv4 (SPW, FEED),
+    - drop variables loaded from columns with all items set to empty array,
+    - transpose the dimensions frequency,receptor
+    - fix dimension names (and order) when needed.
+
+    Parameters
+    ----------
+    generic_sys_cal_xds : xr.Dataset
+        generic dataset read from an MSv2 SYSCAL subtable
+
+    Returns
+    -------
+    generic_sys_cal_xds: xr.Dataset
+        System calibration Xarray Dataset prepared for generic conversion
+        to MSv4.
+    """
+
+    # drop SPW and feed dims
+    generic_sys_cal_xds = generic_sys_cal_xds.isel(SPECTRAL_WINDOW_ID=0, drop=True)
+    generic_sys_cal_xds = generic_sys_cal_xds.isel(FEED_ID=0, drop=True)
+
+    # Often some of the T*_SPECTRUM are present but all the cells are populated
+    # with empty arrays
+    empty_arrays_vars = []
+    for data_var in generic_sys_cal_xds.data_vars:
+        if generic_sys_cal_xds[data_var].size == 0:
+            empty_arrays_vars.append(data_var)
+    if empty_arrays_vars:
+        generic_sys_cal_xds = generic_sys_cal_xds.drop_vars(empty_arrays_vars)
+
+    # Re-arrange receptor and frequency dims depending on input structure
+    if (
+        "receptor" in generic_sys_cal_xds.sizes
+        and "frequency" in generic_sys_cal_xds.sizes
+    ):
+        # dim_3 can be created for example when the T*_SPECTRUM have varying # channels!
+        # more generaly, could transpose with ... to avoid errors with additional spurious dimensions
+        if "dim_3" in generic_sys_cal_xds.dims:
+            generic_sys_cal_xds = generic_sys_cal_xds.drop_dims("dim_3")
+        # From MSv2 tables we get (...,frequency, receptor)
+        #  -> transpose to (...,receptor,frequency) ready for MSv4 sys_cal_xds
+        generic_sys_cal_xds = generic_sys_cal_xds.transpose(
+            "ANTENNA_ID", "TIME", "receptor", "frequency"
+        )
+    elif (
+        "frequency" in generic_sys_cal_xds.sizes
+        and not "dim_3" in generic_sys_cal_xds.sizes
+    ):
+        # because order is (...,frequency,receptor), when frequency is missing
+        # receptor can get wrongly labeled as frequency
+        generic_sys_cal_xds = generic_sys_cal_xds.rename_dims({"frequency": "receptor"})
+    elif (
+        "frequency" not in generic_sys_cal_xds.sizes
+        and "receptor" in generic_sys_cal_xds.sizes
+        and "dim_3" in generic_sys_cal_xds.sizes
+    ):
+        # different *_SPECTRUM array sizes + some empty arrays can create an additional spurious
+        # generic dimension, which should have been "receptor"
+        generic_sys_cal_xds = generic_sys_cal_xds.rename_dims({"receptor": "frequency"})
+        generic_sys_cal_xds = generic_sys_cal_xds.rename_dims({"dim_3": "receptor"})
+        generic_sys_cal_xds = generic_sys_cal_xds.transpose(
+            "ANTENNA_ID", "TIME", "receptor", "frequency"
+        )
+    else:
+        raise RuntimeError(
+            "Cannot understand the arrangement of dimensions of {generic_sys_cal_xds=}"
+        )
+
+    return generic_sys_cal_xds
+
+
+def create_system_calibration_xds(
+    in_file: str,
+    main_xds_frequency: xr.DataArray,
+    ant_xds_name_ids: xr.DataArray,
+    sys_cal_interp_time: Union[xr.DataArray, None] = None,
+):
+    """
+    Creates a system calibration Xarray Dataset from a MSv2 SYSCAL table.
+
+    Parameters
+    ----------
+    in_file: str
+        Input MS name.
+    main_xds_frequency: xr.DataArray
+        frequency array of the main xds (MSv4), containing among other things
+        spectral_window_id and measures metadata
+    ant_xds_name_ids : xr.Dataset
+        antenna_name data array from antenna_xds, with name/id information
+    sys_cal_interp_time: Union[xr.DataArray, None] = None,
+        Time axis to interpolate the data vars to (usually main MSv4 time)
+
+    Returns
+    -------
+    sys_cal_xds: xr.Dataset
+        System calibration Xarray Dataset.
+    """
+
+    spectral_window_id = main_xds_frequency.attrs["spectral_window_id"]
+    try:
+        generic_sys_cal_xds = load_generic_table(
+            in_file,
+            "SYSCAL",
+            rename_ids=subt_rename_ids["SYSCAL"],
+            taql_where=(
+                f" where (SPECTRAL_WINDOW_ID = {spectral_window_id})"
+                f" AND (ANTENNA_ID IN [{','.join(map(str, ant_xds_name_ids.antenna_id.values))}])"
+            ),
+        )
+    except ValueError as _exc:
+        return None
+
+    if not generic_sys_cal_xds.data_vars:
+        # even though SYSCAL is an optional subtable, some write it empty
+        return None
+
+    generic_sys_cal_xds = prepare_generic_sys_cal_xds(generic_sys_cal_xds)
+
+    mandatory_dimensions = ["antenna_name", "time_cal", "receptor_label"]
+    if "frequency" not in generic_sys_cal_xds.sizes:
+        dims_all = mandatory_dimensions
+    else:
+        dims_all = mandatory_dimensions + ["frequency_cal"]
+
+    to_new_data_variables = {
+        "PHASE_DIFF": ["PHASE_DIFFERENCE", ["antenna_name", "time_cal"]],
+        "TCAL": ["TCAL", dims_all],
+        "TCAL_SPECTRUM": ["TCAL", dims_all],
+        "TRX": ["TRX", dims_all],
+        "TRX_SPECTRUM": ["TRX", dims_all],
+        "TSKY": ["TSKY", dims_all],
+        "TSKY_SPECTRUM": ["TSKY", dims_all],
+        "TSYS": ["TSYS", dims_all],
+        "TSYS_SPECTRUM": ["TSYS", dims_all],
+        "TANT": ["TANT", dims_all],
+        "TANT_SPECTRUM": ["TANT", dims_all],
+        "TAN_TSYS": ["TAN_TSYS", dims_all],
+        "TANT_SYS_SPECTRUM": ["TANT_TSYS", dims_all],
+    }
+
+    to_new_coords = {
+        "TIME": ["time_cal", ["time_cal"]],
+        "receptor": ["receptor_label", ["receptor_label"]],
+        "frequency": ["frequency_cal", ["frequency_cal"]],
+    }
+
+    sys_cal_xds = xr.Dataset(attrs={"type": "system_calibration"})
+    coords = {
+        "antenna_name": ant_xds_name_ids.sel(
+            antenna_id=generic_sys_cal_xds["ANTENNA_ID"]
+        ).data,
+        "receptor_label": generic_sys_cal_xds.coords["receptor"].data,
+    }
+    sys_cal_xds = sys_cal_xds.assign_coords(coords)
+    sys_cal_xds = convert_generic_xds_to_xradio_schema(
+        generic_sys_cal_xds, sys_cal_xds, to_new_data_variables, to_new_coords
+    )
+
+    # Add frequency coord and its measures data, if present
+    if "frequency_cal" in dims_all:
+        frequency_coord = {
+            "frequency_cal": generic_sys_cal_xds.coords["frequency"].data
+        }
+        sys_cal_xds = sys_cal_xds.assign_coords(frequency_coord)
+        frequency_measure = {
+            "type": main_xds_frequency.attrs["type"],
+            "units": main_xds_frequency.attrs["units"],
+            "observer": main_xds_frequency.attrs["observer"],
+        }
+        sys_cal_xds.coords["frequency_cal"].attrs.update(frequency_measure)
+
+    if sys_cal_interp_time is not None:
+        sys_cal_xds = interpolate_to_time(
+            sys_cal_xds,
+            sys_cal_interp_time,
+            "system_calibration_xds",
+            time_name="time_cal",
+        )
+
+        time_coord_attrs = {
+            "type": "time",
+            "units": ["s"],
+            "scale": "utc",
+            "format": "unix",
+        }
+        # If interpolating time, rename time_cal => time
+        time_coord = {"time": ("time_cal", sys_cal_interp_time.data)}
+        sys_cal_xds = sys_cal_xds.assign_coords(time_coord)
+        sys_cal_xds.coords["time"].attrs.update(time_coord_attrs)
+        sys_cal_xds = sys_cal_xds.swap_dims({"time_cal": "time"}).drop_vars("time_cal")
+
+    # correct expected types
+    for data_var in sys_cal_xds:
+        if sys_cal_xds.data_vars[data_var].dtype != np.float64:
+            sys_cal_xds[data_var] = sys_cal_xds[data_var].astype(np.float64)
+
+    return sys_cal_xds

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

@@ -23,7 +23,7 @@ subt_rename_ids = {
     "SPECTRAL_WINDOW": {"row": "spectral_window_id", "dim_1": "chan"},
     "SOURCE": {"dim_1": "ra/dec", "dim_2": "line"},
     "STATE": {"row": "state_id"},
-    "SYSCAL": {"dim_1": "channel", "dim_2": "receiver"},
+    "SYSCAL": {"dim_1": "frequency", "dim_2": "receptor"},
     # Would make sense for non-std "WS_NX_STATION_POSITION"
     "WEATHER": {"dim_1": "xyz"},
 }

xradio/{vis/_vis_utils → measurement_set/_utils}/_utils/xds_helper.py

@@ -7,7 +7,7 @@ import xarray as xr
 
 from .cds import CASAVisSet
 from .stokes_types import stokes_types
-from ...._utils.common import get_pad_value
+from xradio._utils.list_and_array import get_pad_value
 
 
 def make_coords(

xradio/{vis/_vis_utils/ms.py → measurement_set/_utils/msv2.py}

@@ -2,14 +2,14 @@ import os
 import toolviper.utils.logger as logger
 from typing import List, Tuple, Union
 
-from ._utils.cds import CASAVisSet
-from ._ms.partitions import (
+from xradio.measurement_set._utils._utils.cds import CASAVisSet
+from xradio.measurement_set._utils._msv2.partitions import (
     finalize_partitions,
     read_ms_ddi_partitions,
     read_ms_scan_subscan_partitions,
 )
-from ._ms.subtables import read_ms_subtables
-from ._utils.xds_helper import vis_xds_packager_cds
+from xradio.measurement_set._utils._msv2.subtables import read_ms_subtables
+from xradio.measurement_set._utils._utils.xds_helper import vis_xds_packager_cds
 
 
 def read_ms(

xradio/{vis/_vis_utils → measurement_set/_utils}/zarr.py

@@ -10,7 +10,7 @@ from ._zarr.read import read_part_keys, read_partitions, read_subtables
 from ._zarr.write import write_metainfo, write_part_keys, write_partitions
 
 
-def is_zarr_vis(inpath: str) -> bool:
+def is_zarr_cor(inpath: str) -> bool:
     """
     Check if a given path has a visibilities dataset in Zarr format
 
@@ -32,7 +32,7 @@ def is_zarr_vis(inpath: str) -> bool:
         return False
 
 
-def read_vis(
+def read_cor(
     inpath: str,
     subtables: bool = True,
     asdm_subtables: bool = False,
@@ -83,7 +83,7 @@ def read_vis(
     return cds
 
 
-def write_vis(
+def write_cor(
     cds: CASAVisSet,
     outpath: str,
     chunks_on_disk: Union[Dict, None] = None,

xradio/{vis → measurement_set}/convert_msv2_to_processing_set.py

@@ -4,8 +4,8 @@ from typing import Dict, Union
 
 import dask
 
-from xradio.vis._vis_utils._ms.partition_queries import create_partitions
-from xradio.vis._vis_utils._ms.conversion import convert_and_write_partition
+from xradio.measurement_set._utils._msv2.partition_queries import create_partitions
+from xradio.measurement_set._utils._msv2.conversion import convert_and_write_partition
 
 
 def convert_msv2_to_processing_set(
@@ -18,6 +18,7 @@ def convert_msv2_to_processing_set(
     pointing_interpolate: bool = False,
     ephemeris_interpolate: bool = False,
     phase_cal_interpolate: bool = False,
+    sys_cal_interpolate: bool = False,
     use_table_iter: bool = False,
     compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
     storage_backend: str = "zarr",
@@ -47,6 +48,10 @@ def convert_msv2_to_processing_set(
         Whether to interpolate the time axis of the pointing sub-dataset to the time axis of the main dataset
     ephemeris_interpolate : bool, optional
         Whether to interpolate the time axis of the ephemeris data variables (of the field_and_source sub-dataset) to the time axis of the main dataset
+    phase_cal_interpolate : bool, optional
+        Whether to interpolate the time axis of the phase calibration data variables to the time axis of the main dataset
+    sys_cal_interpolate : bool, optional
+        Whether to interpolate the time axis of the system calibration data variables (sys_cal_xds) to the time axis of the main dataset
     use_table_iter : bool, optional
         Whether to use the table iterator to read the main table of the MS v2. This should be set to True when reading datasets with large number of rows and few partitions, by default False.
     compressor : numcodecs.abc.Codec, optional
@@ -96,6 +101,7 @@ def convert_msv2_to_processing_set(
                     pointing_interpolate=pointing_interpolate,
                     ephemeris_interpolate=ephemeris_interpolate,
                     phase_cal_interpolate=phase_cal_interpolate,
+                    sys_cal_interpolate=sys_cal_interpolate,
                     compressor=compressor,
                     overwrite=overwrite,
                 )
@@ -114,6 +120,7 @@ def convert_msv2_to_processing_set(
                 pointing_interpolate=pointing_interpolate,
                 ephemeris_interpolate=ephemeris_interpolate,
                 phase_cal_interpolate=phase_cal_interpolate,
+                sys_cal_interpolate=sys_cal_interpolate,
                 compressor=compressor,
                 overwrite=overwrite,
             )