xradio 0.0.48__py3-none-any.whl → 0.0.49__py3-none-any.whl

xradio/measurement_set/_utils/_msv2/msv4_info_dicts.py

@@ -49,28 +49,31 @@ def create_info_dicts(
         line_name = []
 
     info_dicts = {}
-    info_dicts["partition_info"] = {
-        # "spectral_window_id": xds.frequency.attrs["spectral_window_id"],
-        "spectral_window_name": xds.frequency.attrs["spectral_window_name"],
-        # "field_id": to_list(unique_1d(field_id)),
-        "field_name": to_list(np.unique(field_and_source_xds.field_name.values)),
-        "polarization_setup": to_list(xds.polarization.values),
-        "scan_name": to_list(np.unique(partition_info_misc_fields["scan_name"])),
-        "source_name": to_list(np.unique(field_and_source_xds.source_name.values)),
-        # "source_id": to_list(unique_1d(source_id)),
-        "intents": partition_info_misc_fields["intents"].split(","),
-        "taql": partition_info_misc_fields["taql_where"],
-        "line_name": line_name,
-    }
-    if "antenna_name" in partition_info_misc_fields:
-        info_dicts["partition_info"]["antenna_name"] = partition_info_misc_fields[
-            "antenna_name"
-        ]
+    # info_dicts["partition_info"] = {
+    #     # "spectral_window_id": xds.frequency.attrs["spectral_window_id"],
+    #     "spectral_window_name": xds.frequency.attrs["spectral_window_name"],
+    #     # "field_id": to_list(unique_1d(field_id)),
+    #     "field_name": to_list(np.unique(field_and_source_xds.field_name.values)),
+    #     "polarization_setup": to_list(xds.polarization.values),
+    #     "scan_name": to_list(np.unique(partition_info_misc_fields["scan_name"])),
+    #     "source_name": to_list(np.unique(field_and_source_xds.source_name.values)),
+    #     # "source_id": to_list(unique_1d(source_id)),
+    #     "intents": partition_info_misc_fields["intents"].split(","),
+    #     "taql": partition_info_misc_fields["taql_where"],
+    #     "line_name": line_name,
+    # }
+    # if "antenna_name" in partition_info_misc_fields:
+    #     info_dicts["partition_info"]["antenna_name"] = partition_info_misc_fields[
+    #         "antenna_name"
+    #     ]
 
     observation_id = check_if_consistent(
         tb_tool.getcol("OBSERVATION_ID"), "OBSERVATION_ID"
     )
     info_dicts["observation_info"] = create_observation_info(in_file, observation_id)
+    info_dicts["observation_info"]["intents"] = partition_info_misc_fields[
+        "intents"
+    ].split(",")
 
     processor_id = check_if_consistent(tb_tool.getcol("PROCESSOR_ID"), "PROCESSOR_ID")
     info_dicts["processor_info"] = create_processor_info(in_file, processor_id)

xradio/measurement_set/convert_msv2_to_processing_set.py

@@ -18,6 +18,7 @@ def estimate_conversion_memory_and_cores(
     """
     Given an MSv2 and a partition_scheme to use when converting it to MSv4,
     estimates:
+
     - memory (in the sense of the amount expected to be enough to convert)
     - cores (in the sense of the recommended/optimal number of cores to use to convert)
 
@@ -36,7 +37,7 @@ def estimate_conversion_memory_and_cores(
         Partition scheme as used in the function convert_msv2_to_processing_set()
 
     Returns
-    ----------
+    -------
     tuple
         estimated maximum memory required for one partition,
         maximum number of cores it makes sense to use (number of partitions),
@@ -62,7 +63,7 @@ def convert_msv2_to_processing_set(
     use_table_iter: bool = False,
     compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
     storage_backend: str = "zarr",
-    parallel: bool = False,
+    parallel_mode: str = "none",
     overwrite: bool = False,
 ):
     """Convert a Measurement Set v2 into a Processing Set of Measurement Set v4.
@@ -99,14 +100,45 @@ def convert_msv2_to_processing_set(
         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
         The on-disk format to use. "netcdf" is not yet implemented.
-    parallel : bool, optional
-        Makes use of Dask to execute conversion in parallel, by default False.
+    parallel_mode : {"none", "partition", "time"}, optional
+        Choose whether to use Dask to execute conversion in parallel, by default "none" and conversion occurs serially.
+        The option "partition", parallelises the conversion over partitions specified by `partition_scheme`. The option "time" can only be used for phased array interferometers where there are no partitions
+        in the MS v2; instead the MS v2 is parallelised along the time dimension and can be controlled by `main_chunksize`.
     overwrite : bool, optional
         Whether to overwrite an existing processing set, by default False.
     """
 
+    # Create empty data tree
+    import xarray as xr
+
+    ps_dt = xr.DataTree()
+
+    if not str(out_file).endswith("ps.zarr"):
+        out_file += ".ps.zarr"
+
+    print("Output file: ", out_file)
+
+    if overwrite:
+        ps_dt.to_zarr(store=out_file, mode="w")
+    else:
+        ps_dt.to_zarr(store=out_file, mode="w-")
+
+    # Check `parallel_mode` is valid
+    try:
+        assert parallel_mode in ["none", "partition", "time"]
+    except AssertionError:
+        logger.warning(
+            f"`parallel_mode` {parallel_mode} not recognosed. Defauling to 'none'."
+        )
+        parallel_mode = "none"
+
     partitions = create_partitions(in_file, partition_scheme=partition_scheme)
     logger.info("Number of partitions: " + str(len(partitions)))
+    if parallel_mode == "time":
+        assert (
+            len(partitions) == 1
+        ), "MS v2 contains more than one partition. `parallel_mode = 'time'` not valid."
+
     delayed_list = []
 
     for ms_v4_id, partition_info in enumerate(partitions):
@@ -132,7 +164,7 @@ def convert_msv2_to_processing_set(
 
         # prepend '0' to ms_v4_id as needed
         ms_v4_id = f"{ms_v4_id:0>{len(str(len(partitions) - 1))}}"
-        if parallel:
+        if parallel_mode == "partition":
             delayed_list.append(
                 dask.delayed(convert_and_write_partition)(
                     in_file,
@@ -149,6 +181,7 @@ def convert_msv2_to_processing_set(
                     phase_cal_interpolate=phase_cal_interpolate,
                     sys_cal_interpolate=sys_cal_interpolate,
                     compressor=compressor,
+                    parallel_mode=parallel_mode,
                     overwrite=overwrite,
                 )
             )
@@ -168,8 +201,15 @@ def convert_msv2_to_processing_set(
                 phase_cal_interpolate=phase_cal_interpolate,
                 sys_cal_interpolate=sys_cal_interpolate,
                 compressor=compressor,
+                parallel_mode=parallel_mode,
                 overwrite=overwrite,
             )
 
-    if parallel:
+    if parallel_mode == "partition":
         dask.compute(delayed_list)
+
+    import zarr
+
+    root_group = zarr.open(out_file, mode="r+")  # Open in read/write mode
+    root_group.attrs["type"] = "processing_set"  # Replace
+    zarr.convenience.consolidate_metadata(root_group.store)

xradio/measurement_set/load_processing_set.py

@@ -1,80 +1,115 @@
 import os
-from xradio.measurement_set import ProcessingSet
 from typing import Dict, Union
+import dask
+import xarray as xr
+import s3fs
 
 
 def load_processing_set(
     ps_store: str,
-    sel_parms: dict,
-    data_variables: Union[list, None] = None,
+    sel_parms: dict = None,
+    data_group_name: str = None,
+    include_variables: Union[list, None] = None,
+    drop_variables: Union[list, None] = None,
     load_sub_datasets: bool = True,
-) -> ProcessingSet:
+) -> xr.DataTree:
     """Loads a processing set into memory.
 
     Parameters
     ----------
     ps_store : str
         String of the path and name of the processing set. For example '/users/user_1/uid___A002_Xf07bba_Xbe5c_target.lsrk.vis.zarr' for a file stored on a local file system, or 's3://viper-test-data/Antennae_North.cal.lsrk.split.vis.zarr/' for a file in AWS object storage.
-    sel_parms : dict
-        A dictionary where the keys are the names of the ms_xds's and the values are slice_dicts.
+    sel_parms : dict, optional
+        A dictionary where the keys are the names of the ms_xdt's (measurement set xarray data trees) and the values are slice_dicts.
         slice_dicts: A dictionary where the keys are the dimension names and the values are slices.
+
         For example::
 
             {
+
                 'ms_v4_name_1': {'frequency': slice(0, 160, None),'time':slice(0,100)},
                 ...
                 'ms_v4_name_n': {'frequency': slice(0, 160, None),'time':slice(0,100)},
             }
 
-    data_variables : Union[list, None], optional
+        By default None, which loads all ms_xdts.
+    data_group_name : str, optional
+        The name of the data group to select. By default None, which loads all data groups.
+    include_variables : Union[list, None], optional
         The list of data variables to load into memory for example ['VISIBILITY', 'WEIGHT, 'FLAGS']. By default None which will load all data variables into memory.
+    drop_variables : Union[list, None], optional
+        The list of data variables to drop from memory for example ['VISIBILITY', 'WEIGHT, 'FLAGS']. By default None which will not drop any data variables from memory.
     load_sub_datasets : bool, optional
         If true sub-datasets (for example weather_xds, antenna_xds, pointing_xds, system_calibration_xds ...) will be loaded into memory, by default True.
 
     Returns
     -------
-    ProcessingSet
-        In memory representation of processing set (data is represented by Dask.arrays).
+    xarray.DataTree
+        In memory representation of processing set using xr.DataTree.
     """
-    from xradio._utils.zarr.common import _open_dataset, _get_file_system_and_items
-    from xradio.measurement_set import MeasurementSetXds
+    from xradio._utils.zarr.common import _get_file_system_and_items
 
     file_system, ms_store_list = _get_file_system_and_items(ps_store)
 
-    ps = ProcessingSet()
-    for ms_name, ms_xds_isel in sel_parms.items():
-        ms_store = os.path.join(ps_store, ms_name)
-        correlated_store = os.path.join(ms_store, "correlated_xds")
-
-        xds = _open_dataset(
-            correlated_store,
-            file_system,
-            ms_xds_isel,
-            data_variables,
-            load=True,
-        )
-        data_groups = xds.attrs["data_groups"]
-
-        if load_sub_datasets:
-            from xradio.measurement_set.open_processing_set import _open_sub_xds
-
-            sub_xds_dict, field_and_source_xds_dict = _open_sub_xds(
-                ms_store, file_system=file_system, load=True, data_groups=data_groups
+    with dask.config.set(
+        scheduler="synchronous"
+    ):  # serial scheduler, critical so that this can be used within delayed functions.
+        ps_xdt = xr.DataTree()
+
+        if sel_parms:
+            for ms_name, ms_xds_isel in sel_parms.items():
+                ms_store = os.path.join(ps_store, ms_name)
+
+                if isinstance(file_system, s3fs.core.S3FileSystem):
+                    ms_store = s3fs.S3Map(root=ps_store, s3=file_system, check=False)
+
+                if ms_xds_isel:
+                    ms_xdt = (
+                        xr.open_datatree(
+                            ms_store, engine="zarr", drop_variables=drop_variables
+                        )
+                        .isel(ms_xds_isel)
+                        .xr_ms.sel(data_group_name=data_group_name)
+                    )
+                else:
+                    ms_xdt = xr.open_datatree(
+                        ms_store, engine="zarr", drop_variables=drop_variables
+                    ).xr_ms.sel(data_group_name=data_group_name)
+
+                if include_variables is not None:
+                    for data_vars in ms_xdt.ds.data_vars:
+                        if data_vars not in include_variables:
+                            ms_xdt.ds = ms_xdt.ds.drop_vars(data_vars)
+
+                ps_xdt[ms_name] = ms_xdt
+
+            ps_xdt.attrs["type"] = "processing_set"
+        else:
+            ps_xdt = xr.open_datatree(
+                ps_store, engine="zarr", drop_variables=drop_variables
             )
 
-            xds.attrs = {
-                **xds.attrs,
-                **sub_xds_dict,
-            }
-            for data_group_name, data_group_vals in data_groups.items():
+            if (include_variables is not None) or data_group_name:
+                for ms_name, ms_xdt in ps_xdt.items():
+
+                    ms_xdt = ms_xdt.xr_ms.sel(data_group_name=data_group_name)
+
+                    if include_variables is not None:
+                        for data_vars in ms_xdt.ds.data_vars:
+                            if data_vars not in include_variables:
+                                ms_xdt.ds = ms_xdt.ds.drop_vars(data_vars)
+                    ps_xdt[ms_name] = ms_xdt
 
-                xds[data_group_vals["correlated_data"]].attrs[
-                    "field_and_source_xds"
-                ] = field_and_source_xds_dict[data_group_name]
+        if not load_sub_datasets:
+            for ms_xdt in ps_xdt.children.values():
+                ms_xdt_names = list(ms_xdt.keys())
+                for sub_xds_name in ms_xdt_names:
+                    if "xds" in sub_xds_name:
+                        del ms_xdt[sub_xds_name]
 
-        ps[ms_name] = MeasurementSetXds(xds)
+        ps_xdt = ps_xdt.load()
 
-    return ps
+    return ps_xdt
 
 
 class ProcessingSetIterator:
@@ -82,8 +117,10 @@ class ProcessingSetIterator:
         self,
         sel_parms: dict,
         input_data_store: str,
-        input_data: Union[Dict, ProcessingSet, None] = None,
-        data_variables: list = None,
+        input_data: Union[Dict, xr.DataTree, None] = None,
+        data_group_name: str = None,
+        include_variables: Union[list, None] = None,
+        drop_variables: Union[list, None] = None,
         load_sub_datasets: bool = True,
     ):
         """An iterator that will go through a processing set one MS v4 at a time.
@@ -102,10 +139,16 @@ class ProcessingSetIterator:
                 }
         input_data_store : str
             String of the path and name of the processing set. For example '/users/user_1/uid___A002_Xf07bba_Xbe5c_target.lsrk.vis.zarr'.
-        input_data : Union[Dict, processing_set, None], optional
+        input_data : Union[Dict, xr.DataTree, None], optional
             If the processing set is in memory already it can be supplied here. By default None which will make the iterator load data using the supplied input_data_store.
-        data_variables : list, optional
+        data_group_name : str, optional
+            The name of the data group to select. By default None, which loads all data groups.
+        data_group_name : str, optional
+            The name of the data group to select. By default None, which loads all data groups.
+        include_variables : Union[list, None], optional
             The list of data variables to load into memory for example ['VISIBILITY', 'WEIGHT, 'FLAGS']. By default None which will load all data variables into memory.
+        drop_variables : Union[list, None], optional
+            The list of data variables to drop from memory for example ['VISIBILITY', 'WEIGHT, 'FLAGS']. By default None which will not drop any data variables from memory.
         load_sub_datasets : bool, optional
             If true sub-datasets (for example weather_xds, antenna_xds, pointing_xds, system_calibration_xds ...) will be loaded into memory, by default True.
         """
@@ -114,7 +157,9 @@ class ProcessingSetIterator:
         self.input_data_store = input_data_store
         self.sel_parms = sel_parms
         self.xds_name_iter = iter(sel_parms.keys())
-        self.data_variables = data_variables
+        self.data_group_name = data_group_name
+        self.include_variables = include_variables
+        self.drop_variables = drop_variables
         self.load_sub_datasets = load_sub_datasets
 
     def __iter__(self):
@@ -122,20 +167,22 @@ class ProcessingSetIterator:
 
     def __next__(self):
         try:
-            xds_name = next(self.xds_name_iter)
+            sub_xds_name = next(self.xds_name_iter)
         except Exception as e:
             raise StopIteration
 
         if self.input_data is None:
-            slice_description = self.sel_parms[xds_name]
-            ps = load_processing_set(
+            slice_description = self.sel_parms[sub_xds_name]
+            ps_xdt = load_processing_set(
                 ps_store=self.input_data_store,
-                sel_parms={xds_name: slice_description},
-                data_variables=self.data_variables,
+                sel_parms={sub_xds_name: slice_description},
+                data_group_name=self.data_group_name,
+                include_variables=self.include_variables,
+                drop_variables=self.drop_variables,
                 load_sub_datasets=self.load_sub_datasets,
             )
-            xds = ps.get(0)
+            sub_xdt = ps_xdt.get(0)
         else:
-            xds = self.input_data[xds_name]  # In memory
+            sub_xdt = self.input_data[sub_xds_name]  # In memory
 
-        return xds
+        return sub_xdt

xradio/measurement_set/measurement_set_xdt.py

@@ -0,0 +1,197 @@
+import pandas as pd
+from xradio._utils.list_and_array import to_list
+import xarray as xr
+import numpy as np
+import numbers
+import os
+from collections.abc import Mapping, Iterable
+from typing import Any, Union
+
+MS_DATASET_TYPES = {"visibility", "spectrum", "radiometer"}
+
+
+class InvalidAccessorLocation(ValueError):
+    """
+    Raised by MeasurementSetXdt accessor functions called on a wrong DataTree node (not MSv4).
+    """
+
+    pass
+
+
+@xr.register_datatree_accessor("xr_ms")
+class MeasurementSetXdt:
+    """Accessor to the Measurement Set DataTree node. Provides MSv4 specific functionality
+    such as:
+
+        - get_partition_info(): produce an info dict with a general MSv4 description including
+          intents, SPW name, field and source names, etc.
+        - get_field_and_source_xds() to retrieve the field_and_source_xds for a given data
+          group.
+        - sel(): select data by dimension labels, for example by data group and polaritzation
+
+    """
+
+    _xdt: xr.DataTree
+
+    def __init__(self, datatree: xr.DataTree):
+        """
+        Initialize the MeasurementSetXdt instance.
+
+        Parameters
+        ----------
+        datatree: xarray.DataTree
+            The MSv4 DataTree node to construct a MeasurementSetXdt accessor.
+        """
+
+        self._xdt = datatree
+        self.meta = {"summary": {}}
+
+    def sel(
+        self,
+        indexers: Union[Mapping[Any, Any], None] = None,
+        method: Union[str, None] = None,
+        tolerance: Union[int, float, Iterable[Union[int, float]], None] = None,
+        drop: bool = False,
+        **indexers_kwargs: Any,
+    ) -> xr.DataTree:
+        """
+        Select data along dimension(s) by label. Alternative to `xarray.Dataset.sel <https://xarray.pydata.org/en/stable/generated/xarray.Dataset.sel.html>`__ so that a data group can be selected by name by using the `data_group_name` parameter.
+        For more information on data groups see `Data Groups <https://xradio.readthedocs.io/en/latest/measurement_set_overview.html#Data-Groups>`__ section. See `xarray.Dataset.sel <https://xarray.pydata.org/en/stable/generated/xarray.Dataset.sel.html>`__ for parameter descriptions.
+
+        Returns
+        -------
+        xarray.DataTree
+            xarray DataTree with MeasurementSetXdt accessors
+
+        Examples
+        --------
+        >>> # Select data group 'corrected' and polarization 'XX'.
+        >>> selected_ms_xdt = ms_xdt.xr_ms.sel(data_group_name='corrected', polarization='XX')
+
+        >>> # Select data group 'corrected' and polarization 'XX' using a dict.
+        >>> selected_ms_xdt = ms_xdt.xr_ms.sel({'data_group_name':'corrected', 'polarization':'XX')
+        """
+
+        if self._xdt.attrs.get("type") not in MS_DATASET_TYPES:
+            raise InvalidAccessorLocation(f"{self._xdt.path} is not a MSv4 node.")
+
+        assert self._xdt.attrs["type"] in [
+            "visibility",
+            "spectrum",
+            "radiometer",
+        ], "The type of the xdt must be 'visibility', 'spectrum' or 'radiometer'."
+
+        if "data_group_name" in indexers_kwargs:
+            data_group_name = indexers_kwargs["data_group_name"]
+            del indexers_kwargs["data_group_name"]
+        elif (indexers is not None) and ("data_group_name" in indexers):
+            data_group_name = indexers["data_group_name"]
+            del indexers["data_group_name"]
+        else:
+            data_group_name = None
+
+        if data_group_name is not None:
+            sel_data_group_set = set(
+                self._xdt.attrs["data_groups"][data_group_name].values()
+            )
+
+            data_variables_to_drop = []
+            for dg in self._xdt.attrs["data_groups"].values():
+                temp_set = set(dg.values()) - sel_data_group_set
+                data_variables_to_drop.extend(list(temp_set))
+
+            data_variables_to_drop = list(set(data_variables_to_drop))
+
+            sel_ms_xdt = self._xdt
+
+            sel_corr_xds = self._xdt.ds.sel(
+                indexers, method, tolerance, drop, **indexers_kwargs
+            ).drop_vars(data_variables_to_drop)
+
+            sel_ms_xdt.ds = sel_corr_xds
+
+            sel_ms_xdt.attrs["data_groups"] = {
+                data_group_name: self._xdt.attrs["data_groups"][data_group_name]
+            }
+
+            return sel_ms_xdt
+        else:
+            return self._xdt.sel(indexers, method, tolerance, drop, **indexers_kwargs)
+
+    def get_field_and_source_xds(self, data_group_name: str = None) -> xr.Dataset:
+        """Get the field_and_source_xds associated with data group `data_group_name`.
+
+        Parameters
+        ----------
+        data_group_name : str, optional
+            The data group to process. Default is "base" or if not found to first data group.
+
+        Returns
+        -------
+        xarray.Dataset
+            field_and_source_xds associated with the data group.
+        """
+        if self._xdt.attrs.get("type") not in MS_DATASET_TYPES:
+            raise InvalidAccessorLocation(f"{self._xdt.path} is not a MSv4 node.")
+
+        if data_group_name is None:
+            if "base" in self._xdt.attrs["data_groups"].keys():
+                data_group_name = "base"
+            else:
+                data_group_name = list(self._xdt.attrs["data_groups"].keys())[0]
+
+        return self._xdt[f"field_and_source_xds_{data_group_name}"].ds
+
+    def get_partition_info(self, data_group_name: str = None) -> dict:
+        """
+        Generate a partition info dict for an MSv4, with general MSv4 description including
+        information such as field and source names, SPW name, scan name, the intents string,
+        etc.
+
+        The information is gathered from various coordinates, secondary datasets, and info
+        dicts of the MSv4. For example, the SPW name comes from the attributes of the
+        frequency coordinate, whereas field and source related information such as field and
+        source names come from the field_and_source_xds (base) dataset of the MSv4.
+
+        Parameters
+        ----------
+        data_group_name : str, optional
+            The data group to process. Default is "base" or if not found to first data group.
+
+        Returns
+        -------
+        dict
+            Partition info dict for the MSv4
+        """
+        if self._xdt.attrs.get("type") not in MS_DATASET_TYPES:
+            raise InvalidAccessorLocation(
+                f"{self._xdt.path} is not a MSv4 node (type {self._xdt.attrs.get('type')}."
+            )
+
+        if data_group_name is None:
+            if "base" in self._xdt.attrs["data_groups"].keys():
+                data_group_name = "base"
+            else:
+                data_group_name = list(self._xdt.attrs["data_groups"].keys())[0]
+
+        field_and_source_xds = self._xdt.xr_ms.get_field_and_source_xds(data_group_name)
+
+        if "line_name" in field_and_source_xds.coords:
+            line_name = to_list(
+                np.unique(np.ravel(field_and_source_xds.line_name.values))
+            )
+        else:
+            line_name = []
+
+        partition_info = {
+            "spectral_window_name": self._xdt.frequency.attrs["spectral_window_name"],
+            "field_name": to_list(np.unique(field_and_source_xds.field_name.values)),
+            "polarization_setup": to_list(self._xdt.polarization.values),
+            "scan_name": to_list(np.unique(self._xdt.scan_name.values)),
+            "source_name": to_list(np.unique(field_and_source_xds.source_name.values)),
+            "intents": self._xdt.observation_info["intents"],
+            "line_name": line_name,
+            "data_group_name": data_group_name,
+        }
+
+        return partition_info