xradio 0.0.39__py3-none-any.whl → 0.0.41__py3-none-any.whl

xradio/__init__.py

@@ -1,5 +1,5 @@
 import os
-from graphviper.utils.logger import setup_logger
+from toolviper.utils.logger import setup_logger
 
 # _logger_name = "xradio"
 # if os.getenv("VIPER_LOGGER_NAME") != _logger_name:

xradio/_utils/schema.py

@@ -1,12 +1,12 @@
-import graphviper.utils.logger as logger
+import toolviper.utils.logger as logger
 import xarray as xr
 
 
 def convert_generic_xds_to_xradio_schema(
     generic_xds: xr.Dataset,
     msv4_xds: xr.Dataset,
-    to_new_data_variables: dict,
-    to_new_coords: dict,
+    to_new_data_variables: dict[str, list],
+    to_new_coords: dict[str, list],
 ) -> xr.Dataset:
     """Converts a generic xarray Dataset to the xradio schema.
 
@@ -55,6 +55,7 @@ def convert_generic_xds_to_xradio_schema(
         "column_descriptions"
     ]
     coords = {}
+    coord_attrs = {}
 
     name_keys = list(generic_xds.data_vars.keys()) + list(generic_xds.coords.keys())
 
@@ -80,7 +81,14 @@ def convert_generic_xds_to_xradio_schema(
                 new_coord[1],
                 generic_xds[key].data,
             )
+
+            if msv4_measure:
+                coord_attrs[new_coord[0]] = msv4_measure
+
     msv4_xds = msv4_xds.assign_coords(coords)
+    for coord, coord_attrs in coord_attrs.items():
+        msv4_xds.coords[coord].attrs.update(coord_attrs)
+
     return msv4_xds
 
 
@@ -102,6 +110,9 @@ def column_description_casacore_to_msv4_measure(
             casacore_column_description["keywords"]["QuantumUnits"]
         )
 
+        # Beware: casa_ref won't be found in cases such as the custom
+        # 'NRAO_GBT_USER/NRAO_GBT_USER_DIR_REF' in POINTING
+        casa_ref = None
         # Reference frame to convert?
         if "Ref" in msv4_measure_conversion:
             # Find reference frame

xradio/{vis → correlated_data}/__init__.py

@@ -1,11 +1,12 @@
-from .read_processing_set import read_processing_set
+from .processing_set import ProcessingSet
+from .open_processing_set import open_processing_set
 from .load_processing_set import load_processing_set
 from .convert_msv2_to_processing_set import convert_msv2_to_processing_set
 
 from .schema import VisibilityXds
 
 __all__ = [
-    "read_processing_set",
+    "open_processing_set",
     "load_processing_set",
     "convert_msv2_to_processing_set",
     "VisibilityXds",

xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/load_main_table.py

@@ -1,4 +1,4 @@
-import graphviper.utils.logger as logger
+import toolviper.utils.logger as logger
 from typing import Dict, List, Tuple, Union
 
 import pandas as pd
@@ -12,7 +12,7 @@ from .read_main_table import get_partition_ids, redim_id_data_vars, rename_vars
 from .read import add_units_measures, convert_casacore_time, extract_table_attributes
 from .write import revert_time
 from .table_query import open_query, open_table_ro
-from xradio.vis._vis_utils._ms._tables.read_main_table import (
+from xradio.correlated_data._utils._ms._tables.read_main_table import (
     get_baselines,
     get_baseline_indices,
 )

xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/read.py

@@ -1,4 +1,4 @@
-import graphviper.utils.logger as logger
+import toolviper.utils.logger as logger
 import os
 from pathlib import Path
 import re
@@ -20,9 +20,23 @@ MJD_DIF_UNIX = 40587
 
 
 def table_exists(path: str) -> bool:
+    """
+    Whether a casacore table exists on disk (in the casacore.tables.tableexists sense)
+    """
     return tables.tableexists(path)
 
 
+def table_has_column(path: str, column_name: str) -> bool:
+    """
+    Whether a column is present in a casacore table
+    """
+    with open_table_ro(path) as tb_tool:
+        if column_name in tb_tool.colnames():
+            return True
+        else:
+            return False
+
+
 def convert_casacore_time(
     rawtimes: np.ndarray, convert_to_datetime: bool = True
 ) -> np.ndarray:

xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/read_main_table.py

@@ -1,4 +1,4 @@
-import graphviper.utils.logger as logger
+import toolviper.utils.logger as logger
 from typing import Any, Dict, List, Tuple, Union
 
 import dask, dask.array

xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/read_subtables.py

@@ -1,4 +1,4 @@
-import graphviper.utils.logger as logger
+import toolviper.utils.logger as logger
 from pathlib import Path
 from typing import Dict, Tuple, Union
 

xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/_tables/write.py

@@ -1,4 +1,4 @@
-import graphviper.utils.logger as logger, os
+import toolviper.utils.logger as logger, os
 from typing import Tuple
 
 import numpy as np

xradio/{vis/_vis_utils → correlated_data/_utils}/_ms/conversion.py

@@ -1,21 +1,28 @@
+import datetime
+import importlib
 import numcodecs
-import time
-from .._zarr.encoding import add_encoding
-from typing import Dict, Union
-import graphviper.utils.logger as logger
 import os
 import pathlib
+import time
+from typing import Dict, Union
 
 import numpy as np
 import xarray as xr
 
+import toolviper.utils.logger as logger
 from casacore import tables
-from xradio.vis._vis_utils._ms.msv4_sub_xdss import (
+from xradio.correlated_data._utils._ms.msv4_sub_xdss import (
     create_pointing_xds,
+    create_system_calibration_xds,
     create_weather_xds,
 )
-from xradio.vis._vis_utils._ms.create_antenna_xds import create_antenna_xds
-from xradio.vis._vis_utils._ms.create_field_and_source_xds import (
+from .msv4_info_dicts import create_info_dicts
+from xradio.correlated_data._utils._ms.create_antenna_xds import (
+    create_antenna_xds,
+    create_gain_curve_xds,
+    create_phase_calibration_xds,
+)
+from xradio.correlated_data._utils._ms.create_field_and_source_xds import (
     create_field_and_source_xds,
 )
 from xradio._utils.schema import column_description_casacore_to_msv4_measure
@@ -25,6 +32,7 @@ from .msv2_to_msv4_meta import (
     col_dims,
 )
 
+from .._zarr.encoding import add_encoding
 from .subtables import subt_rename_ids
 from ._tables.table_query import open_table_ro, open_query
 from ._tables.read import (
@@ -142,7 +150,7 @@ def mem_chunksize_to_dict_main(chunksize: float, xds: xr.Dataset) -> Dict[str, i
     polarization).
     """
 
-    sizeof_vis = itemsize_vis_spec(xds)
+    sizeof_vis = itemsize_spec(xds)
     size_all_pols = sizeof_vis * xds.sizes["polarization"]
     if size_all_pols / GiBYTES_TO_BYTES > chunksize:
         raise RuntimeError(
@@ -324,7 +332,7 @@ def find_baseline_or_antenna_var(xds: xr.Dataset) -> str:
     return baseline_or_antenna_name
 
 
-def itemsize_vis_spec(xds: xr.Dataset) -> int:
+def itemsize_spec(xds: xr.Dataset) -> int:
     """
     Size in bytes of one visibility (or spectrum) value.
     """
@@ -693,6 +701,7 @@ def convert_and_write_partition(
     pointing_interpolate: bool = False,
     ephemeris_interpolate: bool = False,
     phase_cal_interpolate: bool = False,
+    sys_cal_interpolate: bool = False,
     compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
     storage_backend="zarr",
     overwrite: bool = False,
@@ -705,7 +714,7 @@ def convert_and_write_partition(
         _description_
     out_file : str
         _description_
-    obs_mode : str
+    intents : str
         _description_
     ddi : int, optional
         _description_, by default 0
@@ -723,6 +732,10 @@ def convert_and_write_partition(
         _description_, by default None
     ephemeris_interpolate : bool, optional
         _description_, by default None
+    phase_cal_interpolate : bool, optional
+        _description_, by default None
+    sys_cal_interpolate : bool, optional
+        _description_, by default None
     compressor : numcodecs.abc.Codec, optional
         _description_, by default numcodecs.Zstd(level=2)
     storage_backend : str, optional
@@ -738,7 +751,7 @@ def convert_and_write_partition(
 
     taql_where = create_taql_query(partition_info)
     ddi = partition_info["DATA_DESC_ID"][0]
-    obs_mode = str(partition_info["OBS_MODE"][0])
+    intents = str(partition_info["OBS_MODE"][0])
 
     start = time.time()
     with open_table_ro(in_file) as mtable:
@@ -766,24 +779,32 @@ def convert_and_write_partition(
                 tb_tool.getcol("OBSERVATION_ID"), "OBSERVATION_ID"
             )
 
-            def get_observation_info(in_file, observation_id, obs_mode):
+            def get_observation_info(in_file, observation_id, intents):
                 generic_observation_xds = load_generic_table(
                     in_file,
                     "OBSERVATION",
                     taql_where=f" where (ROWID() IN [{str(observation_id)}])",
                 )
 
-                if obs_mode == "None":
-                    obs_mode = "obs_" + str(observation_id)
+                if intents == "None":
+                    intents = "obs_" + str(observation_id)
 
-                return generic_observation_xds["TELESCOPE_NAME"].values[0], obs_mode
+                return generic_observation_xds["TELESCOPE_NAME"].values[0], intents
 
-            telescope_name, obs_mode = get_observation_info(
-                in_file, observation_id, obs_mode
+            telescope_name, intents = get_observation_info(
+                in_file, observation_id, intents
             )
 
             start = time.time()
-            xds = xr.Dataset()
+            xds = xr.Dataset(
+                attrs={
+                    "creation_date": datetime.datetime.now().isoformat(),
+                    "xradio_version": importlib.metadata.version("xradio"),
+                    "schema_version": "4.0.-9999",
+                    "type": "visibility",
+                }
+            )
+
             # interval = check_if_consistent(tb_tool.getcol("INTERVAL"), "INTERVAL")
             interval = tb_tool.getcol("INTERVAL")
 
@@ -874,22 +895,51 @@ def convert_and_write_partition(
                 antenna_id,
                 feed_id,
                 telescope_name,
+            )
+
+            logger.debug("Time antenna xds  " + str(time.time() - start))
+
+            start = time.time()
+            gain_curve_xds = create_gain_curve_xds(
+                in_file, xds.frequency.attrs["spectral_window_id"], ant_xds
+            )
+            logger.debug("Time gain_curve xds  " + str(time.time() - start))
+
+            start = time.time()
+            phase_calibration_xds = create_phase_calibration_xds(
+                in_file,
+                xds.frequency.attrs["spectral_window_id"],
+                ant_xds,
                 time_min_max,
                 phase_cal_interp_time,
             )
+            logger.debug("Time phase_calibration xds  " + str(time.time() - start))
 
             # Change antenna_ids to antenna_names
             xds = antenna_ids_to_names(xds, ant_xds, is_single_dish)
+            # but before, keep the name-id arrays, we need them for the pointing and weather xds
             ant_xds_name_ids = ant_xds["antenna_name"].set_xindex("antenna_id")
-            ant_xds = ant_xds.drop_vars(
-                "antenna_id"
-            )  # No longer needed after converting to name.
+            ant_xds_station_name_ids = ant_xds["station"].set_xindex("antenna_id")
+            # No longer needed after converting to name.
+            ant_xds = ant_xds.drop_vars("antenna_id")
 
-            logger.debug("Time ant xds  " + str(time.time() - start))
+            # Create system_calibration_xds
+            start = time.time()
+            if sys_cal_interpolate:
+                sys_cal_interp_time = xds.time.values
+            else:
+                sys_cal_interp_time = None
+            system_calibration_xds = create_system_calibration_xds(
+                in_file,
+                xds.frequency,
+                ant_xds_name_ids,
+                sys_cal_interp_time,
+            )
+            logger.debug("Time system_calibation " + str(time.time() - start))
 
             # Create weather_xds
             start = time.time()
-            weather_xds = create_weather_xds(in_file)
+            weather_xds = create_weather_xds(in_file, ant_xds_station_name_ids)
             logger.debug("Time weather " + str(time.time() - start))
 
             # Create pointing_xds
@@ -914,7 +964,6 @@ def convert_and_write_partition(
                 )
 
             start = time.time()
-            xds.attrs["type"] = "visibility"
 
             # Time and frequency should always be increasing
             if len(xds.frequency) > 1 and xds.frequency[1] - xds.frequency[0] < 0:
@@ -943,7 +992,7 @@ def convert_and_write_partition(
             # assert len(col_unique) == 1, col_name + " is not consistent."
             # return col_unique[0]
 
-            field_and_source_xds, source_id, num_lines = create_field_and_source_xds(
+            field_and_source_xds, source_id, _num_lines = create_field_and_source_xds(
                 in_file,
                 field_id,
                 xds.frequency.attrs["spectral_window_id"],
@@ -965,11 +1014,18 @@ def convert_and_write_partition(
                     "FIELD_PHASE_CENTER"
                 ].attrs["frame"]
 
-            if overwrite:
-                mode = "w"
-            else:
-                mode = "w-"
+            partition_info_misc_fields = {
+                "scan_id": scan_id,
+                "intents": intents,
+                "taql_where": taql_where,
+            }
+            info_dicts = create_info_dicts(
+                in_file, xds, field_and_source_xds, partition_info_misc_fields, tb_tool
+            )
+            xds.attrs.update(info_dicts)
 
+            # xds ready, prepare to write
+            start = time.time()
             main_chunksize = parse_chunksize(main_chunksize, "main", xds)
             add_encoding(xds, compressor=compressor, chunks=main_chunksize)
             logger.debug("Time add compressor and chunk " + str(time.time() - start))
@@ -979,31 +1035,10 @@ def convert_and_write_partition(
                 pathlib.Path(in_file).name.replace(".ms", "") + "_" + str(ms_v4_id),
             )
 
-            if "line_name" in field_and_source_xds.coords:
-                line_name = to_list(
-                    unique_1d(np.ravel(field_and_source_xds.line_name.values))
-                )
+            if overwrite:
+                mode = "w"
             else:
-                line_name = []
-
-            xds.attrs["partition_info"] = {
-                # "spectral_window_id": xds.frequency.attrs["spectral_window_id"],
-                "spectral_window_name": xds.frequency.attrs["spectral_window_name"],
-                # "field_id": to_list(unique_1d(field_id)),
-                "field_name": to_list(
-                    np.unique(field_and_source_xds.field_name.values)
-                ),
-                # "source_id": to_list(unique_1d(source_id)),
-                "line_name": line_name,
-                "scan_number": to_list(np.unique(scan_id)),
-                "source_name": to_list(
-                    np.unique(field_and_source_xds.source_name.values)
-                ),
-                "polarization_setup": to_list(xds.polarization.values),
-                "num_lines": num_lines,
-                "obs_mode": obs_mode.split(","),
-                "taql": taql_where,
-            }
+                mode = "w-"
 
             start = time.time()
             if storage_backend == "zarr":
@@ -1022,6 +1057,21 @@ def convert_and_write_partition(
                         store=os.path.join(file_name, "POINTING"), mode=mode
                     )
 
+                if system_calibration_xds:
+                    system_calibration_xds.to_zarr(
+                        store=os.path.join(file_name, "SYSCAL"), mode=mode
+                    )
+
+                if gain_curve_xds:
+                    gain_curve_xds.to_zarr(
+                        store=os.path.join(file_name, "GAIN_CURVE"), mode=mode
+                    )
+
+                if phase_calibration_xds:
+                    phase_calibration_xds.to_zarr(
+                        store=os.path.join(file_name, "PHASE_CAL"), mode=mode
+                    )
+
                 if weather_xds:
                     weather_xds.to_zarr(
                         store=os.path.join(file_name, "WEATHER"), mode=mode
@@ -1080,7 +1130,7 @@ def add_data_groups(xds):
     xds.attrs["data_groups"] = {}
     if "VISIBILITY" in xds:
         xds.attrs["data_groups"]["base"] = {
-            "visibility": "VISIBILITY",
+            "correlated_data": "VISIBILITY",
             "flag": "FLAG",
             "weight": "WEIGHT",
             "uvw": "UVW",
@@ -1088,7 +1138,7 @@ def add_data_groups(xds):
 
     if "VISIBILITY_CORRECTED" in xds:
         xds.attrs["data_groups"]["corrected"] = {
-            "visibility": "VISIBILITY_CORRECTED",
+            "correlated_data": "VISIBILITY_CORRECTED",
             "flag": "FLAG",
             "weight": "WEIGHT",
             "uvw": "UVW",
@@ -1096,7 +1146,7 @@ def add_data_groups(xds):
 
     if "VISIBILITY_MODEL" in xds:
         xds.attrs["data_groups"]["model"] = {
-            "visibility": "VISIBILITY_MODEL",
+            "correlated_data": "VISIBILITY_MODEL",
             "flag": "FLAG",
             "weight": "WEIGHT",
             "uvw": "UVW",
@@ -1105,7 +1155,16 @@ def add_data_groups(xds):
     is_single_dish = False
     if "SPECTRUM" in xds:
         xds.attrs["data_groups"]["base"] = {
-            "spectrum": "SPECTRUM",
+            "correlated_data": "SPECTRUM",
+            "flag": "FLAG",
+            "weight": "WEIGHT",
+            "uvw": "UVW",
+        }
+        is_single_dish = True
+
+    if "SPECTRUM_MODEL" in xds:
+        xds.attrs["data_groups"]["model"] = {
+            "correlated_data": "SPECTRUM_MODEL",
             "flag": "FLAG",
             "weight": "WEIGHT",
             "uvw": "UVW",
@@ -1114,7 +1173,7 @@ def add_data_groups(xds):
 
     if "SPECTRUM_CORRECTED" in xds:
         xds.attrs["data_groups"]["corrected"] = {
-            "spectrum": "SPECTRUM_CORRECTED",
+            "correlated_data": "SPECTRUM_CORRECTED",
             "flag": "FLAG",
             "weight": "WEIGHT",
             "uvw": "UVW",