xradio 0.0.55__py3-none-any.whl → 0.0.58__py3-none-any.whl

xradio/_utils/_casacore/tables.py

@@ -1,5 +1,10 @@
-from casacore import tables
+try:
+    from casacore import tables
+except ImportError:
+    import xradio._utils._casacore.casacore_from_casatools as tables
+
 from contextlib import contextmanager
+import os
 from typing import Dict, Generator
 
 # common casacore table handling code

xradio/_utils/coord_math.py

@@ -63,38 +63,37 @@ def wrap_to_pi(angles):
     return (angles + np.pi) % (2 * np.pi) - np.pi
 
 
-def wrap_to_pi(angles):
-    return (angles + np.pi) % (2 * np.pi) - np.pi
+def convert_to_si_units(xds):
+    import xarray as xr
 
+    with xr.set_options(keep_attrs=True):
+        for data_var in xds.data_vars:
+            if "units" in xds[data_var].attrs:
+                u = xds[data_var].attrs["units"]
 
-def convert_to_si_units(xds):
-    for data_var in xds.data_vars:
-        if "units" in xds[data_var].attrs:
-            for u_i, u in enumerate(xds[data_var].attrs["units"]):
                 if u == "km":
-                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * 1e3
-                    xds[data_var].attrs["units"][u_i] = "m"
+                    xds[data_var] = xds[data_var] * 1e3
+                    xds[data_var].attrs["units"] = "m"
                 if u == "km/s":
-                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * 1e3
-                    xds[data_var].attrs["units"][u_i] = "m/s"
+                    xds[data_var] = xds[data_var] * 1e3
+                    xds[data_var].attrs["units"] = "m/s"
                 if u == "deg":
-                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * np.pi / 180
-                    xds[data_var].attrs["units"][u_i] = "rad"
+                    xds[data_var] = xds[data_var] * np.pi / 180
+                    xds[data_var].attrs["units"] = "rad"
                 if u == "Au" or u == "AU":
-                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * 149597870700
-                    xds[data_var].attrs["units"][u_i] = "m"
+                    xds[data_var] = xds[data_var] * 149597870700
+                    xds[data_var].attrs["units"] = "m"
                 if u == "Au/d" or u == "AU/d":
-                    xds[data_var][..., u_i] = (
-                        xds[data_var][..., u_i] * 149597870700 / 86400
-                    )
-                    xds[data_var].attrs["units"][u_i] = "m/s"
+                    xds[data_var] = xds[data_var] * 149597870700 / 86400
+                    xds[data_var].attrs["units"] = "m/s"
                 if u == "arcsec":
-                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * np.pi / 648000
-                    xds[data_var].attrs["units"][u_i] = "rad"
+                    xds[data_var] = xds[data_var] * np.pi / 648000
+                    xds[data_var].attrs["units"] = "rad"
                 if u == "hPa":
-                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * 100.0
-                    xds[data_var].attrs["units"][u_i] = "Pa"
+                    xds[data_var] = xds[data_var] * 100.0
+                    xds[data_var].attrs["units"] = "Pa"
                 if u == "m-2":
                     # IONOS_ELECTRON sometimes has "m-2" instead of "/m^2"
-                    xds[data_var].attrs["units"][u_i] = "/m^2"
+                    xds[data_var].attrs["units"] = "/m^2"
+
     return xds

xradio/_utils/dict_helpers.py

@@ -1,49 +1,114 @@
 def make_quantity(value, units: str, dims: list = []) -> dict:
     """
     create a quantity dictionary given value and units
+
     Parameters
     ----------
     value : numeric or array of numerics
         Quantity value
     units: str
         Quantity units
+
     Returns
     -------
     dict
     """
-    u = units if isinstance(units, list) else [units]
-    return {"data": value, "dims": dims, "attrs": {"units": u, "type": "quantity"}}
+    return {"data": value, "dims": dims, "attrs": make_quantity_attrs(units)}
+
+
+def ensure_units_are_consistent(units):
+    if isinstance(units, str):
+        return units
+    else:
+        u0 = units[0]
+        for u in units:
+            assert u0 == u, f"Units are not consistent: {u0} != {u}. "
+        return u0
 
 
-def make_frequency_reference_dict(
+def make_quantity_attrs(units: str) -> dict:
+    """
+    Creates the dict of attributes of a quantity
+
+    Parameters
+    ----------
+    units: str
+        Quantity units
+
+    Returns
+    -------
+    dict
+    """
+    return {"units": ensure_units_are_consistent(units), "type": "quantity"}
+
+
+def make_spectral_coord_reference_dict(
     value: float, units: str, observer: str = "lsrk"
 ) -> dict:
-    u = units if isinstance(units, list) else [units]
+    """
+    creates a spectral_coord measure dict given the value, units, and observer
+
+    Parameters
+    ----------
+    value : numeric or array of numerics
+        measure value
+    units : str
+        measure units
+    observer :
+        observer reference frame
+
+    Returns
+    -------
+    dict
+    """
+    u = ensure_units_are_consistent(units)
     return {
-        "attrs": {"units": u, "observer": observer.lower(), "type": "frequency"},
+        "attrs": make_spectral_coord_measure_attrs(
+            u,
+            observer.lower() if observer not in ["TOPO", "BARY", "REST"] else observer,
+        ),
         "data": value,
         "dims": [],
     }
 
 
-def make_skycoord_dict(data: list[float], units: list[str], frame: str) -> dict:
+def make_spectral_coord_measure_attrs(units: str, observer: str = "lsrk") -> dict:
+    """
+    Creates a spectral_coord measure attrs dict given units and observer
+
+    Parameters
+    ----------
+    units: str or list of str
+        Spectral coordinate units
+    observer: str
+        Spectral reference frame
+    Returns
+    -------
+    dict
+        Attrs dict for a spectral_coord measure
+    """
+    u = ensure_units_are_consistent(units)
+    return {"units": u, "observer": observer, "type": "spectral_coord"}
+
+
+def make_skycoord_dict(data: list[float], units: str, frame: str) -> dict:
     return {
         "attrs": {
             "frame": frame.lower(),
             "type": "sky_coord",
-            "units": units,
+            "units": ensure_units_are_consistent(units),
         },
         "data": data,
         "dims": ["l", "m"],
     }
 
 
-def make_time_measure_attrs(units=["s"], scale="utc", time_format="mjd") -> dict:
-    u = units if isinstance(units, list) else [units]
+def make_time_measure_attrs(units="s", scale="utc", time_format="mjd") -> dict:
+    u = ensure_units_are_consistent(units)
     return {"units": u, "scale": scale, "format": time_format, "type": "time"}
 
 
-def make_time_measure_dict(data, units=["s"], scale="utc", time_format="mjd") -> dict:
+def make_time_measure_dict(data, units="s", scale="utc", time_format="mjd") -> dict:
     """
     create a time measure dictionary given value and units
     Parameters
@@ -67,7 +132,7 @@ def make_time_measure_dict(data, units=["s"], scale="utc", time_format="mjd") ->
     return x
 
 
-def make_time_coord_attrs(units=["s"], scale="utc", time_format="mjd") -> dict:
+def make_time_coord_attrs(units="s", scale="utc", time_format="mjd") -> dict:
     """
     create a time measure dictionary given value and units
     Parameters

xradio/_utils/schema.py

@@ -1,6 +1,7 @@
 import toolviper.utils.logger as logger
 import xarray as xr
 from typing import Union
+from xradio._utils.dict_helpers import ensure_units_are_consistent
 
 
 def convert_generic_xds_to_xradio_schema(
@@ -108,7 +109,7 @@ def column_description_casacore_to_msv4_measure(
 
         # Convert type, copy unit
         msv4_measure["type"] = msv4_measure_conversion["type"]
-        msv4_measure["units"] = list(
+        msv4_measure["units"] = ensure_units_are_consistent(
             casacore_column_description["keywords"]["QuantumUnits"]
         )
 
@@ -154,7 +155,9 @@ def column_description_casacore_to_msv4_measure(
     elif "QuantumUnits" in casacore_column_description["keywords"]:
         msv4_measure = {
             "type": "quantity",
-            "units": list(casacore_column_description["keywords"]["QuantumUnits"]),
+            "units": ensure_units_are_consistent(
+                casacore_column_description["keywords"]["QuantumUnits"]
+            ),
         }
 
     return msv4_measure

xradio/_utils/zarr/common.py

@@ -1,6 +1,4 @@
-import copy
 import xarray as xr
-import zarr
 import s3fs
 import os
 from botocore.exceptions import NoCredentialsError
@@ -34,7 +32,7 @@ def _get_file_system_and_items(ps_store: str):
                 for bd in file_system.listdir(ps_store, detail=False)
             ]
 
-        except (NoCredentialsError, PermissionError) as e:
+        except (NoCredentialsError, PermissionError):
             # only public, read-only buckets will be accessible
             # we will want to add messaging and error handling here
             file_system = s3fs.S3FileSystem(anon=True)
@@ -97,73 +95,3 @@ def _open_dataset(
         with dask.config.set(scheduler="synchronous"):
             xds = xds.load()
     return xds
-
-
-# Code to depricate:
-def _get_attrs(zarr_obj):
-    """
-    get attributes of zarr obj (groups or arrays)
-    """
-    return {k: v for k, v in zarr_obj.attrs.asdict().items() if not k.startswith("_NC")}
-
-
-def _load_no_dask_zarr(zarr_name, slice_dict={}):
-    """
-    Alternative to xarray open_zarr where the arrays are not Dask Arrays.
-
-    slice_dict: A dictionary of slice objects for which values to read form a dimension.
-                For example silce_dict={'time':slice(0,10)} would select the first 10 elements in the time dimension.
-                If a dim is not specified all values are retruned.
-    return:
-        xarray.Dataset()
-
-    #Should go into general utils.
-    """
-    # Used by xarray to store array labeling info in zarr meta data.
-    DIMENSION_KEY = "_ARRAY_DIMENSIONS"
-
-    # logger = _get_logger()
-    zarr_group = zarr.open_group(store=zarr_name, mode="r")
-    group_attrs = _get_attrs(zarr_group)
-
-    slice_dict_complete = copy.deepcopy(slice_dict)
-    coords = {}
-    xds = xr.Dataset()
-    for var_name, var in zarr_group.arrays():
-        print("Hallo 3", var_name, var.shape)
-        var_attrs = _get_attrs(var)
-
-        for dim in var_attrs[DIMENSION_KEY]:
-            if dim not in slice_dict_complete:
-                slice_dict_complete[dim] = slice(None)  # No slicing.
-
-        if (var_attrs[DIMENSION_KEY][0] == var_name) and (
-            len(var_attrs[DIMENSION_KEY]) == 1
-        ):
-            coord = var[
-                slice_dict_complete[var_attrs[DIMENSION_KEY][0]]
-            ]  # Dimension coordinates.
-            del var_attrs["_ARRAY_DIMENSIONS"]
-            xds = xds.assign_coords({var_name: coord})
-            xds[var_name].attrs = var_attrs
-        else:
-            # Construct slicing
-            slicing_list = []
-            for dim in var_attrs[DIMENSION_KEY]:
-                slicing_list.append(slice_dict_complete[dim])
-            slicing_tuple = tuple(slicing_list)
-
-            print(var_attrs[DIMENSION_KEY])
-
-            xds[var_name] = xr.DataArray(
-                var[slicing_tuple], dims=var_attrs[DIMENSION_KEY]
-            )
-
-            if "coordinates" in var_attrs:
-                del var_attrs["coordinates"]
-            del var_attrs["_ARRAY_DIMENSIONS"]
-            xds[var_name].attrs = var_attrs
-
-    xds.attrs = group_attrs
-
-    return xds

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,
@@ -33,7 +41,7 @@ from ...._utils._casacore.tables import extract_table_attributes, open_table_ro
 from xradio._utils.coord_math import _deg_to_rad
 from xradio._utils.dict_helpers import (
     _casacore_q_to_xradio_q,
-    make_frequency_reference_dict,
+    make_spectral_coord_reference_dict,
     make_quantity,
     make_skycoord_dict,
     make_time_measure_dict,
@@ -77,34 +85,50 @@ def _casa_image_to_xds_image_attrs(image: casa_image, history: bool = True) -> d
     attrs["units"] = image.unit()
     attrs["telescope"] = {}
     telescope = attrs["telescope"]
+
     for k in ("observer", "obsdate", "telescope", "telescopeposition"):
         if k.startswith("telescope"):
             if k == "telescope":
                 telescope["name"] = coord_dict[k]
             elif k in coord_dict:
                 casa_pos = coord_dict[k]
-                location = {}
-                loc_attrs = {}
-                loc_attrs["type"] = "location"
-                loc_attrs["frame"] = casa_pos["refer"]
-                """
-                if casa_pos["refer"] == "ITRF":
-                    loc_attrs["ellipsoid"] = "GRS80"
-                """
-                loc_attrs["units"] = [
-                    casa_pos["m0"]["unit"],
-                    casa_pos["m1"]["unit"],
-                    casa_pos["m2"]["unit"],
-                ]
-                loc_attrs["coordinate_system"] = "geocentric"
-                loc_attrs["origin_object_name"] = "earth"
-                location["attrs"] = loc_attrs
-                location["data"] = [
-                    casa_pos["m0"]["value"],
-                    casa_pos["m1"]["value"],
-                    casa_pos["m2"]["value"],
-                ]
-                telescope["location"] = location
+                telescope["direction"] = {
+                    "attrs": {
+                        "coordinate_system": "geocentric",
+                        "frame": casa_pos["refer"],
+                        "origin_object_name": "earth",
+                        "type": "location",
+                        "units": casa_pos["m0"]["unit"],
+                    },
+                    "data": [casa_pos["m0"]["value"], casa_pos["m1"]["value"]],
+                    "dims": ["ellipsoid_dir_label"],
+                    "coords": {
+                        "ellipsoid_dir_label": {
+                            "dims": ["ellipsoid_dir_label"],
+                            "data": ["lon", "lat"],
+                        }
+                    },
+                }
+                telescope["distance"] = {
+                    "attrs": {
+                        "coordinate_system": "geocentric",
+                        "frame": casa_pos["refer"],
+                        "origin_object_name": "earth",
+                        "type": "location",
+                        "units": casa_pos["m2"]["unit"],
+                    },
+                    "data": [casa_pos["m2"]["value"]],
+                    "dims": ["ellipsoid_dis_label"],
+                    "coords": {
+                        "ellipsoid_dis_label": {
+                            "dims": ["ellipsoid_dis_label"],
+                            "data": [
+                                "dist",
+                            ],
+                        }
+                    },
+                }
+
                 """
                 del (
                     telescope["position"]["refer"],
@@ -147,7 +171,7 @@ def _casa_image_to_xds_image_attrs(image: casa_image, history: bool = True) -> d
     # associated with it.
     # point_center = coord_dict["pointingcenter"]
     attrs[_pointing_center] = make_skycoord_dict(
-        coord_dict["pointingcenter"]["value"].tolist(), ["rad", "rad"], frame
+        coord_dict["pointingcenter"]["value"].tolist(), "rad", frame
     )
     imageinfo = meta_dict["imageinfo"]
     obj = "objectname"
@@ -213,7 +237,7 @@ def _add_time_attrs(xds: xr.Dataset, coord_dict: dict) -> xr.Dataset:
 
 def _add_vel_attrs(xds: xr.Dataset, coord_dict: dict) -> xr.Dataset:
     vel_coord = xds["velocity"]
-    meta = {"units": ["m/s"]}
+    meta = {"units": "m/s"}
     for k in coord_dict:
         if k.startswith("spectral"):
             sd = coord_dict[k]
@@ -251,7 +275,7 @@ def _casa_image_to_xds_attrs(img_full_path: str) -> dict:
         dir_dict = {}
 
         dir_dict["reference"] = make_skycoord_dict(
-            data=[0.0, 0.0], units=["rad", "rad"], frame=ap_system
+            data=[0.0, 0.0], units="rad", frame=ap_system
         )
         if ap_equinox:
             dir_dict["reference"]["attrs"]["equinox"] = ap_equinox
@@ -526,14 +550,14 @@ def _get_freq_values_attrs(
                     cdelt=wcs["cdelt"],
                 )
                 attrs["rest_frequency"] = make_quantity(sd["restfreq"], "Hz")
-                # attrs["type"] = "frequency"
-                # attrs["units"] = sd["unit"]
-                # attrs["frame"] = sd["system"]
-                attrs["wave_unit"] = [sd["waveUnit"]]
+                attrs["type"] = "spectral_coord"
+                attrs["units"] = sd["unit"]
+                attrs["frame"] = sd["system"]
+                attrs["wave_units"] = sd["waveUnit"]
                 # attrs["crval"] = sd["wcs"]["crval"]
                 # attrs["cdelt"] = sd["wcs"]["cdelt"]
 
-                attrs["reference_value"] = make_frequency_reference_dict(
+                attrs["reference_frequency"] = make_spectral_coord_reference_dict(
                     value=sd["wcs"]["crval"],
                     units=sd["unit"],
                     observer=sd["system"],
@@ -664,7 +688,7 @@ def _get_time_values_attrs(cimage_coord_dict: dict) -> Tuple[List[float], dict]:
     attrs["type"] = "time"
     attrs["scale"] = cimage_coord_dict["obsdate"]["refer"]
     unit = cimage_coord_dict["obsdate"]["m0"]["unit"]
-    attrs["units"] = unit if isinstance(unit, list) else [unit]
+    attrs["units"] = unit
     time_val = cimage_coord_dict["obsdate"]["m0"]["value"]
     attrs["format"] = _get_time_format(time_val, unit)
     return ([time_val], copy.deepcopy(attrs))
@@ -794,7 +818,7 @@ def _get_velocity_values_attrs(
     if not attrs:
         attrs["doppler_type"] = _doppler_types[0]
 
-    attrs["units"] = ["m/s"]
+    attrs["units"] = "m/s"
     attrs["type"] = "doppler"
     return (
         _compute_velocity_values(