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

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/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/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/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.dist-info → xradio-0.0.50.dist-info}/METADATA

@@ -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.dist-info → xradio-0.0.50.dist-info}/RECORD

@@ -27,15 +27,15 @@ xradio/image/_util/_zarr/zarr_low_level.py,sha256=xnYm6EmVbmLxMlOSXH32SABfQBLHfr
 xradio/measurement_set/__init__.py,sha256=Vrr1Py50TvbzeZ_VMCswYNz0Wcccbf-iJDj4ArlfcJ0,870
 xradio/measurement_set/convert_msv2_to_processing_set.py,sha256=uLZjXplVPXa0XnNa-Fty85k_-fsw6ZC98Hfiwd1WF-U,9704
 xradio/measurement_set/load_processing_set.py,sha256=8EPApyGy0Tmzu6Seeby7dKxvtxtAFA585kK50DYVHas,8164
-xradio/measurement_set/measurement_set_xdt.py,sha256=t9CKwgZlnogE6KoGQSVDwHel9FzYiOlA0ldoRyt5jZo,7636
+xradio/measurement_set/measurement_set_xdt.py,sha256=kN337gyn7Q8nF4ENy292PYsmBJJLu5ozhJ3FyT5BcVo,11986
 xradio/measurement_set/open_processing_set.py,sha256=kMODJmXT2KU12L6Y2NdTV8shvLGb5PgLIOqJgMCzlHI,5308
-xradio/measurement_set/processing_set_xdt.py,sha256=fIsfhhnBg5uq7KU0dDRJpyP12jpHo6f4ppbRJTqUAAc,59342
-xradio/measurement_set/schema.py,sha256=ppH-iibynbCw3kp_Ms_scpCanJNKnBHYdXC4GMwbLAU,85646
+xradio/measurement_set/processing_set_xdt.py,sha256=LhonH1v9knkcFg5wYC1KHBB5pqhH0xPH6EkIms3-0xM,59353
+xradio/measurement_set/schema.py,sha256=6vQEuI81V9mQ5y-Az310NfwgbdOEXkJDYQfIoZuMPdc,85757
 xradio/measurement_set/_utils/__init__.py,sha256=XE-h1yMfr6tVD6gdUwXO1CVq5SQ6kD_oj-e5TFwslds,97
 xradio/measurement_set/_utils/msv2.py,sha256=7hnZMFoQ-s1g0ATjEupLvtdqQCdroPv-Rl5OwjqXjh8,4430
 xradio/measurement_set/_utils/zarr.py,sha256=ehXlu0Xh_UZ5Xm2RnHCxESsRZ26c3DQAO5rqMK5MwTk,3947
 xradio/measurement_set/_utils/_msv2/chunks.py,sha256=JTPk3il6fk570BjWZMoOAtsbvnLmqPcBv9EPY6A2yOs,2964
-xradio/measurement_set/_utils/_msv2/conversion.py,sha256=KH6DQ9HKeAaIkXBjDNKKnUuQI7xp647ucLqu9m_TsFE,52931
+xradio/measurement_set/_utils/_msv2/conversion.py,sha256=ERpoGBpi2x2qgX-NSql0LtpWE-uy_PG8uTvE_lJcQuY,52997
 xradio/measurement_set/_utils/_msv2/create_antenna_xds.py,sha256=qLUDxbkJBOaD7EaVx7ufiU0CL5f8VVxK-923-j4XpXc,17758
 xradio/measurement_set/_utils/_msv2/create_field_and_source_xds.py,sha256=iyR5LzCupA6ZJmIOjIylwCQRI1EU4JEofBL7sbDWzwA,37069
 xradio/measurement_set/_utils/_msv2/descr.py,sha256=PGY39PYQj0K4th5RUv0jOWszcHlZDt6VQRTOuntCeYI,5213
@@ -70,8 +70,8 @@ xradio/schema/metamodel.py,sha256=WjtW7pAVzcjLRWifRH3sQoOiN6TV810hARpOIz1M_gw,38
 xradio/schema/typing.py,sha256=8-o6fZd99kJ4FVdgBYRTIRJ-wDqpcUNXzCTfJvl3TIw,10439
 xradio/sphinx/__init__.py,sha256=VGY-7Ty3q67qpnBee0-znbiJ-Iy0F93UO--IpjEdHXc,380
 xradio/sphinx/schema_table.py,sha256=uq33habbAbReqnEG6ASKSd4UOMZGUzA3qoTX45rq84U,12373
-xradio-0.0.49.dist-info/licenses/LICENSE.txt,sha256=9CYIJt7riOXo9AD0eXBZviLxo_HebD-2JJI8oiWtzfg,1807
-xradio-0.0.49.dist-info/METADATA,sha256=-rwVHT0ItsA4EO7VmVQ0gmCmVrUpgOvufSxEhcIPfbc,4441
-xradio-0.0.49.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xradio-0.0.49.dist-info/top_level.txt,sha256=dQu27fGBZJ2Yk-gW5XeD-dZ76Xa4Xcvk60Vz-dwXp7k,7
-xradio-0.0.49.dist-info/RECORD,,
+xradio-0.0.50.dist-info/licenses/LICENSE.txt,sha256=9CYIJt7riOXo9AD0eXBZviLxo_HebD-2JJI8oiWtzfg,1807
+xradio-0.0.50.dist-info/METADATA,sha256=XpYcb4kWmCYjPkjybQhk-fJxKyRwM2W0wjyOMfr_s74,4471
+xradio-0.0.50.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xradio-0.0.50.dist-info/top_level.txt,sha256=dQu27fGBZJ2Yk-gW5XeD-dZ76Xa4Xcvk60Vz-dwXp7k,7
+xradio-0.0.50.dist-info/RECORD,,