xradio 0.0.29__tar.gz → 0.0.30__tar.gz

{xradio-0.0.29/src/xradio.egg-info → xradio-0.0.30}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.29
+Version: 0.0.30
 Summary:  Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License

{xradio-0.0.29 → xradio-0.0.30}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "xradio"
-version = "0.0.29"
+version = "0.0.30"
 description = " Xarray Radio Astronomy Data IO"
 authors = [
     {name = "Jan-Willem Steeb", email="jsteeb@nrao.edu"},

{xradio-0.0.29 → xradio-0.0.30}/src/xradio/vis/_processing_set.py

@@ -6,12 +6,6 @@ class processing_set(dict):
         super().__init__(*args, **kwargs)
         self.meta = {"summary": {}}
 
-    #     generate_meta(self)
-
-    # def generate_meta(self):
-    #     self.meta['summary'] = {"base": _summary(self)}
-    #     self.meta['max_dims'] = _get_ps_max_dims(self)
-
     def summary(self, data_group="base"):
         if data_group in self.meta["summary"]:
             return self.meta["summary"][data_group]
@@ -36,28 +30,34 @@ class processing_set(dict):
     def _summary(self, data_group="base"):
         summary_data = {
             "name": [],
-            "ddi": [],
             "intent": [],
-            "field_id": [],
+            "shape": [],
+            "polarization": [],
+            "spw_id": [],
             "field_name": [],
+            "field_id": [],
+            "field_coords": [],
             "start_frequency": [],
             "end_frequency": [],
-            "shape": [],
-            "field_coords": [],
         }
         from astropy.coordinates import SkyCoord
         import astropy.units as u
 
         for key, value in self.items():
             summary_data["name"].append(key)
-            summary_data["ddi"].append(value.attrs["ddi"])
-            summary_data["intent"].append(value.attrs["intent"])
+            summary_data["intent"].append(value.attrs["partition_info"]["intent"])
+            summary_data["spw_id"].append(
+                value.attrs["partition_info"]["spectral_window_id"]
+            )
+            summary_data["polarization"].append(value.polarization.values)
 
             if "visibility" in value.attrs["data_groups"][data_group]:
                 data_name = value.attrs["data_groups"][data_group]["visibility"]
+                center_name = "FIELD_PHASE_CENTER"
 
             if "spectrum" in value.attrs["data_groups"][data_group]:
                 data_name = value.attrs["data_groups"][data_group]["spectrum"]
+                center_name = "FIELD_REFERENCE_CENTER"
 
             summary_data["shape"].append(value[data_name].shape)
 
@@ -72,13 +72,11 @@ class processing_set(dict):
                 summary_data["field_coords"].append("Ephemeris")
             else:
                 ra_dec_rad = (
-                    value[data_name]
-                    .attrs["field_and_source_xds"]["FIELD_PHASE_CENTER"]
-                    .values
+                    value[data_name].attrs["field_and_source_xds"][center_name].values
                 )
                 frame = (
                     value[data_name]
-                    .attrs["field_and_source_xds"]["FIELD_PHASE_CENTER"]
+                    .attrs["field_and_source_xds"][center_name]
                     .attrs["frame"]
                     .lower()
                 )
@@ -90,8 +88,8 @@ class processing_set(dict):
                 summary_data["field_coords"].append(
                     [
                         frame,
-                        coord.ra.to_string(unit=u.hour),
-                        coord.dec.to_string(unit=u.deg),
+                        coord.ra.to_string(unit=u.hour, precision=2),
+                        coord.dec.to_string(unit=u.deg, precision=2),
                     ]
                 )
 
@@ -131,3 +129,36 @@ class processing_set(dict):
 
     def get(self, id):
         return self[list(self.keys())[id]]
+
+    def sel(self, **kwargs):
+        import numpy as np
+
+        summary_table = self.summary()
+        for key, value in kwargs.items():
+            if isinstance(value, list) or isinstance(value, np.ndarray):
+                summary_table = summary_table[summary_table[key].isin(value)]
+            elif isinstance(value, slice):
+                summary_table = summary_table[
+                    summary_table[key].between(value.start, value.stop)
+                ]
+            else:
+                summary_table = summary_table[summary_table[key] == value]
+
+        sub_ps = processing_set()
+        for key, val in self.items():
+            if key in summary_table["name"].values:
+                sub_ps[key] = val
+
+        return sub_ps
+
+    def ms_sel(self, **kwargs):
+        sub_ps = processing_set()
+        for key, val in self.items():
+            sub_ps[key] = val.sel(kwargs)
+        return sub_ps
+
+    def ms_isel(self, **kwargs):
+        sub_ps = processing_set()
+        for key, val in self.items():
+            sub_ps[key] = val.isel(kwargs)
+        return sub_ps

{xradio-0.0.29 → xradio-0.0.30}/src/xradio/vis/_vis_utils/_ms/conversion.py

@@ -754,9 +754,6 @@ def convert_and_write_partition(
 
             start = time.time()
 
-            xds.attrs["intent"] = intent
-            xds.attrs["ddi"] = ddi
-
             # Time and frequency should always be increasing
             if len(xds.frequency) > 1 and xds.frequency[1] - xds.frequency[0] < 0:
                 xds = xds.sel(frequency=slice(None, None, -1))
@@ -765,41 +762,7 @@ def convert_and_write_partition(
                 xds = xds.sel(time=slice(None, None, -1))
 
             # Add data_groups and field_info
-            xds.attrs["data_groups"] = {}
-            if "VISIBILITY" in xds:
-                xds.attrs["data_groups"]["base"] = {
-                    "visibility": "VISIBILITY",
-                    "flag": "FLAG",
-                    "weight": "WEIGHT",
-                    "uvw": "UVW",
-                }
-
-            if "VISIBILITY_CORRECTED" in xds:
-                xds.attrs["data_groups"]["corrected"] = {
-                    "visibility": "VISIBILITY_CORRECTED",
-                    "flag": "FLAG",
-                    "weight": "WEIGHT",
-                    "uvw": "UVW",
-                }
-
-            is_single_dish = False
-            if "SPECTRUM" in xds:
-                xds.attrs["data_groups"]["base"] = {
-                    "spectrum": "SPECTRUM",
-                    "flag": "FLAG",
-                    "weight": "WEIGHT",
-                    "uvw": "UVW",
-                }
-                is_single_dish = True
-
-            if "SPECTRUM_CORRECTED" in xds:
-                xds.attrs["data_groups"]["corrected"] = {
-                    "spectrum": "SPECTRUM_CORRECTED",
-                    "flag": "FLAG",
-                    "weight": "WEIGHT",
-                    "uvw": "UVW",
-                }
-                is_single_dish = True
+            xds, is_single_dish = add_data_groups(xds)
 
             # Create field_and_source_xds (combines field, source and ephemeris data into one super dataset)
             start = time.time()
@@ -910,3 +873,51 @@ def convert_and_write_partition(
             logger.debug("Write data  " + str(time.time() - start))
 
     # logger.info("Saved ms_v4 " + file_name + " in " + str(time.time() - start_with) + "s")
+
+
+def add_data_groups(xds):
+    xds.attrs["data_groups"] = {}
+    if "VISIBILITY" in xds:
+        xds.attrs["data_groups"]["base"] = {
+            "visibility": "VISIBILITY",
+            "flag": "FLAG",
+            "weight": "WEIGHT",
+            "uvw": "UVW",
+        }
+
+    if "VISIBILITY_CORRECTED" in xds:
+        xds.attrs["data_groups"]["corrected"] = {
+            "visibility": "VISIBILITY_CORRECTED",
+            "flag": "FLAG",
+            "weight": "WEIGHT",
+            "uvw": "UVW",
+        }
+
+    if "VISIBILITY_MODEL" in xds:
+        xds.attrs["data_groups"]["model"] = {
+            "visibility": "VISIBILITY_MODEL",
+            "flag": "FLAG",
+            "weight": "WEIGHT",
+            "uvw": "UVW",
+        }
+
+    is_single_dish = False
+    if "SPECTRUM" in xds:
+        xds.attrs["data_groups"]["base"] = {
+            "spectrum": "SPECTRUM",
+            "flag": "FLAG",
+            "weight": "WEIGHT",
+            "uvw": "UVW",
+        }
+        is_single_dish = True
+
+    if "SPECTRUM_CORRECTED" in xds:
+        xds.attrs["data_groups"]["corrected"] = {
+            "spectrum": "SPECTRUM_CORRECTED",
+            "flag": "FLAG",
+            "weight": "WEIGHT",
+            "uvw": "UVW",
+        }
+        is_single_dish = True
+
+    return xds, is_single_dish

{xradio-0.0.29 → xradio-0.0.30}/src/xradio/vis/_vis_utils/_ms/msv2_to_msv4_meta.py

@@ -4,6 +4,7 @@ col_to_data_variable_names = {
     "FLOAT_DATA": "SPECTRUM",
     "DATA": "VISIBILITY",
     "CORRECTED_DATA": "VISIBILITY_CORRECTED",
+    "MODEL_DATA": "VISIBILITY_MODEL",
     "WEIGHT_SPECTRUM": "WEIGHT",
     "WEIGHT": "WEIGHT",
     "FLAG": "FLAG",
@@ -14,6 +15,7 @@ col_to_data_variable_names = {
 col_dims = {
     "DATA": ("time", "baseline_id", "frequency", "polarization"),
     "CORRECTED_DATA": ("time", "baseline_id", "frequency", "polarization"),
+    "MODEL_DATA": ("time", "baseline_id", "frequency", "polarization"),
     "WEIGHT_SPECTRUM": ("time", "baseline_id", "frequency", "polarization"),
     "WEIGHT": ("time", "baseline_id", "frequency", "polarization"),
     "FLAG": ("time", "baseline_id", "frequency", "polarization"),

{xradio-0.0.29 → xradio-0.0.30}/src/xradio/vis/convert_msv2_to_processing_set.py

@@ -107,66 +107,3 @@ def convert_msv2_to_processing_set(
 
     if parallel:
         dask.compute(delayed_list)
-
-    # delayed_list = []
-    # ms_v4_id = 0
-    # for idx, pair in partition_enumerated_product:
-    #     ddi, state_id, field_id, scan_id = pair
-    #     # logger.debug(
-    #     #     "DDI " + str(ddi) + ", STATE " + str(state_id) + ", FIELD " + str(field_id) + ", SCAN " + str(scan_id)
-    #     # )
-
-    #     # if scan_id == 67: #67
-    #     #     logger.debug(
-    #     #     "DDI " + str(ddi) + ", STATE " + str(state_id) + ", FIELD " + str(field_id) + ", SCAN " + str(scan_id)
-    #     #     )
-    #     if partition_scheme == "ddi_intent_field":
-    #         intent = intents[idx[1]]
-    #     else:
-    #         intent = intents[idx[1]] + "_" + str(state_id)
-
-    #     if parallel:
-    #         delayed_list.append(
-    #             dask.delayed(convert_and_write_partition)(
-    #                 in_file,
-    #                 out_file,
-    #                 intent,
-    #                 ms_v4_id,
-    #                 ddi,
-    #                 state_id,
-    #                 field_id,
-    #                 scan_id,
-    #                 partition_scheme,
-    #                 main_chunksize=main_chunksize,
-    #                 with_pointing=with_pointing,
-    #                 pointing_chunksize=pointing_chunksize,
-    #                 pointing_interpolate=pointing_interpolate,
-    #                 ephemeris_interpolate=ephemeris_interpolate,
-    #                 compressor=compressor,
-    #                 overwrite=overwrite,
-    #             )
-    #         )
-    #     else:
-    #         convert_and_write_partition(
-    #             in_file,
-    #             out_file,
-    #             intent,
-    #             ms_v4_id,
-    #             ddi,
-    #             state_id,
-    #             field_id,
-    #             scan_id,
-    #             partition_scheme,
-    #             main_chunksize=main_chunksize,
-    #             with_pointing=with_pointing,
-    #             pointing_chunksize=pointing_chunksize,
-    #             pointing_interpolate=pointing_interpolate,
-    #             ephemeris_interpolate=ephemeris_interpolate,
-    #             compressor=compressor,
-    #             storage_backend=storage_backend,
-    #             overwrite=overwrite,
-    #         )
-    #     ms_v4_id = ms_v4_id + 1
-
-    # if parallel:
-    #     dask.compute(delayed_list)

{xradio-0.0.29 → xradio-0.0.30}/src/xradio/vis/read_processing_set.py

@@ -37,7 +37,9 @@ def read_processing_set(
         xds = _open_dataset(ms_main_store, file_system)
         data_groups = xds.attrs["data_groups"]
 
-        if (intents is None) or (xds.attrs["intent"] in intents):
+        print(xds)
+
+        if (intents is None) or (xds.attrs["partition_info"]["intent"] in intents):
             sub_xds_dict, field_and_source_xds_dict = _read_sub_xds(
                 ms_store, file_system=file_system, data_groups=data_groups
             )

{xradio-0.0.29 → xradio-0.0.30/src/xradio.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xradio
-Version: 0.0.29
+Version: 0.0.30
 Summary:  Xarray Radio Astronomy Data IO
 Author-email: Jan-Willem Steeb <jsteeb@nrao.edu>
 License: BSD 3-Clause License