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

xradio/measurement_set/_utils/_msv2/_tables/write_exp_api.py

@@ -1,388 +0,0 @@
-import os, time
-from typing import List, Optional, Union
-
-import dask
-import numpy as np
-import xarray as xr
-
-from ..._utils.xds_helper import flatten_xds, calc_optimal_ms_chunk_shape
-from .write import write_generic_table, write_main_table_slice
-from .write import create_table, revert_time
-
-try:
-    from casacore import tables
-except ImportError:
-    import xradio._utils._casacore.casacore_from_casatools as tables
-
-
-# TODO: this should be consolidated with the equivalent in read_main_table,
-# if we keep this mapping
-rename_to_msv2_cols = {
-    "antenna1_id": "ANTENNA1",
-    "antenna2_id": "ANTENNA2",
-    "feed1_id": "FEED1",
-    "feed2_id": "FEED2",
-    # optional cols:
-    # "WEIGHT": "WEIGHT_SPECTRUM",
-    "VIS_CORRECTED": "CORRECTED_DATA",
-    "VIS": "DATA",
-    "VIS_MODEL": "MODEL_DATA",
-    "AUTOCORR": "FLOAT_DATA",
-}
-# cols added in xds not in MSv2
-cols_not_in_msv2 = ["baseline_ant1_id", "baseline_ant2_id"]
-
-
-def cols_from_xds_to_ms(cols: List[str]) -> List[str]:
-    """
-    Translates between lowercase/uppercase convention
-    Rename some MS_colum_names <-> xds_data_var_names
-    Excludes the pointing_ vars that are in the xds but should not be written to MS
-    """
-    return {
-        rename_to_msv2_cols.get(col, col).upper(): col
-        for col in cols
-        if (col and col not in cols_not_in_msv2 and not col.startswith("pointing_"))
-    }
-
-
-def write_ms(
-    mxds: xr.Dataset,
-    outfile: str,
-    infile: str = None,
-    subtables: bool = False,
-    modcols: Union[List[str], None] = None,
-    verbose: bool = False,
-    execute: bool = True,
-) -> Optional[list]:
-    """
-    Write ms format xds contents back to casacore MS (CTDS - casacore Table Data System) format on disk
-
-    Parameters
-    ----------
-    mxds : xr.Dataset,
-        Source multi-xarray dataset (originally created by read_ms)
-    outfile : str
-        Destination filename
-    infile : Union[str, None] (Default value = None)
-        Source filename to copy subtables from. Generally faster than reading/writing through mxds via the subtables parameter. Default None
-        does not copy subtables to output.
-    subtables : bool (Default value = False)
-        Also write subtables from mxds. Default of False only writes mxds attributes that begin with xdsN to the MS main table.
-        Setting to True will write all other mxds attributes to subtables of the main table.  This is probably going to be SLOW!
-        Use infile instead whenever possible.
-    modcols : Union[List[str], None] (Default value = None)
-        List of strings indicating what column(s) were modified (aka xds data_vars). Different logic can be applied to speed up processing when
-        a data_var has not been modified from the input. Default None assumes everything has been modified (SLOW)
-    verbose : bool (Default value = False)
-        Whether or not to print output progress. Since writes will typically execute the DAG, if something is
-        going to go wrong, it will be here.  Default False
-    execute : bool (Default value = True)
-        Whether or not to actually execute the DAG, or just return it with write steps appended. Default True will execute it
-
-    Returns
-    -------
-    Optional[list]
-        delayed write functions
-    """
-    outfile = os.path.expanduser(outfile)
-    if verbose:
-        print("initializing output...")
-    start = time.time()
-
-    xds_list = [flatten_xds(xds) for _key, xds in mxds.partitions.items()]
-
-    cols = cols_from_xds_to_ms(
-        list(set([dv for dx in xds_list for dv in dx.data_vars]))
-    )
-    if modcols is None:
-        modcols = cols
-
-    # create an empty main table with enough space for all desired xds partitions
-    # the first selected xds partition will be passed to create_table to provide a definition of columns and table keywords
-    # we first need to add in additional keywords for the selected subtables that will be written as well
-    max_rows = np.sum([dx.row.shape[0] for dx in xds_list])
-    create_table(
-        outfile, xds_list[0], max_rows=max_rows, infile=infile, cols=cols, generic=False
-    )
-
-    # start a list of dask delayed writes to disk (to be executed later)
-    # the SPECTRAL_WINDOW, POLARIZATION, and DATA_DESCRIPTION tables must always be present and will always be written
-    delayed_writes = [
-        dask.delayed(write_generic_table)(
-            mxds.metainfo["spectral_window"], outfile, "SPECTRAL_WINDOW", cols=None
-        )
-    ]
-    delayed_writes += [
-        dask.delayed(write_generic_table)(
-            mxds.metainfo["polarization"], outfile, "POLARIZATION", cols=None
-        )
-    ]
-    # should data_description be kept somewhere (in attrs?) or rebuilt?
-    # delayed_writes += [
-    #     dask.delayed(write_generic_table)(
-    #         mxds.metainfo["data_description"], outfile, "DATA_DESCRIPTION", cols=None
-    #     )
-    # ]
-    if subtables:  # also write the rest of the subtables
-        for subtable in list(mxds.attrs.keys()):
-            if (
-                subtable.startswith("xds")
-                or (subtable in ["spectral_window", "polarization", "data_description"])
-                or not isinstance(subtable, xr.Dataset)
-            ):
-                continue
-
-            if verbose:
-                print("writing subtable %s..." % subtable)
-            delayed_writes += [
-                dask.delayed(write_generic_table)(
-                    mxds.attrs[subtable], outfile, subtable, cols=None
-                )
-            ]
-
-    ddi_row_start = 0  # output rows will be ordered by DDI
-    for xds in xds_list:
-        txds = xds.copy().unify_chunks()
-        # TODO: carry over or rebuild?
-        ddi = 0  # txds.data_desc_id[:1].values[0]
-
-        # serial write entire DDI column first so subsequent delayed writes can find their spot
-        if verbose:
-            print("setting up DDI %i..." % ddi)
-
-        # write each chunk of each modified data_var, triggering the DAG along the way
-        for col in modcols:
-            if col not in txds:
-                continue  # this can happen with bad_cols, should still be created in create_table()
-
-            if col in cols_not_in_msv2:
-                continue
-
-            chunks = txds[col].chunks
-            dims = txds[col].dims
-            for d0 in range(len(chunks[0])):
-                d0start = ([0] + list(np.cumsum(chunks[0][:-1])))[d0]
-
-                for d1 in range(len(chunks[1]) if len(chunks) > 1 else 1):
-                    d1start = (
-                        ([0] + list(np.cumsum(chunks[1][:-1])))[d1]
-                        if len(chunks) > 1
-                        else 0
-                    )
-
-                    for d2 in range(len(chunks[2]) if len(chunks) > 2 else 1):
-                        d2start = (
-                            ([0] + list(np.cumsum(chunks[2][:-1])))[d2]
-                            if len(chunks) > 2
-                            else 0
-                        )
-
-                        starts = [d0start, d1start, d2start]
-                        lengths = [
-                            chunks[0][d0],
-                            (chunks[1][d1] if len(chunks) > 1 else 0),
-                            (chunks[2][d2] if len(chunks) > 2 else 0),
-                        ]
-                        slices = [
-                            slice(starts[0], starts[0] + lengths[0]),
-                            slice(starts[1], starts[1] + lengths[1]),
-                            slice(starts[2], starts[2] + lengths[2]),
-                        ]
-                        txda = txds[col].isel(
-                            dict(zip(dims, slices)), missing_dims="ignore"
-                        )
-                        starts[0] = starts[0] + ddi_row_start  # offset to end of table
-                        delayed_writes += [
-                            dask.delayed(write_main_table_slice)(
-                                txda,
-                                outfile,
-                                ddi=ddi,
-                                col=col,
-                                full_shape=txds[col].shape[1:],
-                                starts=starts,
-                            )
-                        ]
-
-        # now write remaining data_vars from the xds that weren't modified
-        # this can be done faster by collapsing the chunking to maximum size (minimum #) possible
-        max_chunk_size = np.prod(
-            [txds.chunks[kk][0] for kk in txds.chunks if kk in ["row", "freq", "pol"]]
-        )
-        for col in list(np.setdiff1d(list(cols), modcols)):
-            if col not in txds:
-                continue  # this can happen with bad_cols, should still be created in create_table()
-
-            if col in cols_not_in_msv2:
-                continue
-
-            col_chunk_size = np.prod([kk[0] for kk in txds[col].chunks])
-            if max_chunk_size <= 0:
-                max_chunk_size = 19200
-            if col_chunk_size <= 0:
-                col_rows = max_chunk_size
-            else:
-                col_rows = (
-                    int(np.ceil(max_chunk_size / col_chunk_size))
-                    * txds[col].chunks[0][0]
-                )
-            for rr in range(0, txds[col].row.shape[0], col_rows):
-                txda = txds[col].isel(row=slice(rr, rr + col_rows))
-                delayed_writes += [
-                    dask.delayed(write_main_table_slice)(
-                        txda,
-                        outfile,
-                        ddi=ddi,
-                        col=rename_to_msv2_cols.get(col, col).upper(),
-                        full_shape=txda.shape[1:],
-                        starts=(rr + ddi_row_start,) + (0,) * (len(txda.shape) - 1),
-                    )
-                ]
-
-        ddi_row_start += txds.row.shape[0]  # next xds will be appended after this one
-
-    if execute:
-        if verbose:
-            print("triggering DAG...")
-        zs = dask.compute(delayed_writes)
-        if verbose:
-            print(
-                "execution time %0.2f sec. Compute result len: %d"
-                % ((time.time() - start), len(zs))
-            )
-    else:
-        if verbose:
-            print("returning delayed task list")
-        return delayed_writes
-
-
-def write_ms_serial(
-    mxds: xr.Dataset,
-    outfile: str,
-    infile: str = None,
-    subtables: bool = False,
-    verbose: bool = False,
-    execute: bool = True,
-    memory_available_in_bytes: int = 500000000000,
-):
-    """
-    Write ms format xds contents back to casacore table format on disk
-
-    Parameters
-    ----------
-    mxds : xr.Dataset
-        Source multi-xarray dataset (originally created by read_ms)
-    outfile : str
-        Destination filename
-    infile : str (Default value = None)
-        Source filename to copy subtables from. Generally faster than reading/writing through mxds via the subtables parameter. Default None
-        does not copy subtables to output.
-    subtables : bool (Default value = False)
-        Also write subtables from mxds. Default of False only writes mxds attributes that begin with xdsN to the MS main table.
-        Setting to True will write all other mxds attributes to subtables of the main table.  This is probably going to be SLOW!
-        Use infile instead whenever possible.
-    verbose : bool (Default value = False)
-        Whether or not to print output progress. Since writes will typically execute the DAG, if something is
-        going to go wrong, it will be here.  Default False
-
-    execute : bool (Default value = True)
-        Whether or not to actually execute the DAG, or just return it with write steps appended. Default True will execute it
-    memory_available_in_bytes : (Default value = 500000000000)
-
-    Returns
-    -------
-
-    """
-
-    print("*********************")
-    outfile = os.path.expanduser(outfile)
-    if verbose:
-        print("initializing output...")
-    # start = time.time()
-
-    xds_list = [flatten_xds(xds) for _key, xds in mxds.partitions.items()]
-    cols = list(set([dv for dx in xds_list for dv in dx.data_vars]))
-    cols = cols_from_xds_to_ms(list(np.atleast_1d(cols)))
-
-    # create an empty main table with enough space for all desired xds partitions
-    # the first selected xds partition will be passed to create_table to provide a definition of columns and table keywords
-    # we first need to add in additional keywords for the selected subtables that will be written as well
-    max_rows = np.sum([dx.row.shape[0] for dx in xds_list])
-    create_table(
-        outfile, xds_list[0], max_rows=max_rows, infile=infile, cols=cols, generic=False
-    )
-
-    # start a list of dask delayed writes to disk (to be executed later)
-    # the SPECTRAL_WINDOW, POLARIZATION, and DATA_DESCRIPTION tables must always be present and will always be written
-    write_generic_table(
-        mxds.metainfo["spectral_window"], outfile, "SPECTRAL_WINDOW", cols=None
-    )
-    write_generic_table(
-        mxds.metainfo["polarization"], outfile, "POLARIZATION", cols=None
-    )
-    # should data_description be kept somewhere (in attrs?) or rebuilt?
-    # write_generic_table(mxds.metainfo.data_description, outfile, "DATA_DESCRIPTION", cols=None)
-
-    if subtables:  # also write the rest of the subtables
-        # for subtable in list(mxds.attrs.keys()):
-        #'OBSERVATION','STATE'
-        # ['FEED','OBSERVATION','FIELD','ANTENNA','HISTORY','STATE']
-        # ['FEED','FIELD','ANTENNA','HISTORY']
-        # ,'FIELD','ANTENNA'
-        # for subtable in ['OBSERVATION']:
-        for subtable in list(mxds.metainfo.keys()):
-            if subtable.startswith("xds") or (
-                subtable in ["spectral_window", "polarization", "data_description"]
-            ):
-                continue
-            if verbose:
-                print("writing subtable %s..." % subtable)
-            # print(subtable)
-            # print(mxds.attrs[subtable])
-            try:
-                write_generic_table(
-                    mxds.metainfo[subtable], outfile, subtable.upper(), cols=None
-                )
-            except (RuntimeError, KeyError) as exc:
-                print(f"Exception writing subtable {subtable}: {exc}")
-
-    part_key0 = next(iter(mxds.partitions))
-    vis_data_shape = mxds.partitions[part_key0].VIS.shape
-    rows_chunk_size = calc_optimal_ms_chunk_shape(
-        memory_available_in_bytes, vis_data_shape, 16, "DATA"
-    )
-
-    # print(rows_chunk_size)
-    # rows_chunk_size = 200000000
-    # write each chunk of each modified data_var, triggering the DAG along the way
-    tbs = tables.table(
-        outfile, readonly=False, lockoptions={"option": "permanentwait"}, ack=True
-    )
-
-    start_main = time.time()
-    for col, var_name in cols.items():
-        xda = mxds.partitions[part_key0][var_name]
-        # print(col,xda.dtype)
-
-        for start_row in np.arange(0, vis_data_shape[0], rows_chunk_size):
-            end_row = start_row + rows_chunk_size
-            if end_row > vis_data_shape[0]:
-                end_row = vis_data_shape[0]
-
-            # start = time.time()
-            values = xda[start_row:end_row,].compute().values
-            if xda.dtype == "datetime64[ns]":
-                values = revert_time(values)
-            # print('1. Time', time.time()-start, values.shape)
-
-            # start = time.time()
-            try:
-                tbs.putcol(col, values, start_row, len(values))
-                # print('2. Time', time.time()-start)
-            except RuntimeError as exc:
-                print(f"Exception writing main table column {col}: {exc}")
-
-    print("3. Time", time.time() - start_main)
-
-    tbs.unlock()
-    tbs.close()

xradio/measurement_set/_utils/_msv2/chunks.py

@@ -1,115 +0,0 @@
-from pathlib import Path
-from typing import Dict, Tuple
-
-import xarray as xr
-
-
-from .subtables import subt_rename_ids
-from ._tables.read import load_generic_table
-from ._tables.read_subtables import read_delayed_pointing_table
-
-
-def read_spw_ddi_ant_pol(inpath: str) -> Tuple[xr.Dataset]:
-    """
-    Reads the four metainfo subtables needed to load data chunks into xdss.
-
-    Parameters
-    ----------
-    inpath : str
-        MS path (main table)
-
-    Returns
-    -------
-    Tuple[xr.Dataset]
-        tuple with antenna, ddi, spw, and polarization setup subtables info
-    """
-    spw_xds = load_generic_table(
-        inpath,
-        "SPECTRAL_WINDOW",
-        rename_ids=subt_rename_ids["SPECTRAL_WINDOW"],
-    )
-    ddi_xds = load_generic_table(inpath, "DATA_DESCRIPTION")
-    ant_xds = load_generic_table(
-        inpath, "ANTENNA", rename_ids=subt_rename_ids["ANTENNA"]
-    )
-    pol_xds = load_generic_table(
-        inpath, "POLARIZATION", rename_ids=subt_rename_ids["POLARIZATION"]
-    )
-    return ant_xds, ddi_xds, spw_xds, pol_xds
-
-
-def finalize_chunks(
-    infile: str, chunks: Dict[str, xr.Dataset], chunk_spec: Dict[str, slice]
-) -> Dict[Tuple[int, int], xr.Dataset]:
-    """
-    Adds pointing variables to a dictionary of chunk xdss. This is
-    intended to be added after reading chunks from an MS main table.
-
-    Parameters
-    ----------
-    infile : str
-        MS path (main table)
-    chunks : Dict[str, xr.Dataset]
-        chunk xdss
-    chunk_spec : Dict[str, slice]
-        specification of chunk to load
-
-    Returns
-    -------
-    Dict[Tuple[int, int], xr.Dataset]
-        dictionary of chunk xdss where every xds now has pointing
-        data variables
-    """
-    pnt_name = "POINTING"
-    pnt_path = Path(infile, pnt_name)
-    if "time" in chunk_spec:
-        time_slice = chunk_spec["time"]
-    else:
-        time_slice = None
-    pnt_xds = read_delayed_pointing_table(
-        str(pnt_path),
-        rename_ids=subt_rename_ids.get(pnt_name, None),
-        time_slice=time_slice,
-    )
-
-    if "time" not in pnt_xds.dims:
-        return xr.Dataset()
-
-    pnt_xds = pnt_xds.compute()
-
-    pnt_chunks = {
-        key: finalize_chunk_xds(infile, xds, pnt_xds)
-        for _idx, (key, xds) in enumerate(chunks.items())
-    }
-
-    return pnt_chunks
-
-
-def finalize_chunk_xds(
-    infile: str, chunk_xds: xr.Dataset, pointing_xds: xr.Dataset
-) -> xr.Dataset:
-    """
-    Adds pointing variables to one chunk xds.
-
-    Parameters
-    ----------
-    infile : str
-        MS path (main table)
-    xds_chunk : xr.Dataset
-        chunks xds
-    pointing_xds : xr.Dataset
-        pointing (sub)table xds
-
-    Returns
-    -------
-    xr.Dataset
-        chunk xds with pointing data variables interpolated form
-        the pointing (sub)table
-    """
-
-    interp_pnt = pointing_xds.interp(time=chunk_xds.time, method="nearest")
-
-    for var in interp_pnt.data_vars:
-        chunk_xds[f"pointing_{var}"] = interp_pnt[var]
-
-    return chunk_xds

xradio/measurement_set/_utils/_msv2/descr.py

@@ -1,165 +0,0 @@
-import os
-from typing import Dict, Union
-
-import numpy as np
-import pandas as pd
-import xarray as xr
-
-from ._tables.read import load_generic_table, read_flat_col_chunk
-from ._tables.table_query import open_query, open_table_ro
-from xradio._utils.list_and_array import unique_1d
-
-
-def describe_ms(
-    infile: str, mode: str = "summary", rowmap: Union[dict, None] = None
-) -> Union[pd.DataFrame, Dict]:
-    """
-    Summarize the contents of an MS directory in casacore table format
-
-    Parameters
-    ----------
-    infile : str
-        input MS filename
-    mode : str (Default value = "summary")
-        type of information returned ('summary', 'flat', 'expanded').
-        'summary' returns a pandas dataframe that is nice for displaying in notebooks
-        etc. 'flat' returns a list of tuples of (ddi, row, chan, pol). 'expanded'
-        returns a list of tuples of (ddi, time, baseline, chan, pol). These latter two
-        are good for trying to determine chunk size for read_ms(expand=True/False). (Default value = "summary")
-    rowmap : Union[dict, None] (Default value = None)
-        dict of DDI to tuple of (row indices, channel indices). Returned
-        by ms_selection function. Default None ignores selections
-
-    Returns
-    -------
-    Union[pd.DataFrame, Dict]
-        summary as a pd dataframe
-    """
-    infile = os.path.expanduser(infile)  # does nothing if $HOME is unknown
-    if not os.path.isdir(infile):
-        raise ValueError(f"invalid input filename to describe_ms: {infile}")
-    if mode not in [
-        "summary",
-        "flat",
-        "expanded",
-    ]:
-        raise ValueError("invalid mode, must be summary, flat or expanded")
-
-    ddi_xds = load_generic_table(infile, "DATA_DESCRIPTION")
-    ddis = list(ddi_xds.row.values) if rowmap is None else list(rowmap.keys())
-    summary: Union[pd.DataFrame, Dict] = []
-    if mode == "summary":
-        summary = pd.DataFrame([])
-
-    all_sdf = []
-    with open_table_ro(infile) as tb_tool:
-        for ddi in ddis:
-            taql = f"select * from $tb_tool where DATA_DESC_ID = {ddi}"
-            with open_query(tb_tool, taql) as query_per_ddi:
-                sdf = populate_ms_descr(
-                    infile, mode, query_per_ddi, summary, ddi, ddi_xds
-                )
-                all_sdf.append(sdf)
-
-    if mode == "summary":
-        summary = pd.DataFrame(all_sdf)
-        summary = summary.set_index("ddi").sort_index()
-    else:
-        summary = dict(summary)
-    return summary
-
-
-def populate_ms_descr(
-    infile: str,
-    mode: str,
-    query_per_ddi,
-    summary: dict,
-    ddi: int,
-    ddi_xds: xr.Dataset,
-    rowmap: Union[Dict, None] = None,
-) -> pd.DataFrame:
-    """
-    Adds information from the time and baseline (antenna1+antenna2)
-    columns as well as channel and polarizations, based on a taql
-    query.
-
-    Parameters
-    ----------
-    infile : str
-        input table/MS path
-    mode : str
-        mode (as in describe_ms())
-    query_per_ddi :
-        a TaQL query with data per individual DDI
-    summary : Dict
-        summary dict being populated
-    ddi_xds : xr.Dataset
-        final summary object being populated from the invividual sdf's
-
-    Returns
-    -------
-    pd.DataFrame
-    """
-    spw_ids = ddi_xds.SPECTRAL_WINDOW_ID.values
-    pol_ids = ddi_xds.POLARIZATION_ID.values
-    sdf = {
-        "ddi": ddi,
-        "spw_id": spw_ids[ddi],
-        "pol_id": pol_ids[ddi],
-        "rows": query_per_ddi.nrows(),
-    }
-
-    # figure out characteristics of main table from select subtables (must all be present)
-    spw_xds = load_generic_table(infile, "SPECTRAL_WINDOW")
-    pol_xds = load_generic_table(infile, "POLARIZATION")
-
-    if mode in ["expanded", "summary"]:
-        times = (
-            query_per_ddi.getcol("TIME")
-            if rowmap is None
-            else read_flat_col_chunk(infile, "TIME", (1,), rowmap[ddi][0], 0, 0)
-        )
-        baselines = [
-            (
-                query_per_ddi.getcol(rr)[:, None]
-                if rowmap is None
-                else read_flat_col_chunk(infile, rr, (1,), rowmap[ddi][0], 0, 0)
-            )
-            for rr in ["ANTENNA1", "ANTENNA2"]
-        ]
-        sdf.update(
-            {
-                "times": len(unique_1d(times)),
-                "baselines": len(np.unique(np.hstack(baselines), axis=0)),
-            }
-        )
-
-    chans = spw_xds.NUM_CHAN.values
-    pols = pol_xds.NUM_CORR.values
-    sdf.update(
-        {
-            "chans": (
-                chans[spw_ids[ddi]]
-                if (rowmap is None) or (rowmap[ddi][1] is None)
-                else len(rowmap[ddi][1])
-            ),
-            "pols": pols[pol_ids[ddi]],
-        }
-    )
-    sdf["size_MB"] = np.ceil(
-        (sdf["rows"] * sdf["chans"] * sdf["pols"] * 10) / 1024**2
-    ).astype(int)
-
-    if rowmap is not None:
-        sdf["rows"] = len(rowmap[ddi][0])
-
-    if mode == "summary":
-        summary = pd.concat(
-            [summary, pd.DataFrame(sdf, index=[str(ddi)])], axis=0, sort=False
-        )
-    elif mode == "flat":
-        summary += [(ddi, (sdf["rows"], sdf["chans"], sdf["pols"]))]
-    else:
-        summary += [(ddi, sdf["times"], sdf["baselines"], sdf["chans"], sdf["pols"])]
-
-    return sdf

xradio/measurement_set/_utils/_msv2/msv2_msv3.py

@@ -1,7 +0,0 @@
-ignore_msv2_cols = [
-    "FLAG_ROW",
-    "FLAG_CATEGORY",
-    "TIME_EXTRA_PREC",
-    "DATA_DESC_ID",
-    "SIGMA",
-]