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

xradio/measurement_set/_utils/_msv2/partition_queries.py

@@ -1,5 +1,8 @@
 import itertools
+import time
 import toolviper.utils.logger as logger
+import os
+import pandas as pd
 
 import numpy as np
 
@@ -17,6 +20,9 @@ def enumerated_product(*args):
     )
 
 
+import pickle, gzip
+
+
 def create_partitions(in_file: str, partition_scheme: list) -> list[dict]:
     """Create a list of dictionaries with the partition information.
 
@@ -34,38 +40,53 @@ def create_partitions(in_file: str, partition_scheme: list) -> list[dict]:
     list
         list of dictionaries with the partition information.
     """
-    # vla_otf (bool, optional):  The partioning of VLA OTF (on the fly) mosaics needs a special partitioning scheme. Defaults to False.
 
-    # Create partition table
-    import pandas as pd
-    import os
+    ### Test new implementation without
+    # Always start with these (if available); then extend with user scheme.
+    partition_scheme = [
+        "DATA_DESC_ID",
+        "OBS_MODE",
+        "OBSERVATION_ID",
+        "EPHEMERIS_ID",
+    ] + list(partition_scheme)
 
-    partition_scheme = ["DATA_DESC_ID", "OBS_MODE", "OBSERVATION_ID"] + partition_scheme
+    # partition_scheme = ["DATA_DESC_ID", "OBS_MODE"] + list(
+    #     partition_scheme
+    # )
 
-    # Open MSv2 tables and add columns to partition table (par_df):
-    par_df = pd.DataFrame()
+    t0 = time.time()
+    # --------- Load base columns from MAIN table ----------
     main_tb = tables.table(
         in_file, readonly=True, lockoptions={"option": "usernoread"}, ack=False
     )
-    par_df["DATA_DESC_ID"] = main_tb.getcol("DATA_DESC_ID")
-    par_df["FIELD_ID"] = main_tb.getcol("FIELD_ID")
-    par_df["SCAN_NUMBER"] = main_tb.getcol("SCAN_NUMBER")
-    par_df["STATE_ID"] = main_tb.getcol("STATE_ID")
-    par_df["OBSERVATION_ID"] = main_tb.getcol("OBSERVATION_ID")
-    par_df["ANTENNA1"] = main_tb.getcol("ANTENNA1")
-    par_df = par_df.drop_duplicates()
 
+    # Build minimal DF once. Pull only columns we may need.
+    # Add columns here if you expect to aggregate them per-partition.
+    base_cols = {
+        "DATA_DESC_ID": main_tb.getcol("DATA_DESC_ID"),
+        "FIELD_ID": main_tb.getcol("FIELD_ID"),
+        "SCAN_NUMBER": main_tb.getcol("SCAN_NUMBER"),
+        "STATE_ID": main_tb.getcol("STATE_ID"),
+        "OBSERVATION_ID": main_tb.getcol("OBSERVATION_ID"),
+        "ANTENNA1": main_tb.getcol("ANTENNA1"),
+    }
+    par_df = pd.DataFrame(base_cols).drop_duplicates()
+    logger.debug(
+        f"Loaded MAIN columns in {time.time() - t0:.2f}s "
+        f"({len(par_df):,} unique MAIN rows)"
+    )
+
+    # --------- Optional SOURCE/STATE derived columns ----------
+    # SOURCE_ID (via FIELD table)
+    t1 = time.time()
+    source_id_added = False
     field_tb = tables.table(
         os.path.join(in_file, "FIELD"),
         readonly=True,
         lockoptions={"option": "usernoread"},
         ack=False,
     )
-    # if vla_otf:
-    #     par_df["FIELD_NAME"] = np.array(field_tb.getcol("NAME"))[par_df["FIELD_ID"]]
-
-    # Get source ids if available from source table.
-    if table_exists(os.path.join(os.path.join(in_file, "SOURCE"))):
+    if table_exists(os.path.join(in_file, "SOURCE")):
         source_tb = tables.table(
             os.path.join(in_file, "SOURCE"),
             readonly=True,
@@ -73,13 +94,31 @@ def create_partitions(in_file: str, partition_scheme: list) -> list[dict]:
             ack=False,
         )
         if source_tb.nrows() != 0:
-            par_df["SOURCE_ID"] = field_tb.getcol("SOURCE_ID")[par_df["FIELD_ID"]]
-            # if vla_otf:
-            #     par_df["SOURCE_NAME"] = np.array(source_tb.getcol("NAME"))[
-            #         par_df["SOURCE_ID"]
-            #     ]
+            # Map SOURCE_ID via FIELD_ID
+            field_source = np.asarray(field_tb.getcol("SOURCE_ID"))
+            par_df["SOURCE_ID"] = field_source[par_df["FIELD_ID"]]
+            source_id_added = True
+    logger.debug(
+        f"SOURCE processing in {time.time() - t1:.2f}s "
+        f"(added SOURCE_ID={source_id_added})"
+    )
+
+    if "EPHEMERIS_ID" in field_tb.colnames():
+        ephemeris_id_added = False
+        if field_tb.nrows() != 0:
+            # Map EPHEMERIS_ID via FIELD_ID
+            field_ephemeris = np.asarray(field_tb.getcol("EPHEMERIS_ID"))
+            par_df["EPHEMERIS_ID"] = field_ephemeris[par_df["FIELD_ID"]]
+            ephemeris_id_added = True
+        logger.debug(
+            f"EPHEMERIS processing in {time.time() - t1:.2f}s "
+            f"(added EPHEMERIS_ID={ephemeris_id_added})"
+        )
 
-    # Get intents and subscan numbers if available from state table.
+    # OBS_MODE & SUB_SCAN_NUMBER (via STATE table)
+    t2 = time.time()
+    obs_mode_added = False
+    sub_scan_added = False
     if table_exists(os.path.join(in_file, "STATE")):
         state_tb = tables.table(
             os.path.join(in_file, "STATE"),
@@ -88,30 +127,36 @@ def create_partitions(in_file: str, partition_scheme: list) -> list[dict]:
             ack=False,
         )
         if state_tb.nrows() != 0:
-            # print('state_tb',state_tb.nrows(),state_tb)
-            par_df["OBS_MODE"] = np.array(state_tb.getcol("OBS_MODE"))[
-                par_df["STATE_ID"]
-            ]
-            par_df["SUB_SCAN_NUMBER"] = state_tb.getcol("SUB_SCAN")[par_df["STATE_ID"]]
+            state_obs_mode = np.asarray(state_tb.getcol("OBS_MODE"))
+            state_sub_scan = np.asarray(state_tb.getcol("SUB_SCAN"))
+            # Index by STATE_ID into STATE columns
+            par_df["OBS_MODE"] = state_obs_mode[par_df["STATE_ID"]]
+            par_df["SUB_SCAN_NUMBER"] = state_sub_scan[par_df["STATE_ID"]]
+            obs_mode_added = True
+            sub_scan_added = True
         else:
-            par_df.drop(["STATE_ID"], axis=1)
+            # If STATE empty, drop STATE_ID (it cannot partition anything)
+            if "STATE_ID" in par_df.columns:
+                par_df.drop(columns=["STATE_ID"], inplace=True)
 
-    # Check if all partition scheme criteria are present in the partition table.
-    partition_scheme_updated = []
-    partition_criteria = {}
-    for par in partition_scheme:
-        if par in par_df.columns:
-            partition_criteria[par] = par_df[par].unique()
-            partition_scheme_updated.append(par)
-    logger.info(f"Partition scheme that will be used: {partition_scheme_updated}")
+            if "SUB_SCAN_NUMBER" in par_df.columns:
+                par_df.drop(columns=["SUB_SCAN_NUMBER"], inplace=True)
 
-    # Make all possible combinations of the partition criteria.
-    enumerated_partitions = enumerated_product(*list(partition_criteria.values()))
+    logger.debug(
+        f"STATE processing in {time.time() - t2:.2f}s "
+        f"(OBS_MODE={obs_mode_added}, SUB_SCAN_NUMBER={sub_scan_added})"
+    )
 
-    # print('par_df',par_df)
+    # --------- Decide which partition keys are actually available ----------
+    t3 = time.time()
+    partition_scheme_updated = [k for k in partition_scheme if k in par_df.columns]
+    logger.info(f"Updated partition scheme used: {partition_scheme_updated}")
 
-    # Create a list of dictionaries with the partition information. This will be used to query the MSv2 main table.
-    partitions = []
+    # If none of the requested keys exist, there is a single partition of "everything"
+    if not partition_scheme_updated:
+        partition_scheme_updated = []
+
+    # These are the axes we report per partition (present => aggregate unique values)
     partition_axis_names = [
         "DATA_DESC_ID",
         "OBSERVATION_ID",
@@ -121,30 +166,172 @@ def create_partitions(in_file: str, partition_scheme: list) -> list[dict]:
         "SOURCE_ID",
         "OBS_MODE",
         "SUB_SCAN_NUMBER",
+        "EPHEMERIS_ID",
     ]
+    # Only include ANTENNA1 if user asked for it (keeps output size down)
     if "ANTENNA1" in partition_scheme:
         partition_axis_names.append("ANTENNA1")
 
-    for idx, pair in enumerated_partitions:
-        query = ""
-        for i, par in enumerate(partition_scheme_updated):
-            if isinstance(pair[i], str):
-                query = query + f'{par} == "{pair[i]}" and '
+    # --------- Group only by realized partitions (no Cartesian product!) ----------
+    # observed=True speeds up if categorical; here it’s harmless. sort=False keeps source order.
+    if partition_scheme_updated:
+        grp = par_df.groupby(partition_scheme_updated, sort=False, observed=False)
+        groups_iter = grp
+    else:
+        # Single group: everything
+        groups_iter = [(None, par_df)]
+
+    partitions = []
+    # Fast aggregation: use NumPy for uniques to avoid pandas overhead in the tight loop.
+    for _, gdf in groups_iter:
+        part = {}
+        for name in partition_axis_names:
+            if name in gdf.columns:
+                # Return Python lists to match your prior structure (can be np.ndarray if preferred)
+                part[name] = np.unique(gdf[name].to_numpy()).tolist()
             else:
-                query = query + f"{par} == {pair[i]} and "
-        query = query[:-4]  # remove last and
-        sub_par_df = par_df.query(query).drop_duplicates()
+                part[name] = [None]
+        partitions.append(part)
 
-        if sub_par_df.shape[0] != 0:
-            partition_info = {}
+    logger.debug(
+        f"Partition build in {time.time() - t3:.2f}s; total {len(partitions):,} partitions"
+    )
+    logger.debug(f"Total create_partitions time: {time.time() - t0:.2f}s")
 
-            # FIELD_NAME	SOURCE_NAME
-            for col_name in partition_axis_names:
-                if col_name in sub_par_df.columns:
-                    partition_info[col_name] = sub_par_df[col_name].unique()
-                else:
-                    partition_info[col_name] = [None]
+    # # with gzip.open("partition_original_small.pkl.gz", "wb") as f:
+    # #     pickle.dump(partitions, f, protocol=pickle.HIGHEST_PROTOCOL)
 
-            partitions.append(partition_info)
+    # #partitions[1]["DATA_DESC_ID"] = [999]  # make a change to test comparison
+    # #org_partitions = load_dict_list("partition_original_small.pkl.gz")
+    # org_partitions = load_dict_list("partition_original.pkl.gz")
 
     return partitions
+
+
+from typing import Any, List, Dict
+
+
+def save_dict_list(filename: str, data: List[Dict[str, Any]]) -> None:
+    """
+    Save a list of dictionaries containing NumPy arrays (or other objects)
+    to a compressed pickle file.
+    """
+    with gzip.open(filename, "wb") as f:
+        pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+
+def load_dict_list(filename: str) -> List[Dict[str, Any]]:
+    """
+    Load a list of dictionaries containing NumPy arrays (or other objects)
+    from a compressed pickle file.
+    """
+    with gzip.open(filename, "rb") as f:
+        return pickle.load(f)
+
+
+def dict_list_equal(a: List[Dict[str, Any]], b: List[Dict[str, Any]]) -> bool:
+    """
+    Compare two lists of dictionaries to ensure they are exactly the same.
+    NumPy arrays are compared with array_equal, other objects with ==.
+    """
+    if len(a) != len(b):
+        return False
+
+    for d1, d2 in zip(a, b):
+        if d1.keys() != d2.keys():
+            return False
+        for k in d1:
+            v1, v2 = d1[k], d2[k]
+            if isinstance(v1, np.ndarray) and isinstance(v2, np.ndarray):
+                if not np.array_equal(v1, v2):
+                    return False
+            else:
+                if v1 != v2:
+                    return False
+    return True
+
+
+from typing import Iterable, Mapping, Tuple, List, Dict, Any, Set
+import numpy as np
+
+
+def _to_python_scalar(x: Any) -> Any:
+    """Convert NumPy scalars to Python scalars; leave others unchanged."""
+    if isinstance(x, np.generic):
+        return x.item()
+    return x
+
+
+def _to_hashable_value_list(v: Any) -> Tuple[Any, ...]:
+    """
+    Normalize a dict value (often list/np.ndarray) into a sorted, hashable tuple.
+    - Accepts list/tuple/np.ndarray/scalars/None.
+    - Treats None as a value.
+    - Sorts with a stable key that stringifies items to avoid dtype hiccups.
+    """
+    if isinstance(v, np.ndarray):
+        v = v.tolist()
+    if v is None or isinstance(v, (str, bytes)):
+        # Treat a bare scalar as a single-element collection for consistency.
+        v = [v]
+    elif not isinstance(v, (list, tuple)):
+        v = [v]
+
+    py_vals = [_to_python_scalar(x) for x in v]
+    # Sort by (type name, repr) to keep mixed types stable if present
+    return tuple(sorted(py_vals, key=lambda x: (type(x).__name__, repr(x))))
+
+
+def _canon_partition(
+    d: Mapping[str, Any], ignore_keys: Iterable[str] = ()
+) -> Tuple[Tuple[str, Tuple[Any, ...]], ...]:
+    """
+    Canonicalize a partition dict into a hashable, order-insensitive representation.
+    - Drops keys in ignore_keys.
+    - Converts each value collection to a sorted tuple.
+    - Sorts keys.
+    """
+    ign: Set[str] = set(ignore_keys)
+    items = []
+    for k, v in d.items():
+        if k in ign:
+            continue
+        items.append((k, _to_hashable_value_list(v)))
+    items.sort(key=lambda kv: kv[0])
+    return tuple(items)
+
+
+def compare_partitions_subset(
+    new_partitions: List[Dict[str, Any]],
+    original_partitions: List[Dict[str, Any]],
+    ignore_keys: Iterable[str] = (),
+) -> Tuple[bool, List[Dict[str, Any]]]:
+    """
+    Check that every partition in `new_partitions` also appears in `original_partitions`,
+    ignoring ordering (of partitions and of values within each key).
+
+    Parameters
+    ----------
+    new_partitions : list of dict
+        Partitions produced by the optimized/new code.
+    original_partitions : list of dict
+        Partitions produced by the original code (the reference).
+    ignore_keys : iterable of str, optional
+        Keys to ignore when comparing partitions (e.g., timestamps or debug fields).
+
+    Returns
+    -------
+    (ok, missing)
+        ok : bool
+            True if every new partition is found in the original set.
+        missing : list of dict
+            The list of partitions (from `new_partitions`) that were NOT found in `original_partitions`,
+            useful for debugging diffs.
+    """
+    orig_set = {_canon_partition(p, ignore_keys) for p in original_partitions}
+    missing = []
+    for p in new_partitions:
+        cp = _canon_partition(p, ignore_keys)
+        if cp not in orig_set:
+            missing.append(p)
+    return (len(missing) == 0, missing)

xradio/measurement_set/convert_msv2_to_processing_set.py

@@ -1,10 +1,13 @@
 import toolviper.utils.logger as logger
 import numcodecs
-from typing import Dict, Union
+from typing import Dict, Union, Literal
+import time
 
 import dask
 
-from xradio.measurement_set._utils._msv2.partition_queries import create_partitions
+from xradio.measurement_set._utils._msv2.partition_queries import (
+    create_partitions,
+)
 from xradio.measurement_set._utils._msv2.conversion import (
     convert_and_write_partition,
     estimate_memory_and_cores_for_partitions,
@@ -13,7 +16,7 @@ from xradio.measurement_set._utils._msv2.conversion import (
 
 def estimate_conversion_memory_and_cores(
     in_file: str,
-    partition_scheme: list = ["FIELD_ID"],
+    partition_scheme: list = [],
 ) -> tuple[float, int, int]:
     """
     Given an MSv2 and a partition_scheme to use when converting it to MSv4,
@@ -52,7 +55,7 @@ def estimate_conversion_memory_and_cores(
 def convert_msv2_to_processing_set(
     in_file: str,
     out_file: str,
-    partition_scheme: list = ["FIELD_ID"],
+    partition_scheme: list = [],
     main_chunksize: Union[Dict, float, None] = None,
     with_pointing: bool = True,
     pointing_chunksize: Union[Dict, float, None] = None,
@@ -63,8 +66,8 @@ def convert_msv2_to_processing_set(
     use_table_iter: bool = False,
     compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
     add_reshaping_indices: bool = False,
-    storage_backend: str = "zarr",
-    parallel_mode: str = "none",
+    storage_backend: Literal["zarr", "netcdf"] = "zarr",
+    parallel_mode: Literal["none", "partition", "time"] = "none",
     overwrite: bool = False,
 ):
     """Convert a Measurement Set v2 into a Processing Set of Measurement Set v4.
@@ -80,7 +83,7 @@ def convert_msv2_to_processing_set(
         In addition to data description and polarization setup a finer partitioning is possible by specifying a list of partitioning keys. Any combination of the following keys are possible:
         "FIELD_ID", "SCAN_NUMBER", "STATE_ID", "SOURCE_ID", "SUB_SCAN_NUMBER", "ANTENNA1".
         "ANTENNA1" is intended as a single-dish specific partitioning option.
-        For mosaics where the phase center is rapidly changing (such as VLA on the fly mosaics) partition_scheme should be set to an empty list []. By default, ["FIELD_ID"].
+        For mosaics where the phase center is rapidly changing (such as VLA on the fly mosaics) partition_scheme should be set to an empty list []. By default, [].
     main_chunksize : Union[Dict, float, None], optional
         Defines the chunk size of the main dataset. If given as a dictionary, defines the sizes of several dimensions, and acceptable keys are "time", "baseline_id", "antenna_id", "frequency", "polarization". If given as a float, gives the size of a chunk in GiB. By default, None.
     with_pointing : bool, optional
@@ -101,9 +104,9 @@ def convert_msv2_to_processing_set(
         The Blosc compressor to use when saving the converted data to disk using Zarr, by default numcodecs.Zstd(level=2).
     add_reshaping_indices : bool, optional
         Whether to add the tidxs, bidxs and row_id variables to each partition of the main dataset. These can be used to reshape the data back to the original ordering in the MS v2. This is mainly intended for testing and debugging, by default False.
-    storage_backend : {"zarr", "netcdf"}, optional
+    storage_backend : Literal["zarr", "netcdf"], optional
         The on-disk format to use. "netcdf" is not yet implemented.
-    parallel_mode : {"none", "partition", "time"}, optional
+    parallel_mode : Literal["none", "partition", "time"], optional
         Choose whether to use Dask to execute conversion in parallel, by default "none" and conversion occurs serially.
         The option "partition", parallelises the conversion over partitions specified by `partition_scheme`. The option "time" can only be used for phased array interferometers where there are no partitions
         in the MS v2; instead the MS v2 is parallelised along the time dimension and can be controlled by `main_chunksize`.
@@ -134,6 +137,7 @@ def convert_msv2_to_processing_set(
         parallel_mode = "none"
 
     partitions = create_partitions(in_file, partition_scheme=partition_scheme)
+
     logger.info("Number of partitions: " + str(len(partitions)))
     if parallel_mode == "time":
         assert (
@@ -143,7 +147,6 @@ def convert_msv2_to_processing_set(
     delayed_list = []
 
     for ms_v4_id, partition_info in enumerate(partitions):
-        # print(ms_v4_id,len(partition_info['FIELD_ID']))
 
         logger.info(
             "OBSERVATION_ID "
@@ -156,6 +159,11 @@ def convert_msv2_to_processing_set(
             + str(partition_info["FIELD_ID"])
             + ", SCAN "
             + str(partition_info["SCAN_NUMBER"])
+            + (
+                ", EPHEMERIS " + str(partition_info["EPHEMERIS_ID"])
+                if "EPHEMERIS_ID" in partition_info
+                else ""
+            )
             + (
                 ", ANTENNA " + str(partition_info["ANTENNA1"])
                 if "ANTENNA1" in partition_info
@@ -188,6 +196,7 @@ def convert_msv2_to_processing_set(
                 )
             )
         else:
+            start_time = time.time()
             convert_and_write_partition(
                 in_file,
                 out_file,
@@ -207,6 +216,10 @@ def convert_msv2_to_processing_set(
                 parallel_mode=parallel_mode,
                 overwrite=overwrite,
             )
+            end_time = time.time()
+            logger.debug(
+                f"Time to convert partition {ms_v4_id}: {end_time - start_time:.2f} seconds"
+            )
 
     if parallel_mode == "partition":
         dask.compute(delayed_list)

xradio/measurement_set/measurement_set_xdt.py

@@ -203,19 +203,26 @@ class MeasurementSetXdt:
         else:
             line_name = []
 
-        if "spectral_window_intent" not in self._xdt.frequency.attrs:
+        if "spectral_window_intents" not in self._xdt.frequency.attrs:
             spw_intent = "UNSPECIFIED"
         else:
-            spw_intent = self._xdt.frequency.attrs["spectral_window_intent"]
+            spw_intent = self._xdt.frequency.attrs["spectral_window_intents"]
+
+        if "intents" in self._xdt.observation_info:
+            scan_intents = self._xdt.observation_info["intents"]
+        else:
+            scan_intents = self._xdt.scan_name.attrs.get(
+                "scan_intents", ["UNSPECIFIED"]
+            )
 
         partition_info = {
             "spectral_window_name": self._xdt.frequency.attrs["spectral_window_name"],
-            "spectral_window_intent": spw_intent,
+            "spectral_window_intents": spw_intent,
             "field_name": to_list(np.unique(field_and_source_xds.field_name.values)),
             "polarization_setup": to_list(self._xdt.polarization.values),
             "scan_name": to_list(np.unique(self._xdt.scan_name.values)),
             "source_name": to_list(np.unique(field_and_source_xds.source_name.values)),
-            "intents": self._xdt.observation_info["intents"],
+            "scan_intents": scan_intents,
             "line_name": line_name,
             "data_group_name": data_group_name,
         }

xradio/measurement_set/open_processing_set.py

@@ -5,7 +5,7 @@ import xarray as xr
 
 def open_processing_set(
     ps_store: str,
-    intents: list = None,
+    scan_intents: list | None = None,
 ) -> xr.DataTree:
     """Creates a lazy representation of a Processing Set (only meta-data is loaded into memory).
 
@@ -13,9 +13,9 @@ def open_processing_set(
     ----------
     ps_store : str
         String of the path and name of the processing set. For example '/users/user_1/uid___A002_Xf07bba_Xbe5c_target.lsrk.vis.zarr'.
-    intents : list, optional
-        A list of intents to be opened for example ['OBSERVE_TARGET#ON_SOURCE']. The intents in a processing_set_xdt can be seen by calling processing_set_xdt.ps.summary().
-        By default None, which will include all intents.
+    scan_intents : str | None, optional
+        A list of scan_intents to be opened for example ['OBSERVE_TARGET#ON_SOURCE']. The scan_intents in a processing_set_xdt can be seen by calling processing_set_xdt.ps.summary().
+        By default None, which will include all scan_intents.
 
     Returns
     -------
@@ -34,10 +34,10 @@ def open_processing_set(
 
     # Future work is to add ASDM backend
 
-    if intents is None:
+    if scan_intents is None:
         return ps_xdt
     else:
-        return ps_xdt.xr_ps.query(intents=intents)
+        return ps_xdt.xr_ps.query(scan_intents=scan_intents)
 
 
 # def open_processing_set(