xradio 0.0.60__py3-none-any.whl → 1.0.1__py3-none-any.whl

xradio/_utils/list_and_array.py

@@ -73,7 +73,9 @@ def to_np_array(x):
     return np.array([x])
 
 
-def check_if_consistent(array: np.ndarray, array_name: str) -> np.ndarray:
+def check_if_consistent(
+    array: np.ndarray, array_name: str, err_msg: str = ""
+) -> np.ndarray:
     """_summary_
 
     Parameters
@@ -92,7 +94,7 @@ def check_if_consistent(array: np.ndarray, array_name: str) -> np.ndarray:
         return array.item()
 
     array_unique = unique_1d(array)
-    assert len(array_unique) == 1, array_name + " is not consistent."
+    assert len(array_unique) == 1, array_name + " is not consistent." + err_msg
     return array_unique[0]
 
 

xradio/image/_util/_casacore/xds_to_casacore.py

@@ -5,6 +5,7 @@ import dask.array as da
 import numpy as np
 import xarray as xr
 from astropy.coordinates import Angle
+from astropy import units as apu
 
 try:
     from casacore import tables
@@ -93,10 +94,9 @@ def _compute_spectral_dict(xds: xr.Dataset) -> dict:
     spec["system"] = xds.frequency.attrs["reference_frequency"]["attrs"][
         "observer"
     ].upper()
-    u = xds.frequency.attrs["reference_frequency"]["attrs"]["units"]
-    spec["unit"] = u
+    spec["unit"] = xds.frequency.attrs["reference_frequency"]["attrs"]["units"]
     spec["velType"] = _doppler_types.index(xds.velocity.attrs["doppler_type"])
-    u = xds.velocity.attrs["units"]
+    # u = xds.velocity.attrs["units"]
     spec["version"] = 2
     # vel unit is a list[str] in the xds but needs to be a str in the casa image
     spec["velUnit"] = xds.velocity.attrs["units"]
@@ -106,7 +106,14 @@ def _compute_spectral_dict(xds: xr.Dataset) -> dict:
     wcs["ctype"] = "FREQ"
     wcs["pc"] = 1.0
     wcs["crval"] = float(xds.frequency.attrs["reference_frequency"]["data"])
-    wcs["cdelt"] = float(xds.frequency.values[1] - xds.frequency.values[0])
+    if len(xds.frequency.values) > 1:
+        wcs["cdelt"] = float(xds.frequency.values[1] - xds.frequency.values[0])
+    else:
+        # TODO this is just a temporary fix, likely schema will be updated to include chan widths
+        myu = apu.Unit(spec["unit"])
+        mydel = 1.8 * apu.GHz
+        my_del_converted = mydel.to(spec["unit"])
+        wcs["cdelt"] = my_del_converted.value
     wcs["crpix"] = float((wcs["crval"] - xds.frequency.values[0]) / wcs["cdelt"])
     spec["wcs"] = wcs
     return spec

xradio/image/image.py

@@ -14,8 +14,6 @@ import xarray as xr
 
 # from .._utils.zarr.common import _load_no_dask_zarr
 
-from ._util.casacore import _load_casa_image_block, _xds_to_casa_image
-
 # from ._util.fits import _read_fits_image
 from ._util.image_factory import (
     _make_empty_aperture_image,
@@ -201,6 +199,8 @@ def load_image(infile: str, block_des: dict = None, do_sky_coords=True) -> xr.Da
         # comment next line when done debugging
         # return _load_casa_image_block(infile, selection, do_sky_coords)
         try:
+            from ._util.casacore import _load_casa_image_block
+
             return _load_casa_image_block(infile, selection, do_sky_coords)
         except Exception as e:
             emsgs.append(f"image format appears not to be casacore: {e.args}")
@@ -256,6 +256,8 @@ def write_image(
             )
     my_format = out_format.lower()
     if my_format == "casa":
+        from ._util.casacore import _xds_to_casa_image
+
         _xds_to_casa_image(xds, imagename)
     elif my_format == "zarr":
         _xds_to_zarr(xds, imagename)

xradio/measurement_set/_utils/_msv2/conversion.py

@@ -440,6 +440,7 @@ def create_coordinates(
     baseline_ant1_id: np.ndarray,
     baseline_ant2_id: np.ndarray,
     scan_id: np.ndarray,
+    scan_intents: list[str],
 ) -> tuple[xr.Dataset, int]:
     """
     Creates coordinates of a VisibilityXds/SpectrumXds and assigns them to the input
@@ -464,6 +465,9 @@ def create_coordinates(
         ANTENNA2 ids to be used as coord
     scan_id :
         SCAN_ID values from MSv2, for the scan_name coord
+    scan_intents :
+        list of SCAN_INTENT values from MSv2, for the scan_intents attribute of the
+        scan_name coord
 
     Returns
     -------
@@ -509,6 +513,9 @@ def create_coordinates(
 
     xds = xds.assign_coords(coords)
 
+    ##### Add scan intents attribute to scan_name coord #####
+    xds.scan_name.attrs["scan_intents"] = scan_intents
+
     ###### Create Frequency Coordinate ######
     freq_column_description = spectral_window_xds.attrs["other"]["msv2"]["ctds_attrs"][
         "column_descriptions"
@@ -528,7 +535,7 @@ def create_coordinates(
         spw_name = spw_name + "_" + str(spectral_window_id)
 
     xds.frequency.attrs["spectral_window_name"] = spw_name
-    xds.frequency.attrs["spectral_window_intent"] = "UNSPECIFIED"
+    xds.frequency.attrs["spectral_window_intents"] = ["UNSPECIFIED"]
     msv4_measure = column_description_casacore_to_msv4_measure(
         freq_column_description["REF_FREQUENCY"],
         ref_code=spectral_window_xds["MEAS_FREQ_REF"].data,
@@ -769,15 +776,17 @@ def create_taql_query_where(partition_info: dict):
     taql_where = "WHERE "
     for col_name in main_par_table_cols:
         if col_name in partition_info:
-            taql_where = (
-                taql_where
-                + f"({col_name} IN [{','.join(map(str, partition_info[col_name]))}]) AND"
-            )
-            if col_name == "ANTENNA1":
+
+            if partition_info[col_name][0] is not None:
                 taql_where = (
                     taql_where
-                    + f"(ANTENNA2 IN [{','.join(map(str, partition_info[col_name]))}]) AND"
+                    + f"({col_name} IN [{','.join(map(str, partition_info[col_name]))}]) AND"
                 )
+                if col_name == "ANTENNA1":
+                    taql_where = (
+                        taql_where
+                        + f"(ANTENNA2 IN [{','.join(map(str, partition_info[col_name]))}]) AND"
+                    )
     taql_where = taql_where[:-3]
 
     return taql_where
@@ -933,6 +942,7 @@ def estimate_memory_for_partition(in_file: str, partition: dict) -> float:
 
     taql_partition = create_taql_query_where(partition)
     taql_main = f"select * from $mtable {taql_partition}"
+
     with open_table_ro(in_file) as mtable:
         with open_query(mtable, taql_main) as tb_tool:
             # Do not feel tempted to rely on nrows. nrows tends to underestimate memory when baselines are missing.
@@ -1017,7 +1027,7 @@ def convert_and_write_partition(
         _description_
     out_file : str
         _description_
-    intents : str
+    scan_intents : str
         _description_
     ddi : int, optional
         _description_, by default 0
@@ -1061,7 +1071,7 @@ def convert_and_write_partition(
     taql_where = create_taql_query_where(partition_info)
     table_manager = TableManager(in_file, taql_where)
     ddi = partition_info["DATA_DESC_ID"][0]
-    intents = str(partition_info["OBS_MODE"][0])
+    scan_intents = str(partition_info["OBS_MODE"][0]).split(",")
 
     start = time.time()
     with table_manager.get_table() as tb_tool:
@@ -1085,19 +1095,21 @@ def convert_and_write_partition(
             tb_tool.getcol("OBSERVATION_ID"), "OBSERVATION_ID"
         )
 
-        def get_observation_info(in_file, observation_id, intents):
+        def get_observation_info(in_file, observation_id, scan_intents):
             generic_observation_xds = load_generic_table(
                 in_file,
                 "OBSERVATION",
                 taql_where=f" where (ROWID() IN [{str(observation_id)}])",
             )
 
-            if intents == "None":
-                intents = "obs_" + str(observation_id)
+            if scan_intents == "None":
+                scan_intents = "obs_" + str(observation_id)
 
-            return generic_observation_xds["TELESCOPE_NAME"].values[0], intents
+            return generic_observation_xds["TELESCOPE_NAME"].values[0], scan_intents
 
-        telescope_name, intents = get_observation_info(in_file, observation_id, intents)
+        telescope_name, scan_intents = get_observation_info(
+            in_file, observation_id, scan_intents
+        )
 
         start = time.time()
         xds = xr.Dataset(
@@ -1139,6 +1151,7 @@ def convert_and_write_partition(
             baseline_ant1_id,
             baseline_ant2_id,
             scan_id,
+            scan_intents,
         )
         logger.debug("Time create coordinates " + str(time.time() - start))
 
@@ -1334,7 +1347,6 @@ def convert_and_write_partition(
 
         partition_info_misc_fields = {
             "scan_name": xds.coords["scan_name"].data,
-            "intents": intents,
             "taql_where": taql_where,
         }
         if with_antenna_partitioning:

xradio/measurement_set/_utils/_msv2/create_field_and_source_xds.py

@@ -819,7 +819,9 @@ def extract_field_info_and_check_ephemeris(
     # Need to check if ephemeris_id is present and if ephemeris table is present.
     if "EPHEMERIS_ID" in field_xds:
         # Note: this assumes partition_scheme includes "FIELD_ID"
-        ephemeris_id = check_if_consistent(field_xds.EPHEMERIS_ID, "EPHEMERIS_ID")
+        ephemeris_id = check_if_consistent(
+            field_xds.EPHEMERIS_ID, "EPHEMERIS_ID", taql_where
+        )
 
         if ephemeris_id > -1:
             files = os.listdir(os.path.join(in_file, "FIELD"))

xradio/measurement_set/_utils/_msv2/msv4_info_dicts.py

@@ -1,3 +1,5 @@
+import re
+
 import numpy as np
 import xarray as xr
 
@@ -19,7 +21,7 @@ def create_info_dicts(
     field_and_source_xds: xr.Dataset,
     partition_info_misc_fields: dict,
     tb_tool: tables.table,
-) -> dict:
+) -> dict[str, dict]:
     """
     For an MSv4, produces several info dicts (partition_info, processor_info,
     observation_info). The info dicts are returned in a dictionary that
@@ -49,36 +51,13 @@ def create_info_dicts(
 
     info_dicts = {}
 
-    # if "line_name" in field_and_source_xds.coords:
-    #     line_name = to_list(unique_1d(np.ravel(field_and_source_xds.line_name.values)))
-    # else:
-    #     line_name = []
-
-    # info_dicts["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)),
-    #     "polarization_setup": to_list(xds.polarization.values),
-    #     "scan_name": to_list(np.unique(partition_info_misc_fields["scan_name"])),
-    #     "source_name": to_list(np.unique(field_and_source_xds.source_name.values)),
-    #     # "source_id": to_list(unique_1d(source_id)),
-    #     "intents": partition_info_misc_fields["intents"].split(","),
-    #     "taql": partition_info_misc_fields["taql_where"],
-    #     "line_name": line_name,
-    # }
-    # if "antenna_name" in partition_info_misc_fields:
-    #     info_dicts["partition_info"]["antenna_name"] = partition_info_misc_fields[
-    #         "antenna_name"
-    #     ]
-
     observation_id = check_if_consistent(
         tb_tool.getcol("OBSERVATION_ID"), "OBSERVATION_ID"
     )
     info_dicts["observation_info"] = create_observation_info(in_file, observation_id)
-    info_dicts["observation_info"]["intents"] = partition_info_misc_fields[
-        "intents"
-    ].split(",")
+    # info_dicts["observation_info"]["intents"] = partition_info_misc_fields[
+    #     "intents"
+    # ].split(",")
 
     processor_id = check_if_consistent(tb_tool.getcol("PROCESSOR_ID"), "PROCESSOR_ID")
     info_dicts["processor_info"] = create_processor_info(in_file, processor_id)
@@ -86,9 +65,13 @@ def create_info_dicts(
     return info_dicts
 
 
-def create_observation_info(in_file: str, observation_id: int):
+def create_observation_info(
+    in_file: str, observation_id: int
+) -> dict[str, list[str] | str]:
     """
     Makes a dict with the observation info extracted from the PROCESSOR subtable.
+    When available, it also takes metadata from the ASDM tables (imported 'asis')
+    ASDM_EXECBLOCK and ASDM_SBSUMMARY
 
     Parameters
     ----------
@@ -111,74 +94,238 @@ def create_observation_info(in_file: str, observation_id: int):
     )
 
     observation_info = {
-        "observer": [generic_observation_xds["OBSERVER"].values[0]],
+        "observer": [str(generic_observation_xds["OBSERVER"].values[0])],
         "release_date": str(
             convert_casacore_time(generic_observation_xds["RELEASE_DATE"].values)[0]
         ),
     }
     # could just assume lower:upper case but keeping explicit dict for now
-    mandatory_fields = {"project": "PROJECT"}
-    for field_msv4, row_msv2 in mandatory_fields.items():
-        observation_info[field_msv4] = generic_observation_xds[row_msv2].values[0]
+    mandatory_fields = {"project_UID": "PROJECT", "observing_log": "LOG"}
+    for field_msv4, col_msv2 in mandatory_fields.items():
+        observation_info[field_msv4] = str(generic_observation_xds[col_msv2].values[0])
+
+    execblock_optional_fields = {
+        "execution_block_UID": "execBlockUID",
+        "session_reference_UID": "sessionReference",
+        "observing_log": "observingLog",
+    }
+    execblock_info = try_optional_asdm_asis_table_info(
+        in_file, "ASDM_EXECBLOCK", execblock_optional_fields
+    )
+    observation_info.update(execblock_info)
+
+    sbsummary_optional_fields = {
+        "scheduling_block_UID": "sbSummaryUID",
+    }
+    sbsummary_info = try_optional_asdm_asis_table_info(
+        in_file, "ASDM_SBSUMMARY", sbsummary_optional_fields
+    )
+    observation_info.update(sbsummary_info)
+
+    observation_info = replace_entity_ids(observation_info)
+
+    observation_info = try_find_uids_from_observation_schedule(
+        generic_observation_xds, observation_info
+    )
+
+    return observation_info
+
+
+def try_optional_asdm_asis_table_info(
+    in_file: str, asdm_table_name: str, optional_fields: dict[str, str]
+) -> dict[str, str]:
+    """
+    Tries to find an optional ASDM_* subtable (ASDM_EXECBLOCK, ASDM_SBSUMMARY, etc.),
+    and if available, gets the optional fields requested into a metadata dict. That
+    dict can be used to populate the observation_info dict.
+
+    Parameters
+    ----------
+    in_file: str
+       path to an input MSv2
+    asm_table_name: str
+       name of the "asis" ASDM table to look for.
+    optional_fields: dict[str, str]
+       dictionary of field/column names (as {MSv4_name: MSv2/ASDM_name}
+
+    Returns:
+    --------
+    table_info: dict
+        observation description (partial, some fields) ready for the MSv4
+        observation_info attr
+    """
+    asdm_asis_xds = None
 
-    exec_block_xds = None
     try:
-        exec_block_xds = load_generic_table(in_file, "ASDM_EXECBLOCK")
+        asdm_asis_xds = load_generic_table(in_file, asdm_table_name)
     except ValueError as exc:
         logger.debug(
-            f"Did not find the ASDM_EXECBLOCK subtable, not loading optional fields in observation_info. Exception: {exc}"
+            f"Did not find the {asdm_table_name} subtable, not loading optional fields in observation_info. Exception: {exc}"
         )
-    if exec_block_xds:
-        exec_block_info = extract_exec_block_info(exec_block_xds)
-        observation_info.update(exec_block_info)
 
-    return observation_info
+    if asdm_asis_xds:
+        table_info = extract_optional_fields_asdm_asis_table(
+            asdm_asis_xds, optional_fields
+        )
+    else:
+        table_info = {}
+
+    return table_info
 
 
-def extract_exec_block_info(exec_block_xds: xr.Dataset) -> dict:
+def extract_optional_fields_asdm_asis_table(
+    asdm_asis_xds: xr.Dataset, optional_fields: dict[str, str]
+) -> dict[str, str]:
     """
-    Get the (optional) fields of the observation_info that come from the
-    ASDM_EXECBLOCK subtable.
+    Get the (optional) fields of the observation_info that come from "asis" ASDM
+    tables like the ASDM_EXECBLOCK and ASDM_SBSUMMARY subtables.
 
     Note this does not parse strings like 'session_reference':
-    '<EntityRef entityId="uid://A001/X133d/X169f" partId="X00000000" entityTypeName="OUSStatus"'
-    We might want to simplify that to 'uid://A001/X133d/X169f', but keeping the
-    full string for now, as it has additional information such as the type.
+    '<EntityRef entityId="uid://A001/X133d/X169f" partId="X00000000" entityTypeName="OUSStatus"'.
+    If only the UID is required that needs to be filtered afterwards.
 
     Parameters
     ----------
-    exec_block_xds: xr.Dataset
-        raw xds read from subtable ASDM_EXECBLOCK
+    asdm_asis_xds: xr.Dataset
+        raw xds read from subtable ASDM_*
 
     Returns:
     --------
-    exec_block_info: dict
-        Execution block description ready for the MSv4 observation_info dict
+    info: dict
+        info dict with description from an ASDM_* subtable, ready
+        for the MSv4 observation_info dict
     """
 
-    optional_fields = {
-        "execution_block_id": "execBlockId",
-        "execution_block_number": "execBlockNum",
-        "execution_block_UID": "execBlockUID",
-        "session_reference": "sessionReference",
-        "observing_script": "observingScript",
-        "observing_script_UID": "observingScriptUID",
-        "observing_log": "observingLog",
-    }
-
-    exec_block_info = {}
-    for field_msv4, row_msv2 in optional_fields.items():
-        if row_msv2 in exec_block_xds.data_vars:
-            msv2_value = exec_block_xds[row_msv2].values[0]
+    table_info = {}
+    for field_msv4, col_msv2 in optional_fields.items():
+        if col_msv2 in asdm_asis_xds.data_vars:
+            msv2_value = asdm_asis_xds[col_msv2].values[0]
             if isinstance(msv2_value, np.ndarray):
-                exec_block_info[field_msv4] = ",".join([log for log in msv2_value])
+                table_info[field_msv4] = ",".join([log for log in msv2_value])
             else:
-                exec_block_info[field_msv4] = msv2_value
+                table_info[field_msv4] = msv2_value
+
+    return table_info
+
+
+def try_find_uids_from_observation_schedule(
+    generic_observation_xds: xr.Dataset, observation_info: dict
+) -> dict[str, str]:
+    """
+    This function tries to parse the execution_block_UID and scheduling_block_UID
+    from the SCHEDULE column of the OBSERVATION subtable. If found, and they
+    could not alreadly be loaded from the ASDM_* subtables, adds them to the
+    output observation_info dict.
+
+    Sometimes, even if the ASDM_EXECBLOCK and ASDM_SBSUMMARY are not available to
+    load various ASDM UIDs, we can still find a couple of them in the
+    OBSERVATION/SCHEDULE column (when the MS is imported from an ASDM, by
+    importasdm). The SCHEDULE column can have values like:
 
-    return exec_block_info
+    '[SchedulingBlock uid://A001/X3571/X122, ExecBlock uid://A002/X1003af4/X75a3]'
+
+    Parameters
+    ----------
+    generic_observation_xds: xr.Dataset
+        generic observation dataset from the OBSERVATION subtable
+    observation_info: dict
+        an observation_info being populated
+
+    Returns:
+    --------
+    info: dict
+        info dict with possibly additional UIDs found in the OBSERVATION
+        subtable
+    """
+
+    out_info = dict(observation_info)
+
+    if "SCHEDULE" in generic_observation_xds.data_vars:
+        schedule = generic_observation_xds["SCHEDULE"].values[0]
+        if isinstance(schedule, np.ndarray) and 2 == len(schedule):
+            if "scheduling_block_UID" not in observation_info:
+                scheduling_uid_match = re.search(
+                    "SchedulingBlock ([\\w/:]+)", schedule[0]
+                )
+                if scheduling_uid_match:
+                    out_info["scheduling_block_UID"] = scheduling_uid_match.group(1)
+            if "execution_block_UID" not in observation_info:
+                execution_uid_match = re.search("ExecBlock ([\\w/:]+)", schedule[1])
+                if execution_uid_match:
+                    out_info["execution_block_UID"] = execution_uid_match.group(1)
+
+    return out_info
+
+
+def replace_entity_ids(observation_info: dict) -> dict[str, list[str] | str]:
+    """
+    For several fields of the input dictionary, which are known to be of "UID" type,
+    replace their lengthy XML string with the UID value contained in it. For example, from
+    '<EntityRef entityId="uid://A001/X133d/X169f" partId="X00000000" entityTypeName="OUSStatus">'
+    it takes 'uid://A001/X133d/X169f'.
+
+    The UID values are written in the MSv2 "asis" ASDM_* subtables imported from ASDM tables
+    as the full string of the EntityRef XML elements. This function takes only the entityId
+    ("uid://A00...") from the EntityRef.
+
+
+    Parameters
+    ----------
+    observation_info: dict
+        info dict where some UID fields (as xml element strings) need to be replaced/simplified
+
+    Returns:
+    --------
+    info: dict
+        dictionary as the input where the UIDs have been replaced by their entityId (uid://A00...)
+
+    """
+    out_info = dict(observation_info)
+
+    entity_refs = [
+        "execution_block_UID",
+        "session_reference_UID",
+        "scheduling_block_UID",
+    ]
+    for ref_name in entity_refs:
+        if ref_name in observation_info:
+            out_info[ref_name] = search_entity_id(observation_info[ref_name])
+
+    return out_info
+
+
+def search_entity_id(entity_ref_xml: str) -> str:
+    """
+    Given an EntityRef XML string from an ASDM, like the following
+    examples:
+
+    - example sbSummaryID:
+    '<EntityRef entityId="uid://A001/X133d/X169a" partId="X00000000" entityTypeName="SchedBlock" documentVersion="1"/>'
+
+    - example sessionReferenceUID:
+    '<EntityRef entityId="uid://A001/X133d/X169f" partId="X00000000" entityTypeName="OUSStatus"'
+
+    this funcion takes the "uid://..." value of the entityId.
+
+    Parameters
+    ----------
+    entity_ref_xml: str
+        An EntityRef from an ASDM table (usually ExecBlock or
+        SBSUMMARY) as found in columns like or execBlockUID,
+        sessionReference or sbSummaryUID.
+
+    Returns:
+    --------
+    str
+        the entityId string value of the EntityRef received, or
+        the same string as received if no entityId could be found.
+    """
+    uid_match = re.search('entityId="([\\w/:]+)"', entity_ref_xml)
+    entity_id = uid_match.group(1) if uid_match else entity_ref_xml
+    return entity_id
 
 
-def create_processor_info(in_file: str, processor_id: int):
+def create_processor_info(in_file: str, processor_id: int) -> dict[str, str]:
     """
     Makes a dict with the processor info extracted from the PROCESSOR subtable.