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

xradio/vis/_vis_utils/_ms/partition_queries.py

@@ -14,16 +14,156 @@ from ._tables.read import read_generic_table
 from .subtables import subt_rename_ids
 
 
+def enumerated_product(*args):
+    yield from zip(
+        itertools.product(*(range(len(x)) for x in args)), itertools.product(*args)
+    )
+
+
+def create_partitions(in_file: str, partition_scheme: Union[str, list], vla_otf=False):
+    """Create a list of dictionaries with the partition information.
+
+    Args:
+        in_file (str): Input MSv2 file path.
+        partition_scheme (Union[str, list]): Partition scheme to be used.
+        vla_otf (bool, optional):  The partioning of VLA OTF (on the fly) mosaics needs a special partitioning scheme. Defaults to False.
+
+    Returns:
+        list: list of dictionaries with the partition information.
+    """
+
+    # Create partition table
+    from casacore import tables
+    import numpy as np
+    import xarray as xr
+    import pandas as pd
+    import os
+
+    if partition_scheme == "ddi_intent_field":
+        partition_scheme = ["DATA_DESC_ID", "INTENT", "FIELD_ID"]
+    elif partition_scheme == "ddi_intent_scan":
+        partition_scheme = ["DATA_DESC_ID", "INTENT", "SCAN_NUMBER"]
+
+    # Open MSv2 tables and add columns to partition table (par_df):
+    par_df = pd.DataFrame()
+    main_tb = tables.table(
+        in_file, readonly=True, lockoptions={"option": "usernoread"}, ack=False
+    )
+    par_df["DATA_DESC_ID"] = main_tb.getcol("DATA_DESC_ID")
+    par_df["FIELD_ID"] = main_tb.getcol("FIELD_ID")
+    par_df["SCAN_NUMBER"] = main_tb.getcol("SCAN_NUMBER")
+    par_df["STATE_ID"] = main_tb.getcol("STATE_ID")
+    par_df = par_df.drop_duplicates()
+
+    field_tb = tables.table(
+        os.path.join(in_file, "FIELD"),
+        readonly=True,
+        lockoptions={"option": "usernoread"},
+        ack=False,
+    )
+    if vla_otf:
+        par_df["FIELD_NAME"] = np.array(field_tb.getcol("NAME"))[par_df["FIELD_ID"]]
+
+    # Get source ids if available from source table.
+    if os.path.isdir(os.path.join(os.path.join(in_file, "SOURCE"))):
+        source_tb = tables.table(
+            os.path.join(in_file, "SOURCE"),
+            readonly=True,
+            lockoptions={"option": "usernoread"},
+            ack=False,
+        )
+        if source_tb.nrows() != 0:
+            par_df["SOURCE_ID"] = field_tb.getcol("SOURCE_ID")[par_df["FIELD_ID"]]
+            if vla_otf:
+                par_df["SOURCE_NAME"] = np.array(source_tb.getcol("NAME"))[
+                    par_df["SOURCE_ID"]
+                ]
+
+    # Get intents and subscan numbers if available from state table.
+    if os.path.isdir(os.path.join(in_file, "STATE")):
+        state_tb = tables.table(
+            os.path.join(in_file, "STATE"),
+            readonly=True,
+            lockoptions={"option": "usernoread"},
+            ack=False,
+        )
+        if state_tb.nrows() != 0:
+            # print('state_tb',state_tb.nrows(),state_tb)
+            par_df["INTENT"] = np.array(state_tb.getcol("OBS_MODE"))[par_df["STATE_ID"]]
+            par_df["SUB_SCAN_NUMBER"] = state_tb.getcol("SUB_SCAN")[par_df["STATE_ID"]]
+        else:
+            par_df.drop(["STATE_ID"], axis=1)
+
+    # Check if all partition scheme criteria are present in the partition table.
+    partition_scheme_updated = []
+    partition_criteria = {}
+    for par in partition_scheme:
+        if par in par_df.columns:
+            partition_criteria[par] = par_df[par].unique()
+            partition_scheme_updated.append(par)
+    logger.info(f"Partition scheme that will be used: {partition_scheme_updated}")
+
+    # Make all possible combinations of the partition criteria.
+    enumerated_partitions = enumerated_product(*list(partition_criteria.values()))
+
+    # Create a list of dictionaries with the partition information. This will be used to query the MSv2 main table.
+    partitions = []
+    partition_axis_names = [
+        "DATA_DESC_ID",
+        "FIELD_ID",
+        "SCAN_NUMBER",
+        "STATE_ID",
+        "SOURCE_ID",
+        "INTENT",
+        "SUB_SCAN_NUMBER",
+    ]
+    for idx, pair in enumerated_partitions:
+        query = ""
+        for i, par in enumerate(partition_scheme_updated):
+            if isinstance(pair[i], str):
+                query = query + f'{par} == "{pair[i]}" and '
+            else:
+                query = query + f"{par} == {pair[i]} and "
+        query = query[:-4]  # remove last and
+        sub_par_df = par_df.query(query).drop_duplicates()
+
+        if sub_par_df.shape[0] != 0:
+            partition_info = {}
+
+            # FIELD_NAME	SOURCE_NAME
+            for col_name in partition_axis_names:
+                if col_name in sub_par_df.columns:
+                    partition_info[col_name] = sub_par_df[col_name].unique()
+                else:
+                    partition_info[col_name] = [None]
+
+            partitions.append(partition_info)
+
+    return partitions
+
+
+# Used by code that will be deprecated at some stage.
+
+
 def make_partition_ids_by_ddi_scan(
     infile: str, do_subscans: bool
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """Produces arrays of per-partition ddi, scan, state_id, for when
+    """
+    Produces arrays of per-partition ddi, scan, state_id, for when
     using partiion schemes 'scan' or 'scan/subscan', that is
     partitioning by some variant of (ddi, scan, subscan(state_id))
 
-    :param infile: Path to MS
-    :param do_subscans: also partitioning by subscan, not only scan
-    :return: arrays with indices that define every partition
+    Parameters
+    ----------
+    infile : str
+        Path to MS
+    do_subscans : bool
+        also partitioning by subscan, not only scan
+
+    Returns
+    -------
+    Tuple[np.ndarray, np.ndarray, np.ndarray]
+        arrays with indices that define every partition
     """
     try:
         cctable = None
@@ -66,55 +206,72 @@ def make_partition_ids_by_ddi_scan(
     return data_desc_id, scan_number, state_id
 
 
-def partition_when_empty_state(
-    infile: str,
+def make_partition_ids_by_ddi_intent(
+    infile: str, spw_names: xr.DataArray
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-    """Generate fallback partition ids when trying to partition by
-    'intent' but the STATE table is empty.
+    """
+    Produces arrays of per-partition ddi, scan, state_id, for when
+    using the partition scheme 'intents' (ddi, scan, subscans(state_ids))
 
-    Some MSs have no STATE rows and in the main table STATE_ID==-1
-    (that is not a valid MSv2 but it happens).
+    Parameters
+    ----------
+    infile : str
+        return: arrays with indices that define every partition
+    spw_names: xr.DataArray
 
-    :param infile: Path to the MS
-    :return: same as make_partition_ids_by_ddi_intent but with
-    effectively only ddi indices and other indices set to None ("any
-    IDs found")
+
+    Returns
+    -------
+    Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
+        arrays with indices that define every partition
     """
-    try:
-        main_table = None
+    # TODO: could explore other TAQL alternatives, like
+    # select ... from ::STATE where OBS_MODE = ...
+    #
+    # This will work only if intents are already alphabetically sorted (grouped),
+    # won't work for alternating intents:
+    # taql_intents = "select rowid() as ROWS from $state_tbl GROUPBY OBS_MODE "
 
-        main_table = tables.table(
-            infile, readonly=True, lockoptions={"option": "usernoread"}, ack=False
-        )
-        taql_ddis = "select DISTINCT DATA_DESC_ID from $main_table"
-        with open_query(main_table, taql_ddis) as query_per_intent:
-            # Will take whatever scans given the STATE_IDs and DDIs
-            # scan_number = query_per_intent.getcol("SCAN_NUMBER")
-            distinct_ddis = query_per_intent.getcol("DATA_DESC_ID")
+    with open_table_ro(str(Path(infile, "STATE"))) as state_tbl:
+        distinct_obs_mode = find_distinct_obs_mode(infile, state_tbl)
 
-        logger.debug(
-            f"Producing {len(distinct_ddis)} partitions for ddis: {distinct_ddis}"
-        )
-        nparts = len(distinct_ddis)
+        if distinct_obs_mode is None:
+            return partition_when_empty_state(infile)
 
-    finally:
-        if main_table:
-            main_table.close()
+        with open_table_ro(infile) as main_tbl:
+            (
+                data_desc_id,
+                state_id_partitions,
+                intent_names,
+            ) = make_ddi_state_intent_lists(
+                main_tbl, state_tbl, distinct_obs_mode, spw_names
+            )
 
-    return distinct_ddis, [None] * nparts, [None] * nparts, [""] * nparts
+    # Take whatever scans given by the STATE_IDs and DDIs
+    scan_number = [None] * len(state_id_partitions)
+
+    return data_desc_id, scan_number, state_id_partitions, intent_names
 
 
 def find_distinct_obs_mode(
     infile: str, state_table: tables.table
 ) -> Union[List[str], None]:
-    """Produce a list of unique "scan/subscan" intents.
+    """
+    Produce a list of unique "scan/subscan" intents.
 
-    :param infile: Path to the MS
-    :param state_table: casacore table object to read from
-    :return: List of unique "scan/subscan" intents as given in the
-    OBS_MODE column of the STATE subtable. None if the STATE subtable
-    is empty or there is a problem reading it
+    Parameters
+    ----------
+    infile : str
+        Path to the MS
+    state_table : tables.table
+        casacore table object to read from
 
+    Returns
+    -------
+    Union[List[str], None]
+        List of unique "scan/subscan" intents as given in the
+        OBS_MODE column of the STATE subtable. None if the STATE subtable
+        is empty or there is a problem reading it
     """
     taql_distinct_intents = "select DISTINCT OBS_MODE from $state_table"
     with open_query(state_table, taql_distinct_intents) as query_intents:
@@ -131,58 +288,14 @@ def find_distinct_obs_mode(
         return distinct_obs_mode
 
 
-def filter_intents_per_ddi(
-    ddis: List[int], substr: str, intents: str, spw_name_by_ddi: Dict[int, str]
-) -> List[str]:
-    """For a given pair of:
-    - substring (say 'WVR') associated with a type of intent we want to differentiate
-    - intents string (multiple comma-separated scan/subscan intents)
-    => do: for every DDI passed in the list of ddis, either keep only the
-    intents that have that substring (if there are any) or drop them, depending on
-    whether that substring is present in the SPW name. This is to filter in only
-    the intents that really apply to every DDI/SPW.
-
-    :param ddis: list of ddis for which the intents have to be filtered
-    :param substr: substring to filter by
-    :param intents: string with a comma-separated list of individual
-    scan/subscan intent strings (like scan/subscan intents as stored
-    in the MS STATE/OBS_MODE
-    :param spw_name_by_ddi: SPW names by DDI ID (row index) key
-    :return: list where the intents related to 'substr' have been filtered in our out
-
-    """
-    present = substr in intents
-    # Nothing to effectively filter, full cs-list of intents apply to all DDIs
-    if not present:
-        return [intents] * len(ddis)
-
-    every_intent = intents.split(",")
-    filtered_intents = []
-    for ddi in ddis:
-        spw_name = spw_name_by_ddi[ddi]
-        if not spw_name:
-            # we cannot say / cannot filter
-            filtered_intents.append(intents)
-            continue
-
-        # A not-xor to select/deselect (or keep-only/drop) the intents that apply
-        # to this DDI
-        ddi_intents = [
-            intnt for intnt in every_intent if (substr in intnt) == (substr in spw_name)
-        ]
-        ddi_intents = ",".join(ddi_intents)
-        filtered_intents.append(ddi_intents)
-
-    return filtered_intents
-
-
 def make_ddi_state_intent_lists(
     main_tbl: tables.table,
     state_tbl: tables.table,
     distinct_obs_mode: np.ndarray,
     spw_name_by_ddi: Dict[int, str],
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """Produce arrays of (ddi indices, state indices, intent string)
+    """
+    Produce arrays of (ddi indices, state indices, intent string)
     for every distinct intent string, where every item represents one
     partition of the main table
 
@@ -193,11 +306,21 @@ def make_ddi_state_intent_lists(
     intent is the only kept when the DDI/SPW has WVR in its name). See
     call to filter_intents_per_ddi()
 
-    :param main_tbl: main MS table openend as a casacore.tables.table
-    :state_tbl: STATE subtable openend as a casacore.tables.table
-    :param distinct_obs_mode: list of unique/distinct OBS_MODE strings from the STATE table
-    :return: arrays of (ddi indices, state indices, intent string)
+    Parameters
+    ----------
+    main_tbl : tables.table
+        main MS table openend as a casacore.tables.table
+    state_tbl : tables.table
+        STATE subtable openend as a casacore.tables.table
+    distinct_obs_mode : np.ndarray
+        list of unique/distinct OBS_MODE strings from the STATE table
+    spw_name_by_ddi: Dict[int, str]
 
+
+    Returns
+    -------
+    Tuple[np.ndarray, np.ndarray, np.ndarray]
+        arrays of (ddi indices, state indices, intent string)
     """
     data_desc_id, state_id_partitions, intent_names = [], [], []
     for intent in distinct_obs_mode:
@@ -233,41 +356,106 @@ def make_ddi_state_intent_lists(
     return data_desc_id, state_id_partitions, intent_names
 
 
-def make_partition_ids_by_ddi_intent(
-    infile: str, spw_names: xr.DataArray
-) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-    """Produces arrays of per-partition ddi, scan, state_id, for when
-    using the partition scheme 'intents' (ddi, scan, subscans(state_ids))
+def filter_intents_per_ddi(
+    ddis: List[int], substr: str, intents: str, spw_name_by_ddi: Dict[int, str]
+) -> List[str]:
+    """
+    For a given pair of:
+    - substring (say 'WVR') associated with a type of intent we want to differentiate
+    - intents string (multiple comma-separated scan/subscan intents)
+    => do: for every DDI passed in the list of ddis, either keep only the
+    intents that have that substring (if there are any) or drop them, depending on
+    whether that substring is present in the SPW name. This is to filter in only
+    the intents that really apply to every DDI/SPW.
 
-    :param infile:
-    :return: arrays with indices that define every partition
+    Parameters
+    ----------
+    ddis : List[int]
+        list of ddis for which the intents have to be filtered
+    substr : str
+        substring to filter by
+    intents : str
+        string with a comma-separated list of individual
+        scan/subscan intent strings (like scan/subscan intents as stored
+        in the MS STATE/OBS_MODE
+    spw_name_by_ddi : Dict[int, str]
+        SPW names by DDI ID (row index) key
+
+    Returns
+    -------
+    List[str]
+        list where the intents related to 'substr' have been filtered in our out
     """
-    # TODO: could explore other TAQL alternatives, like
-    # select ... from ::STATE where OBS_MODE = ...
-    #
-    # This will work only if intents are already alphabetically sorted (grouped),
-    # won't work for alternating intents:
-    # taql_intents = "select rowid() as ROWS from $state_tbl GROUPBY OBS_MODE "
+    present = substr in intents
+    # Nothing to effectively filter, full cs-list of intents apply to all DDIs
+    if not present:
+        return [intents] * len(ddis)
 
-    with open_table_ro(str(Path(infile, "STATE"))) as state_tbl:
-        distinct_obs_mode = find_distinct_obs_mode(infile, state_tbl)
+    every_intent = intents.split(",")
+    filtered_intents = []
+    for ddi in ddis:
+        spw_name = spw_name_by_ddi.get(ddi, "")
 
-        if distinct_obs_mode is None:
-            return partition_when_empty_state(infile)
+        if not spw_name:
+            # we cannot say / cannot filter
+            filtered_intents.append(intents)
+            continue
 
-        with open_table_ro(infile) as main_tbl:
-            (
-                data_desc_id,
-                state_id_partitions,
-                intent_names,
-            ) = make_ddi_state_intent_lists(
-                main_tbl, state_tbl, distinct_obs_mode, spw_names
-            )
+        # A not-xor to select/deselect (or keep-only/drop) the intents that apply
+        # to this DDI
+        ddi_intents = [
+            intnt for intnt in every_intent if (substr in intnt) == (substr in spw_name)
+        ]
+        ddi_intents = ",".join(ddi_intents)
+        filtered_intents.append(ddi_intents)
 
-    # Take whatever scans given by the STATE_IDs and DDIs
-    scan_number = [None] * len(state_id_partitions)
+    return filtered_intents
 
-    return data_desc_id, scan_number, state_id_partitions, intent_names
+
+def partition_when_empty_state(
+    infile: str,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Generate fallback partition ids when trying to partition by
+    'intent' but the STATE table is empty.
+
+    Some MSs have no STATE rows and in the main table STATE_ID==-1
+    (that is not a valid MSv2 but it happens).
+
+    Parameters
+    ----------
+    infile : str
+        Path to the MS
+
+    Returns
+    -------
+    Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
+        same as make_partition_ids_by_ddi_intent but with
+        effectively only ddi indices and other indices set to None ("any
+        IDs found")
+    """
+    try:
+        main_table = None
+
+        main_table = tables.table(
+            infile, readonly=True, lockoptions={"option": "usernoread"}, ack=False
+        )
+        taql_ddis = "select DISTINCT DATA_DESC_ID from $main_table"
+        with open_query(main_table, taql_ddis) as query_per_intent:
+            # Will take whatever scans given the STATE_IDs and DDIs
+            # scan_number = query_per_intent.getcol("SCAN_NUMBER")
+            distinct_ddis = query_per_intent.getcol("DATA_DESC_ID")
+
+        logger.debug(
+            f"Producing {len(distinct_ddis)} partitions for ddis: {distinct_ddis}"
+        )
+        nparts = len(distinct_ddis)
+
+    finally:
+        if main_table:
+            main_table.close()
+
+    return distinct_ddis, [None] * nparts, [None] * nparts, [""] * nparts
 
 
 def create_taql_query_and_file_name(out_file, intent, state_ids, field_id, ddi):
@@ -297,7 +485,8 @@ def create_taql_query_and_file_name(out_file, intent, state_ids, field_id, ddi):
 
 
 def get_unqiue_intents(in_file):
-    """_summary_
+    """
+    _summary_
 
     Parameters
     ----------
@@ -334,7 +523,8 @@ def enumerated_product(*args):
 
 
 def create_partition_enumerated_product(in_file: str, partition_scheme: str):
-    """Creates an enumerated_product of the data_desc_ids, state_ids, field_ids in a MS v2 that define the partions in a processing set.
+    """
+    Creates an enumerated_product of the data_desc_ids, state_ids, field_ids in a MS v2 that define the partions in a processing set.
 
     Parameters
     ----------
@@ -365,7 +555,6 @@ def create_partition_enumerated_product(in_file: str, partition_scheme: str):
         intents, state_ids = get_unqiue_intents(in_file)
         field_ids = np.arange(read_generic_table(in_file, "FIELD").sizes["row"])
     else:  # partition_scheme == "ddi_state_field"
-
         if len(state_xds.data_vars) > 0:
             state_ids = [np.arange(state_xds.sizes["row"])]
             intents = state_xds.obs_mode.values

xradio/vis/_vis_utils/_ms/partitions.py

@@ -29,8 +29,19 @@ def make_spw_names_by_ddi(ddi_xds: xr.Dataset, spw_xds: xr.Dataset) -> Dict[int,
 
 
 def split_intents(intents: str):
-    """Make a dict with two scan / subscan levels of intents from an
+    """
+    Make a dict with two scan / subscan levels of intents from an
     intent string from the STATE/OBS_MODE of an MS.
+
+    Parameters
+    ----------
+    intents : str
+        intents "OBS_MODE" string from an MS/STATE row
+
+    Returns
+    -------
+    Dict[str, list]
+        per scan intent list of individual subscan intent strings
     """
     sub_sep = "#"
     if sub_sep not in intents:
@@ -60,13 +71,25 @@ def make_part_key(
     partition_scheme: str,
     intent: str = "",
     scan_state: Union[Tuple, None] = None,
-):
+) -> PartitionKey:
     """
     Makes the key that a partition (sub)xds will have in the partitions dictionary of a cds.
 
-    :param xds: partition xds with data and attrs
-    :param partition_scheme: one of the schemes supported in the read_ms_*_partitions() functions
-    :param scan_state: scan/state ids, required when partition_scheme != 'ddi'
+    Parameters
+    ----------
+    xds : xr.Dataset
+        partition xds with data and attrs
+    partition_scheme : str
+        one of the schemes supported in the read_ms_*_partitions() functions
+    intent : str (Default value = "")
+        partition intent
+    scan_state : Union[Tuple, None] (Default value = None)
+        scan/state ids, required when partition_scheme != 'ddi'
+
+    Returns
+    -------
+    PartitionKey
+        partition key
     """
     spw_id = xds.attrs["partition_ids"]["spw_id"]
     pol_setup_id = xds.attrs["partition_ids"]["pol_setup_id"]
@@ -91,17 +114,28 @@ def read_ms_scan_subscan_partitions(
     expand: bool = False,
     chunks: Union[Tuple[int], List[int], None] = None,
 ) -> Tuple[VisSetPartitions, Dict[str, xr.Dataset], List[str]]:
-    """partitions per scan_number/subscans
+    """
+    partitions per scan_number/subscans
     (main table column SCAN_NUMBER / STATE_ID)
 
-    :param infile: MS path (main table)
-    :param partition_scheme: this functions can do 'intent', 'scan', and 'scan/subscan'
-    :param expand: wether to use (time, baseline) dimensions rather than 1d (row)
-    (only relevant when using the read_flat variant of read functions)
-    :param chunk: Dask chunking as tuple (time, baseline, chan, pol)
-    :return: a dictionary of partitions, a dict of subtable
-    xr.Datasets to use later for metainformation, and a list of the
-    subtables already read
+    Parameters
+    ----------
+    infile : str
+        MS path (main table)
+    partition_scheme : str
+        this functions can do 'intent', 'scan', and 'scan/subscan'
+    expand : bool (Default value = False)
+        wether to use (time, baseline) dimensions rather than 1d (row)
+        (only relevant when using the read_flat variant of read functions)
+    chunk : Union[Tuple[int], List[int], None] (Default value = None)
+        Dask chunking as tuple (time, baseline, chan, pol)
+
+    Returns
+    -------
+    Tuple[VisSetPartitions, Dict[str, xr.Dataset], List[str]]
+        a dictionary of partitions, a dict of subtable
+        xr.Datasets to use later for metainformation, and a list of the
+        subtables already read
     """
 
     spw_xds = read_generic_table(
@@ -140,13 +174,15 @@ def read_ms_scan_subscan_partitions(
         if partition_scheme == "intent":
             intent = distinct_intents[cnt]
             cnt += 1
+        else:
+            intent = ""
 
         if partition_scheme == "scan":
             scan_state = (scan, None)
         else:
             scan_state = (scan, state)
         # experimenting, comparing overheads of expanded vs. flat
-        expanded = True
+        expanded = not expand
         if expanded:
             xds, part_ids, attrs = read_expanded_main_table(
                 infile, ddi, scan_state=scan_state, ignore_msv2_cols=ignore_msv2_cols
@@ -207,11 +243,23 @@ def read_ms_ddi_partitions(
     from the DDIs. First looks into the SPECTRAL_WINDOW, POLARIZATION,
     DATA_DESCRIPTION tables to define the partitions.
 
-    :param expand: redimension (row)->(time,baseline)
-    :param rowmap: to be removed
-    :param chunks: array data chunk sizes
-    :return: dictionary of partitions, dict of subtable xr.Datasets to use later
-    for metainformation, and a list of the subtables already read
+    Parameters
+    ----------
+    infile : str
+        input MS path
+    expand : bool (Default value = False)
+        redimension (row)->(time,baseline)
+    rowmap : Union[dict, None] (Default value = None)
+        to be removed
+    chunks : Union[Tuple[int], List[int], None] (Default value = None)
+        array data chunk sizes
+
+    Returns
+    -------
+    Tuple[VisSetPartitions, Dict[str, xr.Dataset], List[str]]
+        dictionary of partitions, dict of subtable xr.Datasets to use later
+        for metainformation, and a list of the subtables already read
+
     """
     # we need the antenna, spectral window, polarization, and data description tables
     # to define the (sub)datasets (their dims and coords) and to process the main table
@@ -263,7 +311,7 @@ def read_ms_ddi_partitions(
         )
 
         # experimenting, comparing overheads of expanded vs. flat
-        expanded = True
+        expanded = not expand
         if expanded:
             xds, part_ids, attrs = read_expanded_main_table(
                 infile, ddi, ignore_msv2_cols=ignore_msv2_cols
@@ -319,13 +367,22 @@ def read_ms_ddi_partitions(
 def finalize_partitions(
     parts: Dict[str, xr.Dataset], subts: Dict[str, xr.Dataset]
 ) -> Dict[str, xr.Dataset]:
-    """Once the partition datasets and the metainfo/subtable datasets
+    """
+    Once the partition datasets and the metainfo/subtable datasets
     have been read, add to the partitions:
     - pointing variables from the pointing subtable
 
-    :param parts: partitions as xarray datasets, as read from an MS main table
-    :param subts: subtables of an MS read as xarray datasets
-    :return: partitions with additions taken from subtables
+    Parameters
+    ----------
+    parts : Dict[str, xr.Dataset]
+        partitions as xarray datasets, as read from an MS main table
+    subts : Dict[str, xr.Dataset]
+        subtables of an MS read as xarray datasets
+
+    Returns
+    -------
+    Dict[str, xr.Dataset]
+        partitions with additions taken from subtables
     """
     if "pointing" in subts:
         pointing = subts["pointing"]