xradio 0.0.22__tar.gz → 0.0.24__tar.gz

{xradio-0.0.22/src/xradio.egg-info → xradio-0.0.24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.22
+Version: 0.0.24
 Summary: Xarray Radio Astronomy Data IO 
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License

{xradio-0.0.22 → xradio-0.0.24}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "xradio"
-version = "0.0.22"
+version = "0.0.24"
 description = "Xarray Radio Astronomy Data IO "
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},

xradio-0.0.24/src/xradio/_utils/zarr/common.py

@@ -0,0 +1,113 @@
+import copy
+import xarray as xr
+import zarr
+
+
+def _open_dataset(store, xds_isel=None, data_variables=None, load=False):
+    """
+
+    Parameters
+    ----------
+    store : _type_
+        _description_
+    xds_isel : _type_, optional
+        Example {'time':slice(0,10), 'frequency':slice(5,7)}, by default None
+    data_variables : _type_, optional
+        Example ['VISIBILITY','WEIGHT'], by default None
+    load : bool, optional
+        _description_, by default False
+
+    Returns
+    -------
+    _type_
+        _description_
+    """
+
+    import dask
+
+    xds = xr.open_zarr(store)
+
+    if xds_isel is not None:
+        xds = xds.isel(xds_isel)
+
+    if data_variables is not None:
+        xds_sub = xr.Dataset()
+        for dv in data_variables:
+            xds_sub[dv] = xds[dv]
+        xds_sub.attrs = xds.attrs
+        xds = xds_sub
+
+    if load:
+        with dask.config.set(scheduler="synchronous"):
+            xds = xds.load()
+    return xds
+
+
+# Code to depricate:
+def _get_attrs(zarr_obj):
+    """
+    get attributes of zarr obj (groups or arrays)
+    """
+    return {k: v for k, v in zarr_obj.attrs.asdict().items() if not k.startswith("_NC")}
+
+
+def _load_no_dask_zarr(zarr_name, slice_dict={}):
+    """
+    Alternative to xarray open_zarr where the arrays are not Dask Arrays.
+
+    slice_dict: A dictionary of slice objects for which values to read form a dimension.
+                For example silce_dict={'time':slice(0,10)} would select the first 10 elements in the time dimension.
+                If a dim is not specified all values are retruned.
+    return:
+        xarray.Dataset()
+
+    #Should go into general utils.
+    """
+    # Used by xarray to store array labeling info in zarr meta data.
+    DIMENSION_KEY = "_ARRAY_DIMENSIONS"
+
+    # logger = _get_logger()
+    zarr_group = zarr.open_group(store=zarr_name, mode="r")
+    group_attrs = _get_attrs(zarr_group)
+
+    slice_dict_complete = copy.deepcopy(slice_dict)
+    coords = {}
+    xds = xr.Dataset()
+    for var_name, var in zarr_group.arrays():
+        print("Hallo 3", var_name, var.shape)
+        var_attrs = _get_attrs(var)
+
+        for dim in var_attrs[DIMENSION_KEY]:
+            if dim not in slice_dict_complete:
+                slice_dict_complete[dim] = slice(None)  # No slicing.
+
+        if (var_attrs[DIMENSION_KEY][0] == var_name) and (
+            len(var_attrs[DIMENSION_KEY]) == 1
+        ):
+            coord = var[
+                slice_dict_complete[var_attrs[DIMENSION_KEY][0]]
+            ]  # Dimension coordinates.
+            del var_attrs["_ARRAY_DIMENSIONS"]
+            xds = xds.assign_coords({var_name: coord})
+            xds[var_name].attrs = var_attrs
+        else:
+            # Construct slicing
+            slicing_list = []
+            for dim in var_attrs[DIMENSION_KEY]:
+                slicing_list.append(slice_dict_complete[dim])
+            slicing_tuple = tuple(slicing_list)
+
+            print(var_attrs[DIMENSION_KEY])
+
+            xds[var_name] = xr.DataArray(
+                var[slicing_tuple], dims=var_attrs[DIMENSION_KEY]
+            )
+
+            if "coordinates" in var_attrs:
+                del var_attrs["coordinates"]
+            del var_attrs["_ARRAY_DIMENSIONS"]
+            xds[var_name].attrs = var_attrs
+
+    xds.attrs = group_attrs
+
+    return xds

{xradio-0.0.22 → xradio-0.0.24}/src/xradio/image/_util/_zarr/zarr_low_level.py

@@ -70,7 +70,6 @@ image_data_variables_and_dims_single_precision = {
 }
 
 
-
 def pad_array_with_nans(input_array, output_shape, dtype):
     """
     Pad an integer array with NaN values to match the specified output shape.
@@ -93,7 +92,7 @@ def pad_array_with_nans(input_array, output_shape, dtype):
     padded_array[:] = np.nan
 
     # Copy the input array to the appropriate position within the padded array
-    padded_array[: input_shape[0], : input_shape[1]] = input_array
+    padded_array[: input_shape[0], : input_shape[1], : input_shape[2]] = input_array
 
     return padded_array
 
@@ -277,8 +276,7 @@ def create_data_variable_meta_data_on_disk(
     return zarr_meta
 
 
-
-def write_chunk(img_xds,meta,parallel_dims_chunk_id,compressor,image_file):
+def write_chunk(img_xds, meta, parallel_dims_chunk_id, compressor, image_file):
     dims = meta["dims"]
     dtype = meta["dtype"]
     data_varaible_name = meta["name"]
@@ -304,8 +302,6 @@ def write_chunk(img_xds,meta,parallel_dims_chunk_id,compressor,image_file):
 
         write_binary_blob_to_disk(
             array,
-            file_path=os.path.join(
-                image_file, data_varaible_name, chunk_name
-            ),
+            file_path=os.path.join(image_file, data_varaible_name, chunk_name),
             compressor=compressor,
-        )
+        )

xradio-0.0.24/src/xradio/image/_util/zarr.py

@@ -0,0 +1,50 @@
+from ._zarr.xds_to_zarr import _write_zarr
+from ._zarr.xds_from_zarr import _read_zarr
+import numpy as np
+import os
+import xarray as xr
+from ..._utils.zarr.common import _open_dataset
+
+
+def _xds_to_zarr(xds: xr.Dataset, zarr_store: str):
+    _write_zarr(xds, zarr_store)
+
+
+def _xds_from_zarr(
+    zarr_store: str, output: dict = {}, selection: dict = {}
+) -> xr.Dataset:
+    # supported key/values in output are:
+    # "dv"
+    #    what data variables should be returned as.
+    #    "numpy": numpy arrays
+    #    "dask": dask arrays
+    # "coords"
+    #    what coordinates should be returned as
+    #    "numpy": numpy arrays
+    return _read_zarr(zarr_store, output, selection)
+
+
+def _load_image_from_zarr_no_dask(zarr_file: str, selection: dict) -> xr.Dataset:
+    image_xds = _open_dataset(zarr_file, selection, load=True)
+    for h in ["HISTORY", "_attrs_xds_history"]:
+        history = os.sep.join([zarr_file, h])
+        if os.path.isdir(history):
+            image_xds.attrs["history"] = _open_dataset(history, load=True)
+            break
+    _iter_dict(image_xds.attrs)
+    return image_xds
+
+
+def _iter_dict(d: dict) -> None:
+    for k, v in d.items():
+        if isinstance(v, dict):
+            keys = v.keys()
+            if (
+                len(keys) == 3
+                and "_dtype" in keys
+                and "_type" in keys
+                and "_value" in keys
+            ):
+                d[k] = np.array(v["_value"], dtype=v["_dtype"])
+            else:
+                _iter_dict(v)

{xradio-0.0.22 → xradio-0.0.24}/src/xradio/image/image.py

@@ -9,6 +9,8 @@ import copy
 import numpy as np
 import xarray as xr
 
+# from .._utils.zarr.common import _load_no_dask_zarr
+
 from ._util.casacore import _load_casa_image_block, _xds_to_casa_image
 from ._util.fits import _read_fits_image
 from ._util.image_factory import (
@@ -16,7 +18,7 @@ from ._util.image_factory import (
     _make_empty_lmuv_image,
     _make_empty_sky_image,
 )
-from ._util.zarr import _xds_to_zarr, _xds_from_zarr
+from ._util.zarr import _load_image_from_zarr_no_dask, _xds_from_zarr, _xds_to_zarr
 
 warnings.filterwarnings("ignore", category=FutureWarning)
 
@@ -166,9 +168,11 @@ def load_image(infile: str, block_des: dict = {}, do_sky_coords=True) -> xr.Data
         emsgs.append(f'image format appears not to be fits {e.args}')
     """
     # when done debugging, comment out next line
+    # return _load_image_from_zarr_no_dask(infile, block_des)
     # return _xds_from_zarr(infile, {"dv": "numpy"}, selection)
     try:
-        return _xds_from_zarr(infile, {"dv": "numpy", "coords": "numpy"}, selection)
+        return _load_image_from_zarr_no_dask(infile, block_des)
+        # return _xds_from_zarr(infile, {"dv": "numpy", "coords": "numpy"}, selection)
     except Exception as e:
         emsgs.append(f"image format appears not to be zarr {e.args}")
     emsgs.insert(
@@ -177,7 +181,9 @@ def load_image(infile: str, block_des: dict = {}, do_sky_coords=True) -> xr.Data
     raise RuntimeError("\n".join(emsgs))
 
 
-def write_image(xds: xr.Dataset, imagename: str, out_format: str = "casa", overwrite=False) -> None:
+def write_image(
+    xds: xr.Dataset, imagename: str, out_format: str = "casa", overwrite=False
+) -> None:
     """
     Convert an xds image to CASA or zarr image.
     xds : xarray.Dataset
@@ -194,7 +200,8 @@ def write_image(xds: xr.Dataset, imagename: str, out_format: str = "casa", overw
 
     if overwrite:
         import os
-        os.system('rm -rf ' + imagename)
+
+        os.system("rm -rf " + imagename)
 
     if my_format == "casa":
         _xds_to_casa_image(xds, imagename)

xradio-0.0.24/src/xradio/vis/_processing_set.py

@@ -0,0 +1,103 @@
+import pandas as pd
+
+
+class processing_set(dict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.meta = {"summary": {}}
+
+    #     generate_meta(self)
+
+    # def generate_meta(self):
+    #     self.meta['summary'] = {"base": _summary(self)}
+    #     self.meta['max_dims'] = _get_ps_max_dims(self)
+
+    def summary(self, data_group="base"):
+        if data_group in self.meta["summary"]:
+            return self.meta["summary"][data_group]
+        else:
+            self.meta["summary"][data_group] = self._summary(data_group)
+            return self.meta["summary"][data_group]
+
+    def get_ps_max_dims(self):
+        if "max_dims" in self.meta:
+            return self.meta["max_dims"]
+        else:
+            self.meta["max_dims"] = self._get_ps_max_dims()
+            return self.meta["max_dims"]
+
+    def get_ps_freq_axis(self):
+        if "freq_axis" in self.meta:
+            return self.meta["freq_axis"]
+        else:
+            self.meta["freq_axis"] = self._get_ps_freq_axis()
+            return self.meta["freq_axis"]
+
+    def _summary(self, data_group="base"):
+        summary_data = {
+            "name": [],
+            "ddi": [],
+            "intent": [],
+            "field_id": [],
+            "field_name": [],
+            "start_frequency": [],
+            "end_frequency": [],
+            "shape": [],
+        }
+        for key, value in self.items():
+            summary_data["name"].append(key)
+            summary_data["ddi"].append(value.attrs["ddi"])
+            summary_data["intent"].append(value.attrs["intent"])
+
+            if "visibility" in value.attrs["data_groups"][data_group]:
+                data_name = value.attrs["data_groups"][data_group]["visibility"]
+
+            if "spectrum" in value.attrs["data_groups"][data_group]:
+                data_name = value.attrs["data_groups"][data_group]["spectrum"]
+
+            summary_data["shape"].append(value[data_name].shape)
+
+            summary_data["field_id"].append(
+                value[data_name].attrs["field_info"]["field_id"]
+            )
+            summary_data["field_name"].append(
+                value[data_name].attrs["field_info"]["name"]
+            )
+            summary_data["start_frequency"].append(value["frequency"].values[0])
+            summary_data["end_frequency"].append(value["frequency"].values[-1])
+        summary_df = pd.DataFrame(summary_data)
+        return summary_df
+
+    def _get_ps_freq_axis(self):
+        import xarray as xr
+
+        spw_ids = []
+        freq_axis_list = []
+        frame = self.get(0).frequency.attrs["frame"]
+        for ms_xds in self.values():
+            assert (
+                frame == ms_xds.frequency.attrs["frame"]
+            ), "Frequency reference frame not consistent in processing set."
+            if ms_xds.frequency.attrs["spw_id"] not in spw_ids:
+                spw_ids.append(ms_xds.frequency.attrs["spw_id"])
+                freq_axis_list.append(ms_xds.frequency)
+
+        freq_axis = xr.concat(freq_axis_list, dim="frequency").sortby("frequency")
+        return freq_axis
+
+    def _get_ps_max_dims(self):
+        max_dims = None
+        for ms_xds in self.values():
+            if max_dims is None:
+                max_dims = dict(ms_xds.sizes)
+            else:
+                for dim_name, size in ms_xds.sizes.items():
+                    if dim_name in max_dims:
+                        if max_dims[dim_name] < size:
+                            max_dims[dim_name] = size
+                    else:
+                        max_dims[dim_name] = size
+        return max_dims
+
+    def get(self, id):
+        return self[list(self.keys())[id]]

xradio-0.0.24/src/xradio/vis/load_processing_set.py

@@ -0,0 +1,127 @@
+import xarray as xr
+import zarr
+import copy
+import os
+from ._processing_set import processing_set
+from typing import Dict, Union
+
+
+def load_processing_set(
+    ps_store: str,
+    sel_parms: dict,
+    data_variables: Union[list, None] = None,
+    load_sub_datasets: bool = True,
+)->processing_set:
+    """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'.
+    sel_parms : dict
+        A dictionary where the keys are the names of the ms_xds's 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
+        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.
+
+    Returns
+    -------
+    processing_set
+        In memory representation of processing set (data is represented by Dask.arrays). 
+    """    
+    from xradio._utils.zarr.common import _open_dataset
+
+    ps = processing_set()
+    for ms_dir_name, ms_xds_isel in sel_parms.items():
+        xds = _open_dataset(
+            os.path.join(ps_store, ms_dir_name, "MAIN"),
+            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),
+            }
+
+        ps[ms_dir_name] = xds
+    return ps
+
+
+class processing_set_iterator:
+
+    def __init__(
+        self,
+        sel_parms: dict,
+        input_data_store: str,
+        input_data: Union[Dict, processing_set, None] = None,
+        data_variables: list = None,
+        load_sub_datasets: bool = True,
+    ):
+        """An iterator that will go through a processing set one MS v4 at a time.
+
+        Parameters
+        ----------
+        sel_parms : dict
+            A dictionary where the keys are the names of the ms_xds's 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)},
+                }
+        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
+            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
+            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.
+        """        
+
+        self.input_data = input_data
+        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.load_sub_datasets = load_sub_datasets
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            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(
+                ps_store=self.input_data_store,
+                sel_parms={xds_name: slice_description},
+                data_variables=self.data_variables,
+                load_sub_datasets=self.load_sub_datasets,
+            )
+            xds = ps.get(0)
+        else:
+            xds = self.input_data[xds_name]  # In memory
+
+        return xds

xradio-0.0.24/src/xradio/vis/read_processing_set.py

@@ -0,0 +1,85 @@
+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
+
+
+def read_processing_set(
+    ps_store: str, intents: list = None, fields: str = None
+)->processing_set:
+    """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'.
+    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). 
+    """    
+    items = os.listdir(ps_store)
+    ms_xds = xr.Dataset()
+    ps = processing_set()
+    data_group = 'base'
+    for ms_dir_name in items:
+        if "ddi" in ms_dir_name:
+            xds = _open_dataset(os.path.join(ps_store, ms_dir_name, "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
+                ):
+                    xds.attrs = {
+                        **xds.attrs,
+                        **_read_sub_xds(os.path.join(ps_store, ms_dir_name)),
+                    }
+                    ps[ms_dir_name] = xds
+    return ps
+
+
+def _read_sub_xds(ms_store, load=False):
+    sub_xds_dict = {}
+
+    sub_xds = {
+        "antenna_xds": "ANTENNA",
+    }
+    for sub_xds_key, sub_xds_name in sub_xds.items():
+        sub_xds_dict[sub_xds_key] = _open_dataset(
+            os.path.join(ms_store, sub_xds_name), load=load
+        )
+
+    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)
+
+    return sub_xds_dict
+
+
+def _get_data_name(xds, data_group):
+    if "visibility" in xds.attrs["data_groups"][data_group]:
+        data_name = xds.attrs["data_groups"][data_group]["visibility"]
+    elif "spectrum" in xds.attrs["data_groups"][data_group]:
+        data_name = xds.attrs["data_groups"][data_group]["spectrum"]
+    else:
+        error_message = (
+            "No Visibility or Spectrum data variable found in data_group "
+            + data_group
+            + "."
+        )
+        logger.exception(error_message)
+        raise ValueError(error_message)
+    return data_name

{xradio-0.0.22 → xradio-0.0.24/src/xradio.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.22
+Version: 0.0.24
 Summary: Xarray Radio Astronomy Data IO 
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License

{xradio-0.0.22 → xradio-0.0.24}/src/xradio.egg-info/SOURCES.txt

@@ -11,6 +11,8 @@ src/xradio.egg-info/top_level.txt
 src/xradio/_utils/__init__.py
 src/xradio/_utils/common.py
 src/xradio/_utils/_casacore/tables.py
+src/xradio/_utils/zarr/__init__.py
+src/xradio/_utils/zarr/common.py
 src/xradio/image/__init__.py
 src/xradio/image/image.py
 src/xradio/image/_util/__init__.py

xradio-0.0.22/src/xradio/image/_util/zarr.py

@@ -1,21 +0,0 @@
-from ._zarr.xds_to_zarr import _write_zarr
-from ._zarr.xds_from_zarr import _read_zarr
-import xarray as xr
-
-
-def _xds_to_zarr(xds: xr.Dataset, zarr_store: str):
-    _write_zarr(xds, zarr_store)
-
-
-def _xds_from_zarr(
-    zarr_store: str, output: dict = {}, selection: dict = {}
-) -> xr.Dataset:
-    # supported key/values in output are:
-    # "dv"
-    #    what data variables should be returned as.
-    #    "numpy": numpy arrays
-    #    "dask": dask arrays
-    # "coords"
-    #    what coordinates should be returned as
-    #    "numpy": numpy arrays
-    return _read_zarr(zarr_store, output, selection)

xradio-0.0.22/src/xradio/vis/_processing_set.py

@@ -1,45 +0,0 @@
-import pandas as pd
-
-
-class processing_set(dict):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def summary(self, data_group="base"):
-        summary_data = {
-            "name": [],
-            "ddi": [],
-            "intent": [],
-            "field_id": [],
-            "field_name": [],
-            "start_frequency": [],
-            "end_frequency": [],
-        }
-        for key, value in self.items():
-            summary_data["name"].append(key)
-            summary_data["ddi"].append(value.attrs["ddi"])
-            summary_data["intent"].append(value.attrs["intent"])
-
-            if "visibility" in value.attrs["data_groups"][data_group]:
-                data_name = value.attrs["data_groups"][data_group]["visibility"]
-
-            if "spectrum" in value.attrs["data_groups"][data_group]:
-                data_name = value.attrs["data_groups"][data_group]["spectrum"]
-
-            summary_data["field_id"].append(
-                value[data_name].attrs[
-                    "field_info"
-                ]["field_id"]
-            )
-            summary_data["field_name"].append(
-                value[data_name].attrs[
-                    "field_info"
-                ]["name"]
-            )
-            summary_data["start_frequency"].append(value["frequency"].values[0])
-            summary_data["end_frequency"].append(value["frequency"].values[-1])
-        summary_df = pd.DataFrame(summary_data)
-        return summary_df
-
-    def get(self, id):
-        return self[list(self.keys())[id]]

xradio-0.0.22/src/xradio/vis/load_processing_set.py

@@ -1,189 +0,0 @@
-import xarray as xr
-import zarr
-import copy
-import os
-from ._processing_set import processing_set
-
-DIMENSION_KEY = "_ARRAY_DIMENSIONS"  # Used by xarray to store array labeling info in zarr meta data.
-# from xradio._utils._logger import _get_logger
-
-
-def _get_attrs(zarr_obj):
-    """
-    get attributes of zarr obj (groups or arrays)
-    """
-    return {k: v for k, v in zarr_obj.attrs.asdict().items() if not k.startswith("_NC")}
-
-
-def _load_ms_xds(
-    ps_name, ms_xds_name, slice_dict={}, cache_dir=None, chunk_id=None, date_time=""
-):
-    # logger = _get_logger()
-    if cache_dir:
-        xds_cached_name = (
-            os.path.join(cache_dir, ms_xds_name) + "_" + str(chunk_id) + "_" + date_time
-        )
-
-        # Check if already chached:
-        try:
-            ms_xds = _load_ms_xds_core(
-                ms_xds_name=xds_cached_name, slice_dict=slice_dict
-            )
-
-            # logger.debug(ms_xds_name + ' chunk ' + str(slice_dict) + ' was found in cache: ' + xds_cached)
-            found_in_cache = True
-            return xds, found_in_cache
-        except:
-            # logger.debug(xds_cached + ' chunk ' + str(slice_dict) + ' was not found in cache or failed to load. Retrieving chunk from ' + ms_xds_name + ' .')
-            ms_xds = _load_ms_xds_core(
-                ms_xds_name=os.path.join(ps_name, ms_xds_name), slice_dict=slice_dict
-            )
-            write_ms_xds(ms_xds, xds_cached_name)
-
-            found_in_cache = False
-            return xds, found_in_cache
-    else:
-        found_in_cache = None
-        ms_xds = _load_ms_xds_core(
-            ms_xds_name=os.path.join(ps_name, ms_xds_name), slice_dict=slice_dict
-        )
-        return ms_xds, found_in_cache
-
-
-def _write_ms_xds(ms_xds, ms_xds_name):
-    ms_xds_temp = ms_xds
-    xr.Dataset.to_zarr(
-        ms_xds.attrs["ANTENNA"],
-        os.path.join(xds_cached_name, "ANTENNA"),
-        consolidated=True,
-    )
-    ms_xds_temp = ms_xds
-    ms_xds_temp.attrs["ANTENNA"] = {}
-    xr.Dataset.to_zarr(
-        ms_xds_temp, os.path.join(xds_cached_name, "MAIN"), consolidated=True
-    )
-
-
-def _load_ms_xds_core(ms_xds_name, slice_dict):
-    ms_xds = _load_no_dask_zarr(
-        zarr_name=os.path.join(ms_xds_name, "MAIN"), slice_dict=slice_dict
-    )
-    ms_xds.attrs["antenna_xds"] = _load_no_dask_zarr(
-        zarr_name=os.path.join(ms_xds_name, "ANTENNA")
-    )
-    sub_xds = {
-        "antenna_xds": "ANTENNA",
-    }
-    for sub_xds_key, sub_xds_name in sub_xds.items():
-        ms_xds.attrs[sub_xds_key] = _load_no_dask_zarr(
-            zarr_name=os.path.join(ms_xds_name, sub_xds_name)
-        )
-    optional_sub_xds = {
-        "weather_xds": "WEATHER",
-        "pointing_xds": "POINTING",
-    }
-    for sub_xds_key, sub_xds_name in sub_xds.items():
-        sub_xds_path = os.path.join(ms_xds_name, sub_xds_name)
-        if os.path.isdir(sub_xds_path):
-            ms_xds.attrs[sub_xds_key] = _load_no_dask_zarr(
-                zarr_name=os.path.join(ms_xds_name, sub_xds_name)
-            )
-
-    return ms_xds
-
-
-def _load_no_dask_zarr(zarr_name, slice_dict={}):
-    """
-    Alternative to xarray open_zarr where the arrays are not Dask Arrays.
-
-    slice_dict: A dictionary of slice objects for which values to read form a dimension.
-                For example silce_dict={'time':slice(0,10)} would select the first 10 elements in the time dimension.
-                If a dim is not specified all values are retruned.
-    return:
-        xarray.Dataset()
-
-    #Should go into general utils.
-    """
-
-    # logger = _get_logger()
-    zarr_group = zarr.open_group(store=zarr_name, mode="r")
-    group_attrs = _get_attrs(zarr_group)
-
-    slice_dict_complete = copy.deepcopy(slice_dict)
-    coords = {}
-    xds = xr.Dataset()
-    for var_name, var in zarr_group.arrays():
-        var_attrs = _get_attrs(var)
-
-        for dim in var_attrs[DIMENSION_KEY]:
-            if dim not in slice_dict_complete:
-                slice_dict_complete[dim] = slice(None)  # No slicing.
-
-        if (var_attrs[DIMENSION_KEY][0] == var_name) and (
-            len(var_attrs[DIMENSION_KEY]) == 1
-        ):
-            coord = var[
-                slice_dict_complete[var_attrs[DIMENSION_KEY][0]]
-            ]  # Dimension coordinates.
-            del var_attrs["_ARRAY_DIMENSIONS"]
-            xds = xds.assign_coords({var_name: coord})
-            xds[var_name].attrs = var_attrs
-        else:
-            # Construct slicing
-            slicing_list = []
-            for dim in var_attrs[DIMENSION_KEY]:
-                slicing_list.append(slice_dict_complete[dim])
-            slicing_tuple = tuple(slicing_list)
-            xds[var_name] = xr.DataArray(
-                var[slicing_tuple], dims=var_attrs[DIMENSION_KEY]
-            )
-
-            if "coordinates" in var_attrs:
-                del var_attrs["coordinates"]
-            del var_attrs["_ARRAY_DIMENSIONS"]
-            xds[var_name].attrs = var_attrs
-
-    xds.attrs = group_attrs
-
-    return xds
-
-
-def load_processing_set(ps_name, sel_parms):
-    """
-    sel_parms
-        A dictionary where the keys are the names of the ms_xds's and the values are slice_dicts.
-        slice_dicts: A dictionary where the keys are the dimension names and the values are slices.
-    """
-    ps = processing_set()
-    for name_ms_xds, ms_xds_sel_parms in sel_parms.items():
-        ps[name_ms_xds] = _load_ms_xds(ps_name, name_ms_xds, ms_xds_sel_parms)[0]
-    return ps
-
-class processing_set_iterator:
-
-    def __init__(self, data_selection, input_data_store, input_data=None):
-        self.input_data = input_data
-        self.input_data_store = input_data_store
-        self.data_selection = data_selection
-        self.xds_name_iter = iter(data_selection.keys())
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        try:
-            xds_name = next(self.xds_name_iter)
-        except Exception as e:
-            raise StopIteration
-        
-        if self.input_data is None:
-            slice_description = self.data_selection[xds_name]
-            ps = load_processing_set(
-                    ps_name=self.input_data_store,
-                    sel_parms={xds_name: slice_description},
-                )
-            xds = ps.get(0)
-        else:
-            xds = self.input_data[xds_name] #In memory
-
-        return xds

xradio-0.0.22/src/xradio/vis/read_processing_set.py

@@ -1,43 +0,0 @@
-import os
-
-import xarray as xr
-
-from ._processing_set import processing_set
-
-
-def read_processing_set(ps_name, intents=None, data_group='base', fields=None):
-    items = os.listdir(ps_name)
-    ms_xds = xr.Dataset()
-    ps = processing_set()
-    for i in items:
-        if "ddi" in i:
-            xds = xr.open_zarr(ps_name + "/" + i + "/MAIN")
-
-            if (intents is None) or (xds.attrs["intent"] in intents):
-                
-                if "visibility" in xds.attrs["data_groups"][data_group]:
-                    data_name = xds.attrs["data_groups"][data_group]["visibility"]
-
-                if "spectrum" in xds.attrs["data_groups"][data_group]:
-                    data_name = xds.attrs["data_groups"][data_group]["spectrum"]
-
-                if (fields is None) or (xds[data_name].attrs["field_info"]["name"] in fields):
-                    ps[i] = xds
-                    sub_xds = {
-                        "antenna_xds": "ANTENNA",
-                    }
-                    for sub_xds_key, sub_xds_name in sub_xds.items():
-                        ps[i].attrs[sub_xds_key] = xr.open_zarr(
-                            ps_name + "/" + i + "/" + sub_xds_name
-                        )
-
-                    optional_sub_xds = {
-                        "weather_xds": "WEATHER",
-                        "pointing_xds": "POINTING",
-                    }
-                    for sub_xds_key, sub_xds_name in optional_sub_xds.items():
-                        sub_xds_path = ps_name + "/" + i + "/" + sub_xds_name
-                        if os.path.isdir(sub_xds_path):
-                            ps[i].attrs[sub_xds_key] = xr.open_zarr(sub_xds_path)
-
-    return ps