xradio 0.0.25__tar.gz → 0.0.27__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.25
+Version: 0.0.27
 Summary: Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License
@@ -47,6 +47,7 @@ Requires-Dist: prettytable
 Requires-Dist: pytest
 Requires-Dist: pytest-cov
 Requires-Dist: pytest-html
+Requires-Dist: s3fs
 Requires-Dist: scipy
 Requires-Dist: tqdm
 Requires-Dist: xarray

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

@@ -1,6 +1,6 @@
 [project]
 name = "xradio"
-version = "0.0.25"
+version = "0.0.27"
 description = "Xarray Radio Astronomy Data IO"
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},
@@ -22,6 +22,7 @@ dependencies = [
     'pytest',
     'pytest-cov',
     'pytest-html',
+    's3fs',
     'scipy',
     'tqdm',
     'xarray',

{xradio-0.0.25 → xradio-0.0.27}/src/xradio/_utils/zarr/common.py

@@ -1,9 +1,10 @@
 import copy
 import xarray as xr
 import zarr
+import s3fs
 
 
-def _open_dataset(store, xds_isel=None, data_variables=None, load=False):
+def _open_dataset(store, xds_isel=None, data_variables=None, load=False, **kwargs):
     """
 
     Parameters
@@ -25,7 +26,11 @@ def _open_dataset(store, xds_isel=None, data_variables=None, load=False):
 
     import dask
 
-    xds = xr.open_zarr(store)
+    if "s3" in kwargs.keys():
+        mapping = s3fs.S3Map(root=store, s3=kwargs["s3"], check=False)
+        xds = xr.open_zarr(store=mapping)
+    else:
+        xds = xr.open_zarr(store)
 
     if xds_isel is not None:
         xds = xds.isel(xds_isel)

{xradio-0.0.25 → xradio-0.0.27}/src/xradio/image/_util/_casacore/xds_from_casacore.py

@@ -106,7 +106,7 @@ def _add_sky_or_aperture(
         unit = casa_image.unit()
     xda.attrs[_image_type] = image_type
     xda.attrs["units"] = unit
-    name = "sky" if has_sph_dims else "aperture"
+    name = "SKY" if has_sph_dims else "APERTURE"
     xda = xda.rename(name)
     xds[xda.name] = xda
     return xds

{xradio-0.0.25 → xradio-0.0.27}/src/xradio/image/_util/_casacore/xds_to_casacore.py

@@ -8,7 +8,7 @@ from astropy.coordinates import Angle
 from casacore import tables
 
 from .common import _active_mask, _create_new_image, _object_name, _pointing_center
-from ..common import _compute_sky_reference_pixel, _doppler_types
+from ..common import _aperture_or_sky, _compute_sky_reference_pixel, _doppler_types
 from ...._utils._casacore.tables import open_table_rw
 
 
@@ -168,7 +168,7 @@ def _history_from_xds(xds: xr.Dataset, image: str) -> None:
 
 def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
     ii = {}
-    ap_sky = "sky" if "l" in xds.coords else "aperture"
+    ap_sky = _aperture_or_sky(xds)
     ii["image_type"] = (
         xds[ap_sky].attrs["image_type"] if "image_type" in xds[ap_sky].attrs else ""
     )
@@ -208,12 +208,12 @@ def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
 
 
 def _write_casa_data(xds: xr.Dataset, image_full_path: str) -> None:
-    sky_ap = "sky" if "sky" in xds else "aperture"
+    sky_ap = _aperture_or_sky(xds)
     if xds[sky_ap].shape[0] != 1:
         raise RuntimeError("XDS can only be converted if it has exactly one time plane")
     trans_coords = (
         ("frequency", "polarization", "m", "l")
-        if sky_ap == "sky"
+        if sky_ap == "SKY"
         else ("frequency", "polarization", "v", "u")
     )
     casa_image_shape = xds[sky_ap].isel(time=0).transpose(*trans_coords).shape[::-1]
@@ -312,7 +312,7 @@ def _write_initial_image(
 ) -> None:
     if not maskname:
         maskname = ""
-    for dv in ["sky", "aperture"]:
+    for dv in ["SKY", "APERTURE"]:
         if dv in xds.data_vars:
             value = xds[dv][0, 0, 0, 0, 0].values.item()
             if xds[dv][0, 0, 0, 0, 0].values.dtype == "float32":
@@ -352,7 +352,7 @@ def _write_pixels(
     value: xr.DataArray = None,
 ) -> None:
     flip = False
-    if v == "sky" or v == "aperture":
+    if v == "SKY" or v == "APERTURE":
         filename = image_full_path
     else:
         # mask

{xradio-0.0.25 → xradio-0.0.27}/src/xradio/image/_util/_fits/xds_from_fits.py

@@ -685,7 +685,7 @@ def _add_sky_or_aperture(
     unit = helpers["bunit"]
     xda.attrs[_image_type] = image_type
     xda.attrs["units"] = unit
-    name = "sky" if has_sph_dims else "aperture"
+    name = "SKY" if has_sph_dims else "APERTURE"
     xda = xda.rename(name)
     xds[xda.name] = xda
     if helpers["has_mask"]:

{xradio-0.0.25 → xradio-0.0.27}/src/xradio/image/_util/casacore.py

@@ -30,7 +30,7 @@ from ._casacore.xds_to_casacore import (
     _imageinfo_dict_from_xds,
     _write_casa_data,
 )
-from .common import _get_xds_dim_order, _dask_arrayize_dv
+from .common import _aperture_or_sky, _get_xds_dim_order, _dask_arrayize_dv
 
 warnings.filterwarnings("ignore", category=FutureWarning)
 
@@ -104,7 +104,7 @@ def _xds_to_casa_image(xds: xr.Dataset, imagename: str) -> None:
     image_full_path = os.path.expanduser(imagename)
     _write_casa_data(xds, image_full_path)
     # create coordinates
-    ap_sky = "sky" if "sky" in xds.data_vars else "aperture"
+    ap_sky = _aperture_or_sky(xds)
     coord = _coord_dict_from_xds(xds)
     ii = _imageinfo_dict_from_xds(xds)
     units = xds[ap_sky].attrs["units"] if "units" in xds[ap_sky].attrs else None

{xradio-0.0.25 → xradio-0.0.27}/src/xradio/image/_util/common.py

@@ -13,6 +13,10 @@ _doppler_types = ["RADIO", "Z", "RATIO", "BETA", "GAMMA"]
 _image_type = "image_type"
 
 
+def _aperture_or_sky(xds: xr.Dataset) -> str:
+    return "SKY" if "SKY" in xds.data_vars or "l" in xds.coords else "APERTURE"
+
+
 def _get_xds_dim_order(has_sph: bool) -> list:
     dimorder = ["time", "polarization", "frequency"]
     dir_lin = ["l", "m"] if has_sph else ["u", "v"]

{xradio-0.0.25 → xradio-0.0.27}/src/xradio/vis/_vis_utils/_ms/_tables/read_main_table.py

@@ -195,10 +195,8 @@ def read_main_table_chunks(
     n_baseline_chunks = chunks[1]
     # loop over time chunks
     for time_chunk in range(0, n_unique_times, n_time_chunks):
-        time_start = (unique_times[time_chunk] - tol,)
-        time_end = (
-            unique_times[min(n_unique_times, time_chunk + n_time_chunks) - 1] + tol
-        )
+        time_start = unique_times[time_chunk] - tol
+        time_end = unique_times[min(n_unique_times, time_chunk + n_time_chunks) - 1] + tol
 
         # chunk time length
         ctlen = min(n_unique_times, time_chunk + n_time_chunks) - time_chunk

{xradio-0.0.25 → xradio-0.0.27}/src/xradio/vis/load_processing_set.py

@@ -11,13 +11,13 @@ def load_processing_set(
     sel_parms: dict,
     data_variables: Union[list, None] = None,
     load_sub_datasets: bool = True,
-)->processing_set:
+) -> 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'.
+        String of the path and name of the processing set. For example '/users/user_1/uid___A002_Xf07bba_Xbe5c_target.lsrk.vis.zarr' for a file stored on a local file system, or 's3://viper-test-data/Antennae_North.cal.lsrk.split.vis.zarr/' for a file in AWS object storage.
     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.
@@ -37,26 +37,69 @@ def load_processing_set(
     Returns
     -------
     processing_set
-        In memory representation of processing set (data is represented by Dask.arrays). 
-    """    
+        In memory representation of processing set (data is represented by Dask.arrays).
+    """
     from xradio._utils.zarr.common import _open_dataset
+    import s3fs
+    from botocore.exceptions import NoCredentialsError
 
+    s3 = None
     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),
-            }
+
+        # before the _open_dataset call, check if dealing with an S3 bucket URL
+        if ps_store.startswith("s3"):
+            if not ps_store.endswith("/"):
+                # just for consistency, as there is no os.path equivalent in s3fs
+                ps_store = ps_store + "/"
+
+            try:
+                s3 = s3fs.S3FileSystem(anon=False, requester_pays=False)
+
+                main_xds = ps_store + ms_dir_name + "/MAIN"
+                xds = _open_dataset(
+                    main_xds, ms_xds_isel, data_variables, load=True, s3=s3
+                )
+
+                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, s3=s3
+                        ),
+                    }
+
+            except (NoCredentialsError, PermissionError) as e:
+                # only public, read-only buckets will be accessible
+                s3 = s3fs.S3FileSystem(anon=True)
+
+                main_xds = ps_store + ms_dir_name + "/MAIN"
+                xds = _open_dataset(
+                    main_xds, ms_xds_isel, data_variables, load=True, s3=s3
+                )
+
+                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, s3=s3
+                        ),
+                    }
+        else:
+            # fall back to the default case of assuming the files are on local disk
+            main_xds = os.path.join(ps_store, ms_dir_name, "MAIN")
+            xds = _open_dataset(main_xds, 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
@@ -94,7 +137,7 @@ class processing_set_iterator:
             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
@@ -124,4 +167,4 @@ class processing_set_iterator:
         else:
             xds = self.input_data[xds_name]  # In memory
 
-        return xds
+        return xds

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

@@ -0,0 +1,150 @@
+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
+import s3fs
+from botocore.exceptions import NoCredentialsError
+
+
+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).
+    """
+    s3 = None
+    ps_store_is_s3dir = None
+
+    if os.path.isdir(ps_store):
+        ps_store_is_s3dir = False
+        # default to assuming the data are accessible on local file system
+        items = os.listdir(ps_store)
+
+    elif ps_store.startswith("s3"):
+        # only if not found locally, check if dealing with an S3 bucket URL
+        ps_store_is_s3dir = True
+        if not ps_store.endswith("/"):
+            # just for consistency, as there is no os.path equivalent in s3fs
+            ps_store = ps_store + "/"
+
+        try:
+            # initialize the S3 "file system", first attempting to use pre-configured credentials
+            s3 = s3fs.S3FileSystem(anon=False, requester_pays=False)
+
+            items = [bd.split(sep="/")[-1] for bd in s3.listdir(ps_store, detail=False)]
+
+        except (NoCredentialsError, PermissionError) as e:
+            # only public, read-only buckets will be accessible
+            # we will want to add messaging and error handling here
+            s3 = s3fs.S3FileSystem(anon=True)
+
+            items = [bd.split(sep="/")[-1] for bd in s3.listdir(ps_store, detail=False)]
+
+    else:
+        raise (
+            FileNotFoundError,
+            f"Could not find {ps_store} either locally or in the cloud.",
+        )
+
+    ms_xds = xr.Dataset()
+    ps = processing_set()
+    data_group = "base"
+    for ms_dir_name in items:
+        if "ddi" in ms_dir_name:
+            if ps_store_is_s3dir:
+                store_path = ps_store + ms_dir_name
+                store_path_main = store_path + "/MAIN"
+            else:
+                store_path_main = os.path.join(ps_store, ms_dir_name, "MAIN")
+                store_path = os.path.split(store_path_main)[0]
+            if s3 is not None:
+                xds = _open_dataset(store_path_main, s3=s3)
+            else:
+                xds = _open_dataset(store_path_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
+                ):
+                    if s3 is not None:
+                        xds.attrs = {
+                            **xds.attrs,
+                            **_read_sub_xds(store_path, s3=s3),
+                        }
+                        ps[ms_dir_name] = xds
+                    else:
+                        xds.attrs = {
+                            **xds.attrs,
+                            **_read_sub_xds(store_path),
+                        }
+                        ps[ms_dir_name] = xds
+
+    return ps
+
+
+def _read_sub_xds(ms_store, load=False, **kwargs):
+    sub_xds_dict = {}
+
+    sub_xds = {
+        "antenna_xds": "ANTENNA",
+    }
+    for sub_xds_key, sub_xds_name in sub_xds.items():
+        if "s3" in kwargs.keys():
+            joined_store = ms_store + "/" + sub_xds_name
+            sub_xds_dict[sub_xds_key] = _open_dataset(
+                joined_store, load=load, s3=kwargs["s3"]
+            )
+        else:
+            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)
+        elif "s3" in kwargs.keys():
+            joined_store = ms_store + "/" + sub_xds_name
+            if kwargs["s3"].isdir(joined_store):
+                sub_xds_dict[sub_xds_key] = _open_dataset(
+                    joined_store, load=load, s3=kwargs["s3"]
+                )
+
+    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.25 → xradio-0.0.27/src/xradio.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.25
+Version: 0.0.27
 Summary: Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License
@@ -47,6 +47,7 @@ Requires-Dist: prettytable
 Requires-Dist: pytest
 Requires-Dist: pytest-cov
 Requires-Dist: pytest-html
+Requires-Dist: s3fs
 Requires-Dist: scipy
 Requires-Dist: tqdm
 Requires-Dist: xarray

{xradio-0.0.25 → xradio-0.0.27}/src/xradio.egg-info/requires.txt

@@ -10,6 +10,7 @@ prettytable
 pytest
 pytest-cov
 pytest-html
+s3fs
 scipy
 tqdm
 xarray

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

@@ -1,85 +0,0 @@
-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