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

xradio/measurement_set/_utils/_zarr/encoding.py

@@ -4,11 +4,6 @@ def add_encoding(xds, compressor, chunks=None):
 
     chunks = {**dict(xds.sizes), **chunks}  # Add missing sizes if presents.
 
-    encoding = {}
     for da_name in list(xds.data_vars):
-        if chunks:
-            da_chunks = [chunks[dim_name] for dim_name in xds[da_name].sizes]
-            xds[da_name].encoding = {"compressor": compressor, "chunks": da_chunks}
-            # print(xds[da_name].encoding)
-        else:
-            xds[da_name].encoding = {"compressor": compressor}
+        da_chunks = [chunks[dim_name] for dim_name in xds[da_name].sizes]
+        xds[da_name].encoding = {"compressor": compressor, "chunks": da_chunks}

xradio/measurement_set/convert_msv2_to_processing_set.py

@@ -116,8 +116,6 @@ def convert_msv2_to_processing_set(
     if not str(out_file).endswith("ps.zarr"):
         out_file += ".ps.zarr"
 
-    print("Output file: ", out_file)
-
     if overwrite:
         ps_dt.to_zarr(store=out_file, mode="w")
     else:

xradio/measurement_set/load_processing_set.py

@@ -168,7 +168,7 @@ class ProcessingSetIterator:
     def __next__(self):
         try:
             sub_xds_name = next(self.xds_name_iter)
-        except Exception as e:
+        except Exception:
             raise StopIteration
 
         if self.input_data is None:
@@ -181,7 +181,7 @@ class ProcessingSetIterator:
                 drop_variables=self.drop_variables,
                 load_sub_datasets=self.load_sub_datasets,
             )
-            sub_xdt = ps_xdt.get(0)
+            sub_xdt = ps_xdt[sub_xds_name]  # Access by name instead of index
         else:
             sub_xdt = self.input_data[sub_xds_name]  # In memory
 

xradio/measurement_set/measurement_set_xdt.py

@@ -1,12 +1,12 @@
-import pandas as pd
-from xradio._utils.list_and_array import to_list
-import xarray as xr
-import numpy as np
-import numbers
-import os
 from collections.abc import Mapping, Iterable
+import datetime
 from typing import Any, Union
 
+import numpy as np
+import xarray as xr
+
+from xradio._utils.list_and_array import to_list
+
 MS_DATASET_TYPES = {"visibility", "spectrum", "radiometer"}
 
 
@@ -102,7 +102,7 @@ class MeasurementSetXdt:
             data_variables_to_drop = []
             field_and_source_to_drop = []
             for dg_name, dg in self._xdt.attrs["data_groups"].items():
-                print(f"Data group: {dg_name}", dg)
+                # print(f"Data group: {dg_name}", dg)
                 f_and_s = dg["field_and_source"]
                 dg_copy = dg.copy()
                 dg_copy.pop("date", None)
@@ -118,8 +118,8 @@ class MeasurementSetXdt:
 
             sel_ms_xdt = self._xdt
 
-            print("Data variables to drop: ", data_variables_to_drop)
-            print("Field and source to drop: ", field_and_source_to_drop)
+            # print("Data variables to drop: ", data_variables_to_drop)
+            # print("Field and source to drop: ", field_and_source_to_drop)
 
             sel_corr_xds = self._xdt.ds.sel(
                 indexers, method, tolerance, drop, **indexers_kwargs
@@ -152,7 +152,7 @@ class MeasurementSetXdt:
             raise InvalidAccessorLocation(f"{self._xdt.path} is not a MSv4 node.")
 
         if data_group_name is None:
-            if "base" in self._xdt.attrs["data_groups"].keys():
+            if "base" in self._xdt.attrs["data_groups"]:
                 data_group_name = "base"
             else:
                 data_group_name = list(self._xdt.attrs["data_groups"].keys())[0]
@@ -189,7 +189,7 @@ class MeasurementSetXdt:
             )
 
         if data_group_name is None:
-            if "base" in self._xdt.attrs["data_groups"].keys():
+            if "base" in self._xdt.attrs["data_groups"]:
                 data_group_name = "base"
             else:
                 data_group_name = list(self._xdt.attrs["data_groups"].keys())[0]
@@ -205,6 +205,9 @@ class MeasurementSetXdt:
 
         partition_info = {
             "spectral_window_name": self._xdt.frequency.attrs["spectral_window_name"],
+            "spectral_window_intent": self._xdt.frequency.attrs[
+                "spectral_window_intent"
+            ],
             "field_name": to_list(np.unique(field_and_source_xds.field_name.values)),
             "polarization_setup": to_list(self._xdt.polarization.values),
             "scan_name": to_list(np.unique(self._xdt.scan_name.values)),
@@ -228,7 +231,8 @@ class MeasurementSetXdt:
         description: str = None,
         data_group_dv_shared_with: str = None,
     ) -> xr.DataTree:
-        """_summary_
+        """Adds a data group to the MSv4 DataTree, grouping the given data, weight, flag, etc. variables
+        and field_and_source_xds.
 
         Parameters
         ----------
@@ -254,7 +258,7 @@ class MeasurementSetXdt:
         Returns
         -------
         xr.DataTree
-            _description_
+            MSv4 DataTree with the new group added
         """
 
         if data_group_dv_shared_with is None:
@@ -293,14 +297,14 @@ class MeasurementSetXdt:
             ), f"Data variable {uvw} not found in dataset."
 
         if field_and_source_xds is None:
-            field_and_source_xds = default_data_group["field_and_source_xds"]
+            field_and_source_xds = default_data_group["field_and_source"]
         new_data_group["field_and_source"] = field_and_source_xds
         assert (
             field_and_source_xds in self._xdt.children
         ), f"Data variable {field_and_source_xds} not found in dataset."
 
         if date_time is None:
-            date_time = datetime.now().isoformat()
+            date_time = datetime.datetime.now(datetime.timezone.utc).isoformat()
         new_data_group["date"] = date_time
 
         if description is None:
@@ -312,7 +316,7 @@ class MeasurementSetXdt:
         return self._xdt
 
     def _get_default_data_group_name(self):
-        if "base" in self._xdt.attrs["data_groups"].keys():
+        if "base" in self._xdt.attrs["data_groups"]:
             data_group_name = "base"
         else:
             data_group_name = list(self._xdt.attrs["data_groups"].keys())[0]

xradio/measurement_set/open_processing_set.py

@@ -1,6 +1,4 @@
-import os
-import toolviper.utils.logger as logger
-from xradio._utils.zarr.common import _open_dataset, _get_file_system_and_items
+from xradio._utils.zarr.common import _get_file_system_and_items
 import s3fs
 import xarray as xr