xradio 0.0.50__py3-none-any.whl → 0.0.52__py3-none-any.whl

xradio/image/_util/_zarr/xds_from_zarr.py

@@ -12,9 +12,9 @@ from xradio._utils.zarr.common import _get_file_system_and_items
 
 
 def _read_zarr(
-    zarr_store: str, output: dict, selection: dict = {}
+    zarr_store: str, id_dict: dict, selection: dict = {}
 ) -> (xr.Dataset, bool):
-    # supported key/values in output are:
+    # supported key/values in id_dict are:
     # "dv"
     #    what data variables should be returned as.
     #    "numpy": numpy arrays
@@ -29,27 +29,27 @@ def _read_zarr(
     do_dask = False
     do_numpy = False
     do_np_coords = False
-    if "dv" in output:
-        dv = output["dv"]
+    if "dv" in id_dict:
+        dv = id_dict["dv"]
         if dv in ["dask", "numpy"]:
             do_dask = dv == "dask"
             do_numpy = not do_dask
         else:
             raise ValueError(
-                f"Unsupported value {output[dv]} for output[dv]. "
+                f"Unsupported value {id_dict[dv]} for id_dict[dv]. "
                 "Supported values are 'dask' and 'numpy'"
             )
-    if "coords" in output:
-        c = output["coords"]
+    if "coords" in id_dict:
+        c = id_dict["coords"]
         if c == "numpy":
             do_np_coords = True
         else:
             raise ValueError(
-                f"Unexpected value {c} for output[coords]. "
+                f"Unexpected value {c} for id_dict[coords]. "
                 "The supported value is 'numpy'"
             )
     # do not pass selection, because that is only for the top level data vars
-    xds = _decode(xds, zarr_store, output)
+    xds = _decode(xds, zarr_store, id_dict)
     if do_np_coords:
         xds = _coords_to_numpy(xds)
     if do_dask:
@@ -59,11 +59,9 @@ def _read_zarr(
     return xds
 
 
-def _decode(xds: xr.Dataset, zarr_store: str, output: dict) -> (xr.Dataset, bool):
+def _decode(xds: xr.Dataset, zarr_store: str, id_dict: dict) -> xr.Dataset:
     xds.attrs = _decode_dict(xds.attrs, "")
-    sub_xdses = _decode_sub_xdses(zarr_store, output)
-    for k, v in sub_xdses.items():
-        xds.attrs[k] = v
+    _decode_sub_xdses(xds, zarr_store, id_dict)
     return xds
 
 
@@ -85,26 +83,45 @@ def _decode_dict(my_dict: dict, top_key: str) -> dict:
     return my_dict
 
 
-def _decode_sub_xdses(zarr_store: str, output: dict) -> dict:
-    sub_xdses = {}
+def _decode_sub_xdses(xarrayObj, top_dir: str, id_dict: dict) -> None:
+    # FIXME this also needs to support S3
+    # determine immediate subdirs of zarr_store
+    entries = os.scandir(top_dir)
+    for d in entries:
+        path = os.path.join(top_dir, d.name)
+        if os.path.isdir(path):
+            if d.name.startswith(_top_level_sub_xds):
+                ky = d.name[len(_top_level_sub_xds) :]
+                xarrayObj.attrs[ky] = _read_zarr(path, id_dict)
+                # TODO if attrs that are xdses have attrs that are xdses ...
+            else:
+                # descend into the directory
+                _decode_sub_xdses(xarrayObj[d.name], path, id_dict)
 
-    fs, store_contents = _get_file_system_and_items(zarr_store)
 
+"""
+def _decode_sub_xdses(zarr_store: str, id_dict: dict) -> dict:
+    sub_xdses = {}
+    fs, store_contents = _get_file_system_and_items(zarr_store)
     if isinstance(fs, s3fs.core.S3FileSystem):
         # could we just use the items as returned from the helper function..?
         store_tree = fs.walk(zarr_store, topdown=True)
+        # Q: what is prepend_s3 used for? In this version it is defined but not used.
         prepend_s3 = "s3://"
     else:
         store_tree = os.walk(zarr_store, topdown=True)
         prepend_s3 = ""
-
     for root, dirs, files in store_tree:
+        relpath = os.path.relpath(root, zarr_store)
+        print("rpath", relpath)
         for d in dirs:
             if d.startswith(_top_level_sub_xds):
-                xds = _read_zarr(os.sep.join([root, d]), output)
+                xds = _read_zarr(os.sep.join([root, d]), id_dict)
                 # for k, v in xds.data_vars.items():
                 #    xds = xds.drop_vars([k]).assign({k: v.compute()})
-                ky = d[len(_top_level_sub_xds) + 1 :]
+                ky = d[len(_top_level_sub_xds) :]
                 sub_xdses[ky] = xds
-
+    print(f"Sub xdses: {sub_xdses.keys()}")
+    print("return")
     return sub_xdses
+"""

xradio/image/_util/_zarr/xds_to_zarr.py

@@ -1,3 +1,4 @@
+import logging
 import numpy as np
 import xarray as xr
 import os
@@ -6,44 +7,42 @@ from .common import _np_types, _top_level_sub_xds
 
 def _write_zarr(xds: xr.Dataset, zarr_store: str):
     xds_copy = xds.copy(deep=True)
-    xds_copy, xds_dict = _encode(xds_copy)
+    sub_xds_dict = _encode(xds_copy, zarr_store)
     z_obj = xds_copy.to_zarr(store=zarr_store, compute=True)
-    _write_sub_xdses(zarr_store, xds_dict, _top_level_sub_xds)
+    if sub_xds_dict:
+        _write_sub_xdses(sub_xds_dict)
 
 
-def _encode(xds: xr.Dataset):
+def _encode(xds: xr.Dataset, top_path: str) -> dict:
     # encode attrs
-    xds.attrs, xds_dict = _encode_dict(xds.attrs)
-    return xds, xds_dict
+    sub_xds_dict = {}
+    _encode_dict(xds.attrs, top_path, sub_xds_dict)
+    for dv in xds.data_vars:
+        _encode_dict(xds[dv].attrs, os.sep.join([top_path, dv]), sub_xds_dict)
+    logging.debug(f"Encoded sub_xds_dict: {sub_xds_dict}")
+    return sub_xds_dict
 
 
-def _encode_dict(my_dict: dict, top_key="") -> tuple:
-    xds_dict = {}
+def _encode_dict(my_dict: dict, top_path: str, sub_xds_dict) -> tuple:
     del_keys = []
     for k, v in my_dict.items():
         if isinstance(v, dict):
-            z = os.sep.join([top_key, k]) if top_key else k
-            my_dict[k], ret_xds_dict = _encode_dict(v, z)
-            if ret_xds_dict:
-                xds_dict[k] = ret_xds_dict
+            z = os.sep.join([top_path, k])
+            _encode_dict(v, z, sub_xds_dict)
         elif isinstance(v, np.ndarray):
             my_dict[k] = {}
             my_dict[k]["_type"] = "numpy.ndarray"
             my_dict[k]["_value"] = v.tolist()
             my_dict[k]["_dtype"] = str(v.dtype)
         elif isinstance(v, xr.Dataset):
-            xds_dict[k] = v.copy(deep=True)
+            sub_xds_dict[os.sep.join([top_path, f"{_top_level_sub_xds}{k}"])] = v.copy(
+                deep=True
+            )
             del_keys.append(k)
     for k in del_keys:
         del my_dict[k]
-    return my_dict, xds_dict
 
 
-def _write_sub_xdses(zarr_store: str, xds_dict: dict, path: str):
-    for k, v in xds_dict.items():
-        my_path = f"{path}_{k}" if path else f"{k}"
-        if isinstance(v, dict):
-            _write_sub_xdses(zarr_store, xds_dict[k], my_path)
-        elif isinstance(v, xr.Dataset):
-            zs = os.sep.join([zarr_store, my_path])
-            z_obj = v.to_zarr(store=zs, compute=True)
+def _write_sub_xdses(sub_xds: dict):
+    for k, v in sub_xds.items():
+        z_obj = v.to_zarr(store=k, compute=True)

xradio/image/_util/casacore.py

@@ -48,7 +48,9 @@ def _load_casa_image_block(infile: str, block_des: dict, do_sky_coords) -> xr.Da
     block = _get_persistent_block(
         image_full_path, shapes, starts, dimorder, transpose_list, new_axes
     )
-    xds = _add_sky_or_aperture(xds, block, dimorder, image_full_path, ret["sphr_dims"])
+    xds = _add_sky_or_aperture(
+        xds, block, dimorder, image_full_path, ret["sphr_dims"], True
+    )
     mymasks = _get_mask_names(image_full_path)
     for m in mymasks:
         full_path = os.sep.join([image_full_path, m])
@@ -57,7 +59,7 @@ 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, True)
+    xds.attrs = _casa_image_to_xds_attrs(image_full_path)
     mb = _multibeam_array(xds, image_full_path, False)
     if mb is not None:
         selectors = {}
@@ -86,6 +88,7 @@ def _read_casa_image(
         dimorder,
         img_full_path,
         ret["sphr_dims"],
+        history,
     )
     if masks:
         mymasks = _get_mask_names(img_full_path)
@@ -93,7 +96,7 @@ def _read_casa_image(
             ary = _read_image_array(img_full_path, chunks, mask=m, verbose=verbose)
             # 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)
+    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

xradio/measurement_set/_utils/_msv2/conversion.py

@@ -1214,13 +1214,13 @@ def convert_and_write_partition(
         )
         # but before, keep the name-id arrays, we need them for the pointing and weather xds
         ant_xds_name_ids = ant_xds["antenna_name"].set_xindex("antenna_id")
-        ant_xds_station_name_ids = ant_xds["station"].set_xindex("antenna_id")
+        ant_position_xds_with_ids = ant_xds["ANTENNA_POSITION"].set_xindex("antenna_id")
         # No longer needed after converting to name.
         ant_xds = ant_xds.drop_vars("antenna_id")
 
         # Create weather_xds
         start = time.time()
-        weather_xds = create_weather_xds(in_file, ant_xds_station_name_ids)
+        weather_xds = create_weather_xds(in_file, ant_position_xds_with_ids)
         logger.debug("Time weather " + str(time.time() - start))
 
         # Create pointing_xds
@@ -1425,7 +1425,7 @@ def antenna_ids_to_names(
             "antenna_id",
             "antenna_name",
             "mount",
-            "station",
+            "station_name",
         ]
         for unwanted_coord in unwanted_coords_from_ant_xds:
             xds = xds.drop_vars(unwanted_coord)

xradio/measurement_set/_utils/_msv2/create_antenna_xds.py

@@ -109,7 +109,7 @@ def extract_antenna_info(
 
     to_new_coords = {
         "NAME": ["antenna_name", ["antenna_name"]],
-        "STATION": ["station", ["antenna_name"]],
+        "STATION": ["station_name", ["antenna_name"]],
         "MOUNT": ["mount", ["antenna_name"]],
         # "PHASED_ARRAY_ID": ["phased_array_id", ["antenna_name"]],
         "antenna_id": ["antenna_id", ["antenna_name"]],
@@ -158,9 +158,9 @@ def extract_antenna_info(
 
         # None of the native numpy functions work on the github test runner.
         antenna_name = ant_xds["antenna_name"].values
-        station = ant_xds["station"].values
+        station_name = ant_xds["station_name"].values
         antenna_name = np.array(
-            list(map(lambda x, y: x + "_" + y, antenna_name, station))
+            list(map(lambda x, y: x + "_" + y, antenna_name, station_name))
         )
 
         ant_xds["antenna_name"] = xr.DataArray(antenna_name, dims=["antenna_name"])
@@ -376,7 +376,7 @@ def create_gain_curve_xds(
 
     ant_borrowed_coords = {
         "antenna_name": ant_xds.coords["antenna_name"],
-        "station": ant_xds.coords["station"],
+        "station_name": ant_xds.coords["station_name"],
         "mount": ant_xds.coords["mount"],
         "telescope_name": ant_xds.coords["telescope_name"],
         "receptor_label": ant_xds.coords["receptor_label"],
@@ -486,7 +486,7 @@ def create_phase_calibration_xds(
 
     ant_borrowed_coords = {
         "antenna_name": ant_xds.coords["antenna_name"],
-        "station": ant_xds.coords["station"],
+        "station_name": ant_xds.coords["station_name"],
         "mount": ant_xds.coords["mount"],
         "telescope_name": ant_xds.coords["telescope_name"],
         "receptor_label": ant_xds.coords["receptor_label"],

xradio/measurement_set/_utils/_msv2/create_field_and_source_xds.py

@@ -17,7 +17,8 @@ from xradio.measurement_set._utils._msv2._tables.read import (
     make_taql_where_between_min_max,
     load_generic_table,
 )
-from xradio._utils.list_and_array import cast_to_str
+from xradio._utils.list_and_array import cast_to_str, get_pad_value
+
 from xradio._utils.coord_math import (
     convert_to_si_units,
     add_position_offsets,
@@ -226,9 +227,10 @@ def extract_ephemeris_info(
             "type": "location",
             "units": ["deg", "deg", "m"],
             "data": observer_position,
-            "frame": "WGS84",
+            "frame": "ITRS",
             "origin_object_name": "Earth",
             "coordinate_system": ephemeris_meta["obsloc"].lower(),
+            "ellipsoid": "WGS84",
         }
     )  # I think the units are ['deg','deg','m'] and 'WGS84'.
 
@@ -546,10 +548,16 @@ def pad_missing_sources(
         for source_id in unique_source_ids
         if source_id not in source_xds.coords["SOURCE_ID"]
     ]
+    if len(missing_source_ids) < 1:
+        return source_xds
 
     # would like to use the new-ish xr.pad, but it creates issues with indices/coords and is
     # also not free of overheads, as it for example changes all numeric types to float64
-    missing_source_xds = xr.full_like(source_xds.isel(SOURCE_ID=0), fill_value=np.nan)
+    fill_value = {
+        var_name: get_pad_value(var.dtype)
+        for var_name, var in source_xds.data_vars.items()
+    }
+    missing_source_xds = xr.full_like(source_xds.isel(SOURCE_ID=0), fill_value)
     pad_str = "Unknown"
     pad_str_type = "<U9"
     for var in missing_source_xds.data_vars:

xradio/measurement_set/_utils/_msv2/msv4_sub_xdss.py

@@ -172,7 +172,130 @@ def make_taql_where_weather(
     return taql_where
 
 
-def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
+def prepare_generic_weather_xds_and_station_name(
+    generic_weather_xds: xr.Dataset,
+    in_file: str,
+    ant_position_with_ids: xr.DataArray,
+    has_asdm_station_position: bool,
+) -> tuple[xr.Dataset, np.ndarray]:
+    """
+    A generic_weather_xds loaded with load_generic_table() might still need to be reloaded
+    with an additional WHERE condition to constrain the indices of antennas. But this depends on whether
+    ASDM/importasdm extension columns are present or not.
+
+    This also prepares the station_name values:
+    - if has_asdm_station_ids:
+       - tries to find from ASDM_STATION the station names,
+       - otherwise, takes ids (antenna_ids in generic_weather were actually the ASDM_STATION_IDs
+    - else: get the values from antenna_xds (the stations present)
+
+
+    Parameters
+    ----------
+    generic_weather_xds : xr.Dataset
+        generic dataset read from an MSv2 WEATHER subtable
+    in_file : str
+        Input MS name.
+    ant_position_with_ids : xr.DataArray
+        antenna_position data var from the antenna_xds (expected to still include the initial ANTENNA_ID
+        coordinate as well as other coordinates from the antenna_xds)
+    has_asdm_station_position : bool
+        wHether this generic weatehr_xds should be treated as including the nonstandard extensions
+        NS_WX_STATION_ID and NS_WX_STATION_POSITION as created by CASA/importasdm (ALMA and VLA).
+
+    Returns
+    -------
+    (generic_weather_xds, station_name): tuple[[xarray.Dataset, numpy.ndarray]
+        Weather Xarray Dataset prepared for generic conversion to MSv4, values for the station_name coordinate
+    """
+
+    if has_asdm_station_position:
+        asdm_station_path = os.path.join(in_file, "ASDM_STATION")
+        if table_exists(asdm_station_path):
+            asdm_station_xds = load_generic_table(in_file, "ASDM_STATION")
+            station_name = asdm_station_xds.name.values[
+                generic_weather_xds["ANTENNA_ID"].values
+            ]
+        else:
+            # if no info from ASDM_STATION, use the indices from antenna_id which was actually the NS_WX_STATION_ID
+            len_antenna_id = generic_weather_xds.sizes["ANTENNA_ID"]
+            station_name = list(
+                map(
+                    lambda x, y: x + "_" + y,
+                    ["Station"] * len_antenna_id,
+                    generic_weather_xds["ANTENNA_ID"].values.astype(str),
+                )
+            )
+
+    else:
+        taql_where = make_taql_where_weather(in_file, ant_position_with_ids)
+        generic_weather_xds = load_generic_table(
+            in_file,
+            "WEATHER",
+            rename_ids=subt_rename_ids["WEATHER"],
+            taql_where=taql_where,
+        )
+
+        if not generic_weather_xds.data_vars:
+            # for example when the weather subtable only has info for antennas/stations
+            # not present in the MSv4 (no overlap between antennas loaded in ant_xds and weather)
+            return None, None
+
+        stations_present = ant_position_with_ids.sel(
+            antenna_id=generic_weather_xds["ANTENNA_ID"]
+        ).station_name
+        station_name = stations_present.values
+
+    return generic_weather_xds, station_name
+
+
+def finalize_station_position(
+    weather_xds: xr.Dataset, ant_position_with_ids, has_asdm_station_position: bool
+) -> xr.Dataset:
+    """
+    For a STATION_POSITION data var being added to a weather_xds, make sure coordinates and dimensions
+    are conforming to the schema.
+
+    Parameters
+    ----------
+    weather_xds : xr.Dataset
+        weather_xds where we still need to ensure the right coordinates and attributes
+    ant_position_with_ids : xr.DataArray
+        antenna_position data var from the antenna_xds (expected to still include the initial ANTENNA_ID
+        coordinate as well as other coordinates from the antenna_xds)
+    has_asdm_station_position : bool
+        Whether this generic weatehr_xds should be treated as including the nonstandard extensions
+        NS_WX_STATION_ID and NS_WX_STATION_POSITION as created by CASA/importasdm (ALMA and VLA).
+
+    Returns
+    -------
+    weather_xds: xarray.Dataset
+        Weather Xarray Dataset with all coordinates and attributes in STATION_POSITION
+    """
+    if has_asdm_station_position:
+        # STATION_POSITION has been created but needs prooper dimensions and attrs
+        # Drop the time dim
+        weather_xds["STATION_POSITION"] = weather_xds["STATION_POSITION"].sel(
+            time_weather=0, drop=True, method="nearest"
+        )
+        # borrow location frame attributes from antenna position
+        weather_xds["STATION_POSITION"].attrs = ant_position_with_ids.attrs
+    else:
+        # borrow from ant_posision_with_ids but without carrying over other coords
+        weather_xds = weather_xds.assign(
+            {
+                "STATION_POSITION": (
+                    ["station_name", "cartesian_pos_label"],
+                    ant_position_with_ids.values,
+                    ant_position_with_ids.attrs,
+                )
+            }
+        )
+
+    return weather_xds
+
+
+def create_weather_xds(in_file: str, ant_position_with_ids: xr.DataArray):
     """
     Creates a Weather Xarray Dataset from a MS v2 WEATHER table.
 
@@ -180,8 +303,9 @@ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
     ----------
     in_file : str
         Input MS name.
-    ant_xds_station_name_ids : xr.DataArray
-        station name data array from antenna_xds, with name/id information
+    ant_position_with_ids : xr.DataArray
+        antenna_position data var from the antenna_xds (expected to still including the initial ANTENNA_ID coordinate
+        as wellas other coordinates from the antenna_xds)
 
     Returns
     -------
@@ -190,32 +314,32 @@ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
     """
 
     try:
-        taql_where = make_taql_where_weather(in_file, ant_xds_station_name_ids)
         generic_weather_xds = load_generic_table(
             in_file,
             "WEATHER",
             rename_ids=subt_rename_ids["WEATHER"],
-            taql_where=taql_where,
         )
     except ValueError as _exc:
         return None
 
-    if not generic_weather_xds.data_vars:
-        # for example when the weather subtable only has info for antennas/stations
-        # not present in the MSv4 (no overlap between antennas loaded in ant_xds and weather)
+    has_asdm_station_position = (
+        "NS_WX_STATION_POSITION" in generic_weather_xds.data_vars
+    )
+    generic_weather_xds, station_name = prepare_generic_weather_xds_and_station_name(
+        generic_weather_xds, in_file, ant_position_with_ids, has_asdm_station_position
+    )
+    if not generic_weather_xds:
         return None
 
     weather_xds = xr.Dataset(attrs={"type": "weather"})
-    stations_present = ant_xds_station_name_ids.sel(
-        antenna_id=generic_weather_xds["ANTENNA_ID"]
-    )
     coords = {
-        "station_name": stations_present.data,
-        "antenna_name": stations_present.coords["antenna_name"].data,
+        "station_name": station_name,
+        "cartesian_pos_label": ["x", "y", "z"],
     }
     weather_xds = weather_xds.assign_coords(coords)
 
     dims_station_time = ["station_name", "time_weather"]
+    dims_station_time_position = dims_station_time + ["cartesian_pos_label"]
     to_new_data_variables = {
         "H20": ["H2O", dims_station_time],
         "IONOS_ELECTRON": ["IONOS_ELECTRON", dims_station_time],
@@ -226,6 +350,15 @@ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
         "WIND_DIRECTION": ["WIND_DIRECTION", dims_station_time],
         "WIND_SPEED": ["WIND_SPEED", dims_station_time],
     }
+    if has_asdm_station_position:
+        to_new_data_variables.update(
+            {
+                "NS_WX_STATION_POSITION": [
+                    "STATION_POSITION",
+                    dims_station_time_position,
+                ],
+            }
+        )
 
     to_new_coords = {
         "TIME": ["time_weather", ["time_weather"]],
@@ -234,6 +367,9 @@ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
     weather_xds = convert_generic_xds_to_xradio_schema(
         generic_weather_xds, weather_xds, to_new_data_variables, to_new_coords
     )
+    weather_xds = finalize_station_position(
+        weather_xds, ant_position_with_ids, has_asdm_station_position
+    )
 
     # TODO: option to interpolate to main time
 
@@ -256,6 +392,7 @@ def correct_generic_pointing_xds(
     and tries to correct several deviations from the MSv2 specs seen in
     common test data.
     The problems fixed here include wrong dimensions:
+
     - for example transposed dimensions with respect to the MSv2 specs (output
     from CASA simulator),
     - missing/additional unexpected dimensions when some of the columns are
@@ -423,6 +560,7 @@ def prepare_generic_sys_cal_xds(generic_sys_cal_xds: xr.Dataset) -> xr.Dataset:
     sys_cal_xds dataset, as their structure differs in dimensions and order
     of dimensions.
     This function performs various prepareation steps, such as:
+
     - filter out dimensions not neeed for an individual MSv4 (SPW, FEED),
     - drop variables loaded from columns with all items set to empty array,
     - transpose the dimensions frequency,receptor