single beam -> BEAM data_var, beam, mask data_vars to upper case

src/xradio/image/_util/__init__.py

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

src/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"

src/xradio/image/_util/_casacore/xds_from_casacore.py

@@ -37,16 +37,6 @@
     make_skycoord_dict,
 )
 
-"""
-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
-"""
-
-
 def _add_lin_attrs(xds, coord_dict, dir_axes):
     for k in coord_dict:
         if k.startswith("linear"):
@@ -237,12 +227,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
         )
@@ -314,10 +304,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:
@@ -394,14 +380,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:
@@ -595,23 +573,32 @@ 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"]["data"]
-            beam_array[0][p][c][1] = beam_dict["minor"]["data"]
-            beam_array[0][p][c][2] = beam_dict["pa"]["data"]
+    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 = 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":
+        beam_dict = _convert_beam_to_rad(beam)
+        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
 
 
@@ -823,25 +810,22 @@ 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
 

src/xradio/image/_util/_casacore/xds_to_casacore.py

@@ -182,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]},
@@ -204,6 +204,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
         ii["restoringbeam"] = copy.deepcopy(xds.attrs["beam"])
@@ -216,6 +217,7 @@ def _imageinfo_dict_from_xds(xds: xr.Dataset) -> dict:
             del ii["restoringbeam"][k]["data"]
         ii["restoringbeam"]["positionangle"] = copy.deepcopy(ii["restoringbeam"]["pa"])
         del ii["restoringbeam"]["pa"]
+    """
     return ii
 
 

src/xradio/image/_util/_fits/xds_from_fits.py

@@ -54,6 +54,8 @@ def _fits_image_to_xds(
     xds = _add_coord_attrs(xds, helpers)
     if helpers["has_multibeam"]:
         xds = _do_multibeam(xds, img_full_path)
+    elif "beam" in helpers and helpers["beam"] is not None:
+        xds = _add_beam(xds, helpers)
     return xds
 
 
@@ -408,11 +410,11 @@ def _fits_header_to_xds_attrs(hdulist: fits.hdu.hdulist.HDUList) -> dict:
         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
     has_mask = da.any(da.isnan(primary.data)).compute()
-    attrs["active_mask"] = "mask0" if has_mask else None
+    attrs["active_mask"] = "MASK0" if has_mask else None
     helpers["has_mask"] = has_mask
     beam = _beam_attr_from_header(helpers, header)
     if beam != "mb":
-        attrs["beam"] = beam
+        helpers["beam"] = beam
     if "BITPIX" in header:
         v = abs(header["BITPIX"])
         if v == 32:
@@ -635,26 +637,40 @@ def _do_multibeam(xds: xr.Dataset, imname: str) -> xr.Dataset:
             npol = header["NPOL"]
             beam_array = np.zeros([1, nchan, npol, 3])
             data = hdu.data
+            hdulist.close()
             for t in data:
                 beam_array[0, t[3], t[4]] = t[0:3]
             for i in (0, 1, 2):
                 beam_array[:, :, :, i] = (
                     (beam_array[:, :, :, i] * units[i]).to("rad").value
                 )
-            xdb = xr.DataArray(
-                beam_array, dims=["time", "frequency", "polarization", "beam_param"]
-            )
-            xdb = xdb.rename("beam")
-            xdb = xdb.assign_coords(beam_param=["major", "minor", "pa"])
-            xdb.attrs["units"] = "rad"
-            xds["beam"] = xdb
-            return xds
+            return _create_beam_data_var(xds, beam_array)
     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:
+    nchan = xds.sizes["frequency"]
+    npol = xds.sizes["polarization"]
+    beam_array = np.zeros([1, nchan, npol, 3])
+    beam_array[0, :, :, 0] = helpers["beam"]["bmaj"]
+    beam_array[0, :, :, 1] = helpers["beam"]["bmin"]
+    beam_array[0, :, :, 2] = helpers["beam"]["pa"]
+    return _chreate_beam_data_var(xds, beam_array)
+
+
+def _create_beam_data_var(xds: xr.Dataset, beam_array: np.array) -> xr.Dataset:
+    xdb = xr.DataArray(
+        beam_array, dims=["time", "frequency", "polarization", "beam_param"]
+    )
+    xdb = xdb.rename("BEAM")
+    xdb = xdb.assign_coords(beam_param=["major", "minor", "pa"])
+    xdb.attrs["units"] = "rad"
+    xds["BEAM"] = xdb
+    return xds
+
 def _get_uv_values(helpers: dict) -> tuple:
     shape = helpers["shape"]
     ctype = helpers["ctype"]
@@ -698,8 +714,8 @@ def _add_sky_or_aperture(
         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
+        mask = mask.rename("MASK0")
+        xds["MASK0"] = mask
     return xds
 
 

src/xradio/image/_util/casacore.py

@@ -55,15 +55,16 @@ def _load_casa_image_block(infile: str, block_des: dict, do_sky_coords) -> xr.Da
         block = _get_persistent_block(
             full_path, shapes, starts, dimorder, transpose_list, new_axes
         )
-        xds = _add_mask(xds, m, block, dimorder)
+        # data vars are all caps by convention
+        xds = _add_mask(xds, m.upper(), block, dimorder)
     xds.attrs = _casa_image_to_xds_attrs(image_full_path, True)
     mb = _multibeam_array(xds, image_full_path, False)
     if mb is not None:
         selectors = {}
         for k in ("time", "frequency", "polarization"):
             if k in block_des:
                 selectors[k] = block_des[k]
-        xds["beam"] = mb.isel(selectors)
+        xds["BEAM"] = mb.isel(selectors)
     return xds
 
 
@@ -90,11 +91,12 @@ def _read_casa_image(
         mymasks = _get_mask_names(img_full_path)
         for m in mymasks:
             ary = _read_image_array(img_full_path, chunks, mask=m, verbose=verbose)
-            xds = _add_mask(xds, m, ary, dimorder)
+            # data var names are all caps by convention
+            xds = _add_mask(xds, m.upper(), ary, dimorder)
     xds.attrs = _casa_image_to_xds_attrs(img_full_path, history)
     mb = _multibeam_array(xds, img_full_path, True)
     if mb is not None:
-        xds["beam"] = mb
+        xds["BEAM"] = mb
     # xds = _add_coord_attrs(xds, ret["icoords"], ret["dir_axes"])
     xds = _dask_arrayize_dv(xds)
     return xds

src/xradio/image/_util/common.py

@@ -260,21 +260,3 @@ def _l_m_attr_notes() -> Dict[str, str]:
         "See AIPS Memo #27, Section III.",
     }
 
-
-def _set_multibeam_array(xds, beam_ary, units):
-    if "beam" in xds.attrs:
-        if xds.attrs["beam"] is None:
-            del xds.attrs["beam"]
-        else:
-            raise RuntimeError(
-                "Error: xds has beam attr. It must be removed "
-                "before multiple beams can be added"
-            )
-    xdb = xr.DataArray(
-        beam_ary, dims=["time", "frequency", "polarization", "beam_param"]
-    )
-    xdb = xdb.rename("beam")
-    xdb = xdb.assign_coords(beam_param=["major", "minor", "pa"])
-    xdb.attrs["units"] = units
-    xds["beam"] = xdb
-    return xds