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

xradio/image/_util/casacore.py

@@ -9,14 +9,17 @@ import warnings
 from typing import Union
 
 import xarray as xr
+import re
+import dask.array as da
+from xradio._utils.schema import get_data_group_keys
 
 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 (
+
+from xradio.image._util._casacore.xds_from_casacore import (
     _add_mask,
     _add_sky_or_aperture,
     _casa_image_to_xds_attrs,
@@ -28,34 +31,79 @@ from ._casacore.xds_from_casacore import (
     _get_beam,
     _read_image_array,
 )
-from ._casacore.xds_to_casacore import (
+from xradio.image._util._casacore.xds_to_casacore import (
     _coord_dict_from_xds,
     _history_from_xds,
     _imageinfo_dict_from_xds,
     _write_casa_data,
 )
-from .common import _aperture_or_sky, _get_xds_dim_order, _dask_arrayize_dv
+from xradio.image._util.common import (
+    _aperture_or_sky,
+    _get_xds_dim_order,
+    _dask_arrayize_dv,
+)
+from xradio.image._util._casacore.common import _beam_fit_params, _open_image_ro
 
 warnings.filterwarnings("ignore", category=FutureWarning)
 
 
-def _load_casa_image_block(infile: str, block_des: dict, do_sky_coords) -> xr.Dataset:
+def _squeeze_if_needed(ary: da, image_type: str) -> da:
+    if image_type.upper() == "VISIBILITY_NORMALIZATION":
+        shape = ary.shape
+        if shape[3] != 1 or shape[4] != 1:
+            raise ValueError(
+                "VISIBILITY_NORMALIZATION casa image must have l and m of length 1. Found "
+                + [shape[3], shape[4]]
+            )
+        ary = ary.squeeze(axis=(3, 4))
+    return ary
+
+
+def _get_casa_image_metadata(infile: str, do_sky_coords: bool, image_type: str) -> dict:
     image_full_path = os.path.expanduser(infile)
     with _open_image_ro(image_full_path) as casa_image:
         coords = casa_image.coordinates()
         cshape = casa_image.shape()
-    ret = _casa_image_to_xds_coords(image_full_path, False, do_sky_coords)
-    xds = ret["xds"].isel(block_des)
+    ret = _casa_image_to_xds_coords(image_full_path, False, do_sky_coords, image_type)
+    xds = ret["xds"]
+    sphr_dims = ret["sphr_dims"]
     nchan = ret["xds"].sizes["frequency"]
     npol = ret["xds"].sizes["polarization"]
+    dimorder = _get_xds_dim_order(ret["sphr_dims"], image_type)
+    metadata = {
+        "coords": coords,
+        "cshape": cshape,
+        "image_full_path": image_full_path,
+        "nchan": nchan,
+        "npol": npol,
+        "dimorder": dimorder,
+        # "xds_attrs": _casa_image_to_xds_attrs(image_full_path),
+        "sphr_dims": sphr_dims,
+        "xds": xds,
+    }
+    return metadata
+
+
+def _load_casa_image_block(
+    infile: str, block_des: dict, do_sky_coords: bool, image_type: str
+) -> xr.Dataset:
+    md = _get_casa_image_metadata(infile, do_sky_coords, image_type)
+    coords = md["coords"]
+    cshape = md["cshape"]
+    dimorder = md["dimorder"]
+    sphr_dims = md["sphr_dims"]
+    nchan = md["nchan"]
+    npol = md["npol"]
+    xds = md["xds"].isel(block_des)
+    image_full_path = md["image_full_path"]
     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)
     block = _get_persistent_block(
         image_full_path, shapes, starts, dimorder, transpose_list, new_axes
     )
+    block = _squeeze_if_needed(block, image_type)
     xds = _add_sky_or_aperture(
-        xds, block, dimorder, image_full_path, ret["sphr_dims"], True
+        xds, block, dimorder, image_full_path, sphr_dims, False, image_type
     )
     mymasks = _get_mask_names(image_full_path)
     for m in mymasks:
@@ -64,58 +112,168 @@ def _load_casa_image_block(infile: str, block_des: dict, do_sky_coords) -> xr.Da
             full_path, shapes, starts, dimorder, transpose_list, new_axes
         )
         # data vars are all caps by convention
-        xds = _add_mask(xds, m.upper(), block, dimorder)
+        mask_name = re.sub(r"\bMASK(\d+)\b", r"MASK_\1", m.upper())
+        xds = _add_mask(xds, mask_name, block, dimorder)
     xds.attrs = _casa_image_to_xds_attrs(image_full_path)
-    beam = _get_beam(image_full_path, nchan, npol, False)
+    beam = _get_beam(image_full_path, nchan, npol, False, image_type)
+
     if beam is not None:
         selectors = {
             k: block_des[k]
             for k in ("time", "frequency", "polarization")
             if k in block_des
         }
-        xds["BEAM"] = beam.isel(selectors)
+        xds["BEAM_FIT_PARAMS_" + image_type.upper()] = beam.isel(selectors)
+        xds["BEAM_FIT_PARAMS_" + image_type.upper()].attrs["type"] = (
+            "beam_fit_params_" + image_type.lower()
+        )
+        xds[image_type.upper()].attrs[_beam_fit_params] = (
+            "BEAM_FIT_PARAMS_" + image_type.upper()
+        )
     return xds
 
 
-def _read_casa_image(
+def _open_casa_image(
     infile: str,
     chunks: Union[list, dict],
     verbose: bool,
     do_sky_coords: bool,
     masks: bool = True,
-    history: bool = True,
+    history: bool = False,
+    image_type: str = "SKY",
 ) -> xr.Dataset:
-    img_full_path = os.path.expanduser(infile)
-    ret = _casa_image_to_xds_coords(img_full_path, verbose, do_sky_coords)
-    xds = ret["xds"]
-    dimorder = _get_xds_dim_order(ret["sphr_dims"])
+    md = _get_casa_image_metadata(infile, do_sky_coords, image_type)
+    xds = md["xds"]
+    dimorder = md["dimorder"]
+    sphr_dims = md["sphr_dims"]
+    img_full_path = md["image_full_path"]
+    ary = _read_image_array(img_full_path, chunks, verbose=verbose)
+    ary = _squeeze_if_needed(ary, image_type)
     xds = _add_sky_or_aperture(
         xds,
-        _read_image_array(img_full_path, chunks, verbose=verbose),
+        ary,
         dimorder,
         img_full_path,
-        ret["sphr_dims"],
+        sphr_dims,
         history,
+        image_type,
     )
     if masks:
         mymasks = _get_mask_names(img_full_path)
         for m in mymasks:
             ary = _read_image_array(img_full_path, chunks, mask=m, verbose=verbose)
             # data var names are all caps by convention
-            xds = _add_mask(xds, m.upper(), ary, dimorder)
+            mask_name = re.sub(r"\bMASK(\d+)\b", r"MASK_\1", m.upper())
+            xds = _add_mask(xds, mask_name, ary, dimorder)
     xds.attrs = _casa_image_to_xds_attrs(img_full_path)
     beam = _get_beam(
-        img_full_path, xds.sizes["frequency"], xds.sizes["polarization"], True
+        img_full_path,
+        xds.sizes["frequency"],
+        xds.sizes["polarization"],
+        True,
+        image_type,
     )
     if beam is not None:
-        xds["BEAM"] = beam
+        xds["BEAM_FIT_PARAMS_" + image_type.upper()] = beam
+        xds["BEAM_FIT_PARAMS_" + image_type.upper()].attrs["type"] = (
+            "beam_fit_params_" + image_type.lower()
+        )
+        xds[image_type.upper()].attrs[_beam_fit_params] = (
+            "BEAM_FIT_PARAMS_" + image_type.upper()
+        )
+
     # xds = _add_coord_attrs(xds, ret["icoords"], ret["dir_axes"])
     xds = _dask_arrayize_dv(xds)
+
     return xds
 
 
-def _xds_to_casa_image(xds: xr.Dataset, imagename: str) -> None:
-    image_full_path = os.path.expanduser(imagename)
+def _xds_to_multiple_casa_images(xds: xr.Dataset, image_store_name: str) -> None:
+    """Function disentagles xradio xr.Dataset into multiple casa images based on data_groups attribute.
+    An xr.Dataset may contain multiple images (sky, residual, psf, etc) stored under different data variables sharing common coordinates.
+    An addtional complication is that CASA images allow for internal masks and beam fit parameters to be stored alongside the main image data so these also need to be handled.
+    This function creates separate casa images for each image type found in the data_groups attribute of the xr.Dataset.
+
+    Parameters
+    ----------
+    xds : xr.Dataset
+        The xradio xr.Dataset containing multiple images and associated data.
+    image_store_name : str
+        The base name or path for storing the output CASA images.
+        If only one image is written, it will be named image_store_name, esle the images
+        will be named image_store_name.<image_type> where <image_type> is sky, residual, point_spread_function, etc.
+    """
+
+    data_vars_name_set = set(xds.data_vars.keys())
+
+    data_group_keys = list(get_data_group_keys(schema_name="image").keys())
+    internal_image_types_to_exclude = [
+        "flag",
+        "beam_fit_params_sky",
+        "beam_fit_params_point_spread_function",
+    ]
+    n_image_written = 0
+    last_image_written = ""
+    for data_group in xds.attrs["data_groups"].keys():
+        for image_type in data_group_keys:
+            if (image_type in xds.attrs["data_groups"][data_group]) and (
+                image_type not in internal_image_types_to_exclude
+            ):
+                image_name = xds.attrs["data_groups"][data_group][image_type]
+                if image_name in data_vars_name_set:
+                    image_to_write_xds = xr.Dataset()
+                    image_to_write_xds.attrs = xds.attrs.copy()
+
+                    if image_type == "aperture":
+                        image_to_write_xds["APERTURE"] = xds[image_name]
+                    else:
+                        image_to_write_xds["SKY"] = xds[image_name]
+
+                    # This code handles adding internal masks and beam fit params if they exist.
+                    if image_type == "sky":
+                        if (
+                            "beam_fit_params_sky"
+                            in xds.attrs["data_groups"][data_group]
+                        ):
+                            beam_fit_params_name = xds.attrs["data_groups"][data_group][
+                                "beam_fit_params_sky"
+                            ]
+                            image_to_write_xds["BEAM_FIT_PARAMS"] = xds[
+                                beam_fit_params_name
+                            ]
+
+                        if "flag" in xds.attrs["data_groups"][data_group]:
+                            mask_sky_name = xds.attrs["data_groups"][data_group]["flag"]
+                            image_to_write_xds["MASK_0"] = xds[mask_sky_name]
+                            image_to_write_xds["SKY"].attrs["flag"] = "MASK_0"
+
+                    if image_type == "point_spread_function":
+                        if (
+                            "beam_fit_params_point_spread_function"
+                            in xds.attrs["data_groups"][data_group]
+                        ):
+                            beam_fit_params_name = xds.attrs["data_groups"][data_group][
+                                "beam_fit_params_point_spread_function"
+                            ]
+                            image_to_write_xds["BEAM_FIT_PARAMS"] = xds[
+                                beam_fit_params_name
+                            ]
+                    outname = image_store_name + "." + image_type
+                    _xds_to_casa_image(image_to_write_xds, outname)
+                    if not os.path.exists(outname):
+                        raise IOError(f"Failed to write CASA image {outname}")
+                    n_image_written += 1
+                    last_image_written = outname
+                    data_vars_name_set.remove(image_name)
+    if n_image_written == 0:
+        raise ValueError("No valid image types found in xds to write to CASA images.")
+    if n_image_written == 1:
+        # rename the single written image to what the user requested
+        os.rename(last_image_written, image_store_name)
+
+
+def _xds_to_casa_image(xds: xr.Dataset, image_store_name: str) -> None:
+    image_full_path = os.path.expanduser(image_store_name)
     _write_casa_data(xds, image_full_path)
     # create coordinates
     ap_sky = _aperture_or_sky(xds)

xradio/image/_util/common.py

@@ -7,6 +7,7 @@ from typing import Dict, List
 import xarray as xr
 from xradio._utils.coord_math import _deg_to_rad
 from xradio._utils.dict_helpers import make_quantity
+import toolviper.utils.logger as logger
 
 _c = 2.99792458e08 * u.m / u.s
 # OPTICAL = Z
@@ -17,17 +18,18 @@ _doppler_types = [
     "beta",
     "gamma",
 ]
-_image_type = "image_type"
+_image_type = "type"
 
 
 def _aperture_or_sky(xds: xr.Dataset) -> str:
     return "SKY" if "SKY" in xds.data_vars or "l" in xds.coords else "APERTURE"
 
 
-def _get_xds_dim_order(has_sph: bool) -> list:
+def _get_xds_dim_order(has_sph: bool, image_type: str) -> list:
     dimorder = ["time", "frequency", "polarization"]
-    dir_lin = ["l", "m"] if has_sph else ["u", "v"]
-    dimorder.extend(dir_lin)
+    if image_type.upper() != "VISIBILITY_NORMALIZATION":
+        dir_lin = ["l", "m"] if has_sph else ["u", "v"]
+        dimorder.extend(dir_lin)
     return dimorder
 
 
@@ -252,10 +254,10 @@ def _compute_sky_reference_pixel(xds: xr.Dataset) -> np.ndarray:
 
 def _l_m_attr_notes() -> Dict[str, str]:
     return {
-        "l": "l is the angle measured from the phase center to the east. "
-        "So l = x*cdelt, where x is the number of pixels from the phase center. "
+        "l": "l is the angle measured from the reference direction to the east. "
+        "So l = x*cdelt, where x is the number of pixels from the reference direction. "
         "See AIPS Memo #27, Section III.",
-        "m": "m is the angle measured from the phase center to the north. "
-        "So m = y*cdelt, where y is the number of pixels from the phase center. "
+        "m": "m is the angle measured from the reference direction to the north. "
+        "So m = y*cdelt, where y is the number of pixels from the reference direction. "
         "See AIPS Memo #27, Section III.",
     }