xradio 0.0.23__tar.gz → 0.0.25__tar.gz

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.23
-Summary: Xarray Radio Astronomy Data IO 
+Version: 0.0.25
+Summary: Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License
         
@@ -32,7 +32,7 @@ License: BSD 3-Clause License
         OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
         OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
         
-Requires-Python: <3.12,>=3.8
+Requires-Python: <3.12,>=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
 Requires-Dist: astropy

{xradio-0.0.23 → xradio-0.0.25}/pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "xradio"
-version = "0.0.23"
-description = "Xarray Radio Astronomy Data IO "
+version = "0.0.25"
+description = "Xarray Radio Astronomy Data IO"
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},
 ]
 license = {file = "LICENSE.txt"}
 readme = "README.md"
-requires-python = ">= 3.8, < 3.12"
+requires-python = ">= 3.9, < 3.12"
 
 dependencies = [
     'astropy',

{xradio-0.0.23 → xradio-0.0.25}/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.25}/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.
@@ -110,19 +109,55 @@ def write_binary_blob_to_disk(arr, file_path, compressor):
     Returns:
     - None
     """
+    import graphviper.utils.logger as logger
     # Encode the NumPy array using the codec
+    logger.debug('1. Before compressor ' + file_path)
     compressed_arr = compressor.encode(np.ascontiguousarray(arr))
 
+    logger.debug('2. Before makedir')
     # Ensure the directory exists before saving the file
     os.makedirs(os.path.dirname(file_path), exist_ok=True)
 
-    # Save the compressed array to disk
-    with open(file_path, "wb") as file:
-        file.write(compressed_arr)
+    arr_len = len(compressed_arr)
+    logger.debug('3. Before write the len is: ' + str(arr_len))
+    #Save the compressed array to disk
+    # with open(file_path, "wb") as file:
+    #     file.write(compressed_arr)
+    
+    logger.debug('4. Using new writer: ' + str(arr_len))
+    write_to_lustre_chunked(file_path, compressed_arr)
+        
+    # /.lustre/aoc/sciops/pford/CHILES/cube_image/uid___A002_Xee7674_X2844_Cube_3.img.zarr/SKY/0.0.110.0.0
+    # 348192501 bytes
+    # 332.0622453689575 M
+    
+    # from io import BufferedWriter
+    # # Calculate buffer size based on compressed_arr size (adjust multiplier)
+    # buffer_size = min(len(compressed_arr), 1024 * 1024 * 4)  # Max 4 MB buffer
+    # with BufferedWriter(open(file_path, "wb"), buffer_size) as f:
+    #     f.write(compressed_arr)
+    #     f.flush()  # Ensure data gets written to disk 
+    
+    
+    logger.debug('4. Write completed')
 
     # print(f"Compressed array saved to {file_path}")
 
 
+def write_to_lustre_chunked(file_path, compressed_arr, chunk_size=1024 * 1024 * 128):  # 128 MiB chunks
+  """
+  Writes compressed data to a Lustre file path with chunking.
+
+  Args:
+      file_path: Path to the file for writing.
+      compressed_arr: Compressed data array to write.
+      chunk_size: Size of each data chunk in bytes (default: 128 MiB).
+  """
+  with open(file_path, "wb") as f:
+    for i in range(0, len(compressed_arr), chunk_size):
+      chunk = compressed_arr[i:i + chunk_size]
+      f.write(chunk)
+
 def read_binary_blob_from_disk(file_path, compressor, dtype=np.float64):
     """
     Read a compressed binary blob from disk and decode it using Blosc.
@@ -277,8 +312,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 +338,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.25}/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.25/src/xradio/vis/_processing_set.py

@@ -0,0 +1,112 @@
+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": [],
+            "field_coords": []
+        }
+        from astropy.coordinates import SkyCoord
+        import astropy.units as u
+        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])
+
+            ra_dec_rad = value[data_name].attrs["field_info"]['phase_direction']['data']
+            frame = value[data_name].attrs["field_info"]['phase_direction']['attrs']['frame'].lower()
+            coord = SkyCoord(ra=ra_dec_rad[0]*u.rad, dec=ra_dec_rad[1]*u.rad, frame=frame)
+
+            summary_data["field_coords"].append([frame, coord.ra.to_string(unit=u.hour), coord.dec.to_string(unit=u.deg)])
+        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.25/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.25/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.25/src/xradio.egg-info}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.23
-Summary: Xarray Radio Astronomy Data IO 
+Version: 0.0.25
+Summary: Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License
         
@@ -32,7 +32,7 @@ License: BSD 3-Clause License
         OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
         OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
         
-Requires-Python: <3.12,>=3.8
+Requires-Python: <3.12,>=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
 Requires-Dist: astropy

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

@@ -1,84 +0,0 @@
-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 _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_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.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