xradio 0.0.49__tar.gz → 0.0.50__tar.gz

{xradio-0.0.49/src/xradio.egg-info → xradio-0.0.50}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xradio
-Version: 0.0.49
+Version: 0.0.50
 Summary:  Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License
@@ -56,6 +56,7 @@ Requires-Dist: zarr<3,>=2
 Requires-Dist: pyarrow
 Requires-Dist: python_casacore>=3.6.1; sys_platform != "darwin"
 Requires-Dist: typeguard
+Requires-Dist: numcodecs<0.16
 Provides-Extra: interactive
 Requires-Dist: matplotlib; extra == "interactive"
 Requires-Dist: prettytable; extra == "interactive"

{xradio-0.0.49 → xradio-0.0.50}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "xradio"
-version = "0.0.49"
+version = "0.0.50"
 description = " Xarray Radio Astronomy Data IO"
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},
@@ -25,7 +25,8 @@ dependencies = [
     'zarr>=2,<3',
     'pyarrow',
     'python_casacore>=3.6.1; sys_platform != "darwin" ',
-    'typeguard'
+    'typeguard',
+    'numcodecs<0.16'
 ]
 
 [project.optional-dependencies]

{xradio-0.0.49 → xradio-0.0.50}/src/xradio/measurement_set/_utils/_msv2/conversion.py

@@ -1341,7 +1341,7 @@ def convert_and_write_partition(
 
         ms_xdt["/antenna_xds"] = ant_xds
         for group_name in xds.attrs["data_groups"]:
-            ms_xdt["/" + f"field_and_source_xds_{group_name}"] = field_and_source_xds
+            ms_xdt["/" + f"field_and_source_{group_name}_xds"] = field_and_source_xds
 
         if with_pointing and len(pointing_xds.data_vars) > 0:
             ms_xdt["/pointing_xds"] = pointing_xds
@@ -1468,6 +1468,7 @@ def add_group_to_data_groups(
         "correlated_data": correlated_data_name,
         "flag": "FLAG",
         "weight": "WEIGHT",
+        "field_and_source": f"field_and_source_{what_group}_xds",
         "description": f"Data group derived from the data column '{correlated_data_name}' of an MSv2 converted to MSv4",
         "date": datetime.datetime.now(datetime.timezone.utc).isoformat(),
     }

{xradio-0.0.49 → xradio-0.0.50}/src/xradio/measurement_set/measurement_set_xdt.py

@@ -93,17 +93,34 @@ class MeasurementSetXdt:
         if data_group_name is not None:
             sel_data_group_set = set(
                 self._xdt.attrs["data_groups"][data_group_name].values()
-            )
+            ) - set(["date", "description"])
+
+            sel_field_and_source_xds = self._xdt.attrs["data_groups"][data_group_name][
+                "field_and_source"
+            ]
 
             data_variables_to_drop = []
-            for dg in self._xdt.attrs["data_groups"].values():
-                temp_set = set(dg.values()) - sel_data_group_set
+            field_and_source_to_drop = []
+            for dg_name, dg in self._xdt.attrs["data_groups"].items():
+                print(f"Data group: {dg_name}", dg)
+                f_and_s = dg["field_and_source"]
+                dg_copy = dg.copy()
+                dg_copy.pop("date", None)
+                dg_copy.pop("description", None)
+                dg_copy.pop("field_and_source", None)
+                temp_set = set(dg_copy.values()) - sel_data_group_set
                 data_variables_to_drop.extend(list(temp_set))
 
+                if f_and_s != sel_field_and_source_xds:
+                    field_and_source_to_drop.append(f_and_s)
+
             data_variables_to_drop = list(set(data_variables_to_drop))
 
             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)
+
             sel_corr_xds = self._xdt.ds.sel(
                 indexers, method, tolerance, drop, **indexers_kwargs
             ).drop_vars(data_variables_to_drop)
@@ -140,7 +157,10 @@ class MeasurementSetXdt:
             else:
                 data_group_name = list(self._xdt.attrs["data_groups"].keys())[0]
 
-        return self._xdt[f"field_and_source_xds_{data_group_name}"].ds
+        field_and_source_xds_name = self._xdt.attrs["data_groups"][data_group_name][
+            "field_and_source"
+        ]
+        return self._xdt[field_and_source_xds_name].ds
 
     def get_partition_info(self, data_group_name: str = None) -> dict:
         """
@@ -195,3 +215,105 @@ class MeasurementSetXdt:
         }
 
         return partition_info
+
+    def add_data_group(
+        self,
+        new_data_group_name: str,
+        correlated_data: str = None,
+        weight: str = None,
+        flag: str = None,
+        uvw: str = None,
+        field_and_source_xds: str = None,
+        date_time: str = None,
+        description: str = None,
+        data_group_dv_shared_with: str = None,
+    ) -> xr.DataTree:
+        """_summary_
+
+        Parameters
+        ----------
+        new_data_group_name : str
+            _description_
+        correlated_data : str, optional
+            _description_, by default None
+        weights : str, optional
+            _description_, by default None
+        flag : str, optional
+            _description_, by default None
+        uvw : str, optional
+            _description_, by default None
+        field_and_source_xds : str, optional
+            _description_, by default None
+        date_time : str, optional
+            _description_, by default None
+        description : str, optional
+            _description_, by default None
+        data_group_dv_shared_with : str, optional
+            _description_, by default "base"
+
+        Returns
+        -------
+        xr.DataTree
+            _description_
+        """
+
+        if data_group_dv_shared_with is None:
+            data_group_dv_shared_with = self._xdt.xr_ms._get_default_data_group_name()
+        default_data_group = self._xdt.attrs["data_groups"][data_group_dv_shared_with]
+
+        new_data_group = {}
+
+        if correlated_data is None:
+            correlated_data = default_data_group["correlated_data"]
+        new_data_group["correlated_data"] = correlated_data
+        assert (
+            correlated_data in self._xdt.ds.data_vars
+        ), f"Data variable {correlated_data} not found in dataset."
+
+        if weight is None:
+            weight = default_data_group["weight"]
+        new_data_group["weight"] = weight
+        assert (
+            weight in self._xdt.ds.data_vars
+        ), f"Data variable {weight} not found in dataset."
+
+        if flag is None:
+            flag = default_data_group["flag"]
+        new_data_group["flag"] = flag
+        assert (
+            flag in self._xdt.ds.data_vars
+        ), f"Data variable {flag} not found in dataset."
+
+        if self._xdt.attrs["type"] == "visibility":
+            if uvw is None:
+                uvw = default_data_group["uvw"]
+            new_data_group["uvw"] = uvw
+            assert (
+                uvw in self._xdt.ds.data_vars
+            ), 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"]
+        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()
+        new_data_group["date"] = date_time
+
+        if description is None:
+            description = ""
+        new_data_group["description"] = description
+
+        self._xdt.attrs["data_groups"][new_data_group_name] = new_data_group
+
+        return self._xdt
+
+    def _get_default_data_group_name(self):
+        if "base" in self._xdt.attrs["data_groups"].keys():
+            data_group_name = "base"
+        else:
+            data_group_name = list(self._xdt.attrs["data_groups"].keys())[0]
+        return data_group_name

{xradio-0.0.49 → xradio-0.0.50}/src/xradio/measurement_set/processing_set_xdt.py

@@ -193,7 +193,7 @@ class ProcessingSetXdt:
             )
             summary_data["end_frequency"].append(to_list(value["frequency"].values)[-1])
 
-            field_and_source_xds = value["field_and_source_xds_" + data_group]
+            field_and_source_xds = value.xr_ms.get_field_and_source_xds(data_group)
 
             if field_and_source_xds.attrs["type"] == "field_and_source_ephemeris":
                 summary_data["field_coords"].append("Ephemeris")
@@ -349,7 +349,7 @@ class ProcessingSetXdt:
                 "correlated_data"
             ]
 
-            field_and_source_xds = ms_xdt["field_and_source_xds_" + data_group].ds
+            field_and_source_xds = ms_xdt.xr_ms.get_field_and_source_xds(data_group)
 
             if not field_and_source_xds.attrs["type"] == "field_and_source_ephemeris":
 
@@ -450,9 +450,9 @@ class ProcessingSetXdt:
                 "correlated_data"
             ]
 
-            field_and_source_xds = field_and_source_xds = ms_xdt[
-                "field_and_source_xds_" + data_group
-            ].ds
+            field_and_source_xds = field_and_source_xds = (
+                ms_xdt.xr_ms.get_field_and_source_xds(data_group)
+            )
 
             if field_and_source_xds.attrs["type"] == "field_and_source_ephemeris":
 

{xradio-0.0.49 → xradio-0.0.50}/src/xradio/measurement_set/schema.py

@@ -1410,6 +1410,8 @@ class DataGroupDict:
     """ Name of the weight variable of the group, for example 'WEIGHT'. """
     uvw: Optional[str]
     """ Name of the UVW variable of the group, for example 'UVW'. """
+    field_and_source: str
+    """ Name of the field_and_source_xds, for example field_and_source_xds_base. """
     description: str
     """ More details about the data group. """
     date: str

{xradio-0.0.49 → xradio-0.0.50/src/xradio.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xradio
-Version: 0.0.49
+Version: 0.0.50
 Summary:  Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License
@@ -56,6 +56,7 @@ Requires-Dist: zarr<3,>=2
 Requires-Dist: pyarrow
 Requires-Dist: python_casacore>=3.6.1; sys_platform != "darwin"
 Requires-Dist: typeguard
+Requires-Dist: numcodecs<0.16
 Provides-Extra: interactive
 Requires-Dist: matplotlib; extra == "interactive"
 Requires-Dist: prettytable; extra == "interactive"

{xradio-0.0.49 → xradio-0.0.50}/src/xradio.egg-info/requires.txt

@@ -13,6 +13,7 @@ xarray
 zarr<3,>=2
 pyarrow
 typeguard
+numcodecs<0.16
 
 [:sys_platform != "darwin"]
 python_casacore>=3.6.1