xradio 0.0.55__py3-none-any.whl → 0.0.58__py3-none-any.whl

xradio/image/_util/_casacore/xds_to_casacore.py

@@ -5,7 +5,11 @@ import dask.array as da
 import numpy as np
 import xarray as xr
 from astropy.coordinates import Angle
-from casacore import tables
+
+try:
+    from casacore import tables
+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
@@ -28,7 +32,10 @@ def _compute_direction_dict(xds: xr.Dataset) -> dict:
         direction["system"] = "J2000"
     direction["projection"] = xds_dir["projection"]
     direction["projection_parameters"] = xds_dir["projection_parameters"]
-    direction["units"] = xds_dir["reference"]["attrs"]["units"]
+    direction["units"] = [
+        xds_dir["reference"]["attrs"]["units"],
+        xds_dir["reference"]["attrs"]["units"],
+    ]
     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)
@@ -39,7 +46,7 @@ def _compute_direction_dict(xds: xr.Dataset) -> dict:
         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"][0]).deg
+            Angle(str(xds_dir[m]["data"]) + xds_dir[m]["attrs"]["units"]).deg
         )
     return direction
 
@@ -83,20 +90,22 @@ def _compute_spectral_dict(xds: xr.Dataset) -> dict:
     spec["restfreq"] = xds.frequency.attrs["rest_frequency"]["data"]
     # spec["restfreqs"] = copy.deepcopy(xds.frequency.attrs["restfreqs"]["value"])
     spec["restfreqs"] = np.array([spec["restfreq"]])
-    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["system"] = xds.frequency.attrs["reference_frequency"]["attrs"][
+        "observer"
+    ].upper()
+    u = xds.frequency.attrs["reference_frequency"]["attrs"]["units"]
+    spec["unit"] = u
     spec["velType"] = _doppler_types.index(xds.velocity.attrs["doppler_type"])
     u = xds.velocity.attrs["units"]
     spec["version"] = 2
     # vel unit is a list[str] in the xds but needs to be a str in the casa image
-    spec["velUnit"] = xds.velocity.attrs["units"][0]
+    spec["velUnit"] = xds.velocity.attrs["units"]
     # wave unit is a list[str] in the xds but needs to be a str in the casa image
-    spec["waveUnit"] = xds.frequency.attrs["wave_unit"][0]
+    spec["waveUnit"] = xds.frequency.attrs["wave_units"]
     wcs = {}
     wcs["ctype"] = "FREQ"
     wcs["pc"] = 1.0
-    wcs["crval"] = float(xds.frequency.attrs["reference_value"]["data"])
+    wcs["crval"] = float(xds.frequency.attrs["reference_frequency"]["data"])
     wcs["cdelt"] = float(xds.frequency.values[1] - xds.frequency.values[0])
     wcs["crpix"] = float((wcs["crval"] - xds.frequency.values[0]) / wcs["cdelt"])
     spec["wcs"] = wcs
@@ -110,26 +119,46 @@ def _coord_dict_from_xds(xds: xr.Dataset) -> dict:
         tel = xds[sky_ap].attrs["telescope"]
         if "name" in tel:
             coord["telescope"] = xds[sky_ap].attrs["telescope"]["name"]
-        if "location" in tel:
-            xds_telloc = tel["location"]
+
+        if "direction" in tel:
+            xds_telloc = tel["direction"]
             telloc = {}
             telloc["refer"] = xds_telloc["attrs"]["frame"]
             if telloc["refer"] == "GRS80":
                 telloc["refer"] = "ITRF"
-            for i in range(3):
+            for i in range(2):
                 telloc[f"m{i}"] = {
-                    "unit": xds_telloc["attrs"]["units"][i],
+                    "unit": xds_telloc["attrs"]["units"],
                     "value": xds_telloc["data"][i],
                 }
+            telloc[f"m{2}"] = {
+                "unit": tel["distance"]["attrs"]["units"],
+                "value": tel["distance"]["data"][0],
+            }
+
             telloc["type"] = "position"
             coord["telescopeposition"] = telloc
+
+        # if "location" in tel:
+        #     xds_telloc = tel["location"]
+        #     telloc = {}
+        #     telloc["refer"] = xds_telloc["attrs"]["frame"]
+        #     if telloc["refer"] == "GRS80":
+        #         telloc["refer"] = "ITRF"
+        #     for i in range(3):
+        #         telloc[f"m{i}"] = {
+        #             "unit": xds_telloc["attrs"]["units"],
+        #             "value": xds_telloc["data"][i],
+        #         }
+        #     telloc["type"] = "position"
+        #     coord["telescopeposition"] = telloc
     if "observer" in xds[sky_ap].attrs:
         coord["observer"] = xds[sky_ap].attrs["observer"]
     obsdate = {}
     obsdate["refer"] = xds.coords["time"].attrs["scale"]
     obsdate["type"] = "epoch"
     obsdate["m0"] = {}
-    obsdate["m0"]["unit"] = xds.coords["time"].attrs["units"][0]
+    obsdate["m0"]["unit"] = xds.coords["time"].attrs["units"]
     obsdate["m0"]["value"] = float(xds.coords["time"].values[0])
     coord["obsdate"] = obsdate
     if _pointing_center in xds[sky_ap].attrs:
@@ -197,7 +226,7 @@ def _history_from_xds(xds: xr.Dataset, image: str) -> None:
 def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
     ii = {}
     ap_sky = _aperture_or_sky(xds)
-    ii["image_type"] = (
+    ii["imagetype"] = (
         xds[ap_sky].attrs["image_type"] if "image_type" in xds[ap_sky].attrs else ""
     )
     ii["objectname"] = (
@@ -208,6 +237,7 @@ def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
         pp = {}
         pp["nChannels"] = xds.sizes["frequency"]
         pp["nStokes"] = xds.sizes["polarization"]
+
         bu = xds.BEAM.attrs["units"]
         chan = 0
         polarization = 0
@@ -225,6 +255,7 @@ 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
@@ -232,7 +263,7 @@ def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
         for k in ["major", "minor", "pa"]:
             # print("*** ", k, ii["restoringbeam"][k])
             del ii["restoringbeam"][k]["dims"]
-            ii["restoringbeam"][k]["unit"] = ii["restoringbeam"][k]["attrs"]["units"][0]
+            ii["restoringbeam"][k]["units"] = ii["restoringbeam"][k]["attrs"]["units"]
             del ii["restoringbeam"][k]["attrs"]
             ii["restoringbeam"][k]["value"] = ii["restoringbeam"][k]["data"]
             del ii["restoringbeam"][k]["data"]

xradio/image/_util/_fits/xds_from_fits.py

@@ -1,8 +1,26 @@
-import astropy as ap
+import copy
+import re
+from typing import Union
+
+import dask
+import dask.array as da
+import numpy as np
+import psutil
+import xarray as xr
 from astropy import units as u
 from astropy.io import fits
 from astropy.time import Time
-from ..common import (
+
+from xradio._utils.coord_math import _deg_to_rad
+from xradio._utils.dict_helpers import (
+    make_quantity,
+    make_spectral_coord_reference_dict,
+    make_skycoord_dict,
+    make_time_measure_dict,
+)
+
+from xradio.measurement_set._utils._utils.stokes_types import stokes_types
+from xradio.image._util.common import (
     _compute_linear_world_values,
     _compute_velocity_values,
     _compute_world_sph_dims,
@@ -15,37 +33,34 @@ from ..common import (
     _image_type,
     _l_m_attr_notes,
 )
-from xradio._utils.coord_math import _deg_to_rad
-from xradio._utils.dict_helpers import (
-    make_frequency_reference_dict,
-    make_quantity,
-    make_skycoord_dict,
-    make_time_measure_dict,
-)
-import copy
-import dask
-import dask.array as da
-import numpy as np
-import re
-from typing import Union
-import xarray as xr
 
 
 def _fits_image_to_xds(
-    img_full_path: str, chunks: dict, verbose: bool, do_sky_coords: bool
+    img_full_path: str,
+    chunks: dict,
+    verbose: bool,
+    do_sky_coords: bool,
+    compute_mask: bool,
 ) -> dict:
     """
+    compute_mask : bool, optional
+        If True (default), compute and attach valid data masks to the xds.
+        If False, skip mask generation for performance. It is solely the responsibility
+        of the user to ensure downstream apps can handle NaN values; do not
+        ask package developers to add this non-standard behavior.
+
     TODO: complete documentation
     Create an xds without any pixel data from metadata from the specified FITS image
     """
-    # memmap = True allows only part of data to be loaded into memory
     # may also need to pass mode='denywrite'
     # https://stackoverflow.com/questions/35759713/astropy-io-fits-read-row-from-large-fits-file-with-mutliple-hdus
-    hdulist = fits.open(img_full_path, memmap=True)
-    attrs, helpers, header = _fits_header_to_xds_attrs(hdulist)
-    hdulist.close()
-    # avoid keeping reference to mem-mapped fits file
-    del hdulist
+    try:
+        hdulist = fits.open(img_full_path, memmap=True)
+        attrs, helpers, header = _fits_header_to_xds_attrs(hdulist, compute_mask)
+    finally:
+        hdulist.close()
+        # avoid keeping reference to mem-mapped fits file
+        del hdulist
     xds = _create_coords(helpers, header, do_sky_coords)
     sphr_dims = helpers["sphr_dims"]
     ary = _read_image_array(img_full_path, chunks, helpers, verbose)
@@ -83,10 +98,10 @@ def _add_freq_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
         meta["rest_frequency"] = make_quantity(helpers["restfreq"], "Hz")
         meta["rest_frequencies"] = [meta["rest_frequency"]]
         meta["type"] = "frequency"
-        meta["wave_unit"] = ["mm"]
+        meta["wave_units"] = "mm"
         freq_axis = helpers["freq_axis"]
-        meta["reference_value"] = make_frequency_reference_dict(
-            helpers["crval"][freq_axis], ["Hz"], helpers["specsys"]
+        meta["reference_frequency"] = make_spectral_coord_reference_dict(
+            helpers["crval"][freq_axis], "Hz", helpers["specsys"]
         )
         # meta["cdelt"] = helpers["cdelt"][freq_axis]
     if not meta:
@@ -99,7 +114,7 @@ def _add_freq_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
 
 def _add_vel_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
     vel_coord = xds.coords["velocity"]
-    meta = {"units": ["m/s"]}
+    meta = {"units": "m/s"}
     if helpers["has_freq"]:
         meta["doppler_type"] = helpers.get("doppler", "RADIO")
     else:
@@ -156,9 +171,7 @@ 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", "rad"], frame=ref_sys
-    )
+    direction["reference"] = make_skycoord_dict([0.0, 0.0], units="rad", frame=ref_sys)
     dir_axes = helpers["dir_axes"]
     ddata = []
     dunits = []
@@ -236,16 +249,43 @@ def _get_telescope_metadata(helpers: dict, header) -> dict:
         r = np.sqrt(np.sum(xyz * xyz))
         lat = np.arcsin(z / r)
         long = np.arctan2(y, x)
-        tel["location"] = {
+        tel["direction"] = {
             "attrs": {
                 "coordinate_system": "geocentric",
                 # I haven't seen a FITS keyword for reference frame of telescope posiiton
                 "frame": "ITRF",
                 "origin_object_name": "earth",
                 "type": "location",
-                "units": ["rad", "rad", "m"],
+                "units": "rad",
+            },
+            "data": np.array([long, lat]),
+            "dims": ["ellipsoid_dir_label"],
+            "coords": {
+                "ellipsoid_dir_label": {
+                    "dims": ["ellipsoid_dir_label"],
+                    "data": ["lon", "lat"],
+                }
+            },
+        }
+        tel["distance"] = {
+            "attrs": {
+                "coordinate_system": "geocentric",
+                # I haven't seen a FITS keyword for reference frame of telescope posiiton
+                "frame": "ITRF",
+                "origin_object_name": "earth",
+                "type": "location",
+                "units": "m",
+            },
+            "data": np.array([r]),
+            "dims": ["ellipsoid_dis_label"],
+            "coords": {
+                "ellipsoid_dis_label": {
+                    "dims": ["ellipsoid_dis_label"],
+                    "data": [
+                        "dist",
+                    ],
+                }
             },
-            "data": np.array([long, lat, r]),
         }
     return tel
 
@@ -263,9 +303,7 @@ def _compute_pointing_center(helpers: dict, header) -> dict:
     pc_lat = float(header[f"CRVAL{t_axes[1]}"]) * unit[1]
     pc_long = pc_long.to(u.rad).value
     pc_lat = pc_lat.to(u.rad).value
-    return make_skycoord_dict(
-        [pc_long, pc_lat], units=["rad", "rad"], frame=helpers["ref_sys"]
-    )
+    return make_skycoord_dict([pc_long, pc_lat], units="rad", frame=helpers["ref_sys"])
 
 
 def _user_attrs_from_header(header) -> dict:
@@ -364,12 +402,41 @@ def _create_dim_map(helpers: dict, header) -> dict:
     return dim_map
 
 
-def _fits_header_to_xds_attrs(hdulist: fits.hdu.hdulist.HDUList) -> tuple:
+def _fits_header_to_xds_attrs(
+    hdulist: fits.hdu.hdulist.HDUList, compute_mask: bool
+) -> tuple:
+    # First: Guard for unsupported compressed images
+    for i, hdu in enumerate(hdulist):
+        if isinstance(hdu, fits.CompImageHDU):
+            raise RuntimeError(
+                f"HDU {i}, name={hdu.name} is a CompImageHDU, which is not supported "
+                "for memory-mapping. "
+                "Cannot memory-map compressed FITS image (CompImageHDU). "
+                "Workaround: decompress the FITS using tools like `funpack`, `cfitsio`, "
+                "or Astropy's `.scale()`/`.copy()` workflows"
+            )
     primary = None
+    # FIXME beams is set but never actually used in this function. What's up with that?
     beams = None
     for hdu in hdulist:
         if hdu.name == "PRIMARY":
             primary = hdu
+            # Memory map support check
+            # avoid possibly non-existent hdu.scale_type attribute check and check header instead
+            header = hdu.header
+            scale = hdu.header.get("BSCALE", 1.0)
+            zero = hdu.header.get("BZERO", 0.0)
+            if not (scale == 1.0 and zero == 0.0):
+                raise RuntimeError(
+                    "Cannot memory-map scaled FITS data (BSCALE/BZERO set). "
+                    f"BZERO={zero}, BSCALE={scale}. "
+                    "Workaround: remove scaling with Astropy's"
+                    "  `HDU.data = HDU.data * BSCALE + BZERO` and save a new file"
+                )
+            # NOTE: check for primary.data size being too large removed, since
+            # data is read in chunks, so no danger of exhausting memory
+            # NOTE: sanity-check for ndarray type has been removed to avoid
+            # forcing eager memory load of possibly very large data array.
         elif hdu.name == "BEAMS":
             beams = hdu
         else:
@@ -399,13 +466,57 @@ def _fits_header_to_xds_attrs(hdulist: fits.hdu.hdulist.HDUList) -> tuple:
         raise RuntimeError("Could not find both direction axes")
     if dir_axes is not None:
         attrs["direction"] = _xds_direction_attrs_from_header(helpers, header)
-    # FIXME read fits data in chunks in case all data too large to hold in memory
-    helpers["has_mask"] = da.any(da.isnan(primary.data)).compute()
+    helpers["has_mask"] = False
+    if compute_mask:
+        # 🧠 Why the primary.data reference here is Safe (does not cause
+        # an eager read of entire data array)
+        # primary.data is a memory-mapped array (because fits.open(..., memmap=True)
+        # is used upstream)
+        # da.from_array(...) wraps this without reading it immediately
+        # The actual read occurs inside:
+        # .map_blocks(...).any().compute()
+        # ...and that triggers blockwise loading via Dask → safe and parallel
+        # 💡 Gotcha
+        # What would be dangerous:
+        # arr = np.isnan(primary.data).any()
+        # That would pull the whole array into memory. But we're not doing that.
+        data_dask = da.from_array(primary.data, chunks="auto")
+        # The following code black has corner case exposure, although the guard should
+        # eliminate it. But there is a cleaner, dask-y way that should work that we implement
+        # next, with cautions
+        # def chunk_has_nan(block):
+        #     if not isinstance(block, np.ndarray) or block.size == 0:
+        #         return False
+        #    return np.isnan(block).any()
+        # helpers["has_mask"] = data_dask.map_blocks(chunk_has_nan, dtype=bool).any().compute()
+        # ✅ Option: np.isnan(data_dask).any().compute()
+        # 🔒 Pros:
+        # Cleaner and shorter (no custom function)
+        # Handles all chunk shapes robustly — no risk of empty inputs
+        # Uses Dask’s own optimized blockwise operations under the hood
+        # ⚠️ Cons:
+        # Might trigger more eager computation if Dask can't optimize well:
+        # If chunks are misaligned or small, Dask might combine many or materialize more blocks than needed
+        # Especially on large images, it could bump memory pressure slightly
+        # But since we already call .compute(), we will load some block data no matter
+        # what — this just changes how much and how smartly.
+        # ✅ Verdict for compute_mask
+        # Because this is explicitly for computing a global has-NaN flag (not building the
+        # dataset), recommend:
+        # helpers["has_mask"] = np.isnan(data_dask).any().compute()
+        # It's concise, robust to shape edge cases, and still parallelized.
+        # We can always revisit it later if perf becomes a concern — and even then,
+        # it's likely a matter of tuning chunks= manually rather than the expression itself.
+        #
+        # This compute will normally be done in parallel
+        helpers["has_mask"] = np.isnan(data_dask).any().compute()
     beam = _beam_attr_from_header(helpers, header)
     if beam != "mb":
         helpers["beam"] = beam
     if "BITPIX" in header:
         v = abs(header["BITPIX"])
+        if v == 16:
+            helpers["dtype"] = "int16"
         if v == 32:
             helpers["dtype"] = "float32"
         elif v == 64:
@@ -487,8 +598,8 @@ def _create_coords(
                 cdelt=pick(helpers["cdelt"]),
                 cunit=pick(helpers["cunit"]),
             )
+            helpers["cunit"] = my_ret["units"]
             for j, i in enumerate(dir_axes):
-                helpers["cunit"][i] = my_ret["unit"][j]
                 helpers["crval"][i] = my_ret["ref_val"][j]
                 helpers["cdelt"][i] = my_ret["inc"][j]
             coords[my_ret["axis_name"][0]] = (["l", "m"], my_ret["value"][0])
@@ -507,32 +618,7 @@ def _get_time_values(helpers):
 
 
 def _get_pol_values(helpers):
-    # as mapped in casacore Stokes.h
-    stokes_map = [
-        "Undefined",
-        "I",
-        "Q",
-        "U",
-        "V",
-        "RR",
-        "RL",
-        "LR",
-        "LL",
-        "XX",
-        "XY",
-        "YX",
-        "YY",
-        "RX",
-        "RY",
-        "LX",
-        "LY",
-        "XR",
-        "XL",
-        "YR",
-        "YL",
-        "PP",
-        "PQ",
-    ]
+
     idx = helpers["ctype"].index("STOKES")
     if idx >= 0:
         vals = []
@@ -542,7 +628,13 @@ def _get_pol_values(helpers):
         stokes_start_idx = crval - cdelt * crpix
         for i in range(helpers["shape"][idx]):
             stokes_idx = (stokes_start_idx + i) * cdelt
-            vals.append(stokes_map[stokes_idx])
+            if 0 <= stokes_idx < len(stokes_types):
+                # stokes_types provides the index-label mapping from casacore Stokes.h
+                vals.append(stokes_types[stokes_idx])
+            else:
+                raise RuntimeError(
+                    "Can't find the Stokes type using the FITS header index"
+                )
         return vals
     else:
         return ["I"]
@@ -582,9 +674,9 @@ def _get_freq_values(helpers: dict) -> list:
         freq, vel = _freq_from_vel(
             crval, cdelt, crpix, cunit, "Z", helpers["shape"][v_idx], restfreq
         )
-        helpers["velocity"] = vel["value"] * u.Unit(vel["unit"])
-        helpers["crval"][v_idx] = (freq["crval"] * u.Unit(freq["unit"])).to(u.Hz).value
-        helpers["cdelt"][v_idx] = (freq["cdelt"] * u.Unit(freq["unit"])).to(u.Hz).value
+        helpers["velocity"] = vel["value"] * u.Unit(vel["units"])
+        helpers["crval"][v_idx] = (freq["crval"] * u.Unit(freq["units"])).to(u.Hz).value
+        helpers["cdelt"][v_idx] = (freq["cdelt"] * u.Unit(freq["units"])).to(u.Hz).value
         return list(freq["value"])
     else:
         return [1420e6]
@@ -603,6 +695,9 @@ def _get_velocity_values(helpers: dict) -> list:
         return v
 
 
+# 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:
     """Only run if we are sure there are multiple beams"""
     hdulist = fits.open(imname)
@@ -837,12 +932,12 @@ def _get_transpose_list(helpers: dict) -> tuple:
 
 def _read_image_chunk(img_full_path, shapes: tuple, starts: tuple) -> np.ndarray:
     hdulist = fits.open(img_full_path, memmap=True)
-    s = []
-    for start, length in zip(starts, shapes):
-        s.append(slice(start, start + length))
-    t = tuple(s)
-    z = hdulist[0].data[t]
+    hdu = hdulist[0]
+    # Chunk slice
+    slices = tuple(
+        slice(start, start + length) for start, length in zip(starts, shapes)
+    )
+    chunk = hdu.data[slices]
     hdulist.close()
-    # delete to avoid having a reference to a mem-mapped hdulist
     del hdulist
-    return z
+    return chunk

xradio/image/_util/casacore.py

@@ -10,7 +10,11 @@ from typing import Union
 
 import xarray as xr
 
-from casacore import tables
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
+
 from ._casacore.common import _open_image_ro
 from ._casacore.xds_from_casacore import (
     _add_mask,
@@ -42,8 +46,8 @@ def _load_casa_image_block(infile: str, block_des: dict, do_sky_coords) -> xr.Da
         cshape = casa_image.shape()
     ret = _casa_image_to_xds_coords(image_full_path, False, do_sky_coords)
     xds = ret["xds"].isel(block_des)
-    nchan = ret["xds"].dims["frequency"]
-    npol = ret["xds"].dims["polarization"]
+    nchan = ret["xds"].sizes["frequency"]
+    npol = ret["xds"].sizes["polarization"]
     starts, shapes, slices = _get_starts_shapes_slices(block_des, coords, cshape)
     dimorder = _get_xds_dim_order(ret["sphr_dims"])
     transpose_list, new_axes = _get_transpose_list(coords)
@@ -101,7 +105,7 @@ def _read_casa_image(
             xds = _add_mask(xds, m.upper(), ary, dimorder)
     xds.attrs = _casa_image_to_xds_attrs(img_full_path)
     beam = _get_beam(
-        img_full_path, xds.dims["frequency"], xds.dims["polarization"], True
+        img_full_path, xds.sizes["frequency"], xds.sizes["polarization"], True
     )
     if beam is not None:
         xds["BEAM"] = beam
@@ -129,6 +133,7 @@ def _xds_to_casa_image(xds: xr.Dataset, imagename: str) -> None:
         lockoptions={"option": "permanentwait"},
         ack=False,
     )
+
     tb.putkeyword("coords", coord)
     tb.putkeyword("imageinfo", ii)
     if units:

xradio/image/_util/common.py

@@ -110,7 +110,7 @@ def _numpy_arrayize_dv(xds: xr.Dataset) -> xr.Dataset:
 def _default_freq_info() -> dict:
     return {
         "rest_frequency": make_quantity(1420405751.7860003, "Hz"),
-        "type": "frequency",
+        "type": "spectral_coord",
         "frame": "lsrk",
         "units": "Hz",
         "waveUnit": "mm",
@@ -141,7 +141,7 @@ def _freq_from_vel(
     vel = vel * u.Unit(cunit)
     v_dict = {
         "value": vel.value,
-        "unit": cunit,
+        "units": cunit,
         "crval": crval,
         "cdelt": cdelt,
         "crpix": crpix,
@@ -154,7 +154,7 @@ def _freq_from_vel(
         fcdelt = -restfreq / _c / (crval * vel.unit / _c + 1) ** 2 * cdelt * vel.unit
         f_dict = {
             "value": freq.value,
-            "unit": "Hz",
+            "units": "Hz",
             "crval": fcrval.to(u.Hz).value,
             "cdelt": fcdelt.to(u.Hz).value,
             "crpix": crpix,
@@ -180,7 +180,7 @@ def _compute_world_sph_dims(
         "axis_name": [None, None],
         "ref_val": [None, None],
         "inc": [None, None],
-        "unit": ["rad", "rad"],
+        "units": "rad",
         "value": [None, None],
     }
     for i in range(2):

xradio/image/_util/image_factory.py

@@ -6,7 +6,7 @@ from typing import List, Union
 from .common import _c, _compute_world_sph_dims, _l_m_attr_notes
 from xradio._utils.coord_math import _deg_to_rad
 from xradio._utils.dict_helpers import (
-    make_frequency_reference_dict,
+    make_spectral_coord_reference_dict,
     make_quantity,
     make_skycoord_dict,
     make_time_coord_attrs,
@@ -50,24 +50,24 @@ def _add_common_attrs(
     cell_size: Union[List[float], np.ndarray],
     projection: str,
 ) -> xr.Dataset:
-    xds.time.attrs = make_time_coord_attrs(units=["d"], scale="utc", time_format="mjd")
+    xds.time.attrs = make_time_coord_attrs(units="d", scale="utc", time_format="mjd")
     freq_vals = np.array(xds.frequency)
     xds.frequency.attrs = {
         "observer": spectral_reference.lower(),
-        "reference_value": make_frequency_reference_dict(
+        "reference_frequency": make_spectral_coord_reference_dict(
             value=freq_vals[len(freq_vals) // 2].item(),
-            units=["Hz"],
+            units="Hz",
             observer=spectral_reference.lower(),
         ),
         "rest_frequencies": make_quantity(restfreq, "Hz"),
         "rest_frequency": make_quantity(restfreq, "Hz"),
-        "type": "frequency",
-        "units": ["Hz"],
-        "wave_unit": ["mm"],
+        "type": "spectral_coord",
+        "units": "Hz",
+        "wave_units": "mm",
     }
     xds.velocity.attrs = {"doppler_type": "radio", "type": "doppler", "units": "m/s"}
     reference = make_skycoord_dict(
-        data=phase_center, units=["rad", "rad"], frame=direction_reference
+        data=phase_center, units="rad", frame=direction_reference
     )
     reference["attrs"].update({"equinox": "j2000.0"})
     xds.attrs = {