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

xradio/__init__.py

@@ -1,13 +1,13 @@
 import os
 from graphviper.utils.logger import setup_logger
 
-_logger_name = "xradio"
-if os.getenv("VIPER_LOGGER_NAME") != _logger_name:
-    os.environ["VIPER_LOGGER_NAME"] = _logger_name
-    setup_logger(
-        logger_name="xradio",
-        log_to_term=True,
-        log_to_file=False,  # True
-        log_file="xradio-logfile",
-        log_level="DEBUG",
-    )
+# _logger_name = "xradio"
+# if os.getenv("VIPER_LOGGER_NAME") != _logger_name:
+#     os.environ["VIPER_LOGGER_NAME"] = _logger_name
+#     setup_logger(
+#         logger_name="xradio",
+#         log_to_term=True,
+#         log_to_file=False,  # True
+#         log_file="xradio-logfile",
+#         log_level="DEBUG",
+#     )

xradio/_utils/common.py

@@ -1,3 +1,60 @@
 import numpy as np
 
 _deg_to_rad = np.pi / 180
+
+
+def cast_to_str(x):
+    if isinstance(x, list):
+        return x[0]
+    else:
+        return x
+
+
+def convert_to_si_units(xds):
+    for data_var in xds.data_vars:
+        if "units" in xds[data_var].attrs:
+            for u_i, u in enumerate(xds[data_var].attrs["units"]):
+                if u == "km":
+                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * 1e3
+                    xds[data_var].attrs["units"][u_i] = "m"
+                if u == "km/s":
+                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * 1e3
+                    xds[data_var].attrs["units"][u_i] = "m/s"
+                if u == "deg":
+                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * np.pi / 180
+                    xds[data_var].attrs["units"][u_i] = "rad"
+                if u == "Au" or u == "AU":
+                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * 149597870700
+                    xds[data_var].attrs["units"][u_i] = "m"
+                if u == "Au/d" or u == "AU/d":
+                    xds[data_var][..., u_i] = (
+                        xds[data_var][..., u_i] * 149597870700 / 86400
+                    )
+                    xds[data_var].attrs["units"][u_i] = "m/s"
+                if u == "arcsec":
+                    xds[data_var][..., u_i] = xds[data_var][..., u_i] * np.pi / 648000
+                    xds[data_var].attrs["units"][u_i] = "rad"
+    return xds
+
+
+def add_position_offsets(dv_1, dv_2):
+    # Fun with angles: We are adding angles together. We need to make sure that the results are between -pi and pi.
+    new_pos = dv_1 + dv_2
+
+    while np.any(new_pos[:, 0] > np.pi) or np.any(new_pos[:, 0] < -np.pi):
+        new_pos[:, 0] = np.where(
+            new_pos[:, 0] > np.pi, new_pos[:, 0] - 2 * np.pi, new_pos[:, 0]
+        )
+        new_pos[:, 0] = np.where(
+            new_pos[:, 0] < -np.pi, new_pos[:, 0] + 2 * np.pi, new_pos[:, 0]
+        )
+
+    while np.any(new_pos[:, 1] > np.pi / 2) or np.any(new_pos[:, 1] < -np.pi / 2):
+        new_pos[:, 1] = np.where(
+            new_pos[:, 1] > np.pi / 2, new_pos[:, 1] - np.pi, new_pos[:, 1]
+        )
+        new_pos[:, 1] = np.where(
+            new_pos[:, 1] < -np.pi / 2, new_pos[:, 1] + np.pi, new_pos[:, 1]
+        )
+
+    return new_pos

xradio/_utils/{array.py → list_and_array.py}

@@ -1,6 +1,23 @@
 """Contains optimised functions to be used within other modules."""
 
 import numpy as np
+import xarray as xr
+
+
+def to_list(x):
+    if isinstance(x, (list, np.ndarray)):
+        if x.ndim == 0:
+            return [x.item()]
+        return list(x)  # needed for json serialization
+    return [x]
+
+
+def to_np_array(x):
+    if isinstance(x, (list, np.ndarray)):
+        if x.ndim == 0:
+            return np.array([x.item()])
+        return np.array(x)  # needed for json serialization
+    return np.array([x])
 
 
 def check_if_consistent(array: np.ndarray, array_name: str) -> np.ndarray:
@@ -45,6 +62,12 @@ def unique_1d(array: np.ndarray) -> np.ndarray:
         a sorted array of unique values.
 
     """
+    if isinstance(array, xr.core.dataarray.DataArray):
+        array = array.values
+
+    if array.ndim == 0:
+        return np.array([array.item()])
+
     return np.sort(pd.unique(array))
 
 

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,49 +30,65 @@ class processing_set(dict):
     def _summary(self, data_group="base"):
         summary_data = {
             "name": [],
-            "ddi": [],
-            "intent": [],
-            "field_id": [],
+            "obs_mode": [],
+            "shape": [],
+            "polarization": [],
+            "spw_id": [],
+            # "field_id": [],
             "field_name": [],
+            # "source_id": [],
+            "source_name": [],
+            "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["obs_mode"].append(value.attrs["partition_info"]["obs_mode"])
+            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)
 
-            summary_data["field_id"].append(value.attrs["partition_info"]["field_id"])
+            # summary_data["field_id"].append(value.attrs["partition_info"]["field_id"])
+            # summary_data["source_id"].append(value.attrs["partition_info"]["source_id"])
+
             summary_data["field_name"].append(
-                value[data_name].attrs["field_and_source_xds"].attrs["field_name"]
+                value.attrs["partition_info"]["field_name"]
+            )
+            summary_data["source_name"].append(
+                value.attrs["partition_info"]["source_name"]
             )
             summary_data["start_frequency"].append(value["frequency"].values[0])
             summary_data["end_frequency"].append(value["frequency"].values[-1])
 
             if value[data_name].attrs["field_and_source_xds"].is_ephemeris:
                 summary_data["field_coords"].append("Ephemeris")
+            elif (
+                "time"
+                in value[data_name].attrs["field_and_source_xds"][center_name].coords
+            ):
+                summary_data["field_coords"].append("Multi-Phase-Center")
             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 +100,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),
                     ]
                 )
 
@@ -108,8 +118,8 @@ class processing_set(dict):
             assert (
                 frame == ms_xds.frequency.attrs["frame"]
             ), "Frequency reference frame not consistent in processing set."
-            if ms_xds.frequency.attrs["spw_id"] not in spw_ids:
-                spw_ids.append(ms_xds.frequency.attrs["spw_id"])
+            if ms_xds.frequency.attrs["spectral_window_id"] not in spw_ids:
+                spw_ids.append(ms_xds.frequency.attrs["spectral_window_id"])
                 freq_axis_list.append(ms_xds.frequency)
 
         freq_axis = xr.concat(freq_axis_list, dim="frequency").sortby("frequency")
@@ -131,3 +141,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