xradio 0.0.41__py3-none-any.whl → 0.0.42__py3-none-any.whl

xradio/{correlated_data/_utils/_ms → measurement_set/_utils/_msv2}/msv4_sub_xdss.py

@@ -6,6 +6,7 @@ from typing import Tuple, Union
 import numpy as np
 import xarray as xr
 
+from xradio._utils.coord_math import convert_to_si_units
 from xradio._utils.schema import (
     column_description_casacore_to_msv4_measure,
     convert_generic_xds_to_xradio_schema,
@@ -144,7 +145,7 @@ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
     }
     weather_xds = weather_xds.assign_coords(coords)
 
-    dims_station_time = ["station_name", "time"]
+    dims_station_time = ["station_name", "time_weather"]
     to_new_data_variables = {
         "H20": ["H2O", dims_station_time],
         "IONOS_ELECTRON": ["IONOS_ELECTRON", dims_station_time],
@@ -157,13 +158,23 @@ def create_weather_xds(in_file: str, ant_xds_station_name_ids: xr.DataArray):
     }
 
     to_new_coords = {
-        "TIME": ["time", ["time"]],
+        "TIME": ["time_weather", ["time_weather"]],
     }
 
     weather_xds = convert_generic_xds_to_xradio_schema(
         generic_weather_xds, weather_xds, to_new_data_variables, to_new_coords
     )
 
+    # TODO: option to interpolate to main time
+
+    # PRESSURE: hPa in MSv2 specs and some MSs => Pa
+    weather_xds = convert_to_si_units(weather_xds)
+
+    # correct expected types (for example "IONOS_ELECTRON", "PRESSURE" can be float32)
+    for data_var in weather_xds:
+        if weather_xds.data_vars[data_var].dtype != np.float64:
+            weather_xds[data_var] = weather_xds[data_var].astype(np.float64)
+
     return weather_xds
 
 
@@ -199,9 +210,8 @@ def correct_generic_pointing_xds(
 
     correct_pointing_xds = generic_pointing_xds.copy()
 
-    for key in generic_pointing_xds:
-        if key in to_new_data_variables:
-            data_var_name = to_new_data_variables[key]
+    for data_var_name in generic_pointing_xds:
+        if data_var_name in to_new_data_variables:
             # Corrects dim sizes of "empty cell" variables, such as empty DIRECTION, TARGET, etc.
             if (
                 "dim_2" in generic_pointing_xds.sizes
@@ -296,6 +306,8 @@ def create_pointing_xds(
         size = generic_pointing_xds.sizes["n_polynomial"]
         if size == 1:
             generic_pointing_xds = generic_pointing_xds.sel({"n_polynomial": 0})
+        elif size == 0:
+            generic_pointing_xds = generic_pointing_xds.drop_dims("n_polynomial")
 
     time_ant_dims = ["time", "antenna_name"]
     time_ant_dir_dims = time_ant_dims + ["local_sky_dir_label"]
@@ -341,8 +353,8 @@ def prepare_generic_sys_cal_xds(generic_sys_cal_xds: xr.Dataset) -> xr.Dataset:
     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,
-    - fix dimension names when needed.
+    - transpose the dimensions frequency,receptor
+    - fix dimension names (and order) when needed.
 
     Parameters
     ----------
@@ -374,15 +386,38 @@ def prepare_generic_sys_cal_xds(generic_sys_cal_xds: xr.Dataset) -> xr.Dataset:
         "receptor" in generic_sys_cal_xds.sizes
         and "frequency" in generic_sys_cal_xds.sizes
     ):
+        # dim_3 can be created for example when the T*_SPECTRUM have varying # channels!
+        # more generaly, could transpose with ... to avoid errors with additional spurious dimensions
+        if "dim_3" in generic_sys_cal_xds.dims:
+            generic_sys_cal_xds = generic_sys_cal_xds.drop_dims("dim_3")
         # From MSv2 tables we get (...,frequency, receptor)
         #  -> transpose to (...,receptor,frequency) ready for MSv4 sys_cal_xds
         generic_sys_cal_xds = generic_sys_cal_xds.transpose(
             "ANTENNA_ID", "TIME", "receptor", "frequency"
         )
-    else:
+    elif (
+        "frequency" in generic_sys_cal_xds.sizes
+        and not "dim_3" in generic_sys_cal_xds.sizes
+    ):
         # because order is (...,frequency,receptor), when frequency is missing
         # receptor can get wrongly labeled as frequency
         generic_sys_cal_xds = generic_sys_cal_xds.rename_dims({"frequency": "receptor"})
+    elif (
+        "frequency" not in generic_sys_cal_xds.sizes
+        and "receptor" in generic_sys_cal_xds.sizes
+        and "dim_3" in generic_sys_cal_xds.sizes
+    ):
+        # different *_SPECTRUM array sizes + some empty arrays can create an additional spurious
+        # generic dimension, which should have been "receptor"
+        generic_sys_cal_xds = generic_sys_cal_xds.rename_dims({"receptor": "frequency"})
+        generic_sys_cal_xds = generic_sys_cal_xds.rename_dims({"dim_3": "receptor"})
+        generic_sys_cal_xds = generic_sys_cal_xds.transpose(
+            "ANTENNA_ID", "TIME", "receptor", "frequency"
+        )
+    else:
+        raise RuntimeError(
+            "Cannot understand the arrangement of dimensions of {generic_sys_cal_xds=}"
+        )
 
     return generic_sys_cal_xds
 
@@ -462,7 +497,7 @@ def create_system_calibration_xds(
         "frequency": ["frequency_cal", ["frequency_cal"]],
     }
 
-    sys_cal_xds = xr.Dataset(attrs={"type": "sys_cal"})
+    sys_cal_xds = xr.Dataset(attrs={"type": "system_calibration"})
     coords = {
         "antenna_name": ant_xds_name_ids.sel(
             antenna_id=generic_sys_cal_xds["ANTENNA_ID"]
@@ -483,8 +518,7 @@ def create_system_calibration_xds(
         frequency_measure = {
             "type": main_xds_frequency.attrs["type"],
             "units": main_xds_frequency.attrs["units"],
-            "frame": main_xds_frequency.attrs["frame"],
-            "reference_value": main_xds_frequency.attrs["reference_frequency"],
+            "observer": main_xds_frequency.attrs["observer"],
         }
         sys_cal_xds.coords["frequency_cal"].attrs.update(frequency_measure)
 
@@ -499,8 +533,8 @@ def create_system_calibration_xds(
         time_coord_attrs = {
             "type": "time",
             "units": ["s"],
-            "scale": "UTC",
-            "format": "UNIX",
+            "scale": "utc",
+            "format": "unix",
         }
         # If interpolating time, rename time_cal => time
         time_coord = {"time": ("time_cal", sys_cal_interp_time.data)}

xradio/{correlated_data → measurement_set}/_utils/_utils/xds_helper.py

@@ -7,7 +7,7 @@ import xarray as xr
 
 from .cds import CASAVisSet
 from .stokes_types import stokes_types
-from ...._utils.common import get_pad_value
+from xradio._utils.list_and_array import get_pad_value
 
 
 def make_coords(

xradio/{correlated_data/_utils/ms.py → measurement_set/_utils/msv2.py}

@@ -2,14 +2,14 @@ import os
 import toolviper.utils.logger as logger
 from typing import List, Tuple, Union
 
-from ._utils.cds import CASAVisSet
-from ._ms.partitions import (
+from xradio.measurement_set._utils._utils.cds import CASAVisSet
+from xradio.measurement_set._utils._msv2.partitions import (
     finalize_partitions,
     read_ms_ddi_partitions,
     read_ms_scan_subscan_partitions,
 )
-from ._ms.subtables import read_ms_subtables
-from ._utils.xds_helper import vis_xds_packager_cds
+from xradio.measurement_set._utils._msv2.subtables import read_ms_subtables
+from xradio.measurement_set._utils._utils.xds_helper import vis_xds_packager_cds
 
 
 def read_ms(

xradio/{correlated_data → measurement_set}/convert_msv2_to_processing_set.py

@@ -4,8 +4,8 @@ from typing import Dict, Union
 
 import dask
 
-from xradio.correlated_data._utils._ms.partition_queries import create_partitions
-from xradio.correlated_data._utils._ms.conversion import convert_and_write_partition
+from xradio.measurement_set._utils._msv2.partition_queries import create_partitions
+from xradio.measurement_set._utils._msv2.conversion import convert_and_write_partition
 
 
 def convert_msv2_to_processing_set(

xradio/{correlated_data → measurement_set}/load_processing_set.py

@@ -1,5 +1,5 @@
 import os
-from xradio.correlated_data import ProcessingSet
+from xradio.measurement_set import ProcessingSet
 from typing import Dict, Union
 
 
@@ -43,10 +43,10 @@ def load_processing_set(
     ps = ProcessingSet()
     for ms_name, ms_xds_isel in sel_parms.items():
         ms_store = os.path.join(ps_store, ms_name)
-        ms_main_store = os.path.join(ms_store, "MAIN")
+        correlated_store = os.path.join(ms_store, "correlated_xds")
 
         xds = _open_dataset(
-            ms_main_store,
+            correlated_store,
             file_system,
             ms_xds_isel,
             data_variables,
@@ -55,7 +55,7 @@ def load_processing_set(
         data_groups = xds.attrs["data_groups"]
 
         if load_sub_datasets:
-            from xradio.correlated_data.open_processing_set import _open_sub_xds
+            from xradio.measurement_set.open_processing_set import _open_sub_xds
 
             sub_xds_dict, field_and_source_xds_dict = _open_sub_xds(
                 ms_store, file_system=file_system, load=True, data_groups=data_groups
@@ -76,7 +76,7 @@ def load_processing_set(
     return ps
 
 
-class processing_set_iterator:
+class ProcessingSetIterator:
     def __init__(
         self,
         sel_parms: dict,

xradio/measurement_set/measurement_set_xds.py

@@ -0,0 +1,83 @@
+import pandas as pd
+from xradio._utils.list_and_array import to_list
+import xarray as xr
+import numbers
+import os
+
+
+class MeasurementSetXds(xr.Dataset):
+    __slots__ = ()
+
+    def __init__(self, xds):
+        super().__init__(xds.data_vars, xds.coords, xds.attrs)
+
+    def to_store(self, store, **kwargs):
+        """
+        Write the MeasurementSetXds to a Zarr store.
+        Does not write to cloud storage yet.
+
+        Args:
+            store (str): The path to the Zarr store.
+            **kwargs: Additional keyword arguments to be passed to `xarray.Dataset.to_zarr`. See https://docs.xarray.dev/en/latest/generated/xarray.Dataset.to_zarr.html for more information.
+
+        Returns:
+            None
+        """
+
+        copy_cor_xds = self.copy()  # No deep copy
+
+        # Remove field_and_source_xds from all correlated_data (VISIBILITY/SPECTRUM) data variables
+        # and save them as separate zarr files.
+        for data_group_name, data_group in self.attrs["data_groups"].items():
+            del copy_cor_xds[data_group["correlated_data"]].attrs[
+                "field_and_source_xds"
+            ]
+
+            # print("data_group_name", data_group_name)
+            xr.Dataset.to_zarr(
+                self[data_group["correlated_data"]].attrs["field_and_source_xds"],
+                os.path.join(store, "field_and_source_xds_" + data_group_name),
+                **kwargs,
+            )
+
+        # Remove xds attributes from copy_cor_xds and save xds attributes as separate zarr files.
+        for attrs_name in self.attrs:
+            if "xds" in attrs_name:
+                del copy_cor_xds.attrs[attrs_name]
+                xr.Dataset.to_zarr(
+                    self.attrs[attrs_name], os.path.join(store, attrs_name, **kwargs)
+                )
+
+        # Save copy_cor_xds as zarr file.
+        xr.Dataset.to_zarr(
+            copy_cor_xds, os.path.join(store, "correlated_xds"), **kwargs
+        )
+
+    def sel(self, data_group_name=None, **kwargs):
+        """
+        Select data along dimension(s) by label.
+
+        Args:
+            **kwargs: Keyword arguments to be passed to `xarray.Dataset.sel`. See https://xarray.pydata.org/en/stable/generated/xarray.Dataset.sel.html for more information.
+
+        Returns:
+            MeasurementSetXds
+        """
+
+        if data_group_name is not None:
+            sel_data_group_set = set(
+                self.attrs["data_groups"][data_group_name].values()
+            )
+
+            data_variables_to_drop = []
+            for dg in self.attrs["data_groups"].values():
+                temp_set = set(dg.values()) - sel_data_group_set
+                data_variables_to_drop.extend(list(temp_set))
+
+            data_variables_to_drop = list(set(data_variables_to_drop))
+
+            return MeasurementSetXds(
+                super().sel(**kwargs).drop_vars(data_variables_to_drop)
+            )
+        else:
+            return MeasurementSetXds(super().sel(**kwargs))

xradio/{correlated_data → measurement_set}/open_processing_set.py

@@ -1,6 +1,6 @@
 import os
 
-from xradio.correlated_data import ProcessingSet
+from xradio.measurement_set import ProcessingSet
 import toolviper.utils.logger as logger
 from xradio._utils.zarr.common import _open_dataset, _get_file_system_and_items
 import s3fs
@@ -25,6 +25,8 @@ def open_processing_set(
     processing_set
         Lazy representation of processing set (data is represented by Dask.arrays).
     """
+    from xradio.measurement_set import MeasurementSetXds
+
     file_system, ms_store_list = _get_file_system_and_items(ps_store)
 
     ps = ProcessingSet()
@@ -32,9 +34,9 @@ def open_processing_set(
     for ms_name in ms_store_list:
         # try:
         ms_store = os.path.join(ps_store, ms_name)
-        ms_main_store = os.path.join(ms_store, "MAIN")
+        correlated_store = os.path.join(ms_store, "correlated_xds")
 
-        xds = _open_dataset(ms_main_store, file_system)
+        xds = _open_dataset(correlated_store, file_system)
         data_groups = xds.attrs["data_groups"]
 
         if (intents is None) or (
@@ -54,7 +56,7 @@ def open_processing_set(
                     "field_and_source_xds"
                 ] = field_and_source_xds_dict[data_group_name]
 
-            ps[ms_name] = xds
+            ps[ms_name] = MeasurementSetXds(xds)
         # except Exception as e:
         #     logger.warning(f"Could not open {ms_name} due to {e}")
         #     continue
@@ -66,15 +68,6 @@ def _open_sub_xds(ms_store, file_system, data_groups, load=False):
     sub_xds_dict = {}
     field_and_source_xds_dict = {}
 
-    xds_names = {
-        "ANTENNA": "antenna_xds",
-        "POINTING": "pointing_xds",
-        "SYSCAL": "system_calibration_xds",
-        "GAIN_CURVE": "gain_curve_xds",
-        "PHASE_CAL": "phase_calibration_xds",
-        "WEATHER": "weather_xds",
-    }
-
     if isinstance(file_system, s3fs.core.S3FileSystem):
         file_names = [
             bd.split(sep="/")[-1] for bd in file_system.listdir(ms_store, detail=False)
@@ -83,9 +76,9 @@ def _open_sub_xds(ms_store, file_system, data_groups, load=False):
         file_names = file_system.listdir(ms_store)
     file_names = [item for item in file_names if not item.startswith(".")]
 
-    file_names.remove("MAIN")
+    file_names.remove("correlated_xds")
 
-    field_dict = {"FIELD_AND_SOURCE_" + key.upper(): key for key in data_groups.keys()}
+    field_dict = {"field_and_source_xds_" + key: key for key in data_groups.keys()}
 
     # field_and_source_xds_name_start = "FIELD"
     for n in file_names:
@@ -98,7 +91,7 @@ def _open_sub_xds(ms_store, file_system, data_groups, load=False):
         if n in field_dict.keys():
             field_and_source_xds_dict[field_dict[n]] = xds
         else:
-            sub_xds_dict[xds_names[n]] = xds
+            sub_xds_dict[n] = xds
 
     return sub_xds_dict, field_and_source_xds_dict