xradio 0.0.40__py3-none-any.whl → 0.0.42__py3-none-any.whl

xradio/{vis → measurement_set}/load_processing_set.py

@@ -1,5 +1,5 @@
 import os
-from ._processing_set import processing_set
+from xradio.measurement_set import ProcessingSet
 from typing import Dict, Union
 
 
@@ -8,7 +8,7 @@ def load_processing_set(
     sel_parms: dict,
     data_variables: Union[list, None] = None,
     load_sub_datasets: bool = True,
-) -> processing_set:
+) -> ProcessingSet:
     """Loads a processing set into memory.
 
     Parameters
@@ -29,24 +29,24 @@ def load_processing_set(
     data_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.
     load_sub_datasets : bool, optional
-        If true sub-datasets (for example weather_xds, antenna_xds, pointing_xds, ...) will be loaded into memory, by default True.
+        If true sub-datasets (for example weather_xds, antenna_xds, pointing_xds, system_calibration_xds ...) will be loaded into memory, by default True.
 
     Returns
     -------
-    processing_set
+    ProcessingSet
         In memory representation of processing set (data is represented by Dask.arrays).
     """
     from xradio._utils.zarr.common import _open_dataset, _get_file_system_and_items
 
     file_system, ms_store_list = _get_file_system_and_items(ps_store)
 
-    ps = processing_set()
+    ps = ProcessingSet()
     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")
+        correlated_store = os.path.join(ms_store, "correlated_xds")
 
         xds = _open_dataset(
-            ms_main_store,
+            correlated_store,
             file_system,
             ms_xds_isel,
             data_variables,
@@ -55,9 +55,9 @@ def load_processing_set(
         data_groups = xds.attrs["data_groups"]
 
         if load_sub_datasets:
-            from xradio.vis.read_processing_set import _read_sub_xds
+            from xradio.measurement_set.open_processing_set import _open_sub_xds
 
-            sub_xds_dict, field_and_source_xds_dict = _read_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
             )
 
@@ -66,26 +66,22 @@ def load_processing_set(
                 **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]
-                    )
+
+                xds[data_group_vals["correlated_data"]].attrs[
+                    "field_and_source_xds"
+                ] = field_and_source_xds_dict[data_group_name]
 
         ps[ms_name] = xds
 
     return ps
 
 
-class processing_set_iterator:
+class ProcessingSetIterator:
     def __init__(
         self,
         sel_parms: dict,
         input_data_store: str,
-        input_data: Union[Dict, processing_set, None] = None,
+        input_data: Union[Dict, ProcessingSet, None] = None,
         data_variables: list = None,
         load_sub_datasets: bool = True,
     ):
@@ -110,7 +106,7 @@ class processing_set_iterator:
         data_variables : list, 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.
         load_sub_datasets : bool, optional
-            If true sub-datasets (for example weather_xds, antenna_xds, pointing_xds, ...) will be loaded into memory, by default True.
+            If true sub-datasets (for example weather_xds, antenna_xds, pointing_xds, system_calibration_xds ...) will be loaded into memory, by default True.
         """
 
         self.input_data = input_data

xradio/measurement_set/measurement_set_xds.py

@@ -0,0 +1,83 @@
+import pandas as pd
+from xradio._utils.list_and_array import to_list
+import xarray as xr
+import numbers
+import os
+
+
+class MeasurementSetXds(xr.Dataset):
+    __slots__ = ()
+
+    def __init__(self, xds):
+        super().__init__(xds.data_vars, xds.coords, xds.attrs)
+
+    def to_store(self, store, **kwargs):
+        """
+        Write the MeasurementSetXds to a Zarr store.
+        Does not write to cloud storage yet.
+
+        Args:
+            store (str): The path to the Zarr store.
+            **kwargs: Additional keyword arguments to be passed to `xarray.Dataset.to_zarr`. See https://docs.xarray.dev/en/latest/generated/xarray.Dataset.to_zarr.html for more information.
+
+        Returns:
+            None
+        """
+
+        copy_cor_xds = self.copy()  # No deep copy
+
+        # Remove field_and_source_xds from all correlated_data (VISIBILITY/SPECTRUM) data variables
+        # and save them as separate zarr files.
+        for data_group_name, data_group in self.attrs["data_groups"].items():
+            del copy_cor_xds[data_group["correlated_data"]].attrs[
+                "field_and_source_xds"
+            ]
+
+            # print("data_group_name", data_group_name)
+            xr.Dataset.to_zarr(
+                self[data_group["correlated_data"]].attrs["field_and_source_xds"],
+                os.path.join(store, "field_and_source_xds_" + data_group_name),
+                **kwargs,
+            )
+
+        # Remove xds attributes from copy_cor_xds and save xds attributes as separate zarr files.
+        for attrs_name in self.attrs:
+            if "xds" in attrs_name:
+                del copy_cor_xds.attrs[attrs_name]
+                xr.Dataset.to_zarr(
+                    self.attrs[attrs_name], os.path.join(store, attrs_name, **kwargs)
+                )
+
+        # Save copy_cor_xds as zarr file.
+        xr.Dataset.to_zarr(
+            copy_cor_xds, os.path.join(store, "correlated_xds"), **kwargs
+        )
+
+    def sel(self, data_group_name=None, **kwargs):
+        """
+        Select data along dimension(s) by label.
+
+        Args:
+            **kwargs: Keyword arguments to be passed to `xarray.Dataset.sel`. See https://xarray.pydata.org/en/stable/generated/xarray.Dataset.sel.html for more information.
+
+        Returns:
+            MeasurementSetXds
+        """
+
+        if data_group_name is not None:
+            sel_data_group_set = set(
+                self.attrs["data_groups"][data_group_name].values()
+            )
+
+            data_variables_to_drop = []
+            for dg in self.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))
+
+            return MeasurementSetXds(
+                super().sel(**kwargs).drop_vars(data_variables_to_drop)
+            )
+        else:
+            return MeasurementSetXds(super().sel(**kwargs))

xradio/{vis/read_processing_set.py → measurement_set/open_processing_set.py}

@@ -1,46 +1,48 @@
 import os
 
-from ._processing_set import processing_set
+from xradio.measurement_set import ProcessingSet
 import toolviper.utils.logger as logger
 from xradio._utils.zarr.common import _open_dataset, _get_file_system_and_items
 import s3fs
 
 
-def read_processing_set(
+def open_processing_set(
     ps_store: str,
-    obs_modes: list = None,
-) -> processing_set:
+    intents: list = None,
+) -> ProcessingSet:
     """Creates a lazy representation of a Processing Set (only meta-data is loaded 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'.
-    obs_modes : list, optional
-        A list of obs_mode to be read for example ['OBSERVE_TARGET#ON_SOURCE']. The obs_mode in a processing set can be seen by calling processing_set.summary().
-        By default None, which will read all obs_mode.
+    intents : list, optional
+        A list of intents to be open for example ['OBSERVE_TARGET#ON_SOURCE']. The intents in a processing set can be seen by calling processing_set.summary().
+        By default None, which will open all intents.
 
     Returns
     -------
     processing_set
         Lazy representation of processing set (data is represented by Dask.arrays).
     """
+    from xradio.measurement_set import MeasurementSetXds
+
     file_system, ms_store_list = _get_file_system_and_items(ps_store)
 
-    ps = processing_set()
+    ps = ProcessingSet()
     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")
+        correlated_store = os.path.join(ms_store, "correlated_xds")
 
-        xds = _open_dataset(ms_main_store, file_system)
+        xds = _open_dataset(correlated_store, file_system)
         data_groups = xds.attrs["data_groups"]
 
-        if (obs_modes is None) or (
-            bool(set(xds.attrs["partition_info"]["obs_mode"]).intersection(obs_modes))
+        if (intents is None) or (
+            bool(set(xds.attrs["partition_info"]["intents"]).intersection(intents))
         ):
-            sub_xds_dict, field_and_source_xds_dict = _read_sub_xds(
+            sub_xds_dict, field_and_source_xds_dict = _open_sub_xds(
                 ms_store, file_system=file_system, data_groups=data_groups
             )
 
@@ -50,33 +52,22 @@ def read_processing_set(
             }
 
             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
+                xds[data_group_vals["correlated_data"]].attrs[
+                    "field_and_source_xds"
+                ] = field_and_source_xds_dict[data_group_name]
+
+            ps[ms_name] = MeasurementSetXds(xds)
         # except Exception as e:
-        #     logger.warning(f"Could not read {ms_name} due to {e}")
+        #     logger.warning(f"Could not open {ms_name} due to {e}")
         #     continue
 
     return ps
 
 
-def _read_sub_xds(ms_store, file_system, data_groups, load=False):
+def _open_sub_xds(ms_store, file_system, data_groups, load=False):
     sub_xds_dict = {}
     field_and_source_xds_dict = {}
 
-    xds_names = {
-        "ANTENNA": "antenna_xds",
-        "WEATHER": "weather_xds",
-        "POINTING": "pointing_xds",
-    }
-
     if isinstance(file_system, s3fs.core.S3FileSystem):
         file_names = [
             bd.split(sep="/")[-1] for bd in file_system.listdir(ms_store, detail=False)
@@ -85,9 +76,9 @@ def _read_sub_xds(ms_store, file_system, data_groups, load=False):
         file_names = file_system.listdir(ms_store)
     file_names = [item for item in file_names if not item.startswith(".")]
 
-    file_names.remove("MAIN")
+    file_names.remove("correlated_xds")
 
-    field_dict = {"FIELD_AND_SOURCE_" + key.upper(): key for key in data_groups.keys()}
+    field_dict = {"field_and_source_xds_" + key: key for key in data_groups.keys()}
 
     # field_and_source_xds_name_start = "FIELD"
     for n in file_names:
@@ -100,7 +91,7 @@ def _read_sub_xds(ms_store, file_system, data_groups, load=False):
         if n in field_dict.keys():
             field_and_source_xds_dict[field_dict[n]] = xds
         else:
-            sub_xds_dict[xds_names[n]] = xds
+            sub_xds_dict[n] = xds
 
     return sub_xds_dict, field_and_source_xds_dict