xradio 0.0.34__py3-none-any.whl → 0.0.37__py3-none-any.whl

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/_utils/schema.py

@@ -0,0 +1,190 @@
+import graphviper.utils.logger as logger
+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(
+    casacore_column_description, ref_code=None, time_format="UNIX"
+):
+    import numpy as np
+
+    msv4_measure = {}
+    if "MEASINFO" in casacore_column_description["keywords"]:
+        measinfo = casacore_column_description["keywords"]["MEASINFO"]
+
+        # Get conversion information
+        msv4_measure_conversion = casacore_to_msv4_measure_type[measinfo["type"]]
+
+        # Convert type, copy unit
+        msv4_measure["type"] = msv4_measure_conversion["type"]
+        msv4_measure["units"] = list(
+            casacore_column_description["keywords"]["QuantumUnits"]
+        )
+
+        # Reference frame to convert?
+        if "Ref" in msv4_measure_conversion:
+            # Find reference frame
+            if "TabRefCodes" in measinfo:
+                ref_index = np.where(measinfo["TabRefCodes"] == ref_code)[0][0]
+                casa_ref = measinfo["TabRefTypes"][ref_index]
+            elif "Ref" in measinfo:
+                casa_ref = measinfo["Ref"]
+            elif measinfo["type"] == "frequency":
+                # Some MSv2 don't have the "TabRefCodes".
+                ref_index = np.where(casa_frequency_frames_codes == ref_code)[0][0]
+                casa_ref = casa_frequency_frames[ref_index]
+            else:
+                logger.debug(
+                    f"Could not determine {measinfo['type']} measure "
+                    "reference frame!"
+                )
+
+            # Convert into MSv4 representation of reference frame, warn if unknown
+            if casa_ref in msv4_measure_conversion.get("Ref_map", {}):
+                casa_ref = msv4_measure_conversion["Ref_map"][casa_ref]
+            else:
+                logger.debug(
+                    f"Unknown reference frame for {measinfo['type']} "
+                    f"measure, using verbatim: {casa_ref}"
+                )
+
+            msv4_measure[msv4_measure_conversion["Ref"]] = casa_ref
+
+        if msv4_measure["type"] == "time":
+            msv4_measure["format"] = time_format
+    elif "QuantumUnits" in casacore_column_description["keywords"]:
+        msv4_measure = {
+            "type": "quantity",
+            "units": list(casacore_column_description["keywords"]["QuantumUnits"]),
+        }
+
+    return msv4_measure
+
+
+# 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]

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/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/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/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/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/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?