xradio 0.0.27__py3-none-any.whl → 0.0.29__py3-none-any.whl

xradio/vis/_vis_utils/_ms/subtables.py

@@ -13,11 +13,11 @@ from ._tables.read_subtables import read_ephemerides, read_delayed_pointing_tabl
 
 subt_rename_ids = {
     "ANTENNA": {"row": "antenna_id", "dim_1": "xyz"},
-    "FEED": {"dim_1": "xyz", "dim_2": "receptor", "dim_3": "receptor"},
+    "FEED": {"dim_1": "xyz", "dim_2": "receptor", "dim_3": "receptor2"},
     "FIELD": {"row": "field_id", "dim_1": "poly_id", "dim_2": "ra/dec"},
     "FREQ_OFFSET": {"antenna1": "antenna1_id", "antenna2": "antenna2_id"},
     "OBSERVATION": {"row": "observation_id", "dim_1": "start/end"},
-    "POINTING": {"dim_1": "n_polynomial", "dim_2": "ra/dec", "dim_3": "ra/dec"},
+    "POINTING": {"dim_1": "n_polynomial", "dim_3": "dir"},
     "POLARIZATION": {"row": "pol_setup_id", "dim_2": "product_id"},
     "PROCESSOR": {"row": "processor_id"},
     "SPECTRAL_WINDOW": {"row": "spectral_window_id", "dim_1": "chan"},
@@ -35,10 +35,20 @@ def read_ms_subtables(
     """
     Read MSv2 subtables (main table keywords) as xr.Dataset
 
-    :param infile: input MeasurementSet path
-    :param done_subt: Subtables that were already read, to skip them
-    :param asdm_subtables: Whether to also read ASDM_* subtables
-    :return: dict of xarray datasets read from subtables (metadata tables)
+    Parameters
+    ----------
+    infile : str
+        input MeasurementSet path
+    done_subt : List[str]
+        Subtables that were already read, to skip them
+    asdm_subtables : bool (Default value = False)
+        Whether to also read ASDM_* subtables
+
+    Returns
+    -------
+    Dict[str, xr.Dataset]
+        dict of xarray datasets read from subtables (metadata tables)
+
     """
     ignore_msv2_cols_subt = ["FLAG_CMD", "FLAG_ROW", "BEAM_ID"]
     skip_tables = ["SORTED_TABLE", "FLAG_CMD"] + done_subt
@@ -62,9 +72,8 @@ def read_ms_subtables(
 
         if subt_name == "POINTING":
             subt_path = Path(infile, subt_name)
-            xds = read_delayed_pointing_table(
-                str(subt_path), rename_ids=subt_rename_ids.get(subt_name, None)
-            )
+            rename_ids = {"dim_2": "n_polynomial", "dim_3": "dir"}
+            xds = read_delayed_pointing_table(str(subt_path), rename_ids=rename_ids)
         else:
             xds = read_generic_table(
                 infile,
@@ -88,14 +97,23 @@ def read_ms_subtables(
 def add_pointing_to_partition(
     xds_part: xr.Dataset, xds_pointing: xr.Dataset
 ) -> xr.Dataset:
-    """Take pointing variables from a (delayed) pointing dataset and
+    """
+    Take pointing variables from a (delayed) pointing dataset and
     transfer them to a main table partition dataset (interpolating into
     the destination time axis)
 
-    :param xds_part: a partition/sub-xds of the main table
-    :param xds_pointing: the xds read from the pointing subtable
-    :return: partition xds with pointing variables added/interpolated from the
-    pointing_xds into its time axis
+    Parameters
+    ----------
+    xds_part : xr.Dataset
+        a partition/sub-xds of the main table
+    xds_pointing : xr.Dataset
+        the xds read from the pointing subtable
+
+    Returns
+    -------
+    xr.Dataset
+        partition xds with pointing variables added/interpolated from the
+        pointing_xds into its time axis
 
     """
     interp_xds = xds_pointing.interp(time=xds_part.time, method="nearest")

xradio/vis/_vis_utils/_utils/partition_attrs.py

@@ -29,11 +29,19 @@ VisGroup = TypedDict(
 
 
 def make_vis_group_attr(xds: xr.Dataset) -> Dict:
-    """Add an attribute with the initial data/vis groups that have been
+    """
+    Add an attribute with the initial data/vis groups that have been
     read from the MS (DATA / CORRECTED_DATA / MODEL_DATA)
 
-    :param xds: dataset to make the vis_group depending on its data_vars
-    :return: vis_group derived form this dataset
+    Parameters
+    ----------
+    xds : xr.Dataset
+        dataset to make the vis_group depending on its data_vars
+
+    Returns
+    -------
+    Dict
+        vis_group derived form this dataset
     """
     msv2_extended_vis_vars = ["vis", "vis_corrected", "vis_model"]
     msv2_col_names = ["DATA", "CORRECTED_DATA", "MODEL_DATA"]
@@ -87,7 +95,8 @@ def add_partition_attrs(
     part_ids: PartitionIds,
     other_attrs: Dict,
 ) -> xr.Dataset:
-    """add attributes to the xr.Dataset:
+    """
+    add attributes to the xr.Dataset:
     - sub-dict of partition-id related ones
     - sub-dict of data/vis groups
     - sub-dict of attributes coming from the lower level read
@@ -96,13 +105,23 @@ def add_partition_attrs(
     Produces the partition IDs that can be retrieved from the DD subtable and also
     adds the ones passed in part_ids
 
-    :param xds: dataset partition
-    :param ddi: DDI of this partition
-    :param ddi_xds: dataset for the DATA_DESCRIPTION subtable
-    :param part_ids: partition id attrs
-    :param other_attrs: additional attributes produced by the read functions
-    :return: dataset with attributes added
-
+    Parameters
+    ----------
+    xds : xr.Dataset
+        dataset partition
+    ddi : int
+        DDI of this partition
+    ddi_xds : xr.Dataset
+        dataset for the DATA_DESCRIPTION subtable
+    part_ids : PartitionIds
+        partition id attrs
+    other_attrs : Dict
+        additional attributes produced by the read functions
+
+    Returns
+    -------
+    xr.Dataset
+        dataset with attributes added
     """
 
     xds = xds.assign_attrs(

xradio/vis/_vis_utils/_utils/xds_helper.py

@@ -12,7 +12,8 @@ from .stokes_types import stokes_types
 def make_coords(
     xds: xr.Dataset, ddi: int, subtables: Tuple[xr.Dataset, ...]
 ) -> Dict[str, np.ndarray]:
-    """Make the coords to be added to a partition or chunk (besides
+    """
+    Make the coords to be added to a partition or chunk (besides
     the time, baseline) basic structure
 
     Grabs:
@@ -20,7 +21,18 @@ def make_coords(
     - pol idxs from the pol+ddi subtables -> pol names via the stokes_types
     - antenna IDs from antenna subtable
 
-    :param: sub-xds as (ant_xds, ddi_xds, spw_xds, pol_xds)
+    Parameters
+    ----------
+    xds : xr.Dataset
+
+    ddi : int
+
+    subtables: Tuple[xr.Dataset, ...]
+
+
+    Returns
+    -------
+    Dict[str, np.ndarray]
     """
     ant_xds, ddi_xds, spw_xds, pol_xds = subtables
     freq = spw_xds.chan_freq.values[
@@ -46,15 +58,25 @@ def vis_xds_packager_cds(
     subtables: List[Tuple[str, xr.Dataset]],
     partitions: Dict[Any, xr.Dataset],
     descr_add: str = "",
-):
-    """Takes a a list of subtable xds datasets and a dictionary of data
+) -> CASAVisSet:
+    """
+    Takes a a list of subtable xds datasets and a dictionary of data
     partition xds datasets and and packages them as a CASA vis dataset
     (cds)
 
-    :param partitions: data partiions as xds datasets
-    :param subtables: subtables as xds datasets
-    :param descr_add: substring to add to the short descr string of the cds
-    :return: A "cds" - container for the metainfo subtables and data partitions
+    Parameters
+    ----------
+    partitions : List[Tuple[str, xr.Dataset]]
+        data partiions as xds datasets
+    subtables : Dict[Any, xr.Dataset]
+        subtables as xds datasets
+    descr_add : str (Default value = "")
+        substring to add to the short descr string of the cds
+
+    Returns
+    -------
+    CASAVisSet
+        A "cds" - container for the metainfo subtables and data partitions
     """
     vers = version("xradio")
 
@@ -70,14 +92,24 @@ def vis_xds_packager_mxds(
     subtables: List[Tuple[str, xr.Dataset]],
     add_global_coords: bool = True,
 ) -> xr.Dataset:
-    """Takes a dictionary of data partition xds datasets and a list of
+    """
+    Takes a dictionary of data partition xds datasets and a list of
     subtable xds datasets and packages them as a dataset of datasets
     (mxds)
 
-    :param partitions: data partiions as xds datasets
-    :param subtables: subtables as xds datasets
-    :add_global_coords: whether to add coords to the output mxds
-    :return: A "mxds" - xr.dataset of datasets
+    Parameters
+    ----------
+    partitions : Dict[Any, xr.Dataset]
+        data partiions as xds datasets
+    subtables : List[Tuple[str, xr.Dataset]]
+        subtables as xds datasets
+        :add_global_coords: whether to add coords to the output mxds
+    add_global_coords: bool (Default value = True)
+
+    Returns
+    -------
+    xr.Dataset
+        A "mxds" - xr.dataset of datasets
     """
     mxds = xr.Dataset(attrs={"metainfo": subtables, "partitions": partitions})
 
@@ -87,7 +119,7 @@ def vis_xds_packager_mxds(
     return mxds
 
 
-def make_global_coords(mxds: xr.Dataset):
+def make_global_coords(mxds: xr.Dataset) -> Dict[str, xr.DataArray]:
     coords = {}
     metainfo = mxds.attrs["metainfo"]
     if "antenna" in metainfo:
@@ -125,12 +157,24 @@ def make_global_coords(mxds: xr.Dataset):
 def expand_xds(xds: xr.Dataset) -> xr.Dataset:
     """
     expand single (row) dimension of xds to (time, baseline)
+
+    Parameters
+    ----------
+    xds : xr.Dataset
+        "flat" dataset (with row dimension - without (time, baseline) dimensions)
+
+    Returns
+    -------
+    xr.Dataset
+        expanded dataset, with (time, baseline) dimensions
     """
     assert "baseline" not in xds.coords
 
     txds = xds.copy()
     unique_baselines, baselines = np.unique(
-        [txds.antenna1.values, txds.antenna2.values], axis=1, return_inverse=True
+        [txds.baseline_ant1_id.values, txds.baseline_ant2_id.values],
+        axis=1,
+        return_inverse=True,
     )
     txds["baseline"] = xr.DataArray(baselines.astype("int32"), dims=["row"])
 
@@ -148,7 +192,7 @@ def expand_xds(xds: xr.Dataset) -> xr.Dataset:
             if txds[dv].dtype != xds[dv].dtype:
                 txds[dv] = txds[dv].astype(xds[dv].dtype)
     except Exception as exc:
-        print(
+        logger.warning(
             f"WARNING: Cannot expand rows to (time, baseline), "
             f"possibly duplicate values in (time, baseline). Exception: {exc}"
         )
@@ -160,20 +204,47 @@ def expand_xds(xds: xr.Dataset) -> xr.Dataset:
 def flatten_xds(xds: xr.Dataset) -> xr.Dataset:
     """
     flatten (time, baseline) dimensions of xds back to single dimension (row)
+
+    Parameters
+    ----------
+    xds : xr.Dataset
+
+
+    Returns
+    -------
+    xr.Dataset
     """
-    nan_int = np.array([np.nan]).astype("int32")[0]
+    # known invalid cast warning when casting to integer
+    with np.errstate(invalid="ignore"):
+        nan_int = np.array([np.nan]).astype("int32")[0]
     txds = xds.copy()
 
     # flatten the time x baseline dimensions of main table
     if ("time" in xds.sizes) and ("baseline" in xds.sizes):
         txds = xds.stack({"row": ("time", "baseline")}).transpose("row", ...)
+        # compute for issue https://github.com/hainegroup/oceanspy/issues/332
+        # drop=True silently does compute (or at least used to)
         txds = txds.where(
-            (txds.state_id != nan_int) & (txds.field_id != nan_int), drop=True
+            ((txds.state_id != nan_int) & (txds.field_id != nan_int)).compute(),
+            drop=True,
         )  # .unify_chunks()
-        for dv in list(xds.data_vars):
-            txds[dv] = txds[dv].astype(xds[dv].dtype)
 
-    return txds
+        # re-assigning (implicitly dropping index coords) one by one produces
+        # DeprecationWarnings: https://github.com/pydata/xarray/issues/6505
+        astyped_data_vars = dict(xds.data_vars)
+        for dv in list(txds.data_vars):
+            if txds[dv].dtype != xds[dv].dtype:
+                astyped_data_vars[dv] = txds[dv].astype(xds[dv].dtype)
+            else:
+                astyped_data_vars[dv] = txds[dv]
+
+        flat_xds = xr.Dataset(astyped_data_vars, coords=txds.coords, attrs=txds.attrs)
+        flat_xds = flat_xds.reset_index(["time", "baseline"])
+
+    else:
+        flat_xds = txds
+
+    return flat_xds
 
 
 ####################################
@@ -188,21 +259,30 @@ def optimal_chunking(
     Determine the optimal chunk shape for reading an MS or Image based
     on machine resources and intended operations
 
-    :param ndim: number of dimensions to chunk. An MS is 3, an
-    expanded MS is 4. An image could be anywhere from 2 to 5. Not
-    needed if data_shape is given.
-    :param didxs: dimension indices over which subsequent operations
-    will be performed. Values should be less than ndim. Tries to
-    reduce inter-process communication of data contents. Needs to
-    know the shape to do this well. Default None balances chunk size
-    across all dimensions.
-    :param chunk_size: target chunk size ('large', 'small', 'auto').
-    Default 'auto' tries to guess by looking at CPU core count and
-    available memory.
-    :param data_shape: shape of the total MS DDI or Image data. Helps
-    to know. Default None does not optimize based on shape
-
-    :return: optimal chunking for reading the ms (row, chan, pol)
+    Parameters
+    ----------
+    ndim : Union[int, None] = None
+        number of dimensions to chunk. An MS is 3, an
+        expanded MS is 4. An image could be anywhere from 2 to 5. Not
+        needed if data_shape is given.
+    didxs : Union[Tuple[int], List[int], None] = None
+        dimension indices over which subsequent operations
+        will be performed. Values should be less than ndim. Tries to
+        reduce inter-process communication of data contents. Needs to
+        know the shape to do this well. Default None balances chunk size
+        across all dimensions.
+    chunk_size : str (Default value = "auto")
+        target chunk size ('large', 'small', 'auto').
+        Default 'auto' tries to guess by looking at CPU core count and
+        available memory.
+    data_shape : Union[tuple, None] = None
+        shape of the total MS DDI or Image data. Helps
+        to know. Default None does not optimize based on shape
+
+    Returns
+    -------
+    tuple
+        optimal chunking for reading the ms (row, chan, pol)
     """
     assert (ndim is not None) or (
         data_shape is not None
@@ -278,22 +358,33 @@ def calc_optimal_ms_chunk_shape(
     """
     Calculates the max number of rows (1st dim in shape) of a variable
     that can be fit in the memory for a thread.
+
+    Parameters
+    ----------
+    memory_available_in_bytes :
+
+    shape :
+
+    element_size_in_bytes :
+
+    column_name :
+
+
+    Returns
+    -------
+    int
     """
     factor = 0.8  # Account for memory used by other objects in thread.
     # total_mem = np.prod(shape)*element_size_in_bytes
     single_row_mem = np.prod(shape[1:]) * element_size_in_bytes
 
-    try:
-        assert single_row_mem < factor * memory_available_in_bytes
-    except AssertionError as err:
-        logger.exception(
+    if not single_row_mem < factor * memory_available_in_bytes:
+        msg = (
             "Not engough memory in a thread to contain a row of "
-            + column_name
-            + ". Need at least "
-            + str(single_row_mem / factor)
-            + " bytes."
+            f"{column_name}. Need at least {single_row_mem / factor}"
+            " bytes."
         )
-        raise err
+        raise RuntimeError(msg)
 
     rows_chunk_size = int((factor * memory_available_in_bytes) / single_row_mem)
 

xradio/vis/_vis_utils/_zarr/read.py

@@ -11,9 +11,16 @@ def read_part_keys(inpath: str) -> List[Tuple]:
     """
     Reads the partition keys from a Zarr-stored cds.
 
-    :param inpath: path to read from
+    Parameters
+    ----------
+    inpath : str
+        path to read from
+
+    Returns
+    -------
+    List[Tuple]
+        partition keys from a cds
 
-    :return: partition keys from a cds
     """
 
     xds_keys = xr.open_zarr(
@@ -31,9 +38,19 @@ def read_subtables(inpath: str, asdm_subtables: bool) -> Dict[str, xr.Dataset]:
     """
     Reads the metainfo subtables from a Zarr-stored cds.
 
-    :param inpath: path to read from
+    Parameters
+    ----------
+    inpath : str
+        path to read from
+
+    asdm_subtables : bool
+
+
+    Returns
+    -------
+    Dict[str, xr.Dataset]
+        metainfo subtables from a cds
 
-    :return: metainfo subtables from a cds
     """
 
     metainfo = {}
@@ -53,9 +70,18 @@ def read_partitions(inpath: str, part_keys: List[Tuple]) -> Dict[str, xr.Dataset
     """
     Reads all the data partitions a Zarr-stored cds.
 
-    :param inpath: path to read from
+    Parameters
+    ----------
+    inpath : str
+        path to read from
+    part_keys : List[Tuple]
+
+
+    Returns
+    -------
+    Dict[str, xr.Dataset]
+        partitions from a cds
 
-    :return: partitions from a cds
     """
 
     partitions = {}
@@ -79,13 +105,23 @@ def read_xds(
     """
     Read single xds from zarr storage.
 
-    :param inpath: path to read from
-    :param chunks: set chunk size per dimension. Dict is in the form of
-    'dim':chunk_size, for example {'time':100, 'baseline':400, 'chan':32, 'pol':1}.
-    Default None uses the original chunking in the zarr input.
-    :param consolidated: use zarr consolidated metadata.
-    :param overwrite_encoded_chunks: drop the zarr chunks encoded for each variable
-    when a dataset is loaded with specified chunk sizes.
+    Parameters
+    ----------
+    inpath : str
+        path to read from
+    chunks : Union[Dict, None] (Default value = None)
+        set chunk size per dimension. Dict is in the form of
+        'dim':chunk_size, for example {'time':100, 'baseline':400, 'chan':32, 'pol':1}.
+        Default None uses the original chunking in the zarr input.
+    consolidated : boold (Default value = True)
+        use zarr consolidated metadata.
+    overwrite_encoded_chunks : bool (Default value = True)
+        drop the zarr chunks encoded for each variable
+        when a dataset is loaded with specified chunk sizes.
+
+    Returns
+    -------
+    xr.Dataset
     """
 
     xds = xr.open_zarr(
@@ -99,11 +135,11 @@ def read_xds(
 
 
 def read_zarr(
-    infile,
-    sel_xds=None,
-    chunks=None,
-    consolidated=True,
-    overwrite_encoded_chunks=True,
+    infile: str,
+    sel_xds: Union[List, str] = None,
+    chunks: Dict = None,
+    consolidated: bool = True,
+    overwrite_encoded_chunks: bool = True,
     **kwargs,
 ):
     """
@@ -128,11 +164,12 @@ def read_zarr(
     overwrite_encoded_chunks : bool
         drop the zarr chunks encoded for each variable when a dataset is loaded with
         specified chunk sizes.  Default True, only applies when chunks is not None.
+    **kwargs :
+
 
     Returns
     -------
-    xarray.core.dataset.Dataset
-        New xarray Dataset of Visibility data contents
+
     """
 
     if chunks is None:
@@ -178,8 +215,9 @@ def read_zarr(
 
 
 def _fix_dict_for_ms(name, xds):
-    xds.attrs["column_descriptions"] = xds.attrs["column_descriptions"][0]
-    xds.attrs["info"] = xds.attrs["info"][0]
+    # Used to be:
+    # xds.attrs["column_descriptions"] = xds.attrs["column_descriptions"][0]
+    # xds.attrs["info"] = xds.attrs["info"][0]
 
     if "xds" in name:
         xds.column_descriptions["UVW"]["shape"] = np.array(

xradio/vis/_vis_utils/_zarr/write.py

@@ -10,10 +10,21 @@ import zarr
 def write_part_keys(
     partitions: Dict[Any, xr.Dataset], outpath: str, compressor: numcodecs.abc.Codec
 ) -> None:
-    """Writes an xds with the partition keys.
+    """
+    Writes an xds with the partition keys.
+
+    Parameters
+    ----------
+    partitions : Dict[Any, xr.Dataset]
+        partitions from a cds
+    outpath : str
+        path to write a cds
+    compressor : numcodecs.abc.Codec
+        compressor used for the partition keys variable
+
+    Returns
+    -------
 
-    :param partitions: partitions from a cds
-    :param outpath: path to write a cds
     """
 
     spw_ids, pol_setup_ids, intents = map(list, zip(*partitions.keys()))
@@ -47,6 +58,23 @@ def write_metainfo(
 ) -> None:
     """
     Write all metainfo subtables from a cds to zarr storage
+
+    Parameters
+    ----------
+    outpath : str
+
+    metainfo : Dict[str, xr.Dataset]:
+
+    chunks_on_disk : Union[Dict, None] (Default value = None)
+
+    compressor : Union[numcodecs.abc.Codec, None) (Default value = None)
+
+    consolidated : bool (Default value = True)
+
+
+    Returns
+    -------
+
     """
     metadir = Path(outpath, "metainfo")
     os.mkdir(metadir)
@@ -67,6 +95,23 @@ def write_partitions(
 ) -> None:
     """
     Write all data partitions metainfo from a cds to zarr storage
+
+    Parameters
+    ----------
+    outpath : str :
+
+    partitions : Dict[str, xr.Dataset]
+
+    chunks_on_disk : Union[Dict, None] (Default value = None)
+
+    compressor : Union[numcodecs.abc.Codec, None] (Default value = True)
+
+    consolidated: bool (Default value = True)
+
+
+    Returns
+    -------
+
     """
 
     partdir = Path(outpath, "partitions")
@@ -92,11 +137,28 @@ def write_xds_to_zarr(
     """
     Write one xr dataset from a cds (either metainfo or a partition).
 
-    :param xds: cds (sub)dataset
-    :param name: dataset name (example subtable name, or xds{i})
-    :param graph_name: the time taken to execute the graph and save the
-    dataset is measured and saved as an attribute in the zarr file.
-    The graph_name is the label for this timing information.
+    Parameters
+    ----------
+    xds : xr.Dataset
+        cds (sub)dataset
+    name : str
+        dataset name (example subtable name, or xds{i})
+    outpath: str :
+
+    chunks_on_disk : Union[Dict, None] (Default value = None)
+
+    compressor : Union[numcodecs.abc.Codec, None] (Default value = None)
+
+    consolidated : bool (Default value = True)
+
+    graph_name : str
+        the time taken to execute the graph and save the
+        dataset is measured and saved as an attribute in the zarr file.
+        The graph_name is the label for this timing information.
+
+    Returns
+    -------
+
     """
 
     xds_for_disk = xds
@@ -159,8 +221,20 @@ def write_xds_to_zarr(
 
 
 def prepare_attrs_for_zarr(name: str, xds: xr.Dataset) -> xr.Dataset:
-    """Deal with types that cannot be serialized as they are in the
+    """
+    Deal with types that cannot be serialized as they are in the
     cds/xds (ndarray etc.)
+
+    Parameters
+    ----------
+    name : str
+
+    xds : xr.Dataset
+
+
+    Returns
+    -------
+
     """
     ctds_attrs = xds.attrs["other"]["msv2"]["ctds_attrs"]
     col_descrs = ctds_attrs["column_descriptions"]