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

xradio/_utils/_casacore/casacore_from_casatools.py

@@ -29,6 +29,24 @@ except ModuleNotFoundError as exc:
         f"{exc}"
     )
 
+# Valid casacore ImageInfo enum values
+# Reference: https://casacore.github.io/casacore/classcasacore_1_1ImageInfo.html
+_VALID_IMAGE_TYPES = (
+    "Undefined",
+    "Intensity",
+    "Beam",
+    "ColumnDensity",
+    "DepolarizationRatio",
+    "KineticTemperature",
+    "MagneticField",
+    "OpticalDepth",
+    "RotationMeasure",
+    "RotationalTemperature",
+    "SpectralIndex",
+    "Velocity",
+    "VelocityDispersion",
+)
+
 import numpy as np
 import toolviper.utils.logger as logger
 
@@ -166,6 +184,32 @@ def wrap_class_methods(cls: type) -> type:
     return cls
 
 
+def _validate_image_type(value: str) -> str:
+    """Validate and normalize an image type string.
+
+    Parameters
+    ----------
+    value
+        The image type string to validate.
+
+    Returns
+    -------
+    str
+        A valid casacore ImageInfo enum value with proper capitalization.
+        Returns 'Intensity' if the input is not a valid type.
+
+    Notes
+    -----
+    Validation is case-insensitive. The returned string uses the
+    canonical capitalization from the casacore ImageInfo enum.
+    """
+    value_lower = value.lower()
+    for valid_type in _VALID_IMAGE_TYPES:
+        if valid_type.lower() == value_lower:
+            return valid_type
+    return "Intensity"
+
+
 @wrap_class_methods
 class table(casatools.table):
     """A wrapper for the casatools table object.
@@ -513,7 +557,7 @@ class image(casatools.image):
         self,
         imagename,
         axis=0,
-        maskname="mask0",
+        maskname="MASK_0",
         images=(),
         values=None,
         coordsys=None,
@@ -702,18 +746,36 @@ class image(casatools.image):
     def imageinfo(self) -> dict:
         """Retrieve metadata from the image table.
 
-        This method accesses the image table associated with the image name
-        and attempts to retrieve information stored under the 'imageinfo'
-        keyword. If the 'imageinfo' keyword is not found in the table,
-        a default dictionary containing basic information is returned.
+        Accesses the image table and retrieves information stored under the
+        'imageinfo' keyword. The 'imagetype' value is validated against
+        casacore's ImageInfo enumeration values to mimic python-casacore
+        `image.imageinfo()` behavior.
 
         Returns
         -------
         dict
-            A dictionary containing image metadata. This is either the
-            value associated with the 'imageinfo' keyword in the table,
-            or a default dictionary {'imagetype': 'Intensity',
-            'objectname': ''} if the keyword is absent.
+            Image metadata dictionary containing:
+
+            - **imagetype** : str
+                Type of the image, validated against casacore ImageInfo enum.
+                Defaults to 'Intensity' if invalid or missing.
+            - **objectname** : str
+                Name of the observed object.
+
+        Notes
+        -----
+        image.info()['imageinfo'] and image.imageinfo() from python-casacore
+        always returns "imagetype" in a predefined enum value. When the "imageinfo"
+        keyword is missing from the image table, or a non-standard value of "imagetype"
+        (e.g. 'sky') was written into that keyword, image.info() will just return
+        "Intensity" as the imagetype.
+
+        Examples
+        --------
+        >>> img = image('my_image.im')
+        >>> info = img.imageinfo()
+        >>> info['imagetype']
+        'Intensity'
         """
         with table(self._imagename) as tb:
             if "imageinfo" in tb.keywordnames():
@@ -721,6 +783,10 @@ class image(casatools.image):
             else:
                 image_metadata = {"imagetype": "Intensity", "objectname": ""}
 
+        image_metadata["imagetype"] = _validate_image_type(
+            image_metadata.get("imagetype", "Intensity")
+        )
+
         return image_metadata
 
     def datatype(self):

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"