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

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]
@@ -293,14 +293,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 +312,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