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

xradio/vis/_vis_utils/_ms/descr.py

@@ -3,28 +3,37 @@ from typing import Dict, Union
 
 import numpy as np
 import pandas as pd
+import xarray as xr
 
 from ._tables.read import read_generic_table, read_flat_col_chunk
 from ._tables.table_query import open_query, open_table_ro
-from xradio.vis._vis_utils._ms.optimised_functions import unique_1d
+from xradio._utils.array import unique_1d
 
 
 def describe_ms(
-    infile: str, mode="summary", rowmap: Union[dict, None] = None
+    infile: str, mode: str = "summary", rowmap: Union[dict, None] = None
 ) -> Union[pd.DataFrame, Dict]:
     """
     Summarize the contents of an MS directory in casacore table format
 
-    :param infile: input MS filename
-    :param mode: type of information returned ('summary', 'flat', 'expanded').
-    'summary' returns a pandas dataframe that is nice for displaying in notebooks
-    etc. 'flat' returns a list of tuples of (ddi, row, chan, pol). 'expanded'
-    returns a list of tuples of (ddi, time, baseline, chan, pol). These latter two
-    are good for trying to determine chunk size for read_ms(expand=True/False).
-    :param rowmap: dict of DDI to tuple of (row indices, channel indices). Returned
-    by ms_selection function. Default None ignores selections
-
-    :return: summary as a pd dataframe
+    Parameters
+    ----------
+    infile : str
+        input MS filename
+    mode : str (Default value = "summary")
+        type of information returned ('summary', 'flat', 'expanded').
+        'summary' returns a pandas dataframe that is nice for displaying in notebooks
+        etc. 'flat' returns a list of tuples of (ddi, row, chan, pol). 'expanded'
+        returns a list of tuples of (ddi, time, baseline, chan, pol). These latter two
+        are good for trying to determine chunk size for read_ms(expand=True/False). (Default value = "summary")
+    rowmap : Union[dict, None] (Default value = None)
+        dict of DDI to tuple of (row indices, channel indices). Returned
+        by ms_selection function. Default None ignores selections
+
+    Returns
+    -------
+    Union[pd.DataFrame, Dict]
+        summary as a pd dataframe
     """
     infile = os.path.expanduser(infile)  # does nothing if $HOME is unknown
     if not os.path.isdir(infile):
@@ -38,18 +47,22 @@ def describe_ms(
 
     ddi_xds = read_generic_table(infile, "DATA_DESCRIPTION")
     ddis = list(ddi_xds.row.values) if rowmap is None else list(rowmap.keys())
-
     summary: Union[pd.DataFrame, Dict] = []
     if mode == "summary":
         summary = pd.DataFrame([])
 
+    all_sdf = []
     with open_table_ro(infile) as tb_tool:
         for ddi in ddis:
             taql = f"select * from $tb_tool where DATA_DESC_ID = {ddi}"
             with open_query(tb_tool, taql) as query_per_ddi:
-                populate_ms_descr(infile, mode, query_per_ddi, summary, ddi, ddi_xds)
+                sdf = populate_ms_descr(
+                    infile, mode, query_per_ddi, summary, ddi, ddi_xds
+                )
+                all_sdf.append(sdf)
 
     if mode == "summary":
+        summary = pd.DataFrame(all_sdf)
         summary = summary.set_index("ddi").sort_index()
     else:
         summary = dict(summary)
@@ -57,16 +70,35 @@ def describe_ms(
 
 
 def populate_ms_descr(
-    infile, mode, query_per_ddi, summary, ddi, ddi_xds, rowmap=None
+    infile: str,
+    mode: str,
+    query_per_ddi,
+    summary: dict,
+    ddi: int,
+    ddi_xds: xr.Dataset,
+    rowmap: Union[Dict, None] = None,
 ) -> pd.DataFrame:
-    """Adds information from the time and baseline (antenna1+antenna2)
+    """
+    Adds information from the time and baseline (antenna1+antenna2)
     columns as well as channel and polarizations, based on a taql
     query.
 
-    :param mode: mode (as in describe_ms())
-    :param query_per_ddi: a TaQL query with data per individual DDI
-    :param sdf: summary dict being populated
-    :param summary: final summary object being populated from the invividual sdf's
+    Parameters
+    ----------
+    infile : str
+        input table/MS path
+    mode : str
+        mode (as in describe_ms())
+    query_per_ddi :
+        a TaQL query with data per individual DDI
+    summary : Dict
+        summary dict being populated
+    ddi_xds : xr.Dataset
+        final summary object being populated from the invividual sdf's
+
+    Returns
+    -------
+    pd.DataFrame
     """
     spw_ids = ddi_xds.spectral_window_id.values
     pol_ids = ddi_xds.polarization_id.values

xradio/vis/_vis_utils/_ms/msv2_to_msv4_meta.py

@@ -1,3 +1,5 @@
+import graphviper.utils.logger as logger
+
 col_to_data_variable_names = {
     "FLOAT_DATA": "SPECTRUM",
     "DATA": "VISIBILITY",
@@ -26,17 +28,38 @@ col_to_coord_names = {
     "ANTENNA2": "baseline_ant2_id",
 }
 
+# Map casacore measures to astropy
 casacore_to_msv4_measure_type = {
-    "quanta": {"type": "quantity", "Ref": None},
-    "direction": {"type": "sky_coord", "Ref": "frame"},
-    "epoch": {"type": "time", "Ref": "scale"},
-    "frequency": {"type": "spectral_coord", "Ref": "frame"},
-    "position": {"type": "earth_location", "Ref": "ellipsoid"},
-    "uvw": {"type": "uvw", "Ref": "frame"},
+    "quanta": {
+        "type": "quantity",
+    },
+    "direction": {"type": "sky_coord", "Ref": "frame", "Ref_map": {"J2000": "fk5"}},
+    "epoch": {"type": "time", "Ref": "scale", "Ref_map": {"UTC": "utc"}},
+    "frequency": {
+        "type": "spectral_coord",
+        "Ref": "frame",
+        "Ref_map": {
+            "REST": "REST",
+            "LSRK": "LSRK",
+            "LSRD": "LSRD",
+            "BARY": "BARY",
+            "GEO": "GEO",
+            "TOPO": "TOPO",
+            "GALACTO": "GALACTO",
+            "LGROUP": "LGROUP",
+            "CMB": "CMB",
+            "Undefined": "Undefined",
+        },
+    },
+    "position": {
+        "type": "earth_location",
+        "Ref": "ellipsoid",
+        "Ref_map": {"ITRF": "GRS80"},
+    },
+    "uvw": {"type": "uvw", "Ref": "frame", "Ref_map": {"ITRF": "GRS80"}},
+    "radialvelocity": {"type": "quantity"},
 }
 
-casacore_to_msv4_ref = {"J2000": "FK5", "ITRF": "GRS80"}
-
 casa_frequency_frames = [
     "REST",
     "LSRK",
@@ -54,48 +77,60 @@ casa_frequency_frames_codes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 64]
 
 
 def column_description_casacore_to_msv4_measure(
-    casacore_column_description, ref_code=None, time_format="unix"
+    casacore_column_description, ref_code=None, time_format="UNIX"
 ):
     import numpy as np
 
     msv4_measure = {}
     if "MEASINFO" in casacore_column_description["keywords"]:
-        msv4_measure["type"] = casacore_to_msv4_measure_type[
-            casacore_column_description["keywords"]["MEASINFO"]["type"]
-        ]["type"]
+        measinfo = casacore_column_description["keywords"]["MEASINFO"]
+
+        # Get conversion information
+        msv4_measure_conversion = casacore_to_msv4_measure_type[measinfo["type"]]
+
+        # Convert type, copy unit
+        msv4_measure["type"] = msv4_measure_conversion["type"]
         msv4_measure["units"] = list(
             casacore_column_description["keywords"]["QuantumUnits"]
         )
 
-        if "TabRefCodes" in casacore_column_description["keywords"]["MEASINFO"]:
-            ref_index = np.where(
-                casacore_column_description["keywords"]["MEASINFO"]["TabRefCodes"]
-                == ref_code
-            )[0][0]
-            casa_ref = casacore_column_description["keywords"]["MEASINFO"][
-                "TabRefTypes"
-            ][ref_index]
-        else:
-            if "Ref" in casacore_column_description["keywords"]["MEASINFO"]:
-                casa_ref = casacore_column_description["keywords"]["MEASINFO"]["Ref"]
-            elif (
-                casacore_column_description["keywords"]["MEASINFO"]["type"]
-                == "frequency"
-            ):
+        # Reference frame to convert?
+        if "Ref" in msv4_measure_conversion:
+            # Find reference frame
+            if "TabRefCodes" in measinfo:
+                ref_index = np.where(measinfo["TabRefCodes"] == ref_code)[0][0]
+                casa_ref = measinfo["TabRefTypes"][ref_index]
+            elif "Ref" in measinfo:
+                casa_ref = measinfo["Ref"]
+            elif measinfo["type"] == "frequency":
                 # Some MSv2 don't have the "TabRefCodes".
                 ref_index = np.where(casa_frequency_frames_codes == ref_code)[0][0]
                 casa_ref = casa_frequency_frames[ref_index]
+            else:
+                logger.warning(
+                    f"Could not determine {measinfo['type']} measure "
+                    "reference frame!"
+                )
+
+            # Convert into MSv4 representation of reference frame, warn if unknown
+            if casa_ref in msv4_measure_conversion.get("Ref_map", {}):
+                casa_ref = msv4_measure_conversion["Ref_map"][casa_ref]
+            else:
+                logger.warning(
+                    f"Unknown reference frame for {measinfo['type']} "
+                    f"measure, using verbatim: {casa_ref}"
+                )
 
-        if casa_ref in casacore_to_msv4_ref:
-            casa_ref = casacore_to_msv4_ref[casa_ref]
-        msv4_measure[
-            casacore_to_msv4_measure_type[
-                casacore_column_description["keywords"]["MEASINFO"]["type"]
-            ]["Ref"]
-        ] = casa_ref
+            msv4_measure[msv4_measure_conversion["Ref"]] = casa_ref
 
         if msv4_measure["type"] == "time":
-            msv4_measure["format"] = "unix"
+            msv4_measure["format"] = time_format
+    elif "QuantumUnits" in casacore_column_description["keywords"]:
+        msv4_measure = {
+            "type": "quantity",
+            "units": list(casacore_column_description["keywords"]["QuantumUnits"]),
+        }
+
     return msv4_measure
 
 

xradio/vis/_vis_utils/_ms/msv4_infos.py

@@ -1,59 +0,0 @@
-from .msv2_to_msv4_meta import column_description_casacore_to_msv4_measure
-from .subtables import subt_rename_ids
-from ._tables.read import read_generic_table
-
-
-def create_field_info(in_file, field_id):
-    field_xds = read_generic_table(
-        in_file,
-        "FIELD",
-        rename_ids=subt_rename_ids["FIELD"],
-    ).sel(field_id=field_id)
-    # https://stackoverflow.com/questions/53195684/how-to-navigate-a-dict-by-list-of-keys
-
-    field_column_description = field_xds.attrs["other"]["msv2"]["ctds_attrs"][
-        "column_descriptions"
-    ]
-    # ['DELAY_DIR', 'PHASE_DIR', 'REFERENCE_DIR', 'CODE', 'FLAG_ROW', 'NAME', 'NUM_POLY', 'SOURCE_ID', 'TIME']
-
-    msv4_measure = column_description_casacore_to_msv4_measure(
-        field_column_description["REFERENCE_DIR"],
-        ref_code=getattr(field_xds.get("refdir_ref"), "data", None),
-    )
-    delay_dir = {
-        "dims": "",
-        "data": list(field_xds["delay_dir"].data[0, :]),
-        "attrs": msv4_measure,
-    }
-
-    msv4_measure = column_description_casacore_to_msv4_measure(
-        field_column_description["PHASE_DIR"],
-        ref_code=getattr(field_xds.get("phasedir_ref"), "data", None),
-    )
-    phase_dir = {
-        "dims": "",
-        "data": list(field_xds["phase_dir"].data[0, :]),
-        "attrs": msv4_measure,
-    }
-
-    msv4_measure = column_description_casacore_to_msv4_measure(
-        field_column_description["DELAY_DIR"],
-        ref_code=getattr(field_xds.get("delaydir_ref"), "data", None),
-    )
-    reference_dir = {
-        "dims": "",
-        "data": list(field_xds["delay_dir"].data[0, :]),
-        "attrs": msv4_measure,
-    }
-
-    field_info = {
-        "name": str(field_xds["name"].data),
-        "code": str(field_xds["code"].data),
-        "delay_direction": delay_dir,
-        "phase_direction": phase_dir,
-        "reference_direction": reference_dir,
-        "field_id": field_id,
-    }
-    # xds.attrs["field_info"] = field_info
-
-    return field_info

xradio/vis/_vis_utils/_ms/msv4_sub_xdss.py

@@ -1,13 +1,61 @@
+import graphviper.utils.logger as logger
+import time
+from typing import Tuple, Union
+
 import numpy as np
 import xarray as xr
 
 from .msv2_to_msv4_meta import column_description_casacore_to_msv4_measure
 from .subtables import subt_rename_ids
-from ._tables.read import read_generic_table
+from ._tables.read import make_taql_where_between_min_max, read_generic_table
+
+
+def interpolate_to_time(
+    xds: xr.Dataset, interp_time: Union[xr.DataArray, None], message_prefix: str
+) -> xr.Dataset:
+    """
+    Interpolate the time coordinate of the input xarray dataset to the
+    a data array. This can be used for example to interpolate a pointing_xds
+    to the time coord of the (main) MSv4, or similarly the ephemeris
+    data variables of a field_and_source_xds.
+
+    Uses interpolation method "linear", unless the source number of points is
+    1 in which case "nearest" is used, to avoid divide-by-zero issues.
+
+    Parameters:
+    ----------
+    xds : xr.Dataset
+        Xarray dataset to interpolate (presumably a pointing_xds or an xds of
+        ephemeris variables)
+    interp_time : Union[xr.DataArray, None]
+        Time axis to interpolate the dataset to (usually main MSv4 time)
+    message_prefix: str
+        A prefix for info/debug/etc. messages
+
+    Returns:
+    -------
+    interpolated_xds : xr.Dataset
+        xarray dataset with time axis interpolated to interp_time.
+    """
+    if interp_time is not None:
+        points_before = xds.time.size
+        if points_before > 1:
+            method = "linear"
+        else:
+            method = "nearest"
+        xds = xds.interp(time=interp_time, method=method, assume_sorted=True)
+        points_after = xds.time.size
+        logger.debug(
+            f"{message_prefix}: interpolating the time coordinate "
+            f"from {points_before} to {points_after} points"
+        )
+
+    return xds
 
 
 def create_ant_xds(in_file: str):
-    """Creates an Antenna Xarray Dataset from a MS v2 ANTENNA table.
+    """
+    Creates an Antenna Xarray Dataset from a MS v2 ANTENNA table.
 
     Parameters
     ----------
@@ -16,7 +64,7 @@ def create_ant_xds(in_file: str):
 
     Returns
     -------
-    ant_xds : xarray.Dataset
+    xr.Dataset
         Antenna Xarray Dataset.
     """
     # Dictionaries that define the conversion from MSv2 to MSv4:
@@ -90,7 +138,8 @@ def create_ant_xds(in_file: str):
 
 
 def create_weather_xds(in_file: str):
-    """Creates a Weather Xarray Dataset from a MS v2 WEATHER table.
+    """
+    Creates a Weather Xarray Dataset from a MS v2 WEATHER table.
 
     Parameters
     ----------
@@ -99,7 +148,7 @@ def create_weather_xds(in_file: str):
 
     Returns
     -------
-    weather_xds : xarray.Dataset
+    xr.Dataset
         Weather Xarray Dataset.
     """
     # Dictionaries that define the conversion from MSv2 to MSv4:
@@ -209,8 +258,13 @@ def create_weather_xds(in_file: str):
     return weather_xds
 
 
-def create_pointing_xds(in_file: str):
-    """Creates a Pointing Xarray Dataset from an MS v2 POINTING (sub)table.
+def create_pointing_xds(
+    in_file: str,
+    time_min_max: Union[Tuple[np.float64, np.float64], None],
+    interp_time: Union[xr.DataArray, None] = None,
+) -> xr.Dataset:
+    """
+    Creates a Pointing Xarray Dataset from an MS v2 POINTING (sub)table.
 
     WIP: details of a few direction variables (and possibly moving some to attributes) to be
     settled (see MSv4 spreadsheet).
@@ -219,12 +273,17 @@ def create_pointing_xds(in_file: str):
     ----------
     in_file : str
         Input MS name.
+    time_min_max : tuple
+        min / max times values to constrain loading (from the TIME column)
+    interp_time : Union[xr.DataArray, None] (Default value = None)
+        interpolate time to this (presumably main dataset time)
 
     Returns
     -------
-    pointing_xds : xarray.Dataset
-        Pointing Xarray Dataset.
+    xr.Dataset
+         Pointing Xarray dataset
     """
+    start = time.time()
 
     # Dictionaries that define the conversion from MSv2 to MSv4:
     to_new_data_variable_names = {
@@ -257,11 +316,15 @@ def create_pointing_xds(in_file: str):
     to_new_coord_names = {"ra/dec": "direction"}
     coord_dims = {}
 
+    taql_time_range = make_taql_where_between_min_max(
+        time_min_max, in_file, "POINTING", "TIME"
+    )
     # Read POINTING table into a Xarray Dataset.
     generic_pointing_xds = read_generic_table(
         in_file,
         "POINTING",
         rename_ids=subt_rename_ids["POINTING"],
+        taql_where=taql_time_range,
     )
     if not generic_pointing_xds.data_vars:
         # apparently empty MS/POINTING table => produce empty xds
@@ -319,4 +382,8 @@ def create_pointing_xds(in_file: str):
         }
     # TODO: move also source_offset/pointing_offset from data_vars to attrs?
 
+    pointing_xds = interpolate_to_time(pointing_xds, interp_time, "pointing_xds")
+
+    logger.debug(f"create_pointing_xds() execution time {time.time() - start:0.2f} s")
+
     return pointing_xds

xradio/vis/_vis_utils/_ms/optimised_functions.py

@@ -1,46 +0,0 @@
-"""Contains optimised functions to be used within other modules."""
-
-import numpy as np
-import pandas as pd
-
-
-def unique_1d(array: np.ndarray) -> np.ndarray:
-    """Optimised version of np.unique for 1D arrays.
-
-    Args:
-        array (np.ndarray): a 1D array of values.
-
-    Returns:
-        np.ndarray: a sorted array of unique values.
-    """
-    return np.sort(pd.unique(array))
-
-
-def pairing_function(antenna_pairs: np.ndarray) -> np.ndarray:
-    """Pairing function to convert each array pair to a single value.
-
-    This custom pairing function will only work if the maximum value is less
-    than 2**20 and less than 2,048 if using signed 32-bit integers.
-
-    Args:
-        antenna_pairs (np.ndarray): a 2D array containing antenna 1 and antenna
-        2 ids, which forms a baseline.
-
-    Returns:
-        np.ndarray: a 1D array of the paired values.
-    """
-    return antenna_pairs[:, 0] * 2**20 + antenna_pairs[:, 1]
-
-
-def inverse_pairing_function(paired_array: np.ndarray) -> np.ndarray:
-    """Inverse pairing function to convert each paired value to an antenna pair.
-
-    This inverse pairing function is the inverse of the custom pairing function.
-
-    Args:
-        paired_array (np.ndarray): a 1D array of the paired values.
-
-    Returns:
-        np.ndarray: a 2D array containing antenna 1 and antenna 2 ids.
-    """
-    return np.column_stack(np.divmod(paired_array, 2**20))