xradio 0.0.48__py3-none-any.whl → 0.0.49__py3-none-any.whl

xradio/__init__.py

@@ -1,6 +1,7 @@
 import os
 from toolviper.utils.logger import setup_logger
 
+
 # _logger_name = "xradio"
 # if os.getenv("VIPER_LOGGER_NAME") != _logger_name:
 #     os.environ["VIPER_LOGGER_NAME"] = _logger_name

xradio/_utils/dict_helpers.py

@@ -1,4 +1,4 @@
-def make_quantity(value, units: str) -> dict:
+def make_quantity(value, units: str, dims: list = []) -> dict:
     """
     create a quantity dictionary given value and units
     Parameters
@@ -11,4 +11,71 @@ def make_quantity(value, units: str) -> dict:
     -------
     dict
     """
-    return {"value": value, "units": units, "type": "quantity"}
+    u = units if isinstance(units, list) else [units]
+    return {"data": value, "dims": dims, "attrs": {"units": u, "type": "quantity"}}
+
+
+def make_frequency_reference_dict(
+    value: float, units: str, observer: str = "lsrk"
+) -> dict:
+    u = units if isinstance(units, list) else [units]
+    return {
+        "attrs": {"units": u, "observer": observer.lower(), "type": "frequency"},
+        "data": value,
+        "dims": [],
+    }
+
+
+def make_skycoord_dict(data: list[float], units: list[str], frame: str) -> dict:
+    return {
+        "attrs": {
+            "frame": frame.lower(),
+            "type": "sky_coord",
+            "units": units,
+        },
+        "data": data,
+        "dims": ["l", "m"],
+    }
+
+
+def make_time_measure_attrs(units=["s"], scale="utc", time_format="mjd") -> dict:
+    u = units if isinstance(units, list) else [units]
+    return {"units": u, "scale": scale, "format": time_format, "type": "time"}
+
+
+def make_time_coord_attrs(units=["s"], scale="utc", time_format="mjd") -> dict:
+    """
+    create a time measure dictionary given value and units
+    Parameters
+    ----------
+    value : numeric or array of numerics
+        Time value
+    units: str
+        Time units
+    scale: str
+        Time scale
+    time_format: str
+        Time format
+    Returns
+    -------
+    dict
+    """
+    x = make_time_measure_attrs(units, scale, time_format)
+    del x["type"]
+    return x
+
+
+def _casacore_q_to_xradio_q(q: dict) -> dict:
+    """
+    Convert a casacore quantity to an xradio quantity
+    """
+    if isinstance(q, dict):
+        if "value" in q and "unit" in q:
+            return make_quantity(q["value"], q["unit"])
+        else:
+            p = {}
+            for k in q:
+                p[k] = _casacore_q_to_xradio_q(q[k])
+            return p
+    else:
+        raise ValueError(f"Cannot convert {q} to xradio quantity")

xradio/image/_util/__init__.py

@@ -1,3 +0,0 @@
-from .common import _set_multibeam_array
-
-__all__ = ["_set_multibeam_array"]

xradio/image/_util/_casacore/common.py

@@ -15,19 +15,6 @@ def _open_image_ro(infile: str) -> Generator[images.image, None, None]:
         del image
 
 
-"""
-@contextmanager
-def _open_image_rw(
-    infile: str, mask: str, shape: tuple
-) -> Generator[images.image, None, None]:
-    image = images.image(infile, maskname=mask, shape=shape)
-    try:
-        yield image
-    finally:
-        del image
-"""
-
-
 @contextmanager
 def _create_new_image(
     outfile: str, shape: List[int], mask="", value="default"

xradio/image/_util/_casacore/xds_from_casacore.py

@@ -31,16 +31,12 @@ from ..common import (
 )
 from ...._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 make_quantity
-
-"""
-def _add_coord_attrs(xds: xr.Dataset, icoords: dict, dir_axes: list) -> xr.Dataset:
-    _add_time_attrs(xds, icoords)
-    _add_freq_attrs(xds, icoords)
-    xds = _add_vel_attrs(xds, icoords)
-    xds = _add_lin_attrs(xds, icoords, dir_axes)
-    return xds
-"""
+from xradio._utils.dict_helpers import (
+    make_quantity,
+    make_frequency_reference_dict,
+    make_skycoord_dict,
+    _casacore_q_to_xradio_q,
+)
 
 
 def _add_lin_attrs(xds, coord_dict, dir_axes):
@@ -65,11 +61,16 @@ def _add_freq_attrs(xds, coord_dict):
             sd = coord_dict[k]
             meta["rest_frequency"] = make_quantity(sd["restfreq"], "Hz")
             meta["type"] = "frequency"
-            meta["units"] = sd["unit"]
-            meta["frame"] = sd["system"]
+            # meta["units"] = sd["unit"]
+            # meta["frame"] = sd["system"]
             meta["wave_unit"] = sd["waveUnit"]
-            meta["crval"] = sd["wcs"]["crval"]
-            meta["cdelt"] = sd["wcs"]["cdelt"]
+            # meta["crval"] = sd["wcs"]["crval"]
+            # meta["cdelt"] = sd["wcs"]["cdelt"]
+            meta["reference_value"] = make_frequency_reference_dict(
+                value=sd["wcs"]["crval"],
+                units=sd["unit"],
+                observer=sd["system"],
+            )
     if not meta:
         # this is the default frequency information CASA creates
         meta = _default_freq_info()
@@ -108,7 +109,7 @@ def _add_sky_or_aperture(
 
 
 def _get_time_format(value: float, unit: str) -> str:
-    if value >= 40000 and value <= 100000 and unit == "d":
+    if value >= 40000 and value <= 100000 and (unit == "d" or unit == ["d"]):
         return "MJD"
     else:
         return ""
@@ -119,7 +120,7 @@ def _add_time_attrs(xds: xr.Dataset, coord_dict: dict) -> xr.Dataset:
     meta = {}
     meta["type"] = "time"
     meta["scale"] = coord_dict["obsdate"]["refer"]
-    meta["units"] = coord_dict["obsdate"]["m0"]["unit"]
+    meta["units"] = [coord_dict["obsdate"]["m0"]["unit"]]
     meta["format"] = _get_time_format(xds["time"][0], meta["units"])
     xds["time"].attrs = copy.deepcopy(meta)
     # xds['time'] = time_coord
@@ -128,7 +129,7 @@ def _add_time_attrs(xds: xr.Dataset, coord_dict: dict) -> xr.Dataset:
 
 def _add_vel_attrs(xds: xr.Dataset, coord_dict: dict) -> xr.Dataset:
     vel_coord = xds["velocity"]
-    meta = {"units": "m/s"}
+    meta = {"units": ["m/s"]}
     for k in coord_dict:
         if k.startswith("spectral"):
             sd = coord_dict[k]
@@ -164,22 +165,24 @@ def _casa_image_to_xds_attrs(img_full_path: str, history: bool = True) -> dict:
         casa_system = coord_dir_dict[system]
         ap_system, ap_equinox = _convert_direction_system(casa_system, "native")
         dir_dict = {}
-        dir_dict["reference"] = {
-            "frame": ap_system,
-            "type": "sky_coord",
-            "equinox": ap_equinox if ap_equinox else None,
-            "value": [0.0, 0.0],
-            "units": ["rad", "rad"],
-        }
+
+        dir_dict["reference"] = make_skycoord_dict(
+            data=[0.0, 0.0], units=["rad", "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"]["value"][i] = x.value
+            dir_dict["reference"]["data"][i] = x.value
         k = "latpole"
         if k in coord_dir_dict:
-            for j in (k, "longpole"):
-                dir_dict[j] = make_quantity(coord_dir_dict[j] * _deg_to_rad, "rad")
+            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]
@@ -216,7 +219,8 @@ def _casa_image_to_xds_attrs(img_full_path: str, history: bool = True) -> dict:
                 )
         elif k == "obsdate":
             obsdate["scale"] = coord_dict[k]["refer"]
-            obsdate["units"] = coord_dict[k]["m0"]["unit"]
+            myu = coord_dict[k]["m0"]["unit"]
+            obsdate["units"] = myu if isinstance(myu, list) else [myu]
             obsdate["value"] = coord_dict[k]["m0"]["value"]
             obsdate["format"] = _get_time_format(obsdate["value"], obsdate["units"])
             obsdate["type"] = "time"
@@ -225,12 +229,12 @@ def _casa_image_to_xds_attrs(img_full_path: str, history: bool = True) -> dict:
     imageinfo = meta_dict["imageinfo"]
     obj = "objectname"
     attrs[_object_name] = imageinfo[obj] if obj in imageinfo else ""
-    attrs["beam"] = _get_beam(imageinfo)
     attrs["user"] = meta_dict["miscinfo"]
     defmask = "Image_defaultmask"
     with open_table_ro(img_full_path) 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)
+            casa_table.getkeyword(defmask).upper()
             if defmask in casa_table.keywordnames()
             else None
         )
@@ -280,15 +284,15 @@ def _casa_image_to_xds_coords(
     coords = {}
     coord_attrs = {}
     (coords["time"], coord_attrs["time"]) = _get_time_values_attrs(coord_dict)
-    (coords["polarization"], coord_attrs["polarization"]) = _get_pol_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(
+        coord_dict
+    )
     coords["velocity"] = (["frequency"], velocity_vals)
     if len(sphr_dims) > 0:
         crpix = _flatten_list(csys.get_referencepixel())[::-1]
@@ -302,10 +306,6 @@ def _casa_image_to_xds_coords(
                 naxis=shape[idx], crval=0.0, crpix=crpix[idx], cdelt=delta
             )
             coord_attrs[c] = {
-                "crval": 0.0,
-                "cdelt": delta,
-                "units": "rad",
-                "type": "quantity",
                 "note": attr_note[c],
             }
         if do_sky_coords:
@@ -333,6 +333,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"]
     attrs["shape"] = shape
     xds = xr.Dataset(coords=coords)
     for c in coord_attrs.keys():
@@ -348,18 +349,18 @@ def _convert_direction_system(
         if verbose:
             logger.info(
                 f"J2000 found as {which} reference frame in CASA image "
-                'This corresponds to FK5(equinox="J2000") in astropy. '
+                'This corresponds to fk5(equinox="j2000") in astropy. '
                 "Metadata will be written appropriately"
             )
-        return ("FK5", "J2000.0")
+        return ("fk5", "j2000.0")
     elif casa_system == "B1950":
         if verbose:
             logger.info(
                 f"B1950 found as {which} reference frame in CASA image "
-                'This corresponds to FK4(equinox="B1950") in astropy. '
+                'This corresponds to fk4(equinox="b1950") in astropy. '
                 "Metadata will be written appropriately"
             )
-        return ("FK4", "B1950.0")
+        return ("fk4", "b1950.0")
     elif casa_system in ("GALACTIC", "ICRS"):
         return (casa_system.lower(), None)
     else:
@@ -382,14 +383,6 @@ def _flatten_list(list_of_lists: list) -> list:
     return flat
 
 
-def _get_beam(imageinfo: dict):
-    """Returns None if no beam. Multiple beams are handled elsewhere"""
-    k = "restoringbeam"
-    if k in imageinfo and "major" in imageinfo[k]:
-        return _convert_beam_to_rad(imageinfo[k])
-    return None
-
-
 def _get_chunk_list(
     chunks: dict, coords: list, image_shape: Union[list, tuple]
 ) -> tuple:
@@ -510,12 +503,19 @@ def _get_freq_values_attrs(
                     cdelt=wcs["cdelt"],
                 )
                 attrs["rest_frequency"] = make_quantity(sd["restfreq"], "Hz")
-                attrs["type"] = "frequency"
-                attrs["units"] = sd["unit"]
-                attrs["frame"] = sd["system"]
+                # attrs["type"] = "frequency"
+                # attrs["units"] = sd["unit"]
+                # attrs["frame"] = sd["system"]
                 attrs["wave_unit"] = sd["waveUnit"]
-                attrs["crval"] = sd["wcs"]["crval"]
-                attrs["cdelt"] = sd["wcs"]["cdelt"]
+                # attrs["crval"] = sd["wcs"]["crval"]
+                # attrs["cdelt"] = sd["wcs"]["cdelt"]
+
+                attrs["reference_value"] = make_frequency_reference_dict(
+                    value=sd["wcs"]["crval"],
+                    units=sd["unit"],
+                    observer=sd["system"],
+                )
+
                 attrs = copy.deepcopy(attrs)
                 break
     else:
@@ -576,23 +576,33 @@ def _get_mask_names(infile: str) -> list:
     return mymasks
 
 
-def _get_multibeam(imageinfo: dict) -> Union[np.ndarray, None]:
-    """Returns None if the image does not have multiple (per-plane) beams"""
-    p = "perplanebeams"
-    if p not in imageinfo:
+def _get_beam(imageinfo: dict, nchan: int, npol: int) -> Union[np.ndarray, None]:
+    """Returns None if the image has no beam(s)"""
+    x = ["perplanebeams", "restoringbeam"]
+    r = None
+    for z in x:
+        if z in imageinfo:
+            r = z
+            break
+    if r is None:
         return None
-    beam = imageinfo[p]
-    nchan = beam["nChannels"]
-    npol = beam["nStokes"]
-    beam_array = np.zeros([1, npol, nchan, 3])
-    for c in range(nchan):
-        for p in range(npol):
-            k = nchan * p + c
-            b = beam["*" + str(k)]
-            beam_dict = _convert_beam_to_rad(b)
-            beam_array[0][p][c][0] = beam_dict["major"]["value"]
-            beam_array[0][p][c][1] = beam_dict["minor"]["value"]
-            beam_array[0][p][c][2] = beam_dict["pa"]["value"]
+    beam = imageinfo[r]
+    beam_array = np.zeros([1, nchan, npol, 3])
+    if r == "perplanebeams":
+        for c in range(nchan):
+            for p in range(npol):
+                k = nchan * p + c
+                b = _casacore_q_to_xradio_q(beam["*" + str(k)])
+                beam_dict = _convert_beam_to_rad(b)
+                beam_array[0][c][p][0] = beam_dict["major"]["data"]
+                beam_array[0][c][p][1] = beam_dict["minor"]["data"]
+                beam_array[0][c][p][2] = beam_dict["pa"]["data"]
+    elif r == "restoringbeam":
+        b = _casacore_q_to_xradio_q(beam)
+        beam_dict = _convert_beam_to_rad(b)
+        beam_array[0, :, :, 0] = beam_dict["major"]["data"]
+        beam_array[0, :, :, 1] = beam_dict["minor"]["data"]
+        beam_array[0, :, :, 2] = beam_dict["pa"]["data"]
     return beam_array
 
 
@@ -661,7 +671,7 @@ def _get_time_values_attrs(cimage_coord_dict: dict) -> Tuple[List[float], dict]:
     attrs["type"] = "time"
     attrs["scale"] = cimage_coord_dict["obsdate"]["refer"]
     unit = cimage_coord_dict["obsdate"]["m0"]["unit"]
-    attrs["units"] = unit
+    attrs["units"] = unit if isinstance(unit, list) else [unit]
     time_val = cimage_coord_dict["obsdate"]["m0"]["value"]
     attrs["format"] = _get_time_format(time_val, unit)
     return ([time_val], copy.deepcopy(attrs))
@@ -694,15 +704,15 @@ def _get_transpose_list(coords: coordinates.coordinatesystem) -> list:
             not_covered.remove("v")
         elif b.startswith("frequency"):
             # transpose_list[2] = csys['pixelmap1'][0]
-            transpose_list[2] = i
+            transpose_list[1] = i
             not_covered.remove("f")
         elif b.startswith("stok"):
-            transpose_list[1] = i
+            transpose_list[2] = i
             # transpose_list[1] = csys['pixelmap2'][0]
             not_covered.remove("s")
         else:
             raise Exception(f"Unhandled axis name {c}")
-    h = {"l": 3, "m": 4, "u": 3, "v": 4, "f": 2, "s": 1}
+    h = {"l": 3, "m": 4, "u": 3, "v": 4, "f": 1, "s": 2}
     for p in not_covered:
         transpose_list[h[p]] = last_axis
         new_axes.append(last_axis)
@@ -773,7 +783,7 @@ def _get_velocity_values(coord_dict: dict, freq_values: list) -> list:
             restfreq = coord_dict[k]["restfreq"]
             break
     return _compute_velocity_values(
-        restfreq=restfreq, freq_values=freq_values, doppler="RADIO"
+        restfreq=restfreq, freq_values=freq_values, doppler="radio"
     )
 
 
@@ -791,11 +801,11 @@ def _get_velocity_values_attrs(
     if not attrs:
         attrs["doppler_type"] = _doppler_types[0]
 
-    attrs["units"] = "m/s"
+    attrs["units"] = ["m/s"]
     attrs["type"] = "doppler"
     return (
         _compute_velocity_values(
-            restfreq=restfreq, freq_values=freq_values, doppler="RADIO"
+            restfreq=restfreq, freq_values=freq_values, doppler="radio"
         ),
         copy.deepcopy(attrs),
     )
@@ -804,25 +814,20 @@ def _get_velocity_values_attrs(
 def _multibeam_array(
     xds: xr.Dataset, img_full_path: str, as_dask_array: bool
 ) -> Union[xr.DataArray, None]:
-    """This should only be called after the xds.beam attr has been set"""
-    if xds.attrs["beam"] is None:
-        # the image may have multiple beams
-        with _open_image_ro(img_full_path) as casa_image:
-            imageinfo = casa_image.info()["imageinfo"]
-        mb = _get_multibeam(imageinfo)
-        if mb is not None:
-            # multiple beams are stored as a data varialbe, so remove
-            # the beam xds attr
-            del xds.attrs["beam"]
-            if as_dask_array:
-                mb = da.array(mb)
-            xdb = xr.DataArray(
-                mb, dims=["time", "polarization", "frequency", "beam_param"]
-            )
-            xdb = xdb.rename("beam")
-            xdb = xdb.assign_coords(beam_param=["major", "minor", "pa"])
-            xdb.attrs["units"] = "rad"
-            return xdb
+    # the image may have multiple beams
+    with _open_image_ro(img_full_path) as casa_image:
+        imageinfo = casa_image.info()["imageinfo"]
+    mb = _get_beam(
+        imageinfo, nchan=xds.sizes["frequency"], npol=xds.sizes["polarization"]
+    )
+    if mb is not None:
+        if as_dask_array:
+            mb = da.array(mb)
+        xdb = xr.DataArray(mb, dims=["time", "frequency", "polarization", "beam_param"])
+        xdb = xdb.rename("BEAM")
+        xdb = xdb.assign_coords(beam_param=["major", "minor", "pa"])
+        xdb.attrs["units"] = "rad"
+        return xdb
     else:
         return None
 
@@ -1023,7 +1028,7 @@ def _read_image_array(
     :type blc: a type that is convertable to a list via list(blc)
     :param trc: top right corner, given in the axes ordering of the input image.None=>image shape - 1
     :type trc: a type that is convertable to a list via list(trc)
-    :return: Dask array in time, polarization, frequency, l, m order
+    :return: Dask array in time, frequency, polarization, l, m order
     :rtype: dask.array
 
     """

xradio/image/_util/_casacore/xds_to_casacore.py

@@ -19,19 +19,24 @@ def _compute_direction_dict(xds: xr.Dataset) -> dict:
     """
     direction = {}
     xds_dir = xds.attrs["direction"]
-    direction["system"] = xds_dir["reference"]["equinox"]
+    direction["system"] = xds_dir["reference"]["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"] = np.array(xds_dir["reference"]["units"], dtype="<U16")
-    direction["crval"] = np.array(xds_dir["reference"]["value"])
-    direction["cdelt"] = np.array((xds.l.cdelt, xds.m.cdelt))
+    direction["units"] = np.array(xds_dir["reference"]["attrs"]["units"], dtype="<U16")
+    direction["crval"] = np.array(xds_dir["reference"]["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["axes"] = ["Right Ascension", "Declination"]
     direction["conversionSystem"] = direction["system"]
     for s in ["longpole", "latpole"]:
-        # longpole, latpole are numerical values in degrees in casa images
-        direction[s] = Angle(str(xds_dir[s]["value"]) + xds_dir[s]["units"]).deg
+        m = "lonpole" if s == "longpole" else s
+        # lonpole, latpole are numerical values in degrees in casa images
+        direction[s] = Angle(
+            str(xds_dir[m]["data"]) + xds_dir[m]["attrs"]["units"][0]
+        ).deg
     return direction
 
 
@@ -69,20 +74,22 @@ def _compute_spectral_dict(xds: xr.Dataset) -> dict:
     # spec["nativeType"] = _native_types.index(xds.frequency.attrs["native_type"])
     # FREQ
     spec["nativeType"] = 0
-    spec["restfreq"] = xds.frequency.attrs["rest_frequency"]["value"]
+    spec["restfreq"] = xds.frequency.attrs["rest_frequency"]["data"]
     # spec["restfreqs"] = copy.deepcopy(xds.frequency.attrs["restfreqs"]["value"])
     spec["restfreqs"] = [spec["restfreq"]]
-    spec["system"] = xds.frequency.attrs["frame"]
-    spec["unit"] = xds.frequency.attrs["units"]
+    spec["system"] = xds.frequency.attrs["reference_value"]["attrs"]["observer"].upper()
+    u = xds.frequency.attrs["reference_value"]["attrs"]["units"]
+    spec["unit"] = u if isinstance(u, str) else u[0]
     spec["velType"] = _doppler_types.index(xds.velocity.attrs["doppler_type"])
-    spec["velUnit"] = xds.velocity.attrs["units"]
+    u = xds.velocity.attrs["units"]
+    spec["velUnit"] = u if isinstance(u, str) else u[0]
     spec["version"] = 2
     spec["waveUnit"] = xds.frequency.attrs["wave_unit"]
     wcs = {}
     wcs["ctype"] = "FREQ"
     wcs["pc"] = 1.0
-    wcs["crval"] = xds.frequency.attrs["crval"]
-    wcs["cdelt"] = xds.frequency.attrs["cdelt"]
+    wcs["crval"] = xds.frequency.attrs["reference_value"]["data"]
+    wcs["cdelt"] = xds.frequency.values[1] - xds.frequency.values[0]
     wcs["crpix"] = (wcs["crval"] - xds.frequency.values[0]) / wcs["cdelt"]
     spec["wcs"] = wcs
     return spec
@@ -96,7 +103,7 @@ def _coord_dict_from_xds(xds: xr.Dataset) -> dict:
     obsdate["refer"] = xds.coords["time"].attrs["scale"]
     obsdate["type"] = "epoch"
     obsdate["m0"] = {}
-    obsdate["m0"]["unit"] = xds.coords["time"].attrs["units"]
+    obsdate["m0"]["unit"] = xds.coords["time"].attrs["units"][0]
     obsdate["m0"]["value"] = xds.coords["time"].values[0]
     coord["obsdate"] = obsdate
     coord["pointingcenter"] = xds.attrs[_pointing_center].copy()
@@ -137,7 +144,8 @@ def _coord_dict_from_xds(xds: xr.Dataset) -> dict:
     # this probbably needs some verification
     coord["worldreplace0"] = [0.0, 0.0]
     coord["worldreplace1"] = np.array(coord["stokes1"]["crval"])
-    coord["worldreplace2"] = np.array([xds.frequency.attrs["crval"]])
+    # print("spectral", coord["spectral2"])
+    coord["worldreplace2"] = np.array(coord["spectral2"]["wcs"]["crval"])
     return coord
 
 
@@ -174,17 +182,17 @@ def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
         xds[ap_sky].attrs["image_type"] if "image_type" in xds[ap_sky].attrs else ""
     )
     ii["objectname"] = xds.attrs[_object_name]
-    if "beam" in xds.data_vars:
+    if "BEAM" in xds.data_vars:
         # multi beam
         pp = {}
-        pp["nChannels"] = len(xds.frequency)
-        pp["nStokes"] = len(xds.polarization)
-        bu = xds.beam.attrs["units"]
+        pp["nChannels"] = xds.sizes["frequency"]
+        pp["nStokes"] = xds.sizes["polarization"]
+        bu = xds.BEAM.attrs["units"]
         chan = 0
         polarization = 0
-        bv = xds.beam.values
+        bv = xds.BEAM.values
         for i in range(pp["nChannels"] * pp["nStokes"]):
-            bp = bv[0][polarization][chan][:]
+            bp = bv[0][chan][polarization][:]
             b = {
                 "major": {"unit": bu, "value": bp[0]},
                 "minor": {"unit": bu, "value": bp[1]},
@@ -196,15 +204,20 @@ def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
                 chan = 0
                 polarization += 1
         ii["perplanebeams"] = pp
+    """
     elif "beam" in xds.attrs and xds.attrs["beam"]:
         # do nothing if xds.attrs['beam'] is None
         ii["restoringbeam"] = copy.deepcopy(xds.attrs["beam"])
         for k in ["major", "minor", "pa"]:
-            del ii["restoringbeam"][k]["type"]
-            ii["restoringbeam"][k]["unit"] = ii["restoringbeam"][k]["units"]
-            del ii["restoringbeam"][k]["units"]
+            # print("*** ", k, ii["restoringbeam"][k])
+            del ii["restoringbeam"][k]["dims"]
+            ii["restoringbeam"][k]["unit"] = ii["restoringbeam"][k]["attrs"]["units"][0]
+            del ii["restoringbeam"][k]["attrs"]
+            ii["restoringbeam"][k]["value"] = ii["restoringbeam"][k]["data"]
+            del ii["restoringbeam"][k]["data"]
         ii["restoringbeam"]["positionangle"] = copy.deepcopy(ii["restoringbeam"]["pa"])
         del ii["restoringbeam"]["pa"]
+    """
     return ii
 
 
@@ -319,7 +332,6 @@ def _write_initial_image(
             if xds[dv][0, 0, 0, 0, 0].values.dtype == "float32":
                 value = "default"
             break
-    # print(type(value))
     image_full_path = os.path.expanduser(imagename)
     with _create_new_image(
         image_full_path, mask=maskname, shape=image_shape, value=value