From 07fb59c59ab69c7eed54099ff5097fa8cacebb4c Mon Sep 17 00:00:00 2001 From: Dave Mehringer Date: Wed, 26 Feb 2025 16:36:48 -0500 Subject: [PATCH] single beam -> BEAM data_var, beam, mask data_vars to upper case --- src/xradio/image/_util/__init__.py | 3 - src/xradio/image/_util/_casacore/common.py | 13 --- .../_util/_casacore/xds_from_casacore.py | 102 ++++++++---------- .../image/_util/_casacore/xds_to_casacore.py | 14 +-- src/xradio/image/_util/_fits/xds_from_fits.py | 40 ++++--- src/xradio/image/_util/casacore.py | 10 +- src/xradio/image/_util/common.py | 18 ---- tests/unit/test_image.py | 50 +++++++-- 8 files changed, 125 insertions(+), 125 deletions(-) diff --git a/src/xradio/image/_util/__init__.py b/src/xradio/image/_util/__init__.py index 216f28eb..e69de29b 100644 --- a/src/xradio/image/_util/__init__.py +++ b/src/xradio/image/_util/__init__.py @@ -1,3 +0,0 @@ -from .common import _set_multibeam_array - -__all__ = ["_set_multibeam_array"] diff --git a/src/xradio/image/_util/_casacore/common.py b/src/xradio/image/_util/_casacore/common.py index e9763860..f9c2fd78 100644 --- a/src/xradio/image/_util/_casacore/common.py +++ b/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" diff --git a/src/xradio/image/_util/_casacore/xds_from_casacore.py b/src/xradio/image/_util/_casacore/xds_from_casacore.py index cdcc8fdc..861441bf 100644 --- a/src/xradio/image/_util/_casacore/xds_from_casacore.py +++ b/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 diff --git a/src/xradio/image/_util/_casacore/xds_to_casacore.py b/src/xradio/image/_util/_casacore/xds_to_casacore.py index f891e341..8a8f8329 100644 --- a/src/xradio/image/_util/_casacore/xds_to_casacore.py +++ b/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 diff --git a/src/xradio/image/_util/_fits/xds_from_fits.py b/src/xradio/image/_util/_fits/xds_from_fits.py index 8c3fe2b9..aef776ce 100644 --- a/src/xradio/image/_util/_fits/xds_from_fits.py +++ b/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 diff --git a/src/xradio/image/_util/casacore.py b/src/xradio/image/_util/casacore.py index ef30f64f..08b014ee 100644 --- a/src/xradio/image/_util/casacore.py +++ b/src/xradio/image/_util/casacore.py @@ -55,7 +55,8 @@ 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: @@ -63,7 +64,7 @@ def _load_casa_image_block(infile: str, block_des: dict, do_sky_coords) -> xr.Da 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 diff --git a/src/xradio/image/_util/common.py b/src/xradio/image/_util/common.py index a5a8f569..59edbf46 100644 --- a/src/xradio/image/_util/common.py +++ b/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 diff --git a/tests/unit/test_image.py b/tests/unit/test_image.py index c186c3fc..0582d588 100644 --- a/tests/unit/test_image.py +++ b/tests/unit/test_image.py @@ -191,7 +191,6 @@ class xds_from_image_test(ImageBase): "projection": "SIN", } # TODO make a more interesting beam - _exp_attrs["beam"] = None _exp_attrs["obsdate"] = { "type": "time", "scale": "UTC", @@ -201,7 +200,7 @@ class xds_from_image_test(ImageBase): } _exp_attrs["observer"] = "Karl Jansky" _exp_attrs["description"] = None - _exp_attrs["active_mask"] = "mask0" + _exp_attrs["active_mask"] = "MASK0" _exp_attrs["object_name"] = "" _exp_attrs["pointing_center"] = { "value": np.array([6300, -2400]) * np.pi / 180 / 60, @@ -326,15 +325,15 @@ def compare_sky_mask(self, xds: xr.Dataset, fits=False): "Incorrect chunksize", ) self.assertEqual( - xds.mask0.chunksizes["frequency"], (5, 5), "Incorrect chunksize" + xds.MASK0.chunksizes["frequency"], (5, 5), "Incorrect chunksize" ) got_data = da.squeeze(da.transpose(xds[data_variable_name], [1, 2, 4, 3, 0]), 4) - got_mask = da.squeeze(da.transpose(xds.mask0, [1, 2, 4, 3, 0]), 4) + got_mask = da.squeeze(da.transpose(xds.MASK0, [1, 2, 4, 3, 0]), 4) if "sky_array" not in ev: im = casacore.images.image(self.imname()) ev[data_variable_name] = im.getdata() # getmask returns the negated value of the casa image mask, so True - # has the same meaning as it does in xds.mask0 + # has the same meaning as it does in xds.MASK0 ev["mask0"] = im.getmask() ev["sum"] = im.statistics()["sum"][0] if fits: @@ -347,7 +346,7 @@ def compare_sky_mask(self, xds: xr.Dataset, fits=False): (got_data == ev[data_variable_name]).all(), "pixel values incorrect" ) self.assertTrue((got_mask == ev["mask0"]).all(), "mask values incorrect") - got_ma = da.ma.masked_array(xds[data_variable_name], xds.mask0) + got_ma = da.ma.masked_array(xds[data_variable_name], xds.MASK0) self.assertEqual(da.sum(got_ma), ev["sum"], "Incorrect value for sum") self.assertTrue( got_data.dtype == ev[data_variable_name].dtype, @@ -592,7 +591,7 @@ def compare_image_block(self, imagename, zarr=False): "Wrong block SKY array", ) self.assertTrue( - (xds.mask0 == big_xds.mask0[:, 0:1, 0:4, 2:10, 3:15]).all(), + (xds.MASK0 == big_xds.MASK0[:, 0:1, 0:4, 2:10, 3:15]).all(), "Wrong block mask0 array", ) self.dict_equality( @@ -659,7 +658,7 @@ def compare_uv(self, xds: xr.Dataset, image: str) -> None: } x["npix"] = shape[3] if z == "u" else shape[2] self._expec_uv = copy.deepcopy(uv) - expec_coords = set(["time", "polarization", "frequency", "velocity", "u", "v"]) + expec_coords = set(["time", "polarization", "frequency", "velocity", "u", "v", "beam_param"]) self.assertEqual(xds.coords.keys(), expec_coords, "incorrect coordinates") for c in ["u", "v"]: attrs = self._expec_uv[c]["attrs"] @@ -786,6 +785,7 @@ class casacore_to_xds_to_casacore(xds_from_image_test): _outname4_no_sky: str = _outname4 + "_no_sky" _outname5: str = "xds_2_casa_nans_already_masked.im" _outname5_no_sky: str = _outname5 + "_no_sky" + _outname6: str = "single_beam.im" _output_uv: str = "output_uv.im" @classmethod @@ -808,6 +808,7 @@ def tearDownClass(cls): cls._outname4_no_sky, cls._outname5, cls._outname5_no_sky, + cls._outname6, cls._output_uv, ]: if os.path.exists(f): @@ -852,12 +853,15 @@ def test_metadata(self): c2["coordinates"]["spectral2"]["velUnit"] = "km/s" self.dict_equality(c2, c1, "got", "expected") - def test_multibeam(self): + def test_beam(self): """ Verify fix to issue 45 https://github.com/casangi/xradio/issues/45 + irint("*** r", r) """ download(self._imname2), f"failed to download {self._imname2}" + shutil.copytree(self._imname2, self._outname6) + # multibeam image with open_image_ro(self._imname2) as im1: beams1 = im1.imageinfo()["perplanebeams"] for do_sky, outname in zip( @@ -876,6 +880,32 @@ def test_multibeam(self): beam["positionangle"]["value"] *= 180 / np.pi beam["positionangle"]["unit"] = "deg" self.dict_equality(beams1, beams2, "got", "expected") + # convert to single beam image + tb = casacore.tables.table(self._outname6, readonly=False) + beam3 = { + 'major': {'unit': 'arcsec', 'value': 4.0}, + 'minor': {'unit': 'arcsec', 'value': 3.0}, + 'positionangle': {'unit': 'deg', 'value': 5.0} + } + tb.putkeyword( + "imageinfo", { + 'imagetype': 'Intensity', + 'objectname': '', + 'restoringbeam': beam3, + } + ) + xds = read_image(self._outname6) + self.assertFalse("beam" in xds.attrs, "beam should not be in xds.attrs") + expec = np.array([4/180/3600*np.pi, 3/180/3600*np.pi, 5*np.pi/180]) + for i, p in enumerate(["major", "minor", "pa"]): + self.assertTrue( + np.allclose( + xds.BEAM.sel(beam_param=p).values, + expec[i], + ), + f"Incorrect {p} axis", + ) + def test_masking(self): """ @@ -1172,7 +1202,7 @@ def test_multibeam(self): casa_image.tofits(self._outname1) expec = casa_image.imageinfo()["perplanebeams"] xds = read_image(self._outname1) - got = xds.data_vars["beam"] + got = xds.BEAM for p in range(4): for c in range(50): for b in ["major", "minor", "pa"]: