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

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

xradio/image/_util/_fits/xds_from_fits.py

@@ -16,7 +16,11 @@ from ..common import (
     _l_m_attr_notes,
 )
 from xradio._utils.coord_math import _deg_to_rad
-from xradio._utils.dict_helpers import make_quantity
+from xradio._utils.dict_helpers import (
+    make_quantity,
+    make_frequency_reference_dict,
+    make_skycoord_dict,
+)
 import copy
 import dask
 import dask.array as da
@@ -50,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
 
 
@@ -66,6 +72,7 @@ def _add_time_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
     time_coord = xds.coords["time"]
     meta = copy.deepcopy(helpers["obsdate"])
     del meta["value"]
+    # meta["units"] = [ meta["units"] ]
     # meta['format'] = 'MJD'
     # meta['time_scale'] = meta['refer']
     # del meta['refer']
@@ -79,13 +86,16 @@ def _add_freq_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
     meta = {}
     if helpers["has_freq"]:
         meta["rest_frequency"] = make_quantity(helpers["restfreq"], "Hz")
-        meta["frame"] = helpers["specsys"]
-        meta["units"] = "Hz"
+        meta["rest_frequencies"] = [meta["rest_frequency"]]
+        # meta["frame"] = helpers["specsys"]
+        # meta["units"] = "Hz"
         meta["type"] = "frequency"
         meta["wave_unit"] = "mm"
         freq_axis = helpers["freq_axis"]
-        meta["crval"] = helpers["crval"][freq_axis]
-        meta["cdelt"] = helpers["cdelt"][freq_axis]
+        meta["reference_value"] = make_frequency_reference_dict(
+            helpers["crval"][freq_axis], ["Hz"], helpers["specsys"]
+        )
+        # meta["cdelt"] = helpers["cdelt"][freq_axis]
     if not meta:
         # this is the default frequency information CASA creates
         meta = _default_freq_info()
@@ -96,7 +106,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:
@@ -112,10 +122,6 @@ def _add_l_m_attrs(xds: xr.Dataset, helpers: dict) -> xr.Dataset:
     for c in ["l", "m"]:
         if c in xds.coords:
             xds[c].attrs = {
-                "crval": 0.0,
-                "cdelt": helpers[c]["cdelt"],
-                "units": "rad",
-                "type": "quantity",
                 "note": attr_note[c],
             }
     return xds
@@ -167,6 +173,10 @@ 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"] = {
         "type": "sky_coord",
         "frame": ref_sys,
@@ -174,20 +184,39 @@ def _xds_direction_attrs_from_header(helpers: dict, header) -> dict:
         "units": ["rad", "rad"],
         "value": [0.0, 0.0],
     }
+    """
     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")
-        direction["reference"]["value"][i] = x.value
+        ddata.append(x.value)
+        # direction["reference"]["value"][i] = x.value
         x = helpers["cdelt"][i] * u.Unit(_get_unit(helpers["cunit"][i]))
-        x = x.to("rad")
-    direction["latpole"] = make_quantity(header["LATPOLE"] * _deg_to_rad, "rad")
-    direction["longpole"] = make_quantity(header["LONPOLE"] * _deg_to_rad, "rad")
+        dunits.append(x.to("rad"))
+    direction["reference"] = make_skycoord_dict(ddata, units=dunits, frame=ref_sys)
+    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
-            pc[i][j] = header[f"PC{dir_axes[i]+1}_{dir_axes[j]+1}"]
+            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])
@@ -314,9 +343,9 @@ def _beam_attr_from_header(helpers: dict, header) -> Union[dict, str, None]:
     if "BMAJ" in header:
         # single global beam
         beam = {
-            "bmaj": make_quantity(header["BMAJ"], "arcsec"),
-            "bmin": make_quantity(header["BMIN"], "arcsec"),
-            "pa": make_quantity(header["BPA"], "arcsec"),
+            "bmaj": make_quantity(header["BMAJ"], "deg"),
+            "bmin": make_quantity(header["BMIN"], "deg"),
+            "pa": make_quantity(header["BPA"], "deg"),
         }
         return _convert_beam_to_rad(beam)
     elif "CASAMBM" in header and header["CASAMBM"]:
@@ -390,11 +419,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:
@@ -409,7 +438,7 @@ def _fits_header_to_xds_attrs(hdulist: fits.hdu.hdulist.HDUList) -> dict:
     obsdate = {}
     obsdate["type"] = "time"
     obsdate["value"] = Time(header["DATE-OBS"], format="isot").mjd
-    obsdate["units"] = "d"
+    obsdate["units"] = ["d"]
     obsdate["scale"] = header["TIMESYS"]
     obsdate["format"] = "MJD"
     attrs["obsdate"] = obsdate
@@ -453,8 +482,8 @@ def _create_coords(
     helpers["sphr_dims"] = sphr_dims
     coords = {}
     coords["time"] = _get_time_values(helpers)
-    coords["polarization"] = _get_pol_values(helpers)
     coords["frequency"] = _get_freq_values(helpers)
+    coords["polarization"] = _get_pol_values(helpers)
     coords["velocity"] = (["frequency"], _get_velocity_values(helpers))
     if len(sphr_dims) > 0:
         for i, c in enumerate(["l", "m"]):
@@ -497,6 +526,7 @@ def _create_coords(
     else:
         # Fourier image
         coords["u"], coords["v"] = _get_uv_values(helpers)
+    coords["beam_param"] = ["major", "minor", "pa"]
     xds = xr.Dataset(coords=coords)
     return xds
 
@@ -615,28 +645,43 @@ def _do_multibeam(xds: xr.Dataset, imname: str) -> xr.Dataset:
             )
             nchan = header["NCHAN"]
             npol = header["NPOL"]
-            beam_array = np.zeros([1, npol, nchan, 3])
+            beam_array = np.zeros([1, nchan, npol, 3])
             data = hdu.data
+            hdulist.close()
             for t in data:
-                beam_array[0, t[4], t[3]] = t[0:3]
+                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", "polarization", "frequency", "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"]["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)
+
+
+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"]
@@ -680,8 +725,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
 
 
@@ -790,10 +835,10 @@ def _get_transpose_list(helpers: dict) -> tuple:
             or b.startswith("vopt")
             or b.startswith("vrad")
         ):
-            transpose_list[2] = i
+            transpose_list[1] = i
             not_covered.remove("f")
         elif b.startswith("stok"):
-            transpose_list[1] = i
+            transpose_list[2] = i
             not_covered.remove("s")
         else:
             raise RuntimeError(f"Unhandled axis name {c}")

xradio/image/_util/_zarr/zarr_low_level.py

@@ -12,9 +12,9 @@ from numcodecs.compat import (
     ensure_contiguous_ndarray_like,
 )
 
-full_dims_lm = ["time", "polarization", "frequency", "l", "m"]
-full_dims_uv = ["time", "polarization", "frequency", "l", "m"]
-norm_dims = ["polarization", "frequency"]
+full_dims_lm = ["time", "frequency", "polarization", "l", "m"]
+full_dims_uv = ["time", "frequency", "polarization", "l", "m"]
+norm_dims = ["frequency", "polarization"]
 
 image_data_variables_and_dims_double_precision = {
     "aperture": {"dims": full_dims_uv, "dtype": "<c16", "name": "APERTURE"},

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", "polarization", "frequency"):
+        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

xradio/image/_util/common.py

@@ -10,7 +10,13 @@ from xradio._utils.dict_helpers import make_quantity
 
 _c = 2.99792458e08 * u.m / u.s
 # OPTICAL = Z
-_doppler_types = ["RADIO", "Z", "RATIO", "BETA", "GAMMA"]
+_doppler_types = [
+    "radio",
+    "z",
+    "ratio",
+    "beta",
+    "gamma",
+]
 _image_type = "image_type"
 
 
@@ -19,7 +25,7 @@ def _aperture_or_sky(xds: xr.Dataset) -> str:
 
 
 def _get_xds_dim_order(has_sph: bool) -> list:
-    dimorder = ["time", "polarization", "frequency"]
+    dimorder = ["time", "frequency", "polarization"]
     dir_lin = ["l", "m"] if has_sph else ["u", "v"]
     dimorder.extend(dir_lin)
     return dimorder
@@ -34,10 +40,10 @@ def _convert_beam_to_rad(beam: dict) -> dict:
     """
     mybeam = {}
     for k in beam:
-        try:
-            q = u.quantity.Quantity(f"{beam[k]['value']}{beam[k]['unit']}")
-        except:
-            q = u.quantity.Quantity(f"{beam[k]['value']}{beam[k]['units']}")
+        myu = beam[k]["attrs"]["units"]
+        myu = myu[0] if isinstance(myu, list) else myu
+        units = _get_unit(myu)
+        q = u.quantity.Quantity(f"{beam[k]['data']}{units}")
         q = q.to("rad")
         j = "pa" if k == "positionangle" else k
         mybeam[j] = make_quantity(q.value, "rad")
@@ -105,7 +111,7 @@ def _default_freq_info() -> dict:
     return {
         "rest_frequency": make_quantity(1420405751.7860003, "Hz"),
         "type": "frequency",
-        "frame": "LSRK",
+        "frame": "lsrk",
         "units": "Hz",
         "waveUnit": "mm",
         "cdelt": 1000.0,
@@ -253,22 +259,3 @@ def _l_m_attr_notes() -> Dict[str, str]:
         "So m = y*cdelt, where y is the number of pixels from the phase center. "
         "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", "polarization", "frequency", "beam_param"]
-    )
-    xdb = xdb.rename("beam")
-    xdb = xdb.assign_coords(beam_param=["major", "minor", "pa"])
-    xdb.attrs["units"] = units
-    xds["beam"] = xdb
-    return xds