xradio 0.0.33__tar.gz → 0.0.36__tar.gz

{xradio-0.0.33/src/xradio.egg-info → xradio-0.0.36}/PKG-INFO

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

{xradio-0.0.33 → xradio-0.0.36}/pyproject.toml

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

{xradio-0.0.33 → xradio-0.0.36}/src/xradio/_utils/list_and_array.py

@@ -5,10 +5,12 @@ import xarray as xr
 
 
 def to_list(x):
-    if isinstance(x, (list, np.ndarray)):
+    if isinstance(x, np.ndarray):
         if x.ndim == 0:
             return [x.item()]
         return list(x)  # needed for json serialization
+    elif isinstance(x, list):
+        return x
     return [x]
 
 

xradio-0.0.33/src/xradio/vis/_vis_utils/_ms/msv2_to_msv4_meta.py → xradio-0.0.36/src/xradio/_utils/schema.py

@@ -1,81 +1,87 @@
 import graphviper.utils.logger as logger
-
-col_to_data_variable_names = {
-    "FLOAT_DATA": "SPECTRUM",
-    "DATA": "VISIBILITY",
-    "CORRECTED_DATA": "VISIBILITY_CORRECTED",
-    "MODEL_DATA": "VISIBILITY_MODEL",
-    "WEIGHT_SPECTRUM": "WEIGHT",
-    "WEIGHT": "WEIGHT",
-    "FLAG": "FLAG",
-    "UVW": "UVW",
-    "TIME_CENTROID": "TIME_CENTROID",
-    "EXPOSURE": "EFFECTIVE_INTEGRATION_TIME",
-}
-col_dims = {
-    "DATA": ("time", "baseline_id", "frequency", "polarization"),
-    "CORRECTED_DATA": ("time", "baseline_id", "frequency", "polarization"),
-    "MODEL_DATA": ("time", "baseline_id", "frequency", "polarization"),
-    "WEIGHT_SPECTRUM": ("time", "baseline_id", "frequency", "polarization"),
-    "WEIGHT": ("time", "baseline_id", "frequency", "polarization"),
-    "FLAG": ("time", "baseline_id", "frequency", "polarization"),
-    "UVW": ("time", "baseline_id", "uvw_label"),
-    "TIME_CENTROID": ("time", "baseline_id"),
-    "EXPOSURE": ("time", "baseline_id"),
-    "FLOAT_DATA": ("time", "baseline_id", "frequency", "polarization"),
-}
-col_to_coord_names = {
-    "TIME": "time",
-    "ANTENNA1": "baseline_ant1_id",
-    "ANTENNA2": "baseline_ant2_id",
-}
-
-# Map casacore measures to astropy
-casacore_to_msv4_measure_type = {
-    "quanta": {
-        "type": "quantity",
-    },
-    "direction": {"type": "sky_coord", "Ref": "frame", "Ref_map": {"J2000": "fk5"}},
-    "epoch": {"type": "time", "Ref": "scale", "Ref_map": {"UTC": "utc"}},
-    "frequency": {
-        "type": "spectral_coord",
-        "Ref": "frame",
-        "Ref_map": {
-            "REST": "REST",
-            "LSRK": "LSRK",
-            "LSRD": "LSRD",
-            "BARY": "BARY",
-            "GEO": "GEO",
-            "TOPO": "TOPO",
-            "GALACTO": "GALACTO",
-            "LGROUP": "LGROUP",
-            "CMB": "CMB",
-            "Undefined": "Undefined",
-        },
-    },
-    "position": {
-        "type": "earth_location",
-        "Ref": "ellipsoid",
-        "Ref_map": {"ITRF": "GRS80"},
-    },
-    "uvw": {"type": "uvw", "Ref": "frame", "Ref_map": {"ITRF": "GRS80"}},
-    "radialvelocity": {"type": "quantity"},
-}
-
-casa_frequency_frames = [
-    "REST",
-    "LSRK",
-    "LSRD",
-    "BARY",
-    "GEO",
-    "TOPO",
-    "GALACTO",
-    "LGROUP",
-    "CMB",
-    "Undefined",
-]
-
-casa_frequency_frames_codes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 64]
+import xarray as xr
+
+
+def convert_generic_xds_to_xradio_schema(
+    generic_xds: xr.Dataset,
+    msv4_xds: xr.Dataset,
+    to_new_data_variables: dict,
+    to_new_coords: dict,
+) -> xr.Dataset:
+    """Converts a generic xarray Dataset to the xradio schema.
+
+    This function takes a generic xarray Dataset and converts it to an xradio schema
+    represented by the msv4_xds Dataset. It performs the conversion based on the provided
+    mappings in the to_new_data_variables and to_new_coords dictionaries.
+
+    Parameters
+    ----------
+    generic_xds : xr.Dataset
+        The generic xarray Dataset to be converted.
+    msv4_xds : xr.Dataset
+        The xradio schema represented by the msv4_xds Dataset.
+    to_new_data_variables : dict
+        A dictionary mapping the data variables/coordinates in the generic_xds Dataset to the new data variables
+        in the msv4_xds Dataset. The keys are the old data variables/coordinates and the values are a list of the new name and a list of the new dimension names.
+    to_new_coords : dict
+        A dictionary mapping  data variables/coordinates in the generic_xds Dataset to the new coordinates
+        in the msv4_xds Dataset. The keys are the old data variables/coordinates and the values are a list of the new name and a list of the new dimension names.
+
+    Returns
+    -------
+    xr.Dataset
+        The converted xradio schema represented by the msv4_xds Dataset.
+
+    Notes
+    -----
+    Example to_new_data_variables:
+    to_new_data_variables = {
+        "POSITION": ["ANTENNA_POSITION",["name", "cartesian_pos_label"]],
+        "OFFSET": ["ANTENNA_FEED_OFFSET",["name", "cartesian_pos_label"]],
+        "DISH_DIAMETER": ["ANTENNA_DISH_DIAMETER",["name"]],
+    }
+
+    Example to_new_coords:
+    to_new_coords = {
+        "NAME": ["name",["name"]],
+        "STATION": ["station",["name"]],
+        "MOUNT": ["mount",["name"]],
+        "PHASED_ARRAY_ID": ["phased_array_id",["name"]],
+        "antenna_id": ["antenna_id",["name"]],
+    }
+    """
+
+    column_description = generic_xds.attrs["other"]["msv2"]["ctds_attrs"][
+        "column_descriptions"
+    ]
+    coords = {}
+
+    name_keys = list(generic_xds.data_vars.keys()) + list(generic_xds.coords.keys())
+
+    for key in name_keys:
+
+        if key in column_description:
+            msv4_measure = column_description_casacore_to_msv4_measure(
+                column_description[key]
+            )
+        else:
+            msv4_measure = None
+
+        if key in to_new_data_variables:
+            new_dv = to_new_data_variables[key]
+            msv4_xds[new_dv[0]] = xr.DataArray(generic_xds[key].data, dims=new_dv[1])
+
+            if msv4_measure:
+                msv4_xds[new_dv[0]].attrs.update(msv4_measure)
+
+        if key in to_new_coords:
+            new_coord = to_new_coords[key]
+            coords[new_coord[0]] = (
+                new_coord[1],
+                generic_xds[key].data,
+            )
+    msv4_xds = msv4_xds.assign_coords(coords)
+    return msv4_xds
 
 
 def column_description_casacore_to_msv4_measure(
@@ -136,13 +142,49 @@ def column_description_casacore_to_msv4_measure(
     return msv4_measure
 
 
-def create_attribute_metadata(col, main_column_descriptions):
-    attrs_metadata = column_description_casacore_to_msv4_measure(
-        main_column_descriptions[col]
-    )
-    if col in ["DATA", "CORRECTED_DATA", "WEIGHT"]:
-        if not attrs_metadata:
-            attrs_metadata["type"] = "quanta"
-            attrs_metadata["units"] = ["unkown"]
+# Map casacore measures to astropy
+casacore_to_msv4_measure_type = {
+    "quanta": {
+        "type": "quantity",
+    },
+    "direction": {"type": "sky_coord", "Ref": "frame", "Ref_map": {"J2000": "fk5"}},
+    "epoch": {"type": "time", "Ref": "scale", "Ref_map": {"UTC": "utc"}},
+    "frequency": {
+        "type": "spectral_coord",
+        "Ref": "frame",
+        "Ref_map": {
+            "REST": "REST",
+            "LSRK": "LSRK",
+            "LSRD": "LSRD",
+            "BARY": "BARY",
+            "GEO": "GEO",
+            "TOPO": "TOPO",
+            "GALACTO": "GALACTO",
+            "LGROUP": "LGROUP",
+            "CMB": "CMB",
+            "Undefined": "Undefined",
+        },
+    },
+    "position": {
+        "type": "earth_location",
+        "Ref": "ellipsoid",
+        "Ref_map": {"ITRF": "GRS80"},
+    },
+    "uvw": {"type": "uvw", "Ref": "frame", "Ref_map": {"ITRF": "GRS80"}},
+    "radialvelocity": {"type": "quantity"},
+}
+
+casa_frequency_frames = [
+    "REST",
+    "LSRK",
+    "LSRD",
+    "BARY",
+    "GEO",
+    "TOPO",
+    "GALACTO",
+    "LGROUP",
+    "CMB",
+    "Undefined",
+]
 
-    return attrs_metadata
+casa_frequency_frames_codes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 64]

{xradio-0.0.33 → xradio-0.0.36}/src/xradio/_utils/zarr/common.py

@@ -5,11 +5,18 @@ import s3fs
 import os
 from botocore.exceptions import NoCredentialsError
 
+# from xradio.vis._vis_utils._ms.msv2_to_msv4_meta import (
+#     column_description_casacore_to_msv4_measure,
+# )
 
-def _get_ms_stores_and_file_system(ps_store: str):
 
+def _get_file_system_and_items(ps_store: str):
+
+    # default to assuming the data are accessible on local file system
     if os.path.isdir(ps_store):
-        # default to assuming the data are accessible on local file system
+        # handle a common shell convention
+        if ps_store.startswith("~"):
+            ps_store = os.path.expanduser(ps_store)
         items = os.listdir(ps_store)
         file_system = os
 
@@ -36,9 +43,8 @@ def _get_ms_stores_and_file_system(ps_store: str):
                 for bd in file_system.listdir(ps_store, detail=False)
             ]
     else:
-        raise (
-            FileNotFoundError,
-            f"Could not find {ps_store} either locally or in the cloud.",
+        raise FileNotFoundError(
+            f"Could not find {ps_store} either locally or in the cloud."
         )
 
     items = [

{xradio-0.0.33 → xradio-0.0.36}/src/xradio/image/_util/_zarr/xds_from_zarr.py

@@ -5,8 +5,10 @@ import dask.array as da
 import numpy as np
 import os
 import xarray as xr
+import s3fs
 from .common import _np_types, _top_level_sub_xds
 from ..common import _coords_to_numpy, _dask_arrayize_dv, _numpy_arrayize_dv
+from xradio._utils.zarr.common import _get_file_system_and_items
 
 
 def _read_zarr(
@@ -85,8 +87,18 @@ def _decode_dict(my_dict: dict, top_key: str) -> dict:
 
 def _decode_sub_xdses(zarr_store: str, output: dict) -> dict:
     sub_xdses = {}
-    for root, dirs, files in os.walk(zarr_store):
-        # top down walk
+
+    fs, store_contents = _get_file_system_and_items(zarr_store)
+
+    if isinstance(fs, s3fs.core.S3FileSystem):
+        # could we just use the items as returned from the helper function..?
+        store_tree = fs.walk(zarr_store, topdown=True)
+        prepend_s3 = "s3://"
+    else:
+        store_tree = os.walk(zarr_store, topdown=True)
+        prepend_s3 = ""
+
+    for root, dirs, files in store_tree:
         for d in dirs:
             if d.startswith(_top_level_sub_xds):
                 xds = _read_zarr(os.sep.join([root, d]), output)
@@ -94,4 +106,5 @@ def _decode_sub_xdses(zarr_store: str, output: dict) -> dict:
                 #    xds = xds.drop_vars([k]).assign({k: v.compute()})
                 ky = d[len(_top_level_sub_xds) + 1 :]
                 sub_xdses[ky] = xds
+
     return sub_xdses

{xradio-0.0.33 → xradio-0.0.36}/src/xradio/image/_util/_zarr/zarr_low_level.py

@@ -1,6 +1,9 @@
 import os
 import numpy as np
 import json
+import zarr
+import s3fs
+from xradio._utils.zarr.common import _get_file_system_and_items
 
 from numcodecs.compat import (
     ensure_text,
@@ -246,19 +249,29 @@ def write_json_file(data, file_path):
         )
 
 
-def create_data_variable_meta_data_on_disk(
+def create_data_variable_meta_data(
     zarr_group_name, data_variables_and_dims, xds_dims, parallel_coords, compressor
 ):
     zarr_meta = data_variables_and_dims
 
+    fs, items = _get_file_system_and_items(zarr_group_name)
+
     for data_variable_key, dims_dtype_name in data_variables_and_dims.items():
         # print(data_variable_key, dims_dtype_name)
 
         dims = dims_dtype_name["dims"]
         dtype = dims_dtype_name["dtype"]
         data_variable_name = dims_dtype_name["name"]
+
         data_variable_path = os.path.join(zarr_group_name, data_variable_name)
-        os.system("mkdir " + data_variable_path)
+        if isinstance(fs, s3fs.core.S3FileSystem):
+            # N.b.,stateful "folder creation" is not a well defined concept for S3 objects and URIs
+            # see https://github.com/fsspec/s3fs/issues/401
+            # nor is a path specifier (cf. "URI")
+            fs.mkdir(data_variable_path)
+        else:
+            # default to assuming we can use the os module and mkdir system call
+            os.system("mkdir " + data_variable_path)
         # Create .zattrs
         zattrs = {
             "_ARRAY_DIMENSIONS": dims,
@@ -275,7 +288,23 @@ def create_data_variable_meta_data_on_disk(
                 chunks.append(xds_dims[d])
 
         # print(chunks,shape)
-        write_json_file(zattrs, os.path.join(data_variable_path, ".zattrs"))
+        # assuming data_variable_path has been set compatibly
+        zattrs_file = os.path.join(data_variable_path, ".zattrs")
+
+        if isinstance(fs, s3fs.core.S3FileSystem):
+            with fs.open(zattrs_file, "w") as file:
+                json.dump(
+                    zattrs,
+                    file,
+                    indent=4,
+                    sort_keys=True,
+                    ensure_ascii=True,
+                    separators=(",", ": "),
+                    cls=NumberEncoder,
+                )
+        else:
+            # default to assuming we can use primitives
+            write_json_file(zattrs, zattrs_file)
 
         # Create .zarray
         from zarr import n5
@@ -311,36 +340,58 @@ def create_data_variable_meta_data_on_disk(
         zarr_meta[data_variable_key]["chunks"] = chunks
         zarr_meta[data_variable_key]["shape"] = shape
 
-        write_json_file(zarray, os.path.join(data_variable_path, ".zarray"))
+        # again, assuming data_variable_path has been set compatibly
+        zarray_file = os.path.join(data_variable_path, ".zarray")
+
+        if isinstance(fs, s3fs.core.S3FileSystem):
+            with fs.open(zarray_file, "w") as file:
+                json.dump(
+                    zarray,
+                    file,
+                    indent=4,
+                    sort_keys=True,
+                    ensure_ascii=True,
+                    separators=(",", ": "),
+                    cls=NumberEncoder,
+                )
+        else:
+            # default to assuming we can use primitives
+            write_json_file(zarray, zarray_file)
+
     return zarr_meta
 
 
 def write_chunk(img_xds, meta, parallel_dims_chunk_id, compressor, image_file):
     dims = meta["dims"]
     dtype = meta["dtype"]
-    data_varaible_name = meta["name"]
+    data_variable_name = meta["name"]
     chunks = meta["chunks"]
     shape = meta["shape"]
     chunk_name = ""
-    if data_varaible_name in img_xds:
-        for d in img_xds[data_varaible_name].dims:
+    if data_variable_name in img_xds:
+        for d in img_xds[data_variable_name].dims:
             if d in parallel_dims_chunk_id:
                 chunk_name = chunk_name + str(parallel_dims_chunk_id[d]) + "."
             else:
                 chunk_name = chunk_name + "0."
         chunk_name = chunk_name[:-1]
 
-        if list(img_xds[data_varaible_name].shape) != list(chunks):
+        if list(img_xds[data_variable_name].shape) != list(chunks):
             array = pad_array_with_nans(
-                img_xds[data_varaible_name].values,
+                img_xds[data_variable_name].values,
                 output_shape=chunks,
                 dtype=dtype,
             )
         else:
-            array = img_xds[data_varaible_name].values
-
-        write_binary_blob_to_disk(
-            array,
-            file_path=os.path.join(image_file, data_varaible_name, chunk_name),
+            array = img_xds[data_variable_name].values
+
+        z_chunk = zarr.open(
+            os.path.join(image_file, data_variable_name, chunk_name),
+            mode="a",
+            shape=meta["shape"],
+            chunks=meta["chunks"],
+            dtype=meta["dtype"],
             compressor=compressor,
         )
+
+        return z_chunk

{xradio-0.0.33 → xradio-0.0.36}/src/xradio/schema/bases.py

@@ -284,6 +284,22 @@ def _dataset_new(cls, *args, data_vars=None, coords=None, attrs=None, **kwargs):
         val = _np_convert(
             _set_parameter(coords.get(coord.name), mapping.arguments, coord), coord
         )
+        # Determine dimensions / convert to Variable
+        if (
+            val is not None
+            and not isinstance(val, xarray.DataArray)
+            and not isinstance(val, xarray.Variable)
+            and not isinstance(val, tuple)
+        ):
+            default_attrs = {
+                attr.name: attr.default
+                for attr in coord.attributes
+                if attr.default is not None
+            }
+            for dims in coord.dimensions:
+                if len(dims) == len(val.shape):
+                    val = xarray.Variable(dims, val, default_attrs)
+                    break
         if val is not None:
             coords[coord.name] = val
     for data_var in schema.data_vars:
@@ -291,7 +307,9 @@ def _dataset_new(cls, *args, data_vars=None, coords=None, attrs=None, **kwargs):
 
         # Determine dimensions / convert to Variable
         dims = None
-        if isinstance(val, xarray.Variable):
+        if val is None:
+            dims = None
+        elif isinstance(val, xarray.Variable):
             dims = val.dims
         elif isinstance(val, xarray.DataArray):
             val = val.variable
@@ -316,18 +334,29 @@ def _dataset_new(cls, *args, data_vars=None, coords=None, attrs=None, **kwargs):
                     f" expected {' or '.join(options)}!"
                 )
 
+            # Get default attributes
+            default_attrs = {
+                attr.name: attr.default
+                for attr in data_var.attributes
+                if attr.default is not dataclasses.MISSING
+            }
+
             # Replace by variable
-            val = xarray.Variable(dims, val)
+            val = xarray.Variable(dims, val, default_attrs)
 
         # Default coordinates used by this data variable to numpy arange. We
         # can only do this now because we need an example to determine the
         # intended size of the coordinate
-        for coord in schema.coordinates:
-            if coord.name in dims and coords.get(coord.name) is None:
-                dim_ix = dims.index(coord.name)
-                if dim_ix is not None and dim_ix < len(val.shape):
-                    dtype = coord.dtypes[0]
-                    coords[coord.name] = numpy.arange(val.shape[dim_ix], dtype=dtype)
+        if dims is not None:
+            for coord in schema.coordinates:
+                if coord.name in dims and coords.get(coord.name) is None:
+                    dim_ix = dims.index(coord.name)
+                    if dim_ix is not None and dim_ix < len(val.shape):
+                        dtype = coord.dtypes[0]
+                        if numpy.issubdtype(dtype, numpy.number):
+                            coords[coord.name] = numpy.arange(
+                                val.shape[dim_ix], dtype=dtype
+                            )
 
         if val is not None:
             data_vars[data_var.name] = val

{xradio-0.0.33 → xradio-0.0.36}/src/xradio/schema/check.py

@@ -58,7 +58,10 @@ class SchemaIssue:
         err = f"Schema issue with {self.path_str()}: {self.message}"
         if self.expected is not None:
             options = " or ".join(repr(option) for option in self.expected)
-            err += f" (expected: {options} found: {repr(self.found)})"
+            if self.found is not None:
+                err += f" (expected: {options} found: {repr(self.found)})"
+            else:
+                err += f" (expected: {options})"
         return err
 
 
@@ -255,7 +258,7 @@ def check_dimensions(
     if hint_remove and hint_add:
         message = f"Unexpected coordinates, replace {','.join(hint_remove)} by {','.join(hint_add)}?"
     elif hint_remove:
-        message = f"Superflous coordinate {','.join(hint_remove)}?"
+        message = f"Superfluous coordinate {','.join(hint_remove)}?"
     elif hint_add:
         message = f"Missing dimension {','.join(hint_add)}!"
     else:
@@ -381,7 +384,8 @@ def check_data_vars(
                     )
                 else:
                     message = (
-                        f"Required data variable '{data_var_schema.name}' is missing!"
+                        f"Required data variable '{data_var_schema.name}' is missing "
+                        f"(have {','.join(data_vars)})!"
                     )
                 issues.add(
                     SchemaIssue(
@@ -440,6 +444,14 @@ def _check_value(val, ann):
                         )
                     ]
                 )
+            except TypeError as e:
+                return SchemaIssues(
+                    [
+                        SchemaIssue(
+                            path=[], message=str(e), expected=[ann], found=type(val)
+                        )
+                    ]
+                )
 
         if not isinstance(val, xarray.DataArray):
             # Fall through to plain type check

{xradio-0.0.33 → xradio-0.0.36}/src/xradio/schema/dataclass.py

@@ -108,8 +108,8 @@ def extract_xarray_dataclass(klass):
             is_coord = False
         else:
             raise ValueError(
-                f"Unexpected role in '{klass.__name__}',"
-                f" field '{field.name}': {get_role(typ)}"
+                f"Expected field '{field.name}' in '{klass.__name__}' "
+                "to be annotated with either Coord, Data or Attr!"
             )
 
         # Defined using a dataclass, i.e. Coordof/Dataof?