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

xradio/vis/_vis_utils/ms.py

@@ -24,7 +24,8 @@ def read_ms(
     expand: bool = False,
     **kwargs: str,
 ) -> CASAVisSet:
-    """Read a MeasurementSet (MSv2 format) into a next generation CASA
+    """
+    Read a MeasurementSet (MSv2 format) into a next generation CASA
     dataset (visibilities dataset as a set of Xarray datasets).
 
     The MS is partitioned into multiple sub- Xarray datasets (where the data variables are read as
@@ -33,28 +34,37 @@ def read_ms(
     and polarizations) and, subject to experimentation, by scan and subscan. This results in multiple
     partitions as xarray datasets (xds) contained within a main xds (mxds).
 
-    :param infile: Input MS filename
-    :param subtables: Also read and include subtables along with main table selection. Default False will
-    omit subtables (faster)
-    :param asdm_subtables: in addition to MeasurementSet subtables (if enabled), also read extension
-    subtables named "ASDM_*"
-    :param partition_scheme: (experimenting) Whether to partition sub-xds datasets by scan/subscan
-    (in addition to DDI), or other alternative partitioning schemes. Accepted values: 'scan/subscan',
-    'scan', 'ddi', 'intent'. Default: 'intent'
-    :param chunks: Can be used to set a specific chunk shape (with a tuple of ints), or to control the
-    optimization used for automatic chunking (with a list of ints). A tuple of ints in the form of (row,
-    chan, pol) will use a fixed chunk shape. A list or numpy array of ints in the form of [idx1, etc]
-    will trigger auto-chunking optimized for the given indices, with row=0, chan=1, pol=2. Default None
-    uses auto-chunking with a best fit across all dimensions (probably sub-optimal for most cases).
-    :param expand: (to be removed) Whether or not to return the original flat row structure of the MS (False)
-    or expand the rows to time x baseline dimensions (True). Expanding the rows allows for easier indexing
-    and parallelization across time and baseline dimensions, at the cost of some conversion time. Default
-    False
-    :param **kwargs: (to be removed?) Selection parameters from the standard way of making CASA MS
-    selections. Supported keys are: spw, field, scan, baseline, time, scanintent, uvdist, polarization,
-    array, observation.  Values are strings.
-
-    :return: Main xarray dataset of datasets for this visibility dataset
+    Parameters
+    ----------
+    infile : str
+        Input MS filename
+    subtables : bool (Default value = True)
+        Also read and include subtables along with main table selection. Default False will
+        omit subtables (faster)
+    asdm_subtables : bool (Default value = False)
+        in addition to MeasurementSet subtables (if enabled), also read extension
+        subtables named "ASDM_*"
+    partition_scheme : str (Default value = "intent")
+        experimenting) Whether to partition sub-xds datasets by scan/subscan
+        (in addition to DDI), or other alternative partitioning schemes. Accepted values: 'scan/subscan',
+        'scan', 'ddi', 'intent'. Default: 'intent'
+    chunks : Union[Tuple[int], List[int]] (Default value = None)
+        Can be used to set a specific chunk shape (with a tuple of ints), or to control the
+        optimization used for automatic chunking (with a list of ints). A tuple of ints in the form of (row,
+        chan, pol) will use a fixed chunk shape. A list or numpy array of ints in the form of [idx1, etc]
+        will trigger auto-chunking optimized for the given indices, with row=0, chan=1, pol=2. Default None
+        uses auto-chunking with a best fit across all dimensions (probably sub-optimal for most cases).
+    expand : bool (Default value = False)
+        to be removed) Whether or not to return the original flat row structure of the MS (False)
+        or expand the rows to time x baseline dimensions (True). Expanding the rows allows for easier indexing
+        and parallelization across time and baseline dimensions, at the cost of some conversion time.
+    **kwargs: str :
+
+
+    Returns
+    -------
+    CASAVisSet
+        Main xarray dataset of datasets for this visibility dataset
     """
 
     infile = os.path.expanduser(infile)
@@ -104,14 +114,23 @@ def load_vis_chunk(
     block_des: Dict[str, slice],
     partition_key: Tuple[int, int, str],
 ) -> Dict[Tuple[int, int], xr.Dataset]:
-    """Read a chunk of a MeasurementSet (MSv2 format) into an Xarray
+    """
+    Read a chunk of a MeasurementSet (MSv2 format) into an Xarray
     dataset, loading the data in memory.
 
-    :param infile: Input MS filename
-    :param block_des: specification of chunk to load
-
-    :return: Xarray datasets with chunk of visibility data, one per DDI
-    (spw_id, pol_setup_id pair)
+    Parameters
+    ----------
+    infile : str
+        Input MS filename
+    block_des : Dict[str, slice]
+        specification of chunk to load
+    partition_key: partition_key: Tuple[int, int, str]
+
+    Returns
+    -------
+    Dict[Tuple[int, int], xr.Dataset]
+        Xarray datasets with chunk of visibility data, one per DDI
+        (spw_id, pol_setup_id pair)
     """
     infile = os.path.expanduser(infile)
 

xradio/vis/_vis_utils/zarr.py

@@ -11,13 +11,19 @@ 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) -> bool:
+def is_zarr_vis(inpath: str) -> bool:
     """
     Check if a given path has a visibilities dataset in Zarr format
 
-    :param inpath: path to a (possibly) Zarr vis dataset
+    Parameters
+    ----------
+    inpath : str
+        path to a (possibly) Zarr vis dataset
 
-    :return: whether zarr.open can open this path
+    Returns
+    -------
+    bool
+        whether zarr.open can open this path
     """
     try:
         with zarr.open(Path(inpath, "partition_keys"), mode="r"):
@@ -35,11 +41,19 @@ def read_vis(
     """
     Read a CASAVisSet stored in zarr format.
 
-    :param inpath: Input Zarr path
-    :param subtables: Also read and (metainformation) subtables along with main visibilities data.
-    :param asdm_subtables: Also read extension subtables named "ASDM_*"
-
-    :return: Main xarray dataset of datasets for this visibility dataset
+    Parameters
+    ----------
+    inpath : str
+        Input Zarr path
+    subtables : bool (Default value = True)
+        Also read and (metainformation) subtables along with main visibilities data.
+    asdm_subtables : bool (Default value = False)
+        Also read extension subtables named "ASDM_*"
+
+    Returns
+    -------
+    CASAVisSet
+        Main xarray dataset of datasets for this visibility dataset
     """
     inpath = os.path.expanduser(inpath)
     if not os.path.isdir(inpath):
@@ -59,7 +73,7 @@ def read_vis(
     all_time = time.time() - all_start
     logger.info(f"Time to read dataset from_zarr {inpath}: {all_time}")
 
-    vers = xradio.__version__
+    vers = "version-WIP"
     descr_add = "read_vis from zarr"
     cds = CASAVisSet(
         metainfo=metainfo,
@@ -71,25 +85,35 @@ def read_vis(
 
 
 def write_vis(
-    cds,
+    cds: CASAVisSet,
     outpath: str,
     chunks_on_disk: Union[Dict, None] = None,
     compressor: Union[numcodecs.abc.Codec, None] = None,
 ) -> None:
-    """Write CASA vis dataset to zarr format on disk. When
+    """
+    Write CASA vis dataset to zarr format on disk. When
     chunks_on_disk is not specified the chunking in the input dataset
     is used. When chunks_on_disk is specified that dataset is saved
     using that chunking.
 
-    :param cds: CASA visibilities dataset to write to disk
-    :param outpath: output path, generally ends in .zarr
-    :param chunks_on_disk: a dictionary with the chunk size that will
-    be used when writing to disk. For example {'time': 20, 'chan': 6}.
-    If chunks_on_disk is not specified the chunking of dataset will
-    be used.
-    :param compressor: the blosc compressor to use when saving the
-    converted data to disk using zarr. If None the zstd compression
-    algorithm used with compression level 2.
+    Parameters
+    ----------
+    cds : CASAVisSet
+        CASA visibilities dataset to write to disk
+    outpath : str
+        output path, generally ends in .zarr
+    chunks_on_disk : Union[Dict, None] = None (Default value = None)
+        a dictionary with the chunk size that will
+        be used when writing to disk. For example {'time': 20, 'chan': 6}.
+        If chunks_on_disk is not specified the chunking of dataset will
+        be used.
+    compressor : Union[numcodecs.abc.Codec, None] (Default value = None)
+        the blosc compressor to use when saving the
+        converted data to disk using zarr. If None the zstd compression
+        algorithm used with compression level 2.
+
+    Returns
+    -------
     """
 
     if compressor is None:

xradio/vis/convert_msv2_to_processing_set.py

@@ -4,10 +4,7 @@ from typing import Dict, Union
 
 import dask
 
-from xradio.vis._vis_utils._ms.msv2_msv3 import ignore_msv2_cols
-from xradio.vis._vis_utils._ms.partition_queries import (
-    create_partition_enumerated_product,
-)
+from xradio.vis._vis_utils._ms.partition_queries import create_partitions
 from xradio.vis._vis_utils._ms.conversion import convert_and_write_partition
 
 
@@ -15,8 +12,11 @@ def convert_msv2_to_processing_set(
     in_file: str,
     out_file: str,
     partition_scheme: {"ddi_intent_field", "ddi_state_field"} = "ddi_intent_field",
-    main_chunksize: Union[Dict, str, None] = None,
-    pointing_chunksize: Union[Dict, str, None] = None,
+    main_chunksize: Union[Dict, float, None] = None,
+    with_pointing: bool = True,
+    pointing_chunksize: Union[Dict, float, None] = None,
+    pointing_interpolate: bool = False,
+    ephemeris_interpolate: bool = False,
     compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
     storage_backend="zarr",
     parallel: bool = False,
@@ -34,10 +34,16 @@ def convert_msv2_to_processing_set(
         A MS v4 can only contain a single spectral window, polarization setup, intent, and field. Consequently, the MS v2 is partitioned when converting to MS v4.
         The partition_scheme "ddi_intent_field" gives the largest partition that meets the MS v4 specification. The partition_scheme "ddi_state_field" gives a finer granularity where the data is also partitioned by state (the state partitioning will ensure a single intent).
         By default, "ddi_intent_field".
-    main_chunksize : Union[Dict, str, None], optional
-        A dictionary that defines the chunk size of the main dataset. Acceptable keys are "time", "baseline", "antenna", "frequency", "polarization". By default, None.
-    pointing_chunksize : Union[Dict, str, None], optional
-        A dictionary that defines the chunk size of the pointing dataset. Acceptable keys are "time", "antenna", "polarization". By default, None.
+    main_chunksize : Union[Dict, float, None], optional
+        Defines the chunk size of the main dataset. If given as a dictionary, defines the sizes of several dimensions, and acceptable keys are "time", "baseline_id", "antenna_id", "frequency", "polarization". If given as a float, gives the size of a chunk in GiB. By default, None.
+    with_pointing : bool, optional
+        Whether to convert the POINTING subtable into pointing sub-datasets
+    pointing_chunksize : Union[Dict, float, None], optional
+        Defines the chunk size of the pointing dataset. If given as a dictionary, defines the sizes of several dimensions, acceptable keys are "time" and "antenna_id". If given as a float, defines the size of a chunk in GiB. By default, None.
+    pointing_interpolate : bool, optional
+        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
     compressor : numcodecs.abc.Codec, optional
         The Blosc compressor to use when saving the converted data to disk using Zarr, by default numcodecs.Zstd(level=2).
     storage_backend : {"zarr", "netcdf"}, optional
@@ -48,33 +54,36 @@ def convert_msv2_to_processing_set(
         Whether to overwrite an existing processing set, by default False.
     """
 
-    partition_enumerated_product, intents = create_partition_enumerated_product(
-        in_file, partition_scheme
-    )
+    partitions = create_partitions(in_file, partition_scheme=partition_scheme)
+    logger.info("Number of partitions: " + str(len(partitions)))
 
     delayed_list = []
-    for idx, pair in partition_enumerated_product:
-        ddi, state_id, field_id = pair
+    ms_v4_id = 0
+    for partition_info in partitions:
         logger.debug(
-            "DDI " + str(ddi) + ", STATE " + str(state_id) + ", FIELD " + str(field_id)
+            "DDI "
+            + str(partition_info["DATA_DESC_ID"])
+            + ", STATE "
+            + str(partition_info["STATE_ID"])
+            + ", FIELD "
+            + str(partition_info["FIELD_ID"])
+            + ", SCAN "
+            + str(partition_info["SCAN_NUMBER"])
         )
 
-        if partition_scheme == "ddi_intent_field":
-            intent = intents[idx[1]]
-        else:
-            intent = intents[idx[1]] + "_" + str(state_id)
-
         if parallel:
             delayed_list.append(
                 dask.delayed(convert_and_write_partition)(
                     in_file,
                     out_file,
-                    intent,
-                    ddi,
-                    state_id,
-                    field_id,
-                    ignore_msv2_cols=ignore_msv2_cols,
+                    ms_v4_id,
+                    partition_info=partition_info,
+                    partition_scheme=partition_scheme,
                     main_chunksize=main_chunksize,
+                    with_pointing=with_pointing,
+                    pointing_chunksize=pointing_chunksize,
+                    pointing_interpolate=pointing_interpolate,
+                    ephemeris_interpolate=ephemeris_interpolate,
                     compressor=compressor,
                     overwrite=overwrite,
                 )
@@ -83,16 +92,18 @@ def convert_msv2_to_processing_set(
             convert_and_write_partition(
                 in_file,
                 out_file,
-                intent,
-                ddi,
-                state_id,
-                field_id,
-                ignore_msv2_cols=ignore_msv2_cols,
+                ms_v4_id,
+                partition_info=partition_info,
+                partition_scheme=partition_scheme,
                 main_chunksize=main_chunksize,
+                with_pointing=with_pointing,
+                pointing_chunksize=pointing_chunksize,
+                pointing_interpolate=pointing_interpolate,
+                ephemeris_interpolate=ephemeris_interpolate,
                 compressor=compressor,
-                storage_backend=storage_backend,
                 overwrite=overwrite,
             )
+        ms_v4_id = ms_v4_id + 1
 
     if parallel:
         dask.compute(delayed_list)

xradio/vis/load_processing_set.py

@@ -39,74 +39,51 @@ def load_processing_set(
     processing_set
         In memory representation of processing set (data is represented by Dask.arrays).
     """
-    from xradio._utils.zarr.common import _open_dataset
-    import s3fs
-    from botocore.exceptions import NoCredentialsError
+    from xradio._utils.zarr.common import _open_dataset, _get_ms_stores_and_file_system
+
+    file_system, ms_store_list = _get_ms_stores_and_file_system(ps_store)
 
-    s3 = None
     ps = processing_set()
-    for ms_dir_name, ms_xds_isel in sel_parms.items():
-
-        # before the _open_dataset call, check if dealing with an S3 bucket URL
-        if ps_store.startswith("s3"):
-            if not ps_store.endswith("/"):
-                # just for consistency, as there is no os.path equivalent in s3fs
-                ps_store = ps_store + "/"
-
-            try:
-                s3 = s3fs.S3FileSystem(anon=False, requester_pays=False)
-
-                main_xds = ps_store + ms_dir_name + "/MAIN"
-                xds = _open_dataset(
-                    main_xds, ms_xds_isel, data_variables, load=True, s3=s3
-                )
-
-                if load_sub_datasets:
-                    from xradio.vis.read_processing_set import _read_sub_xds
-
-                    xds.attrs = {
-                        **xds.attrs,
-                        **_read_sub_xds(
-                            os.path.join(ps_store, ms_dir_name), load=True, s3=s3
-                        ),
-                    }
-
-            except (NoCredentialsError, PermissionError) as e:
-                # only public, read-only buckets will be accessible
-                s3 = s3fs.S3FileSystem(anon=True)
-
-                main_xds = ps_store + ms_dir_name + "/MAIN"
-                xds = _open_dataset(
-                    main_xds, ms_xds_isel, data_variables, load=True, s3=s3
-                )
-
-                if load_sub_datasets:
-                    from xradio.vis.read_processing_set import _read_sub_xds
-
-                    xds.attrs = {
-                        **xds.attrs,
-                        **_read_sub_xds(
-                            os.path.join(ps_store, ms_dir_name), load=True, s3=s3
-                        ),
-                    }
-        else:
-            # fall back to the default case of assuming the files are on local disk
-            main_xds = os.path.join(ps_store, ms_dir_name, "MAIN")
-            xds = _open_dataset(main_xds, ms_xds_isel, data_variables, load=True)
-            if load_sub_datasets:
-                from xradio.vis.read_processing_set import _read_sub_xds
-
-                xds.attrs = {
-                    **xds.attrs,
-                    **_read_sub_xds(os.path.join(ps_store, ms_dir_name), load=True),
-                }
+    for ms_name, ms_xds_isel in sel_parms.items():
+        ms_store = os.path.join(ps_store, ms_name)
+        ms_main_store = os.path.join(ms_store, "MAIN")
+
+        xds = _open_dataset(
+            ms_main_store,
+            file_system,
+            ms_xds_isel,
+            data_variables,
+            load=True,
+        )
+        data_groups = xds.attrs["data_groups"]
+
+        if load_sub_datasets:
+            from xradio.vis.read_processing_set import _read_sub_xds
+
+            sub_xds_dict, field_and_source_xds_dict = _read_sub_xds(
+                ms_store, file_system=file_system, load=True, data_groups=data_groups
+            )
+
+            xds.attrs = {
+                **xds.attrs,
+                **sub_xds_dict,
+            }
+            for data_group_name, data_group_vals in data_groups.items():
+                if "visibility" in data_group_vals:
+                    xds[data_group_vals["visibility"]].attrs["field_and_source_xds"] = (
+                        field_and_source_xds_dict[data_group_name]
+                    )
+                elif "spectrum" in data_group_vals:
+                    xds[data_group_vals["spectrum"]].attrs["field_and_source_xds"] = (
+                        field_and_source_xds_dict[data_group_name]
+                    )
+
+        ps[ms_name] = xds
 
-        ps[ms_dir_name] = xds
     return ps
 
 
 class processing_set_iterator:
-
     def __init__(
         self,
         sel_parms: dict,

xradio/vis/read_processing_set.py

@@ -2,13 +2,13 @@ import os
 import xarray as xr
 from ._processing_set import processing_set
 import graphviper.utils.logger as logger
-from xradio._utils.zarr.common import _open_dataset
+from xradio._utils.zarr.common import _open_dataset, _get_ms_stores_and_file_system
 import s3fs
-from botocore.exceptions import NoCredentialsError
 
 
 def read_processing_set(
-    ps_store: str, intents: list = None, fields: str = None
+    ps_store: str,
+    intents: list = None,
 ) -> processing_set:
     """Creates a lazy representation of a Processing Set (only meta-data is loaded into memory).
 
@@ -19,119 +19,87 @@ def read_processing_set(
     intents : list, optional
         A list of the intents to be read for example ['OBSERVE_TARGET#ON_SOURCE']. The intents in a processing set can be seem by calling processing_set.summary().
         By default None, which will read all intents.
-    fields : str, optional
-       The list of field names that will be read, by default None which will read all fields.
 
     Returns
     -------
     processing_set
         Lazy representation of processing set (data is represented by Dask.arrays).
     """
-    s3 = None
-    ps_store_is_s3dir = None
+    file_system, ms_store_list = _get_ms_stores_and_file_system(ps_store)
 
-    if os.path.isdir(ps_store):
-        ps_store_is_s3dir = False
-        # default to assuming the data are accessible on local file system
-        items = os.listdir(ps_store)
-
-    elif ps_store.startswith("s3"):
-        # only if not found locally, check if dealing with an S3 bucket URL
-        ps_store_is_s3dir = True
-        if not ps_store.endswith("/"):
-            # just for consistency, as there is no os.path equivalent in s3fs
-            ps_store = ps_store + "/"
-
-        try:
-            # initialize the S3 "file system", first attempting to use pre-configured credentials
-            s3 = s3fs.S3FileSystem(anon=False, requester_pays=False)
-
-            items = [bd.split(sep="/")[-1] for bd in s3.listdir(ps_store, detail=False)]
+    ps = processing_set()
+    data_group = "base"
+    for ms_name in ms_store_list:
+        # try:
+        ms_store = os.path.join(ps_store, ms_name)
+        ms_main_store = os.path.join(ms_store, "MAIN")
 
-        except (NoCredentialsError, PermissionError) as e:
-            # only public, read-only buckets will be accessible
-            # we will want to add messaging and error handling here
-            s3 = s3fs.S3FileSystem(anon=True)
+        xds = _open_dataset(ms_main_store, file_system)
+        data_groups = xds.attrs["data_groups"]
 
-            items = [bd.split(sep="/")[-1] for bd in s3.listdir(ps_store, detail=False)]
+        print(xds)
 
-    else:
-        raise (
-            FileNotFoundError,
-            f"Could not find {ps_store} either locally or in the cloud.",
-        )
+        if (intents is None) or (xds.attrs["partition_info"]["intent"] in intents):
+            sub_xds_dict, field_and_source_xds_dict = _read_sub_xds(
+                ms_store, file_system=file_system, data_groups=data_groups
+            )
 
-    ms_xds = xr.Dataset()
-    ps = processing_set()
-    data_group = "base"
-    for ms_dir_name in items:
-        if "ddi" in ms_dir_name:
-            if ps_store_is_s3dir:
-                store_path = ps_store + ms_dir_name
-                store_path_main = store_path + "/MAIN"
-            else:
-                store_path_main = os.path.join(ps_store, ms_dir_name, "MAIN")
-                store_path = os.path.split(store_path_main)[0]
-            if s3 is not None:
-                xds = _open_dataset(store_path_main, s3=s3)
-            else:
-                xds = _open_dataset(store_path_main)
-
-            if (intents is None) or (xds.attrs["intent"] in intents):
-                data_name = _get_data_name(xds, data_group)
-
-                if (fields is None) or (
-                    xds[data_name].attrs["field_info"]["name"] in fields
-                ):
-                    if s3 is not None:
-                        xds.attrs = {
-                            **xds.attrs,
-                            **_read_sub_xds(store_path, s3=s3),
-                        }
-                        ps[ms_dir_name] = xds
-                    else:
-                        xds.attrs = {
-                            **xds.attrs,
-                            **_read_sub_xds(store_path),
-                        }
-                        ps[ms_dir_name] = xds
+            xds.attrs = {
+                **xds.attrs,
+                **sub_xds_dict,
+            }
+
+            for data_group_name, data_group_vals in data_groups.items():
+                if "visibility" in data_group_vals:
+                    xds[data_group_vals["visibility"]].attrs["field_and_source_xds"] = (
+                        field_and_source_xds_dict[data_group_name]
+                    )
+                elif "spectrum" in data_group_vals:
+                    xds[data_group_vals["spectrum"]].attrs["field_and_source_xds"] = (
+                        field_and_source_xds_dict[data_group_name]
+                    )
+
+            ps[ms_name] = xds
+        # except Exception as e:
+        #     logger.warning(f"Could not read {ms_name} due to {e}")
+        #     continue
 
     return ps
 
 
-def _read_sub_xds(ms_store, load=False, **kwargs):
+def _read_sub_xds(ms_store, file_system, data_groups, load=False):
     sub_xds_dict = {}
+    field_and_source_xds_dict = {}
 
-    sub_xds = {
-        "antenna_xds": "ANTENNA",
+    xds_names = {
+        "ANTENNA": "antenna_xds",
+        "WEATHER": "weather_xds",
+        "POINTING": "pointing_xds",
     }
-    for sub_xds_key, sub_xds_name in sub_xds.items():
-        if "s3" in kwargs.keys():
-            joined_store = ms_store + "/" + sub_xds_name
-            sub_xds_dict[sub_xds_key] = _open_dataset(
-                joined_store, load=load, s3=kwargs["s3"]
-            )
+
+    if isinstance(file_system, s3fs.core.S3FileSystem):
+        file_names = [
+            bd.split(sep="/")[-1] for bd in file_system.listdir(ms_store, detail=False)
+        ]
+    else:
+        file_names = file_system.listdir(ms_store)
+    file_names = [item for item in file_names if not item.startswith(".")]
+
+    file_names.remove("MAIN")
+
+    field_dict = {"FIELD_AND_SOURCE_" + key.upper(): key for key in data_groups.keys()}
+
+    # field_and_source_xds_name_start = "FIELD"
+    for n in file_names:
+        xds = _open_dataset(
+            os.path.join(ms_store, n), load=load, file_system=file_system
+        )
+        if n in field_dict.keys():
+            field_and_source_xds_dict[field_dict[n]] = xds
         else:
-            sub_xds_dict[sub_xds_key] = _open_dataset(
-                os.path.join(ms_store, sub_xds_name), load=load
-            )
+            sub_xds_dict[xds_names[n]] = xds
 
-    optional_sub_xds = {
-        "weather_xds": "WEATHER",
-        "pointing_xds": "POINTING",
-    }
-    for sub_xds_key, sub_xds_name in optional_sub_xds.items():
-        sub_xds_path = os.path.join(ms_store, sub_xds_name)
-        if os.path.isdir(sub_xds_path):
-            sub_xds_dict[sub_xds_key] = _open_dataset(sub_xds_path, load=load)
-        elif "s3" in kwargs.keys():
-            joined_store = ms_store + "/" + sub_xds_name
-            if kwargs["s3"].isdir(joined_store):
-                sub_xds_dict[sub_xds_key] = _open_dataset(
-                    joined_store, load=load, s3=kwargs["s3"]
-                )
-
-    return sub_xds_dict
+    return sub_xds_dict, field_and_source_xds_dict
 
 
 def _get_data_name(xds, data_group):