xradio 0.0.56__tar.gz → 0.0.58__tar.gz

{xradio-0.0.56/src/xradio.egg-info → xradio-0.0.58}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xradio
-Version: 0.0.56
+Version: 0.0.58
 Summary:  Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>, Federico Montesino Pouzols <pouzols@eso.edu>, Dave Mehringer <dmehring@nrao.edu>, Peter Wortmann <peter.wortmann@skao.int>
 License: BSD 3-Clause License
@@ -50,6 +50,7 @@ Requires-Dist: zarr<3,>=2
 Requires-Dist: pyarrow
 Requires-Dist: typeguard
 Requires-Dist: numcodecs<0.16
+Requires-Dist: psutil
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"

{xradio-0.0.56 → xradio-0.0.58}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "xradio"
-version = "v0.0.56"
+version = "v0.0.58"
 description = " Xarray Radio Astronomy Data IO"
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},
@@ -22,7 +22,8 @@ dependencies = [
     'zarr>=2,<3',
     'pyarrow',
     'typeguard',
-    'numcodecs<0.16'
+    'numcodecs<0.16',
+    'psutil' # psutil is needed so large FITS images are not loaded into memory
 ]
 
 # toolviper has the following key dependencies also directly used by XRADIO (https://github.com/casangi/toolviper/blob/main/pyproject.toml):

{xradio-0.0.56 → xradio-0.0.58}/src/xradio/__init__.py

@@ -1,5 +1,5 @@
-import os
-from toolviper.utils.logger import setup_logger
+# import os
+# from toolviper.utils.logger import setup_logger
 
 
 # _logger_name = "xradio"

{xradio-0.0.56 → xradio-0.0.58}/src/xradio/_utils/_casacore/casacore_from_casatools.py

@@ -18,7 +18,17 @@ from typing import Any, Dict, List, Sequence, Union
 # Configure casaconfig settings prior to casatools import
 # this ensures optimal initialization and resource allocation for casatools
 # Also see : https://casadocs.readthedocs.io/en/stable/api/casaconfig.html
-import casaconfig.config
+try:
+    import casaconfig.config
+except ModuleNotFoundError as exc:
+    raise ModuleNotFoundError(
+        f"casaconfig.config cannot be found, probably because the package "
+        "casatools is not available. If you are here that is probably because "
+        "python-casacore is not available either. MSv2 related functionality "
+        "requires either python-casacore or casatools. Import failure details: "
+        f"{exc}"
+    )
+
 import numpy as np
 import toolviper.utils.logger as logger
 
@@ -70,7 +80,7 @@ else:
     casaconfig.config.logfile = "/dev/null"
     casaconfig.config.nologfile = True
 
-import casatools
+import casatools  # noqa: E402 (because of previous config initialization)
 
 casatools.logger.setglobal(True)
 casatools.logger.ompSetNumThreads(1)

{xradio-0.0.56 → xradio-0.0.58}/src/xradio/_utils/_casacore/tables.py

@@ -4,6 +4,7 @@ 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-0.0.56 → xradio-0.0.58}/src/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-0.0.58/src/xradio/_utils/dict_helpers.py

@@ -0,0 +1,170 @@
+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
+    """
+    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_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:
+    """
+    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": make_spectral_coord_measure_attrs(
+            u,
+            observer.lower() if observer not in ["TOPO", "BARY", "REST"] else observer,
+        ),
+        "data": value,
+        "dims": [],
+    }
+
+
+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": ensure_units_are_consistent(units),
+        },
+        "data": data,
+        "dims": ["l", "m"],
+    }
+
+
+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:
+    """
+    create a time measure dictionary given value and units
+    Parameters
+    ----------
+    value : numeric or array of numerics
+        Time value
+    units: str
+        Time units
+    scale: str
+        Time scale
+    time_format: str
+        Time format
+    Returns
+    -------
+    dict
+    """
+    x = {}
+    x["attrs"] = make_time_measure_attrs(units, scale, time_format)
+    x["data"] = data
+    x["dims"] = []
+    return x
+
+
+def make_time_coord_attrs(units="s", scale="utc", time_format="mjd") -> dict:
+    """
+    create a time measure dictionary given value and units
+    Parameters
+    ----------
+    value : numeric or array of numerics
+        Time value
+    units: str
+        Time units
+    scale: str
+        Time scale
+    time_format: str
+        Time format
+    Returns
+    -------
+    dict
+    """
+    x = make_time_measure_attrs(units, scale, time_format)
+    del x["type"]
+    return x
+
+
+def _casacore_q_to_xradio_q(q: dict) -> dict:
+    """
+    Convert a casacore quantity to an xradio quantity
+    """
+    if isinstance(q, dict):
+        if "value" in q and "unit" in q:
+            return make_quantity(q["value"], q["unit"])
+        else:
+            p = {}
+            for k in q:
+                p[k] = _casacore_q_to_xradio_q(q[k])
+            return p
+    else:
+        raise ValueError(f"Cannot convert {q} to xradio quantity")

{xradio-0.0.56 → xradio-0.0.58}/src/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-0.0.56 → xradio-0.0.58}/src/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-0.0.56 → xradio-0.0.58}/src/xradio/image/_util/_casacore/xds_from_casacore.py

@@ -41,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,
@@ -85,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"],
@@ -155,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"
@@ -221,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]
@@ -259,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
@@ -534,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"],
@@ -672,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))
@@ -802,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(