xradio 0.0.54__py3-none-any.whl → 0.0.56__py3-none-any.whl

xradio/_utils/_casacore/tables.py

@@ -1,4 +1,8 @@
-from casacore import tables
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
+
 from contextlib import contextmanager
 from typing import Dict, Generator
 

xradio/image/_util/_casacore/common.py

@@ -1,7 +1,10 @@
-from casacore import images
+try:
+    from casacore import images
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as images
+
 from contextlib import contextmanager
-import numpy as np
-from typing import Dict, Generator, List, Union
+from typing import Generator, List
 
 
 @contextmanager
@@ -12,6 +15,7 @@ def _open_image_ro(infile: str) -> Generator[images.image, None, None]:
     finally:
         # there is no obvious way to close a python-casacore image, so
         # just delete the object to clear it from the table cache
+        # image.unlock() # not necessary
         del image
 
 
@@ -31,6 +35,10 @@ def _create_new_image(
     try:
         yield image
     finally:
+        # Explicitly calling unlock() is important for downstream parallel reads across multiple processes.
+        # Without it, Dask workers may mistakenly consider a freshly written image from another process
+        # as invalid—even if the image directory appears to be fully written and present.
+        image.unlock()
         del image
 
 

xradio/image/_util/_casacore/xds_from_casacore.py

@@ -9,8 +9,15 @@ import toolviper.utils.logger as logger
 import numpy as np
 import xarray as xr
 from astropy import units as u
-from casacore import tables
-from casacore.images import coordinates, image as casa_image
+
+try:
+    from casacore import tables
+    from casacore.images import coordinates, image as casa_image
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
+    import xradio._utils._casacore.casacore_from_casatools as coordinates
+    from xradio._utils._casacore.casacore_from_casatools import image as casa_image
+
 
 from .common import (
     _active_mask,
@@ -18,6 +25,7 @@ from .common import (
     _open_image_ro,
     _pointing_center,
 )
+
 from ..common import (
     _compute_linear_world_values,
     _compute_velocity_values,
@@ -599,36 +607,6 @@ def _get_mask_names(infile: str) -> list:
     return mymasks
 
 
-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[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 = _casacore_q_to_xradio_q(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":
-        b = _casacore_q_to_xradio_q(beam)
-        beam_dict = _convert_beam_to_rad(b)
-        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
-
-
 def _get_persistent_block(
     infile: str,
     shapes: tuple,
@@ -834,25 +812,46 @@ def _get_velocity_values_attrs(
     )
 
 
-def _multibeam_array(
-    xds: xr.Dataset, img_full_path: str, as_dask_array: bool
+def _get_beam(
+    img_full_path: str, nchan: int, npol: int, as_dask_array: bool
 ) -> Union[xr.DataArray, None]:
     # 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:
+    x = ["perplanebeams", "restoringbeam"]
+    r = None
+    for z in x:
+        if z in imageinfo:
+            r = z
+            break
+    if r is None:
         return None
+    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 = _casacore_q_to_xradio_q(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":
+        b = _casacore_q_to_xradio_q(beam)
+        beam_dict = _convert_beam_to_rad(b)
+        beam_array[0, :, :, 0] = beam_dict["major"]["data"]
+        beam_array[0, :, :, 1] = beam_dict["minor"]["data"]
+        beam_array[0, :, :, 2] = beam_dict["pa"]["data"]
+    if as_dask_array:
+        beam_array = da.array(beam_array)
+    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"
+    return xdb
 
 
 ###########################################################################

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
@@ -197,7 +201,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"] = (

xradio/image/_util/_fits/xds_from_fits.py

@@ -1,8 +1,25 @@
-import astropy as ap
+import copy
+import re
+from typing import Union
+
+import dask
+import dask.array as da
+import numpy as np
+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_frequency_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,20 +32,6 @@ 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(
@@ -507,32 +510,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 +520,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"]

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,
@@ -21,7 +25,7 @@ from ._casacore.xds_from_casacore import (
     _get_persistent_block,
     _get_starts_shapes_slices,
     _get_transpose_list,
-    _multibeam_array,
+    _get_beam,
     _read_image_array,
 )
 from ._casacore.xds_to_casacore import (
@@ -42,6 +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"]
     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)
@@ -60,13 +66,14 @@ def _load_casa_image_block(infile: str, block_des: dict, do_sky_coords) -> xr.Da
         # 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)
-    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)
+    beam = _get_beam(image_full_path, nchan, npol, False)
+    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)
     return xds
 
 
@@ -97,9 +104,11 @@ def _read_casa_image(
             # 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)
-    mb = _multibeam_array(xds, img_full_path, True)
-    if mb is not None:
-        xds["BEAM"] = mb
+    beam = _get_beam(
+        img_full_path, xds.dims["frequency"], xds.dims["polarization"], True
+    )
+    if beam is not None:
+        xds["BEAM"] = beam
     # xds = _add_coord_attrs(xds, ret["icoords"], ret["dir_axes"])
     xds = _dask_arrayize_dv(xds)
     return xds

xradio/image/image.py

@@ -159,7 +159,7 @@ def load_image(infile: str, block_des: dict = {}, do_sky_coords=True) -> xr.Data
         do_casa = False
     if do_casa:
         # comment next line when done debugging
-        # return _load_casa_image_block(infile, bd, do_sky_coords)
+        # return _load_casa_image_block(infile, selection, do_sky_coords)
         try:
             return _load_casa_image_block(infile, selection, do_sky_coords)
         except Exception as e:

xradio/measurement_set/_utils/_msv2/_tables/load.py

@@ -2,7 +2,10 @@ from typing import Tuple
 
 import numpy as np
 
-from casacore import tables
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
 from ....._utils.common import get_pad_value
 
 

xradio/measurement_set/_utils/_msv2/_tables/load_main_table.py

@@ -5,7 +5,10 @@ import pandas as pd
 import numpy as np
 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 .load import load_col_chunk
 from .read_main_table import get_partition_ids, redim_id_data_vars, rename_vars

xradio/measurement_set/_utils/_msv2/_tables/read.py

@@ -10,7 +10,11 @@ import pandas as pd
 import xarray as xr
 
 import astropy.units
-from casacore import tables
+
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
 
 from .table_query import open_query, open_table_ro, TableManager
 from xradio._utils.list_and_array import get_pad_value
@@ -308,6 +312,8 @@ def add_units_measures(
                 ):  # Little fix for Meerkat data where the units are a string.
                     cc_units = [cc_units]
 
+                if isinstance(cc_units, np.ndarray):
+                    cc_units = cc_units.tolist()
                 if not isinstance(cc_units, list) or not cc_units:
                     logger.warning(
                         f"Invalid units found for column/variable {col}: {cc_units}"
@@ -1250,22 +1256,20 @@ def read_col_conversion_numpy(
         # Use casacore to get the shape of a row for this column
         #################################################################################
 
-        # getcolshapestring() only works on columns where a row element is an
-        # array ie. fails for TIME
-        # Assumes the RuntimeError is because the column is a scalar
-        try:
+        # getcolshapestring() only works for array-valued columns.
+        # For scalar columns (e.g., EXPOSURE, TIME_CENTROID), it raises a RuntimeError.
+        # So we first check if the column is scalar to avoid that.
+        if tb_tool.isscalarcol(col):
+            extra_dimensions = ()
+        else:
+            # Get the shape string for the first row of the column (e.g., "[4, 2]")
             shape_string = tb_tool.getcolshapestring(col)[0]
-            # Convert `shape_string` into a tuple that numpy understands
+
+            # Convert the shape string into a tuple of integers (e.g., (4, 2)) that numpy
+            # understands.
             extra_dimensions = tuple(
-                [
-                    int(idx)
-                    for idx in shape_string.replace("[", "")
-                    .replace("]", "")
-                    .split(", ")
-                ]
+                int(dim) for dim in shape_string.strip("[]").split(", ")
             )
-        except RuntimeError:
-            extra_dimensions = ()
 
         #################################################################################
 

xradio/measurement_set/_utils/_msv2/_tables/read_main_table.py

@@ -6,7 +6,10 @@ import numpy as np
 import xarray as xr
 import pandas as pd
 
-from casacore import tables
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
 
 from .read import (
     read_flat_col_chunk,

xradio/measurement_set/_utils/_msv2/_tables/read_subtables.py

@@ -7,7 +7,10 @@ import numpy as np
 import pandas as pd
 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 .table_query import open_query, open_table_ro
 from .read import (

xradio/measurement_set/_utils/_msv2/_tables/table_query.py

@@ -1,7 +1,10 @@
 from typing import Generator
 from contextlib import contextmanager
 
-from casacore import tables
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
 
 
 @contextmanager
@@ -17,7 +20,11 @@ def open_table_ro(infile: str) -> Generator[tables.table, None, None]:
 
 @contextmanager
 def open_query(table: tables.table, query: str) -> Generator[tables.table, None, None]:
-    ttq = tables.taql(query)
+
+    if hasattr(tables, "taql"):
+        ttq = tables.taql(query)
+    else:
+        ttq = table.taql(query)
     try:
         yield ttq
     finally:
@@ -43,4 +50,7 @@ class TableManager:
             self.infile, readonly=True, lockoptions={"option": "usernoread"}, ack=False
         ) as mtable:
             query = f"select * from $mtable {self.taql_where}"
-            return tables.taql(query)
+            if hasattr(tables, "taql"):
+                return tables.taql(query)
+            else:
+                return mtable.taql(query)

xradio/measurement_set/_utils/_msv2/_tables/write.py

@@ -4,7 +4,10 @@ from typing import Tuple
 import numpy as np
 import xarray as xr
 
-from casacore import tables
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
 
 
 def revert_time(datetimes: np.ndarray) -> np.ndarray:

xradio/measurement_set/_utils/_msv2/_tables/write_exp_api.py

@@ -9,7 +9,10 @@ from ..._utils.xds_helper import flatten_xds, calc_optimal_ms_chunk_shape
 from .write import write_generic_table, write_main_table_slice
 from .write import create_table, revert_time
 
-from casacore import tables
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
 
 
 # TODO: this should be consolidated with the equivalent in read_main_table,

xradio/measurement_set/_utils/_msv2/conversion.py

@@ -8,9 +8,14 @@ from typing import Dict, Union
 
 import numpy as np
 import xarray as xr
+import traceback
 
 import toolviper.utils.logger as logger
-from casacore import tables
+
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
 
 from xradio.measurement_set._utils._msv2.msv4_sub_xdss import (
     create_pointing_xds,
@@ -648,6 +653,7 @@ def create_data_variables(
                 )
             except Exception as exc:
                 logger.debug(f"Could not load column {col}, exception: {exc}")
+                logger.debug(traceback.format_exc())
 
                 if ("WEIGHT_SPECTRUM" == col) and (
                     "WEIGHT" in col_names

xradio/measurement_set/_utils/_msv2/msv4_info_dicts.py

@@ -1,7 +1,11 @@
 import numpy as np
 import xarray as xr
 
-from casacore import tables
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
+
 import toolviper.utils.logger as logger
 
 from .subtables import subt_rename_ids

xradio/measurement_set/_utils/_msv2/msv4_sub_xdss.py

@@ -828,7 +828,7 @@ def create_phased_array_xds(
     }
     data_vars["COORDINATE_AXES"].attrs = {
         "type": "rotation_matrix",
-        "units": ["undimensioned", "undimensioned", "undimensioned"],
+        "units": ["dimensionless", "dimensionless", "dimensionless"],
     }
     # Remove the "frame" attribute if it exists, because ELEMENT_OFFSET is
     # defined in a station-local frame for which no standard name exists

xradio/measurement_set/_utils/_msv2/partition_queries.py

@@ -6,7 +6,10 @@ from typing import Dict, List, Tuple, Union
 import numpy as np
 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 ._tables.table_query import open_table_ro, open_query
 from ._tables.read import table_exists

xradio/measurement_set/schema.py

@@ -83,8 +83,8 @@ Doppler = Literal["doppler"]
 RotationMatrix = Literal["rotation_matrix"]
 
 # Units of quantities and measures
-UnitsUndimensioned = list[
-    Literal["undimensioned"]
+UnitsDimensionless = list[
+    Literal["dimensionless"]
 ]  # name consistent with casacore measures
 UnitsSeconds = list[Literal["s"]]
 UnitsHertz = list[Literal["Hz"]]
@@ -2261,7 +2261,7 @@ class PhasedArrayCoordinateAxesArray:
 
     data: Data[tuple[AntennaName, CartesianPosLabelLocal, CartesianPosLabel], float]
 
-    units: Attr[UnitsUndimensioned]
+    units: Attr[UnitsDimensionless]
 
     type: Attr[RotationMatrix]
     """ Measure type. Should be ``"rotation_matrix"``."""