xradio 1.0.2__py3-none-any.whl → 1.1.0__py3-none-any.whl

xradio/_utils/_casacore/casacore_from_casatools.py

@@ -513,7 +513,7 @@ class image(casatools.image):
         self,
         imagename,
         axis=0,
-        maskname="mask0",
+        maskname="mask_0",
         images=(),
         values=None,
         coordsys=None,

xradio/_utils/dict_helpers.py

@@ -1,3 +1,6 @@
+from xradio._utils.list_and_array import to_python_type
+
+
 def make_quantity(value, units: str, dims: list = []) -> dict:
     """
     create a quantity dictionary given value and units
@@ -13,7 +16,11 @@ def make_quantity(value, units: str, dims: list = []) -> dict:
     -------
     dict
     """
-    return {"data": value, "dims": dims, "attrs": make_quantity_attrs(units)}
+    return {
+        "data": to_python_type(value),
+        "dims": dims,
+        "attrs": make_quantity_attrs(units),
+    }
 
 
 def ensure_units_are_consistent(units):
@@ -67,7 +74,7 @@ def make_spectral_coord_reference_dict(
             u,
             observer.lower() if observer not in ["TOPO", "BARY", "REST"] else observer,
         ),
-        "data": value,
+        "data": to_python_type(value),
         "dims": [],
     }
 
@@ -98,14 +105,38 @@ def make_skycoord_dict(data: list[float], units: str, frame: str) -> dict:
             "type": "sky_coord",
             "units": ensure_units_are_consistent(units),
         },
-        "data": data,
-        "dims": ["l", "m"],
+        "data": to_python_type(data),
+        "dims": "sky_dir_label",
+        "coords": {"sky_dir_label": {"data": ["ra", "dec"], "dims": "sky_dir_label"}},
+    }
+
+
+def make_direction_location_dict(data: list[float], units: str, frame: str) -> dict:
+    return {
+        "attrs": {
+            "frame": frame.upper(),
+            "type": "location",
+            "units": ensure_units_are_consistent(units),
+        },
+        "data": to_python_type(data),
+        "dims": "ellipsoid_dir_label",
+        "coords": {
+            "ellipsoid_dir_label": {
+                "data": ["lon", "lat"],
+                "dims": "ellipsoid_dir_label",
+            }
+        },
     }
 
 
 def make_time_measure_attrs(units="s", scale="utc", time_format="mjd") -> dict:
     u = ensure_units_are_consistent(units)
-    return {"units": u, "scale": scale, "format": time_format, "type": "time"}
+    return {
+        "units": u,
+        "scale": scale.lower(),
+        "format": time_format.lower(),
+        "type": "time",
+    }
 
 
 def make_time_measure_dict(data, units="s", scale="utc", time_format="mjd") -> dict:
@@ -127,7 +158,7 @@ def make_time_measure_dict(data, units="s", scale="utc", time_format="mjd") -> d
     """
     x = {}
     x["attrs"] = make_time_measure_attrs(units, scale, time_format)
-    x["data"] = data
+    x["data"] = to_python_type(data)
     x["dims"] = []
     return x
 
@@ -149,7 +180,7 @@ def make_time_coord_attrs(units="s", scale="utc", time_format="mjd") -> dict:
     -------
     dict
     """
-    x = make_time_measure_attrs(units, scale, time_format)
+    x = make_time_measure_attrs(units, scale.lower(), time_format.lower())
     del x["type"]
     return x
 

xradio/_utils/list_and_array.py

@@ -53,11 +53,34 @@ def get_pad_value(col_dtype: np.dtype) -> object:
         )
 
 
+def to_python_type(x):
+    """
+    Convert any NumPy scalar, array, or nested structure to native Python types.
+
+    - np.float32, np.float64 → float
+    - np.int32, np.int64 → int
+    - np.bool_ → bool
+    - np.ndarray → list of Python-native types
+    - nested containers (list/tuple/dict) are handled recursively
+    """
+    if isinstance(x, np.generic):  # covers all numpy scalar types
+        return x.item()
+    elif isinstance(x, np.ndarray):
+        return x.tolist()
+    elif isinstance(x, (list, tuple)):
+        return type(x)(to_python_type(v) for v in x)
+    elif isinstance(x, dict):
+        return {k: to_python_type(v) for k, v in x.items()}
+    else:
+        return x
+
+
 def to_list(x):
     if isinstance(x, np.ndarray):
-        if x.ndim == 0:
-            return [x.item()]
-        return list(x)  # needed for json serialization
+        z = x.astype(float)
+        if z.ndim == 0:
+            return [z.item()]
+        return list(z)  # needed for json serialization
     elif isinstance(x, list):
         return x
     return [x]

xradio/_utils/schema.py

@@ -224,3 +224,47 @@ casa_frequency_frames = [
 ]
 
 casa_frequency_frames_codes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 64]
+
+
+def get_data_group_keys(schema_name: str) -> list[tuple[str, bool]]:
+    """Return (name, is_optional) pairs for data group keys for the given schema.
+
+    Parameters
+    ----------
+    schema_name : str
+        The name of the schema to retrieve data group keys for. Supported values are "msv4" and "image".
+
+    Returns
+    -------
+    list[tuple[str, bool]]
+        A list of tuples (key_name, is_optional) for the specified schema.
+
+    Raises
+    ------
+    ValueError
+        If the schema name is unknown.
+    """
+    from typing import get_type_hints, get_origin, get_args, Union
+
+    if schema_name == "msv4":
+        from xradio.measurement_set.schema import DataGroupDict
+
+        cls = DataGroupDict
+    elif schema_name == "image":
+        from xradio.image.schema import DataGroupDict
+
+        cls = DataGroupDict
+    else:
+        raise ValueError(f"Unknown schema name: {schema_name}")
+
+    annotations = get_type_hints(cls)
+    keys_with_optional = {}
+    for name, anno in annotations.items():
+        origin = get_origin(anno)
+        is_optional = False
+        if origin is Union:
+            args = get_args(anno)
+            is_optional = type(None) in args
+        keys_with_optional[name] = is_optional
+
+    return keys_with_optional

xradio/_utils/xarray_helpers.py

@@ -0,0 +1,63 @@
+import xarray as xr
+from xradio._utils.schema import get_data_group_keys
+from collections.abc import Mapping, Iterable
+from typing import Any, Union
+
+
+def get_data_group_name(
+    xdx: Union[xr.Dataset, xr.DataTree], data_group_name: str = None
+) -> str:
+
+    if data_group_name is None:
+        if "base" in xdx.attrs["data_groups"]:
+            data_group_name = "base"
+        else:
+            data_group_name = list(xdx.attrs["data_groups"].keys())[0]
+
+    return data_group_name
+
+
+def create_new_data_group(
+    xdx: Union[xr.Dataset, xr.DataTree],
+    schema_name: str,
+    new_data_group_name: str,
+    data_group: dict,
+    data_group_dv_shared_with: str = None,
+) -> xr.DataTree:
+    """Adds a data group to Xarray Data Structure (Dataset or DataTree).
+
+    Parameters
+    ----------
+    new_data_group_name : str
+        The name of the new data group to add.
+    data_group : dict
+        A dictionary containing the data group variables and their attributes.
+    data_group_dv_shared_with : str, optional
+        The name of the data group to share data variables with, by default "base"
+
+    Returns
+    -------
+    xr.DataTree
+        MSv4 DataTree with the new group added
+    """
+
+    data_group_dv_shared_with = get_data_group_name(xdx, data_group_dv_shared_with)
+
+    default_data_group = xdx.attrs["data_groups"][data_group_dv_shared_with]
+
+    new_data_group = {}
+
+    data_group_keys = get_data_group_keys(schema_name)
+
+    for key, optional in data_group_keys.items():
+        if key in data_group:
+            new_data_group[key] = data_group[key]
+        else:
+            if key not in default_data_group and not optional:
+                raise ValueError(
+                    f"Data group key '{key}' is required but not provided and not present in shared data group '{data_group_dv_shared_with}'."
+                )
+            elif key in default_data_group:
+                new_data_group[key] = default_data_group[key]
+
+    return new_data_group_name, new_data_group

xradio/_utils/zarr/common.py

@@ -1,7 +1,5 @@
 import xarray as xr
-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,
@@ -10,6 +8,9 @@ from botocore.exceptions import NoCredentialsError
 
 def _get_file_system_and_items(ps_store: str):
 
+    import s3fs
+    from botocore.exceptions import NoCredentialsError
+
     # default to assuming the data are accessible on local file system
     if os.path.isdir(ps_store):
         # handle a common shell convention
@@ -74,6 +75,7 @@ def _open_dataset(
     """
 
     import dask
+    import s3fs
 
     if isinstance(file_system, s3fs.core.S3FileSystem):
         mapping = s3fs.S3Map(root=store, s3=file_system, check=False)

xradio/image/__init__.py

@@ -5,15 +5,17 @@ from .image import (
     make_empty_aperture_image,
     make_empty_lmuv_image,
     make_empty_sky_image,
-    read_image,
+    open_image,
     write_image,
 )
 
+from .image_xds import ImageXds
+
 __all__ = [
     "load_image",
     "make_empty_aperture_image",
     "make_empty_lmuv_image",
     "make_empty_sky_image",
-    "read_image",
+    "open_image",
     "write_image",
 ]

xradio/image/_util/_casacore/common.py

@@ -42,7 +42,8 @@ def _create_new_image(
         del image
 
 
-_active_mask = "active_mask"
+_image_flag = "flag"
+_beam_fit_params = "beam_fit_params"
 # _native_types = ["FREQ", "VRAD", "VOPT", "BETA", "WAVE", "AWAV"]
 _object_name = "object_name"
 _pointing_center = "pointing_center"

xradio/image/_util/_casacore/xds_from_casacore.py

@@ -9,6 +9,7 @@ import toolviper.utils.logger as logger
 import numpy as np
 import xarray as xr
 from astropy import units as u
+from xradio._utils.list_and_array import to_python_type
 
 try:
     from casacore import tables
@@ -19,14 +20,15 @@ except ImportError:
     from xradio._utils._casacore.casacore_from_casatools import image as casa_image
 
 
-from .common import (
-    _active_mask,
+from xradio.image._util._casacore.common import (
+    _image_flag,
+    _beam_fit_params,
     _object_name,
     _open_image_ro,
     _pointing_center,
 )
 
-from ..common import (
+from xradio.image._util.common import (
     _compute_linear_world_values,
     _compute_velocity_values,
     _compute_world_sph_dims,
@@ -37,10 +39,11 @@ from ..common import (
     _image_type,
     _l_m_attr_notes,
 )
-from ...._utils._casacore.tables import extract_table_attributes, open_table_ro
+from xradio._utils._casacore.tables import extract_table_attributes, open_table_ro
 from xradio._utils.coord_math import _deg_to_rad
 from xradio._utils.dict_helpers import (
     _casacore_q_to_xradio_q,
+    make_direction_location_dict,
     make_spectral_coord_reference_dict,
     make_quantity,
     make_skycoord_dict,
@@ -68,13 +71,15 @@ def _add_mask(
     xda = xr.DataArray(ary, dims=dimorder)
     # True pixels are good in numpy masked arrays
     xda = da.logical_not(xda)
-    xda.attrs["image_type"] = "Mask"
+    xda.attrs["type"] = "flag"
     xda = xda.rename(name)
     xds[xda.name] = xda
     return xds
 
 
-def _casa_image_to_xds_image_attrs(image: casa_image, history: bool = True) -> dict:
+def _casa_image_to_xds_image_attrs(
+    image: casa_image, history: bool = False, image_type: str = "SKY"
+) -> dict:
     """
     get the image attributes from the casacoreimage object
     """
@@ -177,20 +182,27 @@ def _casa_image_to_xds_image_attrs(image: casa_image, history: bool = True) -> d
     obj = "objectname"
     attrs[_object_name] = imageinfo[obj] if obj in imageinfo else ""
     attrs["user"] = meta_dict["miscinfo"]
+    """
     defmask = "Image_defaultmask"
     with open_table_ro(image.name()) as casa_table:
         # the actual mask is a data var and data var names are all caps by convention
-        attrs[_active_mask] = (
-            casa_table.getkeyword(defmask).upper()
-            if defmask in casa_table.keywordnames()
-            else None
-        )
+        import re
+
+        if defmask in casa_table.keywordnames():
+            am = casa_table.getkeyword(defmask).upper()
+            am = re.sub(r"\bMASK(\d+)\b", r"MASK_\1", am)
+        else:
+            am = None
+        attrs[_image_flag] = "FLAG_" + image_type.upper()
+    """
     attrs["description"] = None
-    # if also loading history, put it as another xds in the image attrs
+    # Store history as a dict (not xr.Dataset) for Xarray compatibility
     if history:
         htable = os.sep.join([os.path.abspath(image.name()), "logtable"])
         if os.path.isdir(htable):
-            attrs["history"] = read_generic_table(htable)
+            history_xds = read_generic_table(htable)
+            # Convert xr.Dataset to dict for serialization compatibility
+            attrs["history"] = history_xds.to_dict()
         else:
             logger.warning(
                 f"Unable to find history table {htable}. History will not be included"
@@ -205,13 +217,14 @@ def _add_sky_or_aperture(
     img_full_path: str,
     has_sph_dims: bool,
     history: bool,
+    image_type: str = "SKY",
 ) -> xr.Dataset:
     xda = xr.DataArray(ary, dims=dimorder).astype(ary.dtype)
     with _open_image_ro(img_full_path) as casa_image:
-        xda.attrs = _casa_image_to_xds_image_attrs(casa_image, history)
+        xda.attrs = _casa_image_to_xds_image_attrs(casa_image, history, image_type)
     # xds.attrs = attrs
-    name = "SKY" if has_sph_dims else "APERTURE"
-    xda = xda.rename(name)
+    # name = "SKY" if has_sph_dims else "APERTURE"
+    xda = xda.rename(image_type)
     xds[xda.name] = xda
     return xds
 
@@ -257,6 +270,8 @@ def _casa_image_to_xds_attrs(img_full_path: str) -> dict:
     """
     with _open_image_ro(img_full_path) as casa_image:
         meta_dict = casa_image.info()
+    # print("meta_dict:", meta_dict)
+    # print("***********")
     coord_dict = copy.deepcopy(meta_dict["coordinates"])
     attrs = {}
     dir_key = None
@@ -272,31 +287,67 @@ def _casa_image_to_xds_attrs(img_full_path: str) -> dict:
             raise RuntimeError("No direction reference frame found")
         casa_system = coord_dir_dict[system]
         ap_system, ap_equinox = _convert_direction_system(casa_system, "native")
-        dir_dict = {}
 
-        dir_dict["reference"] = make_skycoord_dict(
-            data=[0.0, 0.0], units="rad", frame=ap_system
+        coordinate_system_info = {}
+
+        unit0 = u.Unit(_get_unit(coord_dir_dict["units"][0]))
+        unit1 = u.Unit(_get_unit(coord_dir_dict["units"][1]))
+        ra = float((coord_dir_dict["crval"][0] * unit0).to("rad").value)
+        dec = float((coord_dir_dict["crval"][1] * unit1).to("rad").value)
+        coordinate_system_info["reference_direction"] = make_skycoord_dict(
+            data=[ra, dec], units="rad", frame=ap_system
         )
         if ap_equinox:
-            dir_dict["reference"]["attrs"]["equinox"] = ap_equinox
-        for i in range(2):
-            unit = u.Unit(_get_unit(coord_dir_dict["units"][i]))
-            q = coord_dir_dict["crval"][i] * unit
-            x = q.to("rad")
-            dir_dict["reference"]["data"][i] = x.value
-        k = "latpole"
-        if k in coord_dir_dict:
-            for j in (k, "lonpole"):
-                m = "longpole" if j == "lonpole" else j
-                dir_dict[j] = make_quantity(
-                    value=coord_dir_dict[m] * _deg_to_rad, units="rad", dims=["l", "m"]
-                )
-        for j in ("pc", "projection_parameters", "projection"):
-            if j in coord_dir_dict:
-                dir_dict[j] = coord_dir_dict[j]
-        attrs["direction"] = dir_dict
+            coordinate_system_info["reference_direction"]["attrs"][
+                "equinox"
+            ] = ap_equinox
+
+        pol_dir = [-1, coord_dir_dict["latpole"] * _deg_to_rad]
+        if "lonpole" in coord_dir_dict:
+            pol_dir[0] = coord_dir_dict["lonpole"] * _deg_to_rad
+        else:
+            pol_dir[0] = coord_dir_dict["longpole"] * _deg_to_rad
+
+        coordinate_system_info["native_pole_direction"] = make_direction_location_dict(
+            pol_dir, "rad", "native_projection"
+        )
+
+        coordinate_system_info["projection"] = coord_dir_dict["projection"]
+        coordinate_system_info["projection_parameters"] = to_python_type(
+            coord_dir_dict["projection_parameters"]
+        )
+        coordinate_system_info["pixel_coordinate_transformation_matrix"] = (
+            to_python_type(coord_dir_dict["pc"])
+        )
+
+        attrs["coordinate_system_info"] = coordinate_system_info
     return copy.deepcopy(attrs)
 
+    #     dir_dict = {}
+
+    #     dir_dict["reference"] = make_skycoord_dict(
+    #         data=[0.0, 0.0], units="rad", frame=ap_system
+    #     )
+    #     if ap_equinox:
+    #         dir_dict["reference"]["attrs"]["equinox"] = ap_equinox
+    #     for i in range(2):
+    #         unit = u.Unit(_get_unit(coord_dir_dict["units"][i]))
+    #         q = coord_dir_dict["crval"][i] * unit
+    #         x = q.to("rad")
+    #         dir_dict["reference"]["data"][i] = float(x.value)
+    #     k = "latpole"
+    #     if k in coord_dir_dict:
+    #         for j in (k, "lonpole"):
+    #             m = "longpole" if j == "lonpole" else j
+    #             dir_dict[j] = make_quantity(
+    #                 value=coord_dir_dict[m] * _deg_to_rad, units="rad", dims=["l", "m"]
+    #             )
+    #     for j in ("pc", "projection_parameters", "projection"):
+    #         if j in coord_dir_dict:
+    #             dir_dict[j] = coord_dir_dict[j]
+    #     attrs["direction"] = dir_dict
+    # return copy.deepcopy(attrs)
+
 
 def _casa_image_to_xds_coords(
     img_full_path: str, verbose: bool, do_sky_coords: bool
@@ -330,14 +381,12 @@ def _casa_image_to_xds_coords(
     attrs["sphr_dims"] = sphr_dims
     coords = {}
     coord_attrs = {}
-    (coords["time"], coord_attrs["time"]) = _get_time_values_attrs(coord_dict)
-    (coords["frequency"], coord_attrs["frequency"]) = _get_freq_values_attrs(
-        csys, shape
-    )
-    (velocity_vals, coord_attrs["velocity"]) = _get_velocity_values_attrs(
+    coords["time"], coord_attrs["time"] = _get_time_values_attrs(coord_dict)
+    coords["frequency"], coord_attrs["frequency"] = _get_freq_values_attrs(csys, shape)
+    velocity_vals, coord_attrs["velocity"] = _get_velocity_values_attrs(
         coord_dict, coords["frequency"]
     )
-    (coords["polarization"], coord_attrs["polarization"]) = _get_pol_values_attrs(
+    coords["polarization"], coord_attrs["polarization"] = _get_pol_values_attrs(
         coord_dict
     )
     coords["velocity"] = (["frequency"], velocity_vals)
@@ -380,7 +429,7 @@ def _casa_image_to_xds_coords(
         ret = _get_uv_values_attrs(coord_dict, axis_names, shape)
         for z in ["u", "v"]:
             coords[z], coord_attrs[z] = ret[z]
-    coords["beam_param"] = ["major", "minor", "pa"]
+    coords["beam_params_label"] = ["major", "minor", "pa"]
     attrs["shape"] = shape
     xds = xr.Dataset(coords=coords)
     for c in coord_attrs.keys():
@@ -390,7 +439,7 @@ def _casa_image_to_xds_coords(
 
 
 def _convert_direction_system(
-    casa_system: str, which: str, verbose: bool = True
+    casa_system: str, which: str, verbose: bool = False
 ) -> tuple:
     if casa_system == "J2000":
         if verbose:
@@ -634,6 +683,7 @@ def _get_persistent_block(
     block = _read_image_chunk(infile, shapes, starts)
     block = np.expand_dims(block, new_axes)
     block = block.transpose(transpose_list)
+    block = da.from_array(block, chunks=block.shape)
     block = xr.DataArray(block, dims=dimorder)
     return block
 
@@ -686,11 +736,11 @@ def _get_starts_shapes_slices(
 def _get_time_values_attrs(cimage_coord_dict: dict) -> Tuple[List[float], dict]:
     attrs = {}
     attrs["type"] = "time"
-    attrs["scale"] = cimage_coord_dict["obsdate"]["refer"]
+    attrs["scale"] = cimage_coord_dict["obsdate"]["refer"].lower()
     unit = cimage_coord_dict["obsdate"]["m0"]["unit"]
     attrs["units"] = unit
     time_val = cimage_coord_dict["obsdate"]["m0"]["value"]
-    attrs["format"] = _get_time_format(time_val, unit)
+    attrs["format"] = _get_time_format(time_val, unit).lower()
     return ([time_val], copy.deepcopy(attrs))
 
 
@@ -829,7 +879,11 @@ def _get_velocity_values_attrs(
 
 
 def _get_beam(
-    img_full_path: str, nchan: int, npol: int, as_dask_array: bool
+    img_full_path: str,
+    nchan: int,
+    npol: int,
+    as_dask_array: bool,
+    image_type: str = "SKY",
 ) -> Union[xr.DataArray, None]:
     # the image may have multiple beams
     with _open_image_ro(img_full_path) as casa_image:
@@ -862,10 +916,10 @@ def _get_beam(
     if as_dask_array:
         beam_array = da.array(beam_array)
     xdb = xr.DataArray(
-        beam_array, dims=["time", "frequency", "polarization", "beam_param"]
+        beam_array, dims=["time", "frequency", "polarization", "beam_params_label"]
     )
-    xdb = xdb.rename("BEAM")
-    xdb = xdb.assign_coords(beam_param=["major", "minor", "pa"])
+    xdb = xdb.rename("BEAM_FIT_PARAMS_" + image_type.upper())
+    xdb = xdb.assign_coords(beam_params_label=["major", "minor", "pa"])
     xdb.attrs["units"] = "rad"
     return xdb
 
@@ -1054,7 +1108,7 @@ def _read_image_array(
                    indicating the length of a chunk on that particular axis. If
                    a key is missing, the associated chunk length along that axis
                    is 1. 'l' represents the longitude like dimension, and 'm'
-                   represents the latitude like dimension. For apeature images,
+                   represents the latitude like dimension. For aperture images,
                    'u' may be used in place of 'l', and 'v' in place of 'm'.
     :type chunks: list | dict, required
     :param mask: If specified, this is the associated image mask to read, rather than the actual