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

xradio/{vis → correlated_data}/load_processing_set.py

@@ -1,5 +1,5 @@
 import os
-from ._processing_set import processing_set
+from xradio.correlated_data 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,18 +29,18 @@ 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")
@@ -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.correlated_data.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,14 +66,10 @@ 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
 
@@ -85,7 +81,7 @@ class processing_set_iterator:
         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/{vis/read_processing_set.py → correlated_data/open_processing_set.py}

@@ -1,24 +1,24 @@
 import os
 
-from ._processing_set import processing_set
+from xradio.correlated_data 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
     -------
@@ -27,7 +27,7 @@ def read_processing_set(
     """
     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:
@@ -37,10 +37,10 @@ def read_processing_set(
         xds = _open_dataset(ms_main_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,31 +50,29 @@ 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]
-                    )
+                xds[data_group_vals["correlated_data"]].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}")
+        #     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",
+        "SYSCAL": "system_calibration_xds",
+        "GAIN_CURVE": "gain_curve_xds",
+        "PHASE_CAL": "phase_calibration_xds",
+        "WEATHER": "weather_xds",
     }
 
     if isinstance(file_system, s3fs.core.S3FileSystem):

xradio/{vis/_processing_set.py → correlated_data/processing_set.py}

@@ -3,7 +3,7 @@ from xradio._utils.list_and_array import to_list
 import numbers
 
 
-class processing_set(dict):
+class ProcessingSet(dict):
     """
     A dictionary subclass representing a Processing Set (PS) that is a set of Measurement Sets v4 (MS).
 
@@ -74,7 +74,7 @@ class processing_set(dict):
     def _summary(self, data_group="base"):
         summary_data = {
             "name": [],
-            "obs_mode": [],
+            "intents": [],
             "shape": [],
             "polarization": [],
             "scan_number": [],
@@ -94,7 +94,7 @@ class processing_set(dict):
 
         for key, value in self.items():
             summary_data["name"].append(key)
-            summary_data["obs_mode"].append(value.attrs["partition_info"]["obs_mode"])
+            summary_data["intents"].append(value.attrs["partition_info"]["intents"])
             summary_data["spw_name"].append(
                 value.attrs["partition_info"]["spectral_window_name"]
             )
@@ -102,13 +102,12 @@ class processing_set(dict):
             summary_data["scan_number"].append(
                 value.attrs["partition_info"]["scan_number"]
             )
+            data_name = value.attrs["data_groups"][data_group]["correlated_data"]
 
-            if "visibility" in value.attrs["data_groups"][data_group]:
-                data_name = value.attrs["data_groups"][data_group]["visibility"]
+            if "VISIBILITY" in data_name:
                 center_name = "FIELD_PHASE_CENTER"
 
-            if "spectrum" in value.attrs["data_groups"][data_group]:
-                data_name = value.attrs["data_groups"][data_group]["spectrum"]
+            if "SPECTRUM" in data_name:
                 center_name = "FIELD_REFERENCE_CENTER"
 
             summary_data["shape"].append(value[data_name].shape)
@@ -202,12 +201,12 @@ class processing_set(dict):
         """
         Selects a subset of the Processing Set based on column names and values or a Pandas query.
 
-        The following columns are supported: name, obs_mode, polarization, spw_name, field_name, source_name, field_coords, start_frequency, end_frequency.
+        The following columns are supported: name, intents, polarization, spw_name, field_name, source_name, field_coords, start_frequency, end_frequency.
 
         This function will not apply any selection on the MS data so data will not be dropped for example if a MS has field_name=['field_0','field_10','field_08'] and ps.sel(field_name='field_0') is done the resulting MS will still have field_name=['field_0','field_10','field_08'].
 
         Examples:
-            ps.sel(obs_mode='OBSERVE_TARGET#ON_SOURCE', polarization=['RR', 'LL']) # Select all MSs with obs_mode 'OBSERVE_TARGET#ON_SOURCE' and polarization 'RR' or 'LL'.
+            ps.sel(intents='OBSERVE_TARGET#ON_SOURCE', polarization=['RR', 'LL']) # Select all MSs with intents 'OBSERVE_TARGET#ON_SOURCE' and polarization 'RR' or 'LL'.
             ps.sel(query='start_frequency > 100e9 AND end_frequency < 200e9') # Select all MSs with start_frequency greater than 100 GHz and less than 200 GHz.
 
         Args:
@@ -264,7 +263,7 @@ class processing_set(dict):
         if query is not None:
             summary_table = summary_table.query(query)
 
-        sub_ps = processing_set()
+        sub_ps = ProcessingSet()
         for key, val in self.items():
             if key in summary_table["name"].values:
                 sub_ps[key] = val
@@ -281,7 +280,7 @@ class processing_set(dict):
         Returns:
             processing_set: The subset of the Processing Set.
         """
-        sub_ps = processing_set()
+        sub_ps = ProcessingSet()
         for key, val in self.items():
             sub_ps[key] = val.sel(kwargs)
         return sub_ps
@@ -296,7 +295,7 @@ class processing_set(dict):
         Returns:
             processing_set: The subset of the Processing Set.
         """
-        sub_ps = processing_set()
+        sub_ps = ProcessingSet()
         for key, val in self.items():
             sub_ps[key] = val.isel(kwargs)
         return sub_ps