xradio 0.0.23__tar.gz → 0.0.24__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.23
+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.23 → xradio-0.0.24}/pyproject.toml

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

{xradio-0.0.23 → xradio-0.0.24}/src/xradio/_utils/zarr/common.py

@@ -3,6 +3,47 @@ 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)
@@ -33,6 +74,7 @@ def _load_no_dask_zarr(zarr_name, 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]:
@@ -54,6 +96,9 @@ def _load_no_dask_zarr(zarr_name, slice_dict={}):
             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]
             )

{xradio-0.0.23 → 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.
@@ -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.23 → xradio-0.0.24}/src/xradio/image/_util/zarr.py

@@ -3,7 +3,7 @@ from ._zarr.xds_from_zarr import _read_zarr
 import numpy as np
 import os
 import xarray as xr
-from ..._utils.zarr.common import _load_no_dask_zarr
+from ..._utils.zarr.common import _open_dataset
 
 
 def _xds_to_zarr(xds: xr.Dataset, zarr_store: str):
@@ -25,11 +25,11 @@ def _xds_from_zarr(
 
 
 def _load_image_from_zarr_no_dask(zarr_file: str, selection: dict) -> xr.Dataset:
-    image_xds = _load_no_dask_zarr(zarr_file, selection)
+    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"] = _load_no_dask_zarr(history)
+            image_xds.attrs["history"] = _open_dataset(history, load=True)
             break
     _iter_dict(image_xds.attrs)
     return image_xds

{xradio-0.0.23 → xradio-0.0.24}/src/xradio/vis/_processing_set.py

@@ -4,7 +4,8 @@ import pandas as pd
 class processing_set(dict):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.meta = {'summary':{}}
+        self.meta = {"summary": {}}
+
     #     generate_meta(self)
 
     # def generate_meta(self):
@@ -12,18 +13,25 @@ class processing_set(dict):
     #     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]
+        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]
-        
+            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']
+        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['max_dims'] = self._get_ps_max_dims()
-            return self.meta['max_dims'] 
+            self.meta["freq_axis"] = self._get_ps_freq_axis()
+            return self.meta["freq_axis"]
 
     def _summary(self, data_group="base"):
         summary_data = {
@@ -34,7 +42,7 @@ class processing_set(dict):
             "field_name": [],
             "start_frequency": [],
             "end_frequency": [],
-            "shape": []
+            "shape": [],
         }
         for key, value in self.items():
             summary_data["name"].append(key)
@@ -47,25 +55,36 @@ class processing_set(dict):
             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["shape"].append(value[data_name].shape)
 
             summary_data["field_id"].append(
-                value[data_name].attrs[
-                    "field_info"
-                ]["field_id"]
+                value[data_name].attrs["field_info"]["field_id"]
             )
             summary_data["field_name"].append(
-                value[data_name].attrs[
-                    "field_info"
-                ]["name"]
+                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():
@@ -79,6 +98,6 @@ class processing_set(dict):
                     else:
                         max_dims[dim_name] = size
         return max_dims
-    
+
     def get(self, id):
-        return self[list(self.keys())[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.23 → xradio-0.0.24/src/xradio.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.23
+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.23/src/xradio/vis/load_processing_set.py

@@ -1,127 +0,0 @@
-import xarray as xr
-import zarr
-import copy
-import os
-from ._processing_set import processing_set
-from .._utils.zarr.common import _load_no_dask_zarr
-
-# from xradio._utils._logger import _get_logger
-
-
-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_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.23/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