From 3a095ffc96163b53c966a84d9123fbeb11a3b144 Mon Sep 17 00:00:00 2001 From: Jan-Willem Date: Mon, 6 Apr 2026 13:59:11 -0400 Subject: [PATCH 1/4] Add data load stage to graph. --- src/graphviper/graph_tools/__init__.py | 1 + .../graph_tools/coordinate_utils.py | 91 +++++++++ .../graph_tools/generate_dask_workflow.py | 102 +++++++++- src/graphviper/graph_tools/map.py | 138 ++++++++++++++ tests/test_graph_tools.py | 175 ++++++++++++++++++ 5 files changed, 503 insertions(+), 4 deletions(-) diff --git a/src/graphviper/graph_tools/__init__.py b/src/graphviper/graph_tools/__init__.py index 4314577..c396084 100644 --- a/src/graphviper/graph_tools/__init__.py +++ b/src/graphviper/graph_tools/__init__.py @@ -3,6 +3,7 @@ make_frequency_coord, make_parallel_coord, interpolate_data_coords_onto_parallel_coords, + get_disk_chunk_sizes, ) from .map import map from .reduce import reduce diff --git a/src/graphviper/graph_tools/coordinate_utils.py b/src/graphviper/graph_tools/coordinate_utils.py index 72bd407..e9cac12 100644 --- a/src/graphviper/graph_tools/coordinate_utils.py +++ b/src/graphviper/graph_tools/coordinate_utils.py @@ -711,6 +711,97 @@ def interpolate_data_coords_onto_parallel_coords( return node_task_data_mapping +def get_disk_chunk_sizes( + input_data: Union[Dict, xr.DataTree], + parallel_coords: Dict, +) -> Dict[str, int]: + """Determine the native on-disk chunk size for each parallel coordinate dimension. + + When a Zarr store is opened with ``xarray`` using a dask backend + (``chunks={}``), the resulting dask arrays carry the native zarr chunk sizes. + This function reads those chunk sizes from the coordinate arrays so that the + caller can pass them as ``disk_chunk_sizes`` to + :func:`graphviper.graph_tools.map.map`, enabling the data loading layer to + coalesce I/O at disk-chunk granularity. + + Only the dimensions listed in *parallel_coords* are inspected; non-parallel + dimensions are ignored. + + Parameters + ---------- + input_data : Union[Dict, xr.DataTree] + The opened processing set — either an ``xr.DataTree`` or a plain dict of + ``xr.Dataset``. The datasets must have been opened with a dask backend + (``chunks={}`` or equivalent) so that their coordinates are backed by + dask arrays with chunk information. If a coordinate is numpy-backed (no + chunk information), the full coordinate length is used as the chunk size, + which is equivalent to treating the entire axis as a single disk chunk. + parallel_coords : Dict + The parallel coordinates dict as produced by :func:`make_parallel_coord`. + Only the keys (dimension names) are used. + + Returns + ------- + Dict[str, int] + ``{dim: native_chunk_size}`` for every dimension in *parallel_coords* + that is found in at least one dataset. The returned chunk size is the + **minimum first-chunk size** observed across all datasets, which is the + conservative choice when datasets in the same processing set have + different on-disk chunking for the same dimension. + + Notes + ----- + Zarr stores the chunk size uniformly except for the last chunk, which may be + smaller if the array length is not a multiple of the chunk size. This + function always reads the *first* chunk size (index 0) to get the nominal + chunk size and ignores the trailing remainder chunk. + + Examples + -------- + >>> ps_xdt = open_processing_set(ps_store) + >>> parallel_coords = {"frequency": make_parallel_coord(coord=img_xds.frequency, n_chunks=n_chunks)} + >>> disk_chunk_sizes = get_disk_chunk_sizes(ps_xdt, parallel_coords) + >>> # e.g. {"frequency": 200} + >>> viper_graph = map(..., disk_chunk_sizes=disk_chunk_sizes) + """ + disk_chunk_sizes: Dict[str, int] = {} + + for dim in parallel_coords: + min_chunk_size: Optional[int] = None + + for xds_name, xds in input_data.items(): + # DataTree nodes expose their dataset via .ds; plain dicts yield + # the dataset directly. + ds = xds.ds if isinstance(xds, xr.DataTree) else xds + + if dim not in ds.coords: + continue + + coord_array = ds[dim] + coord_data = coord_array.data # dask array or numpy array + + if hasattr(coord_data, "chunks") and coord_data.chunks: + # Dask array: chunks is a tuple-of-tuples, one tuple per axis. + # For a 1-D coordinate there is exactly one axis (index 0). + first_chunk = int(coord_data.chunks[0][0]) + else: + # Numpy-backed (no dask): treat the whole axis as one chunk. + first_chunk = int(coord_array.size) + + if min_chunk_size is None or first_chunk < min_chunk_size: + min_chunk_size = first_chunk + + if min_chunk_size is not None: + disk_chunk_sizes[dim] = min_chunk_size + else: + logger.warning( + f"get_disk_chunk_sizes: dimension '{dim}' not found in any dataset; " + "skipping." + ) + + return disk_chunk_sizes + + # def interpolate_data_coords_onto_parallel_coords( # parallel_coords: dict, # input_data: Union[Dict, xr.DataTree], diff --git a/src/graphviper/graph_tools/generate_dask_workflow.py b/src/graphviper/graph_tools/generate_dask_workflow.py index 3f8bc9e..2707dda 100644 --- a/src/graphviper/graph_tools/generate_dask_workflow.py +++ b/src/graphviper/graph_tools/generate_dask_workflow.py @@ -1,4 +1,5 @@ import dask +from collections import defaultdict def _tree_combine(list_to_combine, reduce_node_task, input_params): @@ -21,13 +22,106 @@ def _single_node(graph, reduce_node_task, input_params): return dask.delayed(reduce_node_task)(graph, input_params) +def _prepare_task_input(loaded_data, relative_data_selection, input_params): + """Sub-select the task slice from a pre-loaded disk chunk and inject it into + the task parameter dict. + + This is the single per-task node that sits between a shared load node and a + map node. Combining sub-selection and injection into one step avoids the + extra intermediate dask node that would otherwise be required if they were + separate delayed calls. + + The load node is shared across all map tasks that fall within the same + on-disk chunk; this function is NOT shared — one instance runs per map task. + + Parameters + ---------- + loaded_data : dict + ``{xds_name: xarray.Dataset}`` returned by a load node covering a full + disk chunk. + relative_data_selection : dict + ``{xds_name: {dim: slice}}`` with indices relative to the start of the + pre-loaded chunk. ``slice(None)`` entries are skipped. + input_params : dict + Per-task parameter dict produced by :func:`graphviper.graph_tools.map.map`. + + Returns + ------- + dict + Shallow copy of *input_params* with ``"input_data"`` set to the + sub-selected dataset dict. + """ + input_data = {} + for xds_name, xds_isel in relative_data_selection.items(): + if xds_name not in loaded_data: + continue + ds = loaded_data[xds_name] + effective_sel = { + dim: sel for dim, sel in xds_isel.items() if sel != slice(None) + } + if effective_sel: + ds = ds.isel(effective_sel) + input_data[xds_name] = ds + + result = dict(input_params) + result["input_data"] = input_data + return result + + def generate_dask_workflow(viper_graph): dask_graph = [] - for input_params in viper_graph["map"]["input_params"]: - dask_graph.append( - dask.delayed(viper_graph["map"]["node_task"])(dask.delayed(input_params)) - ) + + if "load" in viper_graph: + # ------------------------------------------------------------------ + # Build load nodes — one dask.delayed per unique disk-chunk group. + # Dask will compute each load node once even when multiple map nodes + # depend on it, avoiding redundant disk reads. + # + # Each load node and the prepare/map tasks it feeds are annotated + # with ``viper_load_group`` and ``viper_map_pair`` so that the + # ViperGraphPlugin scheduler plugin can assign priorities that + # (a) minimise the number of concurrently-loaded disk chunks and + # (b) schedule reduction-adjacent task pairs together. + # ------------------------------------------------------------------ + load_fn = viper_graph["load"]["node_task"] + load_nodes = [] + for load_id, lp in enumerate(viper_graph["load"]["input_params"]): + delayed_lp = dask.delayed(lp) + with dask.annotate(viper_load_group=load_id): + load_nodes.append(dask.delayed(load_fn)(delayed_lp)) + + load_node_ids = viper_graph["map"]["load_node_ids"] + relative_data_selections = viper_graph["map"]["relative_data_selections"] + map_fn = viper_graph["map"]["node_task"] + + # Count tasks per load group so pair_id can be computed incrementally. + group_task_counter: dict = defaultdict(int) + + for i, input_params in enumerate(viper_graph["map"]["input_params"]): + load_node_id = load_node_ids[i] + delayed_params = dask.delayed(input_params) + if load_node_id == -1: + # No load node for this task — fall back to per-task loading. + dask_graph.append(dask.delayed(map_fn)(delayed_params)) + else: + load_node = load_nodes[load_node_id] + rel_sel = relative_data_selections[i] + # pair_id groups consecutive tasks within the same load group + # into the pairs that will be combined at the first level of + # the binary tree reduction. + pair_id = group_task_counter[load_node_id] // 2 + group_task_counter[load_node_id] += 1 + with dask.annotate(viper_load_group=load_node_id, viper_map_pair=pair_id): + task_input = dask.delayed(_prepare_task_input)( + load_node, rel_sel, delayed_params + ) + dask_graph.append(dask.delayed(map_fn)(task_input)) + else: + for input_params in viper_graph["map"]["input_params"]: + dask_graph.append( + dask.delayed(viper_graph["map"]["node_task"])(dask.delayed(input_params)) + ) if "reduce" in viper_graph: if viper_graph["reduce"]["mode"] == "tree": diff --git a/src/graphviper/graph_tools/map.py b/src/graphviper/graph_tools/map.py index a0e589a..aab91f5 100644 --- a/src/graphviper/graph_tools/map.py +++ b/src/graphviper/graph_tools/map.py @@ -20,6 +20,9 @@ def map( in_memory_compute: bool = False, client=None, date_time: str = None, + data_loading_task: Union[Callable[..., Any], None] = None, + disk_chunk_sizes: Union[Dict[str, int], None] = None, + load_node_input_params: Union[dict, None] = None, ) -> Dict: """Create a perfectly parallel graph where a node is generated for each item in the :ref:`node_task_data_mapping ` using the function specified in the ``node_task`` parameter. @@ -116,9 +119,144 @@ def map( graph = {"map": {"node_task": node_task, "input_params": input_param_list}} + if data_loading_task is not None and disk_chunk_sizes is not None: + if load_node_input_params is None: + load_node_input_params = {} + load_input_params_list, load_node_ids, relative_data_selections = ( + _build_load_stage(input_param_list, disk_chunk_sizes, load_node_input_params) + ) + graph["load"] = { + "node_task": data_loading_task, + "input_params": load_input_params_list, + } + graph["map"]["load_node_ids"] = load_node_ids + graph["map"]["relative_data_selections"] = relative_data_selections + return graph +def _build_load_stage( + input_param_list: list, + disk_chunk_sizes: Dict[str, int], + load_node_input_params: dict, +): + """Groups mapping tasks by their native on-disk chunk and builds load node parameters. + + For each parallel dimension listed in ``disk_chunk_sizes`` the mapping tasks are + grouped by the disk chunk they fall in. Tasks that span a chunk boundary, or have + no disk-chunked dimension selections, are excluded from the load layer + (``load_node_id == -1``) and will fall back to the normal per-task data loading. + + Parameters + ---------- + input_param_list : list of dict + Per-task parameter dicts produced by :func:`map`. Each dict must contain + ``"input_data_store"`` and ``"data_selection"``. + disk_chunk_sizes : dict + ``{dim: native_chunk_size}`` for every parallel dimension whose on-disk + chunking should be used to coalesce I/O. + load_node_input_params : dict + Extra parameters merged into every load node parameter dict + (e.g. ``processing_set_data_group_name``). + + Returns + ------- + load_input_params_list : list of dict + One parameter dict per unique disk-chunk group (= one load node). + Each dict contains ``"input_data_store"``, ``"data_selection"`` (at the + disk-chunk granularity), and any entries from ``load_node_input_params``. + load_node_ids : list of int + Index into ``load_input_params_list`` for each mapping task, or ``-1`` when + no load node applies to that task. + relative_data_selections : list of dict or None + Per-task data selection expressed as indices *relative to the start of the + pre-loaded disk chunk*, or ``None`` when ``load_node_id == -1``. + """ + disk_chunk_group_to_id: Dict = {} + load_input_params_list: list = [] + load_node_ids: list = [] + relative_data_selections: list = [] + + for task_params in input_param_list: + data_selection = task_params.get("data_selection", {}) + + group_key_parts = [] + disk_level_selection: Dict = {} + relative_sel: Dict = {} + skip = False + + for xds_name, xds_sel in data_selection.items(): + disk_level_selection[xds_name] = {} + relative_sel[xds_name] = {} + + for dim, sel in xds_sel.items(): + if ( + dim in disk_chunk_sizes + and isinstance(sel, slice) + and sel.start is not None + and sel.stop is not None + and sel.start >= 0 + and sel.stop >= 0 + ): + chunk_size = disk_chunk_sizes[dim] + abs_start = sel.start + abs_stop = sel.stop + + start_disk_chunk = abs_start // chunk_size + # abs_stop is exclusive, so the last included index is abs_stop - 1 + end_disk_chunk = max(0, abs_stop - 1) // chunk_size + + if start_disk_chunk != end_disk_chunk: + # Task spans a disk chunk boundary — skip the optimization. + logger.debug( + f"Load layer: task spans disk chunk boundary for dim '{dim}' " + f"(sel={sel}, chunk_size={chunk_size}); falling back to per-task loading." + ) + skip = True + break + + disk_start = start_disk_chunk * chunk_size + disk_stop = disk_start + chunk_size + + disk_level_selection[xds_name][dim] = slice(disk_start, disk_stop) + relative_sel[xds_name][dim] = slice( + abs_start - disk_start, + abs_stop - disk_start, + ) + group_key_parts.append((xds_name, dim, start_disk_chunk)) + else: + # Dim not disk-chunked: load node loads everything; the + # original absolute slice is still valid relative to the + # loaded data (which starts at index 0 for this dim). + disk_level_selection[xds_name][dim] = slice(None) + relative_sel[xds_name][dim] = sel + + if skip: + break + + if skip or not group_key_parts: + load_node_ids.append(-1) + relative_data_selections.append(None) + continue + + group_key = frozenset(group_key_parts) + + if group_key not in disk_chunk_group_to_id: + new_id = len(load_input_params_list) + disk_chunk_group_to_id[group_key] = new_id + load_params: Dict = { + "input_data_store": task_params["input_data_store"], + "data_selection": disk_level_selection, + } + load_params.update(load_node_input_params) + load_input_params_list.append(load_params) + + load_node_ids.append(disk_chunk_group_to_id[group_key]) + relative_data_selections.append(relative_sel) + + return load_input_params_list, load_node_ids, relative_data_selections + + def _select_data(input_data, data_selection): if isinstance(input_data, xr.DataTree): input_data_sel = xr.DataTree() diff --git a/tests/test_graph_tools.py b/tests/test_graph_tools.py index 68aab9a..f95c9f2 100644 --- a/tests/test_graph_tools.py +++ b/tests/test_graph_tools.py @@ -182,9 +182,184 @@ def test_ps_partition(): ) +def test_build_load_stage(): + """Unit test for _build_load_stage: grouping and relative-selection logic.""" + from graphviper.graph_tools.map import _build_load_stage + + # 10 mapping tasks, each covering 1 frequency channel (slices 0-9). + # Disk chunk size = 5 → two load nodes covering [0,5) and [5,10). + input_param_list = [ + { + "input_data_store": "fake_store", + "data_selection": {"xds_0": {"frequency": slice(i, i + 1)}}, + } + for i in range(10) + ] + + load_params_list, load_node_ids, relative_sels = _build_load_stage( + input_param_list, + disk_chunk_sizes={"frequency": 5}, + load_node_input_params={}, + ) + + # Exactly two load nodes (one per disk chunk) + assert len(load_params_list) == 2 + + # Tasks 0-4 → load node 0; tasks 5-9 → load node 1 + assert load_node_ids[:5] == [0, 0, 0, 0, 0] + assert load_node_ids[5:] == [1, 1, 1, 1, 1] + + # Disk-level selections cover the full disk chunk + assert load_params_list[0]["data_selection"]["xds_0"]["frequency"] == slice(0, 5) + assert load_params_list[1]["data_selection"]["xds_0"]["frequency"] == slice(5, 10) + + # Relative selections are expressed relative to the disk chunk start + assert relative_sels[0]["xds_0"]["frequency"] == slice(0, 1) # abs 0→1, disk_start 0 + assert relative_sels[4]["xds_0"]["frequency"] == slice(4, 5) # abs 4→5, disk_start 0 + assert relative_sels[5]["xds_0"]["frequency"] == slice(0, 1) # abs 5→6, disk_start 5 + assert relative_sels[9]["xds_0"]["frequency"] == slice(4, 5) # abs 9→10, disk_start 5 + + +def test_build_load_stage_boundary(): + """Tasks spanning a disk chunk boundary receive load_node_id == -1.""" + from graphviper.graph_tools.map import _build_load_stage + + # Task 0: slice(3, 7) spans chunks 0 and 1 (chunk_size=5) → fallback + # Task 1: slice(0, 3) fits entirely in chunk 0 → gets a load node + input_param_list = [ + { + "input_data_store": "fake_store", + "data_selection": {"xds_0": {"frequency": slice(3, 7)}}, + }, + { + "input_data_store": "fake_store", + "data_selection": {"xds_0": {"frequency": slice(0, 3)}}, + }, + ] + + _, load_node_ids, relative_sels = _build_load_stage( + input_param_list, + disk_chunk_sizes={"frequency": 5}, + load_node_input_params={}, + ) + + assert load_node_ids[0] == -1 # boundary-spanning: no load node + assert relative_sels[0] is None + assert load_node_ids[1] == 0 # fits in one chunk: has a load node + assert relative_sels[1] is not None + + +def test_build_load_stage_extra_params(): + """Extra load_node_input_params are merged into every load node dict.""" + from graphviper.graph_tools.map import _build_load_stage + + input_param_list = [ + { + "input_data_store": "fake_store", + "data_selection": {"xds_0": {"frequency": slice(0, 2)}}, + } + ] + + load_params_list, _, _ = _build_load_stage( + input_param_list, + disk_chunk_sizes={"frequency": 10}, + load_node_input_params={"data_group_name": "corrected"}, + ) + + assert load_params_list[0]["data_group_name"] == "corrected" + + +def test_load_layer_generate_dask_workflow(): + """End-to-end test: load layer pre-loads disk chunks and delivers sub-selected + data to each map task via input_params['input_data'].""" + import numpy as np + import xarray as xr + import dask + + from graphviper.graph_tools.map import map as viper_map + from graphviper.graph_tools.coordinate_utils import ( + interpolate_data_coords_onto_parallel_coords, + make_parallel_coord, + ) + from graphviper.graph_tools.generate_dask_workflow import generate_dask_workflow + + # Synthetic dataset: vis[i] = i for frequency i in 0..11 + n_freq = 12 + freq_vals = np.arange(n_freq, dtype=float) + xds = xr.Dataset( + {"vis": (["frequency"], freq_vals.copy())}, + coords={ + "frequency": xr.DataArray( + freq_vals, + dims=["frequency"], + attrs={"units": "Hz", "type": "spectral_coord", "velocity_frame": "lsrk"}, + ) + }, + ) + input_data = {"xds_0": xds} + + # 6 map tasks of 2 channels each; disk chunk = 4 channels → 3 load nodes + n_map_chunks = 6 + parallel_coords = { + "frequency": make_parallel_coord(coord=xds.frequency, n_chunks=n_map_chunks) + } + node_task_data_mapping = interpolate_data_coords_onto_parallel_coords( + parallel_coords, input_data + ) + + def my_task(input_params): + # With the load layer, input_data is pre-populated. + ps = input_params["input_data"] + assert ps is not None, "Expected pre-loaded data from load layer" + return float(ps["xds_0"]["vis"].values.sum()) + + def my_load_fn(load_params): + sel = { + k: v + for k, v in load_params["data_selection"]["xds_0"].items() + if v != slice(None) + } + ds = input_data["xds_0"].isel(sel).load() + return {"xds_0": ds} + + graph = viper_map( + input_data=input_data, + node_task_data_mapping=node_task_data_mapping, + node_task=my_task, + input_params={"input_data_store": "synthetic"}, + in_memory_compute=False, + data_loading_task=my_load_fn, + disk_chunk_sizes={"frequency": 4}, # 4-channel disk chunks, 2-channel map tasks + ) + + # Verify graph structure + assert "load" in graph + assert len(graph["load"]["input_params"]) == 3 # 12 channels / 4 per chunk + assert len(graph["map"]["load_node_ids"]) == n_map_chunks + assert all(lid >= 0 for lid in graph["map"]["load_node_ids"]) + + # Tasks sharing the same disk chunk should share the same load node id + ids = graph["map"]["load_node_ids"] + assert ids[0] == ids[1] # tasks 0 and 1 share disk chunk 0 (freq 0-3) + assert ids[2] == ids[3] # tasks 2 and 3 share disk chunk 1 (freq 4-7) + assert ids[4] == ids[5] # tasks 4 and 5 share disk chunk 2 (freq 8-11) + assert ids[0] != ids[2] != ids[4] + + dask_graph = generate_dask_workflow(graph) + results = dask.compute(dask_graph)[0] + + # Each task k covers 2 consecutive channels 2k and 2k+1 → sum = 4k+1 + expected_total = sum(range(n_freq)) # 0+1+...+11 = 66 + assert abs(sum(results) - expected_total) < 1e-10 + + if __name__ == "__main__": test_map_reduce() test_ps_partition() + test_build_load_stage() + test_build_load_stage_boundary() + test_build_load_stage_extra_params() + test_load_layer_generate_dask_workflow() """ chunk_indx 0 (0, 0) chunk_indx 1 (0, 1) From 89c4ed7724abe9e975977b41bb77659e1dd43b3d Mon Sep 17 00:00:00 2001 From: Jan-Willem Date: Mon, 6 Apr 2026 15:40:02 -0400 Subject: [PATCH 2/4] Update tutorials with new optional data load layer. --- docs/graph_building_tutorial.ipynb | 431 +++++++++++++++++++---- docs/graph_building_tutorial_image.ipynb | 318 +++++++++++++++-- 2 files changed, 642 insertions(+), 107 deletions(-) diff --git a/docs/graph_building_tutorial.ipynb b/docs/graph_building_tutorial.ipynb index 8d26ece..0150373 100644 --- a/docs/graph_building_tutorial.ipynb +++ b/docs/graph_building_tutorial.ipynb @@ -104,8 +104,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "[\u001b[38;2;128;05;128m2026-03-10 15:46:25,378\u001b[0m] \u001b[38;2;255;160;0m WARNING\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m It is recommended that the local cache directory be set using the \u001b[38;2;50;50;205mdask_local_dir\u001b[0m parameter. \n", - "[\u001b[38;2;128;05;128m2026-03-10 15:46:25,378\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m Running client in synchronous mode. \n" + "[\u001b[38;2;128;05;128m2026-04-06 15:38:38,938\u001b[0m] \u001b[38;2;255;160;0m WARNING\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m It is recommended that the local cache directory be set using the \u001b[38;2;50;50;205mdask_local_dir\u001b[0m parameter. \n", + "[\u001b[38;2;128;05;128m2026-04-06 15:38:38,938\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m Running client in synchronous mode. \n" ] } ], @@ -136,14 +136,14 @@ "name": "stdout", "output_type": "stream", "text": [ - "[\u001b[38;2;128;05;128m2026-03-10 15:46:25,381\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m Initializing download... \n", - "[\u001b[38;2;128;05;128m2026-03-10 15:46:25,382\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m File already exists: /Users/joshua/Development/graphviper/docs/Antennae_North.cal.lsrk.split.ms \n", - "[\u001b[38;2;128;05;128m2026-03-10 15:46:27,886\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m Updated partition scheme used: ['DATA_DESC_ID', 'OBS_MODE', 'OBSERVATION_ID'] \n", - "[\u001b[38;2;128;05;128m2026-03-10 15:46:27,888\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m Number of partitions: 4 \n", - "[\u001b[38;2;128;05;128m2026-03-10 15:46:27,888\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m OBSERVATION_ID [0], DDI [0], STATE [23, 24, 25, 30, 31, 32, 33, 34, 37], FIELD [0, 1, 2], SCAN [9, 17, 21, 25], EPHEMERIS [None] \n", - "[\u001b[38;2;128;05;128m2026-03-10 15:46:28,070\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m OBSERVATION_ID [1], DDI [0], STATE [23, 24, 25, 30, 31, 32, 33, 34, 37], FIELD [0, 1, 2], SCAN [26, 34, 38, 42], EPHEMERIS [None] \n", - "[\u001b[38;2;128;05;128m2026-03-10 15:46:28,209\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m OBSERVATION_ID [2], DDI [0], STATE [32, 33, 34], FIELD [0, 1, 2], SCAN [43], EPHEMERIS [None] \n", - "[\u001b[38;2;128;05;128m2026-03-10 15:46:28,321\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m OBSERVATION_ID [3], DDI [0], STATE [39, 40, 41, 46, 47, 48, 49, 50, 53], FIELD [0, 1, 2], SCAN [48, 56, 60, 64], EPHEMERIS [None] \n" + "[\u001b[38;2;128;05;128m2026-04-06 15:38:38,947\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m Initializing download... \n", + "[\u001b[38;2;128;05;128m2026-04-06 15:38:38,948\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m File already exists: /Users/jsteeb/Dropbox/viper_dev/graphviper/docs/Antennae_North.cal.lsrk.split.ms \n", + "[\u001b[38;2;128;05;128m2026-04-06 15:38:39,446\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m Updated partition scheme used: ['DATA_DESC_ID', 'OBS_MODE', 'OBSERVATION_ID'] \n", + "[\u001b[38;2;128;05;128m2026-04-06 15:38:39,453\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m Number of partitions: 4 \n", + "[\u001b[38;2;128;05;128m2026-04-06 15:38:39,453\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m OBSERVATION_ID [0], DDI [0], STATE [23, 24, 25, 30, 31, 32, 33, 34, 37], FIELD [0, 1, 2], SCAN [9, 17, 21, 25], EPHEMERIS [None] \n", + "[\u001b[38;2;128;05;128m2026-04-06 15:38:39,803\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m OBSERVATION_ID [1], DDI [0], STATE [23, 24, 25, 30, 31, 32, 33, 34, 37], FIELD [0, 1, 2], SCAN [26, 34, 38, 42], EPHEMERIS [None] \n", + "[\u001b[38;2;128;05;128m2026-04-06 15:38:40,021\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m OBSERVATION_ID [2], DDI [0], STATE [32, 33, 34], FIELD [0, 1, 2], SCAN [43], EPHEMERIS [None] \n", + "[\u001b[38;2;128;05;128m2026-04-06 15:38:40,298\u001b[0m] \u001b[38;2;50;50;205m INFO\u001b[0m\u001b[38;2;112;128;144m client: \u001b[0m OBSERVATION_ID [3], DDI [0], STATE [39, 40, 41, 46, 47, 48, 49, 50, 53], FIELD [0, 1, 2], SCAN [48, 56, 60, 64], EPHEMERIS [None] \n" ] } ], @@ -957,11 +957,11 @@ "│ FLAG (time, baseline_id, frequency, polarization) bool 36kB dask.array<chunksize=(50, 45, 3, 2), meta=np.ndarray>\n", "│ TIME_CENTROID (time, baseline_id) float64 18kB dask.array<chunksize=(50, 45), meta=np.ndarray>\n", "│ UVW (time, baseline_id, uvw_label) float64 54kB dask.array<chunksize=(50, 45, 3), meta=np.ndarray>\n", - "│ VISIBILITY (time, baseline_id, frequency, polarization) complex128 576kB dask.array<chunksize=(50, 45, 3, 2), meta=np.ndarray>\n", - "│ WEIGHT (time, baseline_id, frequency, polarization) float64 288kB dask.array<chunksize=(50, 45, 3, 2), meta=np.ndarray>\n", + "│ VISIBILITY (time, baseline_id, frequency, polarization) complex64 288kB dask.array<chunksize=(50, 45, 3, 2), meta=np.ndarray>\n", + "│ WEIGHT (time, baseline_id, frequency, polarization) float32 144kB dask.array<chunksize=(50, 45, 3, 2), meta=np.ndarray>\n", "│ Attributes:\n", - "│ creation_date: 2026-03-10T06:46:27.899663+00:00\n", - "│ creator: {'software_name': 'xradio', 'version': '1.1.2'}\n", + "│ creation_date: 2026-04-06T19:38:39.489741+00:00\n", + "│ creator: {'software_name': 'xradio', 'version': '1.1.3'}\n", "│ data_groups: {'base': {'correlated_data': 'VISIBILITY', 'date': '20...\n", "│ observation_info: {'execution_block_UID': 'uid://A002/X1ff7b0/Xb', 'obse...\n", "│ processor_info: {'sub_type': 'ALMA_CORRELATOR_MODE', 'type': 'CORRELAT...\n", @@ -1013,7 +1013,7 @@ " WIND_DIRECTION (station_name, time_weather) float64 4kB dask.array<chunksize=(2, 259), meta=np.ndarray>\n", " WIND_SPEED (station_name, time_weather) float64 4kB dask.array<chunksize=(2, 259), meta=np.ndarray>\n", " Attributes:\n", - " type: weather