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

xradio/image/_util/_casacore/xds_to_casacore.py

@@ -1,5 +1,6 @@
 import copy
 import os
+import time
 
 import dask.array as da
 import numpy as np
@@ -12,9 +13,19 @@ try:
 except ImportError:
     import xradio._utils._casacore.casacore_from_casatools as tables
 
-from .common import _active_mask, _create_new_image, _object_name, _pointing_center
-from ..common import _aperture_or_sky, _compute_sky_reference_pixel, _doppler_types
-from ...._utils._casacore.tables import open_table_rw
+from xradio.image._util._casacore.common import (
+    _image_flag,
+    _beam_fit_params,
+    _create_new_image,
+    _object_name,
+    _pointing_center,
+)
+from xradio.image._util.common import (
+    _aperture_or_sky,
+    _compute_sky_reference_pixel,
+    _doppler_types,
+)
+from xradio._utils._casacore.tables import open_table_rw
 
 
 def _compute_direction_dict(xds: xr.Dataset) -> dict:
@@ -23,31 +34,35 @@ def _compute_direction_dict(xds: xr.Dataset) -> dict:
     for a CASA image coordinate system
     """
     direction = {}
-    xds_dir = xds.attrs["direction"]
+
+    xds_dir = xds.attrs["coordinate_system_info"]
     direction["_axes_sizes"] = np.array(
         [xds.sizes[dim] for dim in ("l", "m")], dtype=np.int32
     )
     direction["_image_axes"] = np.array([2, 3], dtype=np.int32)
-    direction["system"] = xds_dir["reference"]["attrs"]["equinox"].upper()
+    direction["system"] = xds_dir["reference_direction"]["attrs"]["equinox"].upper()
     if direction["system"] == "J2000.0":
         direction["system"] = "J2000"
     direction["projection"] = xds_dir["projection"]
     direction["projection_parameters"] = xds_dir["projection_parameters"]
     direction["units"] = [
-        xds_dir["reference"]["attrs"]["units"],
-        xds_dir["reference"]["attrs"]["units"],
+        xds_dir["reference_direction"]["attrs"]["units"],
+        xds_dir["reference_direction"]["attrs"]["units"],
     ]
-    direction["crval"] = np.array(xds_dir["reference"]["data"])
+    direction["crval"] = np.array(xds_dir["reference_direction"]["data"])
     direction["cdelt"] = np.array((xds.l[1] - xds.l[0], xds.m[1] - xds.m[0]))
     direction["crpix"] = _compute_sky_reference_pixel(xds)
-    direction["pc"] = np.array(xds_dir["pc"])
+    direction["pc"] = np.array(xds_dir["pixel_coordinate_transformation_matrix"])
     direction["axes"] = ["Right Ascension", "Declination"]
     direction["conversionSystem"] = direction["system"]
-    for s in ["longpole", "latpole"]:
+    for i, s in enumerate(["longpole", "latpole"]):
         m = "lonpole" if s == "longpole" else s
         # lonpole, latpole are numerical values in degrees in casa images
         direction[s] = float(
-            Angle(str(xds_dir[m]["data"]) + xds_dir[m]["attrs"]["units"]).deg
+            Angle(
+                str(xds_dir["native_pole_direction"]["data"][i])
+                + xds_dir["native_pole_direction"]["attrs"]["units"]
+            ).deg
         )
     return direction
 
@@ -205,50 +220,94 @@ def _coord_dict_from_xds(xds: xr.Dataset) -> dict:
 
 
 def _history_from_xds(xds: xr.Dataset, image: str) -> None:
-    nrows = len(xds.history.row) if "row" in xds.data_vars else 0
+    """
+    Write history from xds attributes to CASA image logtable.
+
+    Parameters
+    ----------
+    xds : xr.Dataset
+        Dataset with potential history attribute (stored as dict).
+    image : str
+        Path to the CASA image.
+
+    Notes
+    -----
+    History is stored in data variable attributes (e.g., SKY.attrs['history']),
+    not in the main dataset attributes.
+    """
+    # Check if history exists in data variable attributes (SKY, APERTURE, etc.)
+    history_dict = None
+    for var_name in xds.data_vars:
+        if "history" in xds[var_name].attrs:
+            history_dict = xds[var_name].attrs["history"]
+            break
+
+    # Also check dataset-level attributes as a fallback
+    if history_dict is None and "history" in xds.attrs:
+        history_dict = xds.attrs["history"]
+
+    if history_dict is None or not isinstance(history_dict, dict):
+        return
+
+    # Convert dict back to xr.Dataset to access data
+    try:
+        history_xds = xr.Dataset.from_dict(history_dict)
+    except Exception:
+        # If conversion fails, history dict may be malformed, skip writing
+        return
+
+    # Check for row in both dims and coords (it could be either depending on the xarray version)
+    nrows = 0
+    if "row" in history_xds.dims:
+        nrows = history_xds.sizes["row"]
+    elif "row" in history_xds.coords:
+        nrows = len(history_xds.row)
+    elif "row" in history_xds.data_vars:
+        # row might also be a data variable
+        nrows = len(history_xds.row)
+
     if nrows > 0:
         # TODO need to implement nrows == 0 case
         with open_table_rw(os.sep.join([image, "logtable"])) as tb:
             tb.addrows(nrows + 1)
             for c in ["TIME", "PRIORITY", "MESSAGE", "LOCATION", "OBJECT_ID"]:
-                vals = xds.history[c].values
-                if c == "TIME":
-                    k = time.time() + 40587 * 86400
-                elif c == "PRIORITY":
-                    k = "INFO"
-                elif c == "MESSAGE":
-                    k = (
-                        "Wrote xds to "
-                        + os.path.basename(image)
-                        + " using cngi_io.xds_to_casa_image_2()"
-                    )
-                elif c == "LOCATION":
-                    k = "cngi_io.xds_to_casa_image_2"
-                elif c == "OBJECT_ID":
-                    k = ""
-                vals = np.append(vals, k)
-                tb.putcol(c, vals)
+                if c in history_xds.data_vars or c in history_xds.coords:
+                    vals = history_xds[c].values
+                    if c == "TIME":
+                        k = time.time() + 40587 * 86400
+                    elif c == "PRIORITY":
+                        k = "INFO"
+                    elif c == "MESSAGE":
+                        k = (
+                            "Wrote xds to "
+                            + os.path.basename(image)
+                            + " using cngi_io.xds_to_casa_image_2()"
+                        )
+                    elif c == "LOCATION":
+                        k = "cngi_io.xds_to_casa_image_2"
+                    elif c == "OBJECT_ID":
+                        k = ""
+                    vals = np.append(vals, k)
+                    tb.putcol(c, vals)
 
 
 def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
     ii = {}
     ap_sky = _aperture_or_sky(xds)
-    ii["imagetype"] = (
-        xds[ap_sky].attrs["image_type"] if "image_type" in xds[ap_sky].attrs else ""
-    )
+    ii["imagetype"] = xds[ap_sky].attrs["type"] if "type" in xds[ap_sky].attrs else ""
     ii["objectname"] = (
         xds[ap_sky].attrs[_object_name] if _object_name in xds[ap_sky].attrs else ""
     )
-    if "BEAM" in xds.data_vars:
+    if "BEAM_FIT_PARAMS" in xds.data_vars:
         # multi beam
         pp = {}
         pp["nChannels"] = xds.sizes["frequency"]
         pp["nStokes"] = xds.sizes["polarization"]
 
-        bu = xds.BEAM.attrs["units"]
+        bu = xds.BEAM_FIT_PARAMS.attrs["units"]
         chan = 0
         polarization = 0
-        bv = xds.BEAM.values
+        bv = xds.BEAM_FIT_PARAMS.values
         for i in range(pp["nChannels"] * pp["nStokes"]):
             bp = bv[0][chan][polarization][:]
             b = {
@@ -281,7 +340,9 @@ def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
 
 
 def _write_casa_data(xds: xr.Dataset, image_full_path: str) -> None:
+
     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 = (
@@ -290,9 +351,8 @@ def _write_casa_data(xds: xr.Dataset, image_full_path: str) -> None:
         else ("frequency", "polarization", "v", "u")
     )
     casa_image_shape = xds[sky_ap].isel(time=0).transpose(*trans_coords).shape[::-1]
-    active_mask = (
-        xds[sky_ap].attrs["active_mask"] if _active_mask in xds[sky_ap].attrs else ""
-    )
+    flag = xds[sky_ap].attrs[_image_flag] if _image_flag in xds[sky_ap].attrs else ""
+    # TODO this all needs to be refactored to use flag rather than mask
     masks = []
     masks_rec = {}
     mask_rec = {
@@ -311,10 +371,14 @@ def _write_casa_data(xds: xr.Dataset, image_full_path: str) -> None:
     }
     for m in xds.data_vars:
         attrs = xds[m].attrs
-        if "image_type" in attrs and attrs["image_type"] == "Mask":
+        if "type" in attrs and attrs["type"] in ("mask", "flag"):
             masks_rec[m] = mask_rec
             masks_rec[m]["mask"] = f"Table: {os.sep.join([image_full_path, m])}"
             masks.append(m)
+    if flag and flag in xds.data_vars and flag not in masks:
+        masks_rec[flag] = mask_rec
+        masks_rec[flag]["mask"] = f"Table: {os.sep.join([image_full_path, flag])}"
+        masks.append(flag)
     myvars = [sky_ap]
     myvars.extend(masks)
     # xr.apply_ufunc seems to like stripping attributes from xds and its coordinates,
@@ -327,15 +391,15 @@ def _write_casa_data(xds: xr.Dataset, image_full_path: str) -> None:
     do_mask_nans = False
     there_are_nans = nan_mask.any()
     if there_are_nans:
-        has_active_mask = bool(active_mask)
-        if not has_active_mask:
+        has_flag = bool(flag)
+        if not has_flag:
             do_mask_nans = True
         else:
-            notted_active_mask = xr.apply_ufunc(
-                da.logical_not, xds[active_mask].copy(deep=True), dask="allowed"
+            notted_flag = xr.apply_ufunc(
+                da.logical_not, xds[flag].copy(deep=True), dask="allowed"
             )
             some_nans_are_not_already_masked = xr.apply_ufunc(
-                da.logical_and, nan_mask, notted_active_mask, dask="allowed"
+                da.logical_and, nan_mask, notted_flag, dask="allowed"
             )
             do_mask_nans = some_nans_are_not_already_masked.any()
     if do_mask_nans:
@@ -350,8 +414,8 @@ def _write_casa_data(xds: xr.Dataset, image_full_path: str) -> None:
         ] = f"Table: {os.sep.join([image_full_path, mask_name])}"
         masks.append(mask_name)
         arr_masks[mask_name] = nan_mask
-        if active_mask:
-            mask_name = f"mask_xds_nans_or_{active_mask}"
+        if flag:
+            mask_name = f"mask_xds_nans_or_{flag}"
             i = 0
             while mask_name in masks:
                 mask_name = f"mask_xds_nans{i}"
@@ -365,21 +429,22 @@ def _write_casa_data(xds: xr.Dataset, image_full_path: str) -> None:
             # apparent reason, so make a copy of the xds to run it on
             arr_masks[mask_name] = xr.apply_ufunc(
                 da.logical_or,
-                xds[active_mask].copy(deep=True),
+                xds[flag].copy(deep=True),
                 nan_mask,
                 dask="allowed",
             )
-        active_mask = mask_name
+        flag = mask_name
     data_type = "complex" if "u" in xds.coords else "float"
-    _write_initial_image(xds, image_full_path, active_mask, casa_image_shape[::-1])
+
+    _write_initial_image(xds, image_full_path, flag, casa_image_shape[::-1])
     for v in myvars:
-        _write_pixels(v, active_mask, image_full_path, xds)
+        _write_pixels(v, flag, image_full_path, xds)
     for name, v in arr_masks.items():
-        _write_pixels(name, active_mask, image_full_path, xds, v)
+        _write_pixels(name, flag, image_full_path, xds, v)
     if masks:
         with open_table_rw(image_full_path) as tb:
             tb.putkeyword("masks", masks_rec)
-            tb.putkeyword("Image_defaultmask", active_mask)
+            tb.putkeyword("Image_defaultmask", flag)
 
 
 def _write_initial_image(

xradio/image/_util/_fits/xds_from_fits.py

@@ -13,12 +13,14 @@ from astropy.time import Time
 
 from xradio._utils.coord_math import _deg_to_rad
 from xradio._utils.dict_helpers import (
+    make_direction_location_dict,
     make_quantity,
     make_spectral_coord_reference_dict,
     make_skycoord_dict,
     make_time_measure_dict,
 )
 
+from xradio._utils.list_and_array import to_python_type
 from xradio.measurement_set._utils._utils.stokes_types import stokes_types
 from xradio.image._util.common import (
     _compute_linear_world_values,
@@ -34,6 +36,11 @@ from xradio.image._util.common import (
     _l_m_attr_notes,
 )
 
+from xradio.image._util._casacore.common import (
+    _image_flag,
+    _beam_fit_params,
+)
+
 
 def _fits_image_to_xds(
     img_full_path: str,
@@ -41,6 +48,7 @@ def _fits_image_to_xds(
     verbose: bool,
     do_sky_coords: bool,
     compute_mask: bool,
+    image_type: str = "SKY",
 ) -> dict:
     """
     compute_mask : bool, optional
@@ -65,13 +73,27 @@ def _fits_image_to_xds(
     sphr_dims = helpers["sphr_dims"]
     ary = _read_image_array(img_full_path, chunks, helpers, verbose)
     dim_order = _get_xds_dim_order(sphr_dims)
-    xds = _add_sky_or_aperture(xds, ary, dim_order, header, helpers, sphr_dims)
+    xds = _add_sky_or_aperture(
+        xds, ary, dim_order, header, helpers, sphr_dims, image_type
+    )
     xds.attrs = attrs
     xds = _add_coord_attrs(xds, helpers)
     if helpers["has_multibeam"]:
-        xds = _do_multibeam(xds, img_full_path)
+        xds = _do_multibeam(xds, img_full_path, image_type=image_type)
+        xds[image_type.upper()].attrs[_beam_fit_params] = (
+            "BEAM_FIT_PARAMS_" + image_type.upper()
+        )
+        xds["BEAM_FIT_PARAMS_" + image_type.upper()].attrs["type"] = (
+            "beam_fit_params_" + image_type.lower()
+        )
     elif "beam" in helpers and helpers["beam"] is not None:
-        xds = _add_beam(xds, helpers)
+        xds = _add_beam(xds, helpers, image_type)
+        xds[image_type.upper()].attrs[_beam_fit_params] = (
+            "BEAM_FIT_PARAMS_" + image_type.upper()
+        )
+        xds["BEAM_FIT_PARAMS_" + image_type.upper()].attrs["type"] = (
+            "beam_fit_params_" + image_type.lower()
+        )
     return xds
 
 
@@ -116,7 +138,7 @@ def _add_vel_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
     vel_coord = xds.coords["velocity"]
     meta = {"units": "m/s"}
     if helpers["has_freq"]:
-        meta["doppler_type"] = helpers.get("doppler", "RADIO")
+        meta["doppler_type"] = helpers.get("doppler", "radio")
     else:
         meta["doppler_type"] = _doppler_types[0]
     meta["type"] = "doppler"
@@ -151,7 +173,7 @@ def _is_freq_like(v: str) -> bool:
     return v.startswith("FREQ") or v == "VOPT" or v == "VRAD"
 
 
-def _xds_direction_attrs_from_header(helpers: dict, header) -> dict:
+def _xds_coordinate_system_info_attrs_from_header(helpers: dict, header) -> dict:
     # helpers is modified in place, headers is not modified
     t_axes = helpers["t_axes"]
     p0 = header[f"CTYPE{t_axes[0]}"][-3:]
@@ -161,8 +183,8 @@ def _xds_direction_attrs_from_header(helpers: dict, header) -> dict:
             f"Projections for direction axes ({p0}, {p1}) differ, but they "
             "must be the same"
         )
-    direction = {}
-    direction["projection"] = p0
+    coordinate_system_info = {}
+    coordinate_system_info["projection"] = p0
     helpers["projection"] = p0
     ref_sys = header["RADESYS"]
     ref_eqx = None if ref_sys.upper() == "ICRS" else header["EQUINOX"]
@@ -171,7 +193,9 @@ def _xds_direction_attrs_from_header(helpers: dict, header) -> dict:
     helpers["ref_sys"] = ref_sys
     helpers["ref_eqx"] = ref_eqx
     # fits does not support conversion frames
-    direction["reference"] = make_skycoord_dict([0.0, 0.0], units="rad", frame=ref_sys)
+    coordinate_system_info["reference_direction"] = make_skycoord_dict(
+        [0.0, 0.0], units="rad", frame=ref_sys
+    )
     dir_axes = helpers["dir_axes"]
     ddata = []
     dunits = []
@@ -182,15 +206,20 @@ def _xds_direction_attrs_from_header(helpers: dict, header) -> dict:
         # direction["reference"]["value"][i] = x.value
         x = helpers["cdelt"][i] * u.Unit(_get_unit(helpers["cunit"][i]))
         dunits.append("rad")
-    direction["reference"] = make_skycoord_dict(ddata, units=dunits, frame=ref_sys)
-    if ref_eqx is not None:
-        direction["reference"]["attrs"]["equinox"] = ref_eqx.lower()
-    direction["latpole"] = make_quantity(
-        header["LATPOLE"] * _deg_to_rad, "rad", dims=["l", "m"]
+    coordinate_system_info["reference_direction"] = make_skycoord_dict(
+        ddata, units=dunits, frame=ref_sys
     )
-    direction["lonpole"] = make_quantity(
-        header["LONPOLE"] * _deg_to_rad, "rad", dims=["l", "m"]
+    if ref_eqx is not None:
+        coordinate_system_info["reference_direction"]["attrs"][
+            "equinox"
+        ] = ref_eqx.lower()
+
+    coordinate_system_info["native_pole_direction"] = make_direction_location_dict(
+        [header["LONPOLE"] * _deg_to_rad, header["LATPOLE"] * _deg_to_rad],
+        units="rad",
+        frame="native_projection",
     )
+
     pc = np.zeros([2, 2])
     for i in (0, 1):
         for j in (0, 1):
@@ -205,10 +234,61 @@ def _xds_direction_attrs_from_header(helpers: dict, header) -> dict:
                         f"Could not find PC{dir_axes[i]+1}_{dir_axes[j]+1} or "
                         f"PC0{dir_axes[i]+1}_0{dir_axes[j]+1} in FITS header"
                     )
-    direction["pc"] = pc
+    coordinate_system_info["pixel_coordinate_transformation_matrix"] = to_python_type(
+        pc
+    )
     # Is there really no fits header parameter for projection_parameters?
-    direction["projection_parameters"] = np.array([0.0, 0.0])
-    return direction
+    coordinate_system_info["projection_parameters"] = [0.0, 0.0]
+    return coordinate_system_info
+
+    # direction = {}
+    # direction["projection"] = p0
+    # helpers["projection"] = p0
+    # ref_sys = header["RADESYS"]
+    # ref_eqx = None if ref_sys.upper() == "ICRS" else header["EQUINOX"]
+    # if ref_sys == "FK5" and ref_eqx == 2000:
+    #     ref_eqx = "J2000.0"
+    # helpers["ref_sys"] = ref_sys
+    # helpers["ref_eqx"] = ref_eqx
+    # # fits does not support conversion frames
+    # direction["reference"] = make_skycoord_dict([0.0, 0.0], units="rad", frame=ref_sys)
+    # dir_axes = helpers["dir_axes"]
+    # ddata = []
+    # dunits = []
+    # for i in dir_axes:
+    #     x = helpers["crval"][i] * u.Unit(_get_unit(helpers["cunit"][i]))
+    #     x = x.to("rad")
+    #     ddata.append(x.value)
+    #     # direction["reference"]["value"][i] = x.value
+    #     x = helpers["cdelt"][i] * u.Unit(_get_unit(helpers["cunit"][i]))
+    #     dunits.append("rad")
+    # direction["reference"] = make_skycoord_dict(ddata, units=dunits, frame=ref_sys)
+    # if ref_eqx is not None:
+    #     direction["reference"]["attrs"]["equinox"] = ref_eqx.lower()
+    # direction["latpole"] = make_quantity(
+    #     header["LATPOLE"] * _deg_to_rad, "rad", dims=["l", "m"]
+    # )
+    # direction["lonpole"] = make_quantity(
+    #     header["LONPOLE"] * _deg_to_rad, "rad", dims=["l", "m"]
+    # )
+    # pc = np.zeros([2, 2])
+    # for i in (0, 1):
+    #     for j in (0, 1):
+    #         # dir_axes are now 0-based, but fits needs 1-based
+    #         try:
+    #             pc[i][j] = header[f"PC{dir_axes[i]+1}_{dir_axes[j]+1}"]
+    #         except KeyError:
+    #             try:
+    #                 pc[i][j] = header[f"PC0{dir_axes[i]+1}_0{dir_axes[j]+1}"]
+    #             except KeyError:
+    #                 raise RuntimeError(
+    #                     f"Could not find PC{dir_axes[i]+1}_{dir_axes[j]+1} or "
+    #                     f"PC0{dir_axes[i]+1}_0{dir_axes[j]+1} in FITS header"
+    #                 )
+    # direction["pc"] = pc
+    # # Is there really no fits header parameter for projection_parameters?
+    # direction["projection_parameters"] = np.array([0.0, 0.0])
+    # return direction
 
 
 def _fits_header_c_values_to_metadata(helpers: dict, header) -> None:
@@ -258,7 +338,7 @@ def _get_telescope_metadata(helpers: dict, header) -> dict:
                 "type": "location",
                 "units": "rad",
             },
-            "data": np.array([long, lat]),
+            "data": [long, lat],
             "dims": ["ellipsoid_dir_label"],
             "coords": {
                 "ellipsoid_dir_label": {
@@ -276,7 +356,7 @@ def _get_telescope_metadata(helpers: dict, header) -> dict:
                 "type": "location",
                 "units": "m",
             },
-            "data": np.array([r]),
+            "data": [r],
             "dims": ["ellipsoid_dis_label"],
             "coords": {
                 "ellipsoid_dis_label": {
@@ -465,7 +545,9 @@ def _fits_header_to_xds_attrs(
     else:
         raise RuntimeError("Could not find both direction axes")
     if dir_axes is not None:
-        attrs["direction"] = _xds_direction_attrs_from_header(helpers, header)
+        attrs["coordinate_system_info"] = _xds_coordinate_system_info_attrs_from_header(
+            helpers, header
+        )
     helpers["has_mask"] = False
     if compute_mask:
         # 🧠 Why the primary.data reference here is Safe (does not cause
@@ -527,7 +609,7 @@ def _fits_header_to_xds_attrs(
         data=Time(header["DATE-OBS"], format="isot").mjd,
         units=["d"],
         scale=header["TIMESYS"],
-        time_format="MJD",
+        time_format="mjd",
     )
 
     # TODO complete _make_history_xds when spec has been finalized
@@ -608,7 +690,8 @@ def _create_coords(
     else:
         # Fourier image
         coords["u"], coords["v"] = _get_uv_values(helpers)
-    coords["beam_param"] = ["major", "minor", "pa"]
+    coords["beam_params_label"] = ["major", "minor", "pa"]
+
     xds = xr.Dataset(coords=coords)
     return xds
 
@@ -698,7 +781,7 @@ def _get_velocity_values(helpers: dict) -> list:
 # FIXME change namee, even if there is only a single beam, we make a
 # multi beam array using it. If we have a beam, it will always be
 # "mutltibeam" is name is redundant and confusing
-def _do_multibeam(xds: xr.Dataset, imname: str) -> xr.Dataset:
+def _do_multibeam(xds: xr.Dataset, imname: str, image_type: str = "SKY") -> xr.Dataset:
     """Only run if we are sure there are multiple beams"""
     hdulist = fits.open(imname)
     for hdu in hdulist:
@@ -720,31 +803,33 @@ def _do_multibeam(xds: xr.Dataset, imname: str) -> xr.Dataset:
                 beam_array[:, :, :, i] = (
                     (beam_array[:, :, :, i] * units[i]).to("rad").value
                 )
-            return _create_beam_data_var(xds, beam_array)
+            return _create_beam_data_var(xds, beam_array, image_type)
     raise RuntimeError(
         "It looks like there should be a BEAMS table but no "
         "such table found in FITS file"
     )
 
 
-def _add_beam(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
+def _add_beam(xds: xr.Dataset, helpers: dict, image_type: str = "SKY") -> xr.Dataset:
     nchan = xds.sizes["frequency"]
     npol = xds.sizes["polarization"]
     beam_array = np.zeros([1, nchan, npol, 3])
     beam_array[0, :, :, 0] = helpers["beam"]["bmaj"]["data"]
     beam_array[0, :, :, 1] = helpers["beam"]["bmin"]["data"]
     beam_array[0, :, :, 2] = helpers["beam"]["pa"]["data"]
-    return _create_beam_data_var(xds, beam_array)
+    return _create_beam_data_var(xds, beam_array, image_type)
 
 
-def _create_beam_data_var(xds: xr.Dataset, beam_array: np.array) -> xr.Dataset:
+def _create_beam_data_var(
+    xds: xr.Dataset, beam_array: np.array, image_type: str = "SKY"
+) -> xr.Dataset:
     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"
-    xds["BEAM"] = xdb
+    xds["BEAM_FIT_PARAMS_" + image_type.upper()] = xdb
     return xds
 
 
@@ -779,7 +864,9 @@ def _add_sky_or_aperture(
     header,
     helpers: dict,
     has_sph_dims: bool,
+    image_type: str = "SKY",
 ) -> xr.Dataset:
+    # TODO add code to recognize aperture images and set image_type accordingly
     xda = xr.DataArray(ary, dims=dim_order)
     for h, a in zip(
         ["BUNIT", "BTYPE", "OBJECT", "OBSERVER"],
@@ -792,16 +879,17 @@ def _add_sky_or_aperture(
     xda.attrs["telescope"] = _get_telescope_metadata(helpers, header)
     xda.attrs["description"] = None
     xda.attrs["user"] = _user_attrs_from_header(header)
-    name = "SKY" if has_sph_dims else "APERTURE"
+    # name = "SKY" if has_sph_dims else "APERTURE"
+    name = image_type
     xda = xda.rename(name)
     xds[xda.name] = xda
     if helpers["has_mask"]:
         pp = da if type(xda[0].data) == dask.array.core.Array else np
         mask = pp.isnan(xda)
         mask.attrs = {}
-        mask = mask.rename("MASK0")
-        xds["MASK0"] = mask
-        xda.attrs["active_mask"] = "MASK0"
+        mask = mask.rename("FLAG_" + name.upper())
+        xds["FLAG_" + name.upper()] = mask
+        xds[name.upper()].attrs[_image_flag] = "FLAG_" + name.upper()
     xda = xda.rename(name)
     xds[xda.name] = xda
     return xds

xradio/image/_util/_zarr/common.py

@@ -1,6 +1,5 @@
 import numpy as np
 
-
 _np_types = {
     "complex128": np.complex128,
     "complex64": np.complex64,