xradio 0.0.29__py3-none-any.whl → 0.0.30__py3-none-any.whl

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

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

@@ -29,10 +29,10 @@ xradio/schema/dataclass.py,sha256=LyYsHdorNrzqX1NFBr01vBM0vaBA7XY_TIkZXphMdns,12
 xradio/schema/metamodel.py,sha256=WjtW7pAVzcjLRWifRH3sQoOiN6TV810hARpOIz1M_gw,3845
 xradio/schema/typing.py,sha256=8-o6fZd99kJ4FVdgBYRTIRJ-wDqpcUNXzCTfJvl3TIw,10439
 xradio/vis/__init__.py,sha256=0HFHuUezO-2eVn16IKAzVQCqL896Spaj_QIYjVY9Aq4,401
-xradio/vis/_processing_set.py,sha256=k7oCbGO4ukG6PsCDhD5TQo_zZwcpYEGzjosM1NiK0ZE,4733
-xradio/vis/convert_msv2_to_processing_set.py,sha256=C_mBcS6UYrSbYGHAENLyYjRTjaC4RxMqD-PtEOE2PJc,7398
+xradio/vis/_processing_set.py,sha256=ZmrQms_ETDdvunGs8dGtdAKI0hm759LG-E6jjaCGiYI,5906
+xradio/vis/convert_msv2_to_processing_set.py,sha256=nT_19bFZk4NixAZq9Zy8n7HIUAfDz9i_iomL7d_lxYM,5068
 xradio/vis/load_processing_set.py,sha256=xREcp9G1Vywbb9Pc8_XUcbFypt_JAXNAIG5xEKOBEQ8,5807
-xradio/vis/read_processing_set.py,sha256=v0IeOCP-Xyus-cM3ztyhn-asl9vfjrhMTd8_JDc9l2c,4029
+xradio/vis/read_processing_set.py,sha256=NTgcncF2knAPJVbD_1dKFuLDowJ6AM1kAIbwPdSQBTI,4067
 xradio/vis/schema.py,sha256=mYXHWuzcaTo1PD7A98BqcfMobI74UjklJaAp0J4wQTM,23577
 xradio/vis/vis_io.py,sha256=Onv0tBcix06Dc2b9UjXQzcTVlTvykG49KPO36tf5WKI,5714
 xradio/vis/_vis_utils/__init__.py,sha256=Scu6rKJ2SpO8aG7F-xdTZcYfyWx0viV8gFh8E8ur_gI,93
@@ -40,10 +40,10 @@ xradio/vis/_vis_utils/ms.py,sha256=JcaFxidq-iL7tIbyv6F4qnMTeHediaumP_Y3RTAJvLw,5
 xradio/vis/_vis_utils/ms_column_descriptions_dicts.py,sha256=BV2SyiAmIB8P8n-u8osKlKYKiOq8fmiFsB9xSxlkTrM,39906
 xradio/vis/_vis_utils/zarr.py,sha256=WlFPpcRZz2qXVyzrQ7MqmndhBm5NJUZZZaThDvx2Nw8,3962
 xradio/vis/_vis_utils/_ms/chunks.py,sha256=XiQWyzqxPSmcsb1RhKAuAO55YKPmBDnbUmfymLeEIFo,5414
-xradio/vis/_vis_utils/_ms/conversion.py,sha256=ILZbiW7udxEg7uPTRy8NA8-hsKiGxL7wR25ZsFSE8rk,32237
+xradio/vis/_vis_utils/_ms/conversion.py,sha256=f9-2WaxjYushC-5eNYGienBW3iXRKvKFgBKw4KYcZ2A,32242
 xradio/vis/_vis_utils/_ms/descr.py,sha256=ymavAEnFUc6zIcncP_-72y8zoOYk7Sx-euPYWXkUSE0,5206
 xradio/vis/_vis_utils/_ms/msv2_msv3.py,sha256=9AKs2HWly7Ivv_Cjr11dIPGmm33_rtSBoGF9wN5ZwEQ,116
-xradio/vis/_vis_utils/_ms/msv2_to_msv4_meta.py,sha256=b0Dx8WSTDExE_83BYVTr7I5BTV0HnnTfDMGmF89USKA,4862
+xradio/vis/_vis_utils/_ms/msv2_to_msv4_meta.py,sha256=JWv8ggK3b3tBzxgbuQJ-eG0QFOI-Jnbp7fveee2FuDw,4972
 xradio/vis/_vis_utils/_ms/msv4_infos.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xradio/vis/_vis_utils/_ms/msv4_sub_xdss.py,sha256=L5rg6TMkgLA_PvbOuu3khGcoFjGUHUp1E_1v-IrLWII,13411
 xradio/vis/_vis_utils/_ms/optimised_functions.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -66,8 +66,8 @@ xradio/vis/_vis_utils/_utils/xds_helper.py,sha256=ggv3ArWMN75tklA5v0WxjfBM6rSvUB
 xradio/vis/_vis_utils/_zarr/encoding.py,sha256=GENIlThV6a9CUCL6gIGlu9c6NR3OFWNos6mpxZjMwDc,536
 xradio/vis/_vis_utils/_zarr/read.py,sha256=9UKX_kKztE5jhsriBPnLhWtbptk5DRSVUEHT8V5qh8w,7597
 xradio/vis/_vis_utils/_zarr/write.py,sha256=vebYh-rAKfY3Feuv6H400s-2UuciPtfZMIlMtb23YSk,10044
-xradio-0.0.29.dist-info/LICENSE.txt,sha256=9CYIJt7riOXo9AD0eXBZviLxo_HebD-2JJI8oiWtzfg,1807
-xradio-0.0.29.dist-info/METADATA,sha256=HZaK9u171WZYlDQDXHeJ6cDffkdsotRtRnBp27UJlpg,4449
-xradio-0.0.29.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-xradio-0.0.29.dist-info/top_level.txt,sha256=dQu27fGBZJ2Yk-gW5XeD-dZ76Xa4Xcvk60Vz-dwXp7k,7
-xradio-0.0.29.dist-info/RECORD,,
+xradio-0.0.30.dist-info/LICENSE.txt,sha256=9CYIJt7riOXo9AD0eXBZviLxo_HebD-2JJI8oiWtzfg,1807
+xradio-0.0.30.dist-info/METADATA,sha256=ln-6BSbfL6qcPWrskhLb4sMhvYBZ2TReO5kUBmTKTP0,4449
+xradio-0.0.30.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
+xradio-0.0.30.dist-info/top_level.txt,sha256=dQu27fGBZJ2Yk-gW5XeD-dZ76Xa4Xcvk60Vz-dwXp7k,7
+xradio-0.0.30.dist-info/RECORD,,