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

xradio/vis/_vis_utils/_ms/_tables/read_main_table.py

@@ -17,7 +17,7 @@ from .read import (
 )
 
 from .table_query import open_table_ro, open_query
-from xradio.vis._vis_utils._ms.optimised_functions import (
+from xradio._utils.array import (
     unique_1d,
     pairing_function,
     inverse_pairing_function,
@@ -38,14 +38,22 @@ rename_msv2_cols = {
 
 
 def rename_vars(mvars: Dict[str, xr.DataArray]) -> Dict[str, xr.DataArray]:
-    """Apply rename rules. Also preserve ordering of data_vars
+    """
+    Apply rename rules. Also preserve ordering of data_vars
 
     Note: not using xr.DataArray.rename because we have optional
     column renames and rename complains if some of the names passed
     are not present in the dataset
 
-    :param mvars: dictionary of data_vars to be used to create an xr.Dataset
-    :return: similar dictionary after applying MSv2 => MSv3/ngCASA renaming rules
+    Parameters
+    ----------
+    mvars : Dict[str, xr.DataArray]
+        dictionary of data_vars to be used to create an xr.Dataset
+
+    Returns
+    -------
+    Dict[str, xr.DataArray]
+        similar dictionary after applying MSv2 => MSv3/ngCASA renaming rules
     """
     renamed = {
         rename_msv2_cols[name] if name in rename_msv2_cols else name: var
@@ -62,8 +70,16 @@ def redim_id_data_vars(mvars: Dict[str, xr.DataArray]) -> Dict[str, xr.DataArray
     The antenna id data vars:
      From MS (antenna1_id(time, baseline), antenna2_id(time,baseline)
      To cds (baseline_ant1_id(baseline), baseline_ant2_id(baseline)
-    :param mvars: data variables being prepared for a partition xds
-    :return: data variables with the ant id ones modified to cds type
+
+    Parameters
+    ----------
+    mvars : Dict[str, xr.DataArray]
+        data variables being prepared for a partition xds
+
+    Returns
+    -------
+    Dict[str, xr.DataArray]
+        data variables with the ant id ones modified to cds type
     """
     # Vars to drop baseline dim
     var_names = [
@@ -74,23 +90,34 @@ def redim_id_data_vars(mvars: Dict[str, xr.DataArray]) -> Dict[str, xr.DataArray
         "state_id",
     ]
     for vname in var_names:
-        mvars[vname] = mvars[vname].sel(baseline=0, drop=True)
+        if "baseline" in mvars[vname].coords:
+            mvars[vname] = mvars[vname].sel(baseline=0, drop=True)
 
     for idx in ["1", "2"]:
         new_name = f"baseline_ant{idx}_id"
         mvars[new_name] = mvars.pop(f"antenna{idx}_id")
-        mvars[new_name] = mvars[new_name].sel(time=0, drop=True)
+        if "time" in mvars[new_name].coords:
+            mvars[new_name] = mvars[new_name].sel(time=0, drop=True)
 
     return mvars
 
 
 def get_partition_ids(mtable: tables.table, taql_where: str) -> Dict:
-    """Get some of the partition IDs that we have to retrieve from some
+    """
+    Get some of the partition IDs that we have to retrieve from some
     of the top level ID/sorting cols of the main table of the MS.
 
-    :param mtable: MS main table
-    :param taql_where: where part that defines the partition in TaQL
-    :return: ids of array, observation, and processor
+    Parameters
+    ----------
+    mtable : tables.table
+        MS main table
+    taql_where : str
+        where part that defines the partition in TaQL
+
+    Returns
+    -------
+    Dict
+        ids of array, observation, and processor
     """
 
     taql_ids = f"select DISTINCT ARRAY_ID, OBSERVATION_ID, PROCESSOR_ID from $mtable {taql_where}"
@@ -130,6 +157,23 @@ def read_expanded_main_table(
     This is the expanded version (time, baseline) dims.
 
     Chunk tuple: (time, baseline, freq, pol)
+
+    Parameters
+    ----------
+    infile : str
+
+    ddi : int  (Default value = 0)
+
+    scan_state : Union[Tuple[int, int], None] (Default value = None)
+
+    ignore_msv2_cols: Union[list, None] (Default value = None)
+
+    chunks: Tuple[int, ...] (Default value = (400, 200, 100, 2))
+
+
+    Returns
+    -------
+    Tuple[xr.Dataset, Dict[str, Any], Dict[str, Any]]
     """
     if ignore_msv2_cols is None:
         ignore_msv2_cols = []
@@ -175,6 +219,23 @@ def read_main_table_chunks(
     """
     Iterates through the time,baseline chunks and reads slices from
     all the data columns.
+
+    Parameters
+    ----------
+    infile : str
+
+    tb_tool : tables.table
+
+    taql_where : str
+
+    ignore_msv2_cols : Union[list, None] (Default value = None)
+
+    chunks: Tuple[int, ...] (Default value = (400, 200, 100, 2))
+
+
+    Returns
+    -------
+    Tuple[xr.Dataset, Dict[str, Any]]
     """
     baselines = get_baselines(tb_tool)
 
@@ -196,7 +257,9 @@ def read_main_table_chunks(
     # loop over time chunks
     for time_chunk in range(0, n_unique_times, n_time_chunks):
         time_start = unique_times[time_chunk] - tol
-        time_end = unique_times[min(n_unique_times, time_chunk + n_time_chunks) - 1] + tol
+        time_end = (
+            unique_times[min(n_unique_times, time_chunk + n_time_chunks) - 1] + tol
+        )
 
         # chunk time length
         ctlen = min(n_unique_times, time_chunk + n_time_chunks) - time_chunk
@@ -244,6 +307,7 @@ def read_main_table_chunks(
 
     dims = ["time", "baseline", "freq", "pol"]
     mvars = concat_tvars_to_mvars(dims, tvars, pol_cnt, chan_cnt)
+
     mcoords = {
         "time": xr.DataArray(convert_casacore_time(unique_times), dims=["time"]),
         "baseline": xr.DataArray(np.arange(n_baselines), dims=["baseline"]),
@@ -277,16 +341,21 @@ def get_utimes_tol(mtable: tables.table, taql_where: str) -> Tuple[np.ndarray, f
 
 
 def get_baselines(tb_tool: tables.table) -> np.ndarray:
-    """Gets the unique baselines from antenna 1 and antenna 2 ids.
+    """
+    Gets the unique baselines from antenna 1 and antenna 2 ids.
 
     Uses a pairing function and inverse pairing function to decrease the
     computation time of finding unique values.
 
-    Args:
-        tb_tool (tables.table): MeasurementSet table to get the antenna ids.
+    Parameters
+    ----------
+    tb_tool : tables.table
+        MeasurementSet table to get the antenna ids.
 
-    Returns:
-        unique_baselines (np.ndarray): a 2D array of unique antenna pairs
+    Returns
+    -------
+    unique_baselines : np.ndarray
+        a 2D array of unique antenna pairs
         (baselines) from the MeasurementSet table provided.
     """
     ant1, ant2 = tb_tool.getcol("ANTENNA1", 0, -1), tb_tool.getcol("ANTENNA2", 0, -1)
@@ -310,19 +379,23 @@ def get_baselines(tb_tool: tables.table) -> np.ndarray:
 def get_baseline_indices(
     unique_baselines: np.ndarray, baseline_set: np.ndarray
 ) -> np.ndarray:
-    """Finds the baseline indices of a set of baselines using the unique baselines.
+    """
+    Finds the baseline indices of a set of baselines using the unique baselines.
 
     Uses a pairing function to reduce the number of values so it's more
     efficient to find the indices.
 
-    Args:
-        unique_baselines (np.ndarray): a 2D array of unique antenna pairs
-        (baselines).
-        baseline_set (np.ndarray): a 2D array of antenna pairs (baselines). This
-        array may contain duplicates.
-
-    Returns:
-        baseline_indices (np.ndarray): the indices of the baseline set that
+    Parameters
+    ----------
+    unique_baselines : np.ndarray
+        a 2D array of unique antenna pairs (baselines).
+    baseline_set : np.ndarray
+        a 2D array of antenna pairs (baselines). This array may contain duplicates.
+
+    Returns
+    -------
+    baseline_indices : np.ndarray
+        the indices of the baseline set that
         correspond to the unique baselines.
     """
     unique_baselines_paired = pairing_function(unique_baselines)
@@ -343,12 +416,31 @@ def read_all_cols_bvars(
     tb_tool: tables.table,
     chunks: Tuple[int, ...],
     chan_cnt: int,
-    ignore_msv2_cols,
+    ignore_msv2_cols: bool,
     delayed_params: Tuple,
     bvars: Dict[str, xr.DataArray],
 ) -> None:
     """
     Loops over each column and create delayed dask arrays
+
+    Parameters
+    ----------
+    tb_tool : tables.table
+
+    chunks : Tuple[int, ...]
+
+    chan_cnt : int
+
+    ignore_msv2_cols : bool
+
+    delayed_params : Tuple
+
+    bvars : Dict[str, xr.DataArray]
+
+
+    Returns
+    -------
+
     """
 
     col_names = tb_tool.colnames()
@@ -428,6 +520,17 @@ def concat_bvars_update_tvars(
     concats all the dask chunks from each baseline. This is intended to
     be called iteratively, for every time chunk iteration, once all the
     baseline chunks have been read.
+
+    Parameters
+    ----------
+    bvars: Dict[str, xr.DataArray]
+
+    tvars: Dict[str, xr.DataArray]
+
+
+    Returns
+    -------
+
     """
     for kk in bvars.keys():
         if len(bvars[kk]) == 0:
@@ -444,12 +547,21 @@ def concat_tvars_to_mvars(
     Concat into a single dask array all the dask arrays from each time
     chunk to make the final arrays of the xds.
 
-    :param dims: dimension names
-    :param tvars: variables as lists of dask arrays per time chunk
-    :param pol_cnt: len of pol axis/dim
-    :param chan_cnt: len of freq axis/dim (chan indices)
-
-    :return: variables as concated dask arrays
+    Parameters
+    ----------
+    dims : List[str]
+        dimension names
+    tvars : Dict[str, xr.DataArray]
+        variables as lists of dask arrays per time chunk
+    pol_cnt : int
+        len of pol axis/dim
+    chan_cnt : int
+        len of freq axis/dim (chan indices)
+
+    Returns
+    -------
+    Dict[str, xr.DataArray]
+        variables as concated dask arrays
     """
 
     mvars = {}
@@ -492,19 +604,38 @@ def read_flat_main_table(
     features may be missing and/or flaky.
 
     Chunk tuple: (row, freq, pol)
+
+    Parameters
+    ----------
+    infile : str
+
+    ddi : Union[int, None] (Default value = None)
+
+    scan_state : Union[Tuple[int, int], None] (Default value = None)
+
+    rowidxs : np.ndarray (Default value = None)
+
+    ignore_msv2_cols : Union[List, None] (Default value = None)
+
+    chunks : Tuple[int, ...] (Default value = (22000, 512, 2))
+
+    Returns
+    -------
+    Tuple[xr.Dataset, Dict[str, Any], Dict[str, Any]]
     """
     taql_where = f"where DATA_DESC_ID = {ddi}"
     if scan_state:
         # TODO: support additional intent/scan/subscan conditions if
         # we keep this read_flat functionality
-        _scans, states = scan_state
+        scans, states = scan_state
         # get row indices relative to full main table
-        if states:
-            taql_where += (
-                f" AND SCAN_NUMBER = {scan_state[0]} AND STATE_ID = {scan_state[1]}"
-            )
-        else:
-            taql_where += f" AND SCAN_NUMBER = {scan_state[0]}"
+        if type(states) == np.ndarray:
+            state_ids_or = " OR STATE_ID = ".join(np.char.mod("%d", states))
+            taql_where += f" AND (STATE_ID = {state_ids_or})"
+        elif states is not None:
+            taql_where += f" AND STATE_ID = {states}"
+        elif scans is not None:
+            taql_where += f" AND SCAN_NUMBER = {scans}"
 
     mtable = tables.table(
         infile, readonly=True, lockoptions={"option": "usernoread"}, ack=False
@@ -515,11 +646,11 @@ def read_flat_main_table(
         taql_rowid = f"select rowid() as ROWS from $mtable {taql_where}"
         with open_query(mtable, taql_rowid) as query_rows:
             rowidxs = query_rows.getcol("ROWS")
-            mtable.close()
 
     nrows = len(rowidxs)
     if nrows == 0:
-        return xr.Dataset()
+        mtable.close()
+        return xr.Dataset(), {}, {}
 
     part_ids = get_partition_ids(mtable, taql_where)
 
@@ -537,7 +668,7 @@ def read_flat_main_table(
                 (col, query_cols.getcol(col, 0, 1))
                 for col in cols
                 if (col not in ignore)
-                and (ignore_msv2_cols and (col not in ignore_msv2_cols))
+                and not (ignore_msv2_cols and col in ignore_msv2_cols)
             ]
         )
         chan_cnt, pol_cnt = [
@@ -635,7 +766,7 @@ def read_flat_main_table(
             )
 
     mvars["time"] = xr.DataArray(
-        convert_casacore_time(mvars["TIME"].values), dims=["row"]
+        convert_casacore_time(mvars["time"].values), dims=["row"]
     ).chunk({"row": chunks[0]})
 
     # add xds global attributes

xradio/vis/_vis_utils/_ms/_tables/read_subtables.py

@@ -19,7 +19,7 @@ from .read import (
     read_generic_table,
 )
 from .write import revert_time
-from xradio.vis._vis_utils._ms.optimised_functions import unique_1d
+from xradio._utils.array import unique_1d
 
 
 def read_ephemerides(
@@ -28,8 +28,15 @@ def read_ephemerides(
     """
     Read ephemerides info from MSv2 FIELD/EPHEMi_....tab subtables
 
-    :param infile: path to MS
-    :return: ephemerides xds with metainfo as in the MSv3/EPHEMERIDES subtable
+    Parameters
+    ----------
+    infile : str
+        path to MS
+
+    Returns
+    -------
+    Union[xr.Dataset, None]
+        ephemerides xds with metainfo as in the MSv3/EPHEMERIDES subtable
     """
     field_subt = Path(infile, "FIELD")
     subdirs = [
@@ -42,7 +49,7 @@ def read_ephemerides(
         logger.debug(f"Reading ephemerides info from: FIELD / {sdir.name}")
         # One "EPHEM_*.tab" (each with a difference ephemeris_id) to concatenate
         ephem.append(
-            read_generic_table(infile, str(Path("field", sdir)), timecols=["MJD"])
+            read_generic_table(infile, str(Path(*sdir.parts[-2:])), timecols=["MJD"])
         )
 
     if ephem:
@@ -62,9 +69,21 @@ def read_delayed_pointing_table(
     """
     Read MS pointing subtable in delayed arrays into an xr.Dataset
 
-    :param infile: path to pointing table
-    :rename_ids: dict with dimension renaming mapping
-    :param chunks: chunks for the arrays. Chunks tuple: time, antenna, data_vars_dim_1, data_vars_dim_2
+    Parameters
+    ----------
+    infile : str
+        path to pointing table
+    rename_ids : Dict[str, str] (Default value = None)
+        dict with dimension renaming mapping
+    chunks : Tuple (Default value = (10000, 100, 2, 20))
+        chunks for the arrays. Chunks tuple: time, antenna, data_vars_dim_1, data_vars_dim_2
+    time_slice: slice
+        time bounds
+
+    Returns
+    -------
+    xr.Dataset
+        pointing dataset
     """
 
     with open_table_ro(infile) as mtable:
@@ -124,18 +143,25 @@ def read_delayed_pointing_table(
     return xds
 
 
-def normalize_time_slice(mtable: tables.table, time_slice: slice):
+def normalize_time_slice(mtable: tables.table, time_slice: slice) -> slice:
     """
     If we get indices, produce the TIME column time value for the
     start/top indices. If we get timestamps, convert them to casacore
     refeference.
 
-    :param mtable: a casacore table from which we are reading a TIME
-    column
-    :param time_slice: slice giving start/stop time. Can be given as
-    integer indices or as timestamps (Xarray / pandas reference)
-
-    :return: a (start, stop) slice with times in casacore ref frame
+    Parameters
+    ----------
+    mtable : tables.table
+        a casacore table from which we are reading a TIME
+        column
+    time_slice : slice
+        slice giving start/stop time. Can be given as
+        integer indices or as timestamps (Xarray / pandas reference)
+
+    Returns
+    -------
+    slice
+        a (start, stop) slice with times in casacore ref frame
     """
     if type(time_slice.start) == pd.Timestamp and type(time_slice.stop) == pd.Timestamp:
         # Add tol?
@@ -189,12 +215,25 @@ def read_delayed_pointing_times(
     Read pointing table in delayed time / antenna chunks. Loops over
     time chunks
 
-    :param infile: path to pointing table
-    :param antennas: antenna ids
-    :param utimes: unique times from table
-    :param chunks: chunks for the arrays
-    :param query_all: table to read columns
-    :return: dictionary of columns=>variables (read as dask.delayed)
+    Parameters
+    ----------
+    infile : str
+        path to pointing table
+    antennas : np.ndarray
+        antenna ids
+    utimes : np.ndarray
+        unique times from table
+    chunks : tuple
+        chunks for the arrays
+    query_all : tables.table
+        table to read columns
+    time_slice: slice :
+        time bounds
+
+    Returns
+    -------
+    Dict[str, xr.DataArray]
+        dictionary of columns=>variables (read as dask.delayed)
     """
 
     antenna_chunks = range(0, len(antennas), chunks[1])
@@ -232,7 +271,7 @@ def read_delayed_pointing_chunks(
     infile: str,
     antennas: np.ndarray,
     chunks: tuple,
-    utimes: np.array,
+    utimes: np.ndarray,
     tc: int,
     tb_tool: tables.table,
 ) -> Dict[str, xr.DataArray]:
@@ -240,13 +279,25 @@ def read_delayed_pointing_chunks(
     For one time chunk, read the baseline/antenna chunks. Loops over
     antenna_id chunks and reads all columns as dask.delayed calls.
 
-    :param infile: path to pointing table
-    :param antennas: antenna ids
-    :param chunks: chunks for the arrays
-    :param utimes: unique times from table
-    :param tc: time index
-    :param tb_tool: table to read columns
-    :return: dictionary of columns=>variables (read as dask.delayed)
+    Parameters
+    ----------
+    infile : str
+        path to pointing table
+    antennas : np.ndarray
+        antenna ids
+    chunks : tuple
+        chunks for the arrays
+    utimes : np.ndarray
+        unique times from table
+    tc : int
+        time index
+    tb_tool : tables.table
+        table to read columns
+
+    Returns
+    -------
+    Dict[str, xr.DataArray]
+        dictionary of columns=>variables (read as dask.delayed)
     """
 
     # add a tol around the time ranges returned by taql, for the next taql queries