Open in Colab: https://colab.research.google.com/github/casangi/graph_viper/blob/master/docs/graph_building_tutorial_image.ipynb

GraphVIPER Tutorial: Image

[1]:

import os
import xradio

from importlib.metadata import version

try:
    import graphviper

    print("GraphVIPER version", version("graphviper"), "already installed.")
except ImportError as e:
    print(e)
    print("Installing GraphVIPER")

    os.system("pip install graphviper")



    print("GraphVIPER version", version("graphviper"), " installed.")

GraphVIPER version 0.0.39 already installed.

Download dataset

[2]:

import toolviper

toolviper.utils.data.download(file="demo_simulated.im")

[2026-04-06 15:37:46,961]     INFO   toolviper:  Initializing download...

/Users/jsteeb/miniforge3/envs/zinc/lib/python3.13/site-packages/rich/live.py:260: UserWarning: install "ipywidgets"
for Jupyter support
  warnings.warn('install "ipywidgets" for Jupyter support')

Setup Dask Cluster

To simplify things we are going to start of by just using a single thread (everything will run in serial).

[3]:

import dask

dask.config.set(scheduler="synchronous")

[3]:

<dask.config.set at 0x13abed310>

Inspect Image Dataset

[4]:

from xradio.image import load_image, open_image, write_image

img_xds = open_image(store="demo_simulated.im", chunks={"l": 40, "m": 20, "freq": 5})
img_xds

Successful readonly open of default-locked table demo_simulated.im: 1 columns, 1 rows

[4]:

<xarray.Dataset> Size: 20MB
Dimensions:              (time: 1, frequency: 50, polarization: 4, l: 200,
                          m: 100, beam_params_label: 3)
Coordinates:
  * time                 (time) float64 8B 5.154e+04
  * frequency            (frequency) float64 400B 1.415e+09 ... 1.415e+09
    velocity             (frequency) float64 400B 1.146e+06 ... 1.136e+06
  * polarization         (polarization) <U1 16B 'I' 'Q' 'U' 'V'
  * l                    (l) float64 2kB 0.02909 0.0288 ... -0.02851 -0.0288
  * m                    (m) float64 800B -0.01454 -0.01425 ... 0.01396 0.01425
    right_ascension      (l, m) float64 160kB 0.0291 0.0291 ... 6.254 6.254
    declination          (l, m) float64 160kB -0.01454 -0.01425 ... 0.01425
  * beam_params_label    (beam_params_label) <U5 60B 'major' 'minor' 'pa'
Data variables:
    SKY                  (time, frequency, polarization, l, m) float32 16MB dask.array<chunksize=(1, 50, 4, 40, 20), meta=np.ndarray>
    BEAM_FIT_PARAMS_SKY  (time, frequency, polarization, beam_params_label) float64 5kB dask.array<chunksize=(1, 50, 4, 3), meta=np.ndarray>
    FLAG_SKY             (time, frequency, polarization, l, m) bool 4MB dask.array<chunksize=(1, 50, 4, 40, 20), meta=np.ndarray>
Attributes:
    coordinate_system_info:  {'reference_direction': {'attrs': {'frame': 'fk5...
    type:                    image_dataset
    data_groups:             {'base': {'sky': 'SKY', 'beam_fit_params_sky': '...

Create Parallel Coordinates, Map

[5]:

import dask

from toolviper.utils.display import DataDict
from graphviper.graph_tools.coordinate_utils import make_parallel_coord
from graphviper.graph_tools.map import map

from xradio.image import load_image


input_parms = {}
parallel_coords = {}
parallel_coords["frequency"] = make_parallel_coord(coord=img_xds.frequency, n_chunks=6)
sel_parms = {}
input_parms["input_data_store"] = "demo_simulated.im"

input_data = {"img": img_xds}

from graphviper.graph_tools.coordinate_utils import (
    interpolate_data_coords_onto_parallel_coords,
)

node_task_data_mapping = interpolate_data_coords_onto_parallel_coords(
    parallel_coords, input_data
)

DataDict.html(node_task_data_mapping)


def my_func(input_parms):
    if input_parms["input_data"] is None:
        img_xds = load_image(
            input_parms["input_data_store"],
            block_des=input_parms["data_selection"]["img"],
        )
    else:
        img_xds = input_parms["input_data"]["img"]

    toolviper.utils.display.DataDict.html(input_parms)

    print("****")


graph = map(
    input_data=input_data,
    node_task_data_mapping=node_task_data_mapping,
    node_task=my_func,
    input_params=input_parms,
    in_memory_compute=True,
)
dask_graph = dask.compute(graph)

Data Loading Layer

When the map partition size is smaller than the native on-disk zarr chunk size, each map task would otherwise decompress the same chunk independently. The data loading layer inserts shared load nodes that read each chunk once and distribute pre-cut slices to the map tasks:

zarr store  (e.g. 10 channels per compressed chunk)
     │
 [load node]  ← decompresses once, shared across all tasks in the chunk
    │   │
[prepare] [prepare]  ← sub-selects each task's slice in memory
    │         │
[map task] [map task]  ← receives input_params["input_data"], no disk I/O

Detect native on-disk chunk sizes

Open the image with native zarr chunking (chunks={}) so that the coordinate arrays carry dask chunk metadata, then call get_disk_chunk_sizes.

[6]:

from xradio.image import open_image
from graphviper.graph_tools.coordinate_utils import (
    make_parallel_coord,
    interpolate_data_coords_onto_parallel_coords,
    get_disk_chunk_sizes,
)
from graphviper.graph_tools.map import map
from graphviper.graph_tools.generate_dask_workflow import generate_dask_workflow
from graphviper.graph_tools import reduce

# Open with native zarr chunks so get_disk_chunk_sizes can read chunk metadata.
img_xds_native = open_image(store="demo_simulated.im", chunks={})

input_data = {"img": img_xds_native}

# 10 map tasks across 50 frequency channels → 5 channels per task.
parallel_coords = {"frequency": make_parallel_coord(coord=img_xds_native.frequency, n_chunks=10)}

disk_chunk_sizes = get_disk_chunk_sizes(input_data, parallel_coords)
print("Native on-disk chunk sizes:", disk_chunk_sizes)

Successful readonly open of default-locked table demo_simulated.im: 1 columns, 1 rows
Native on-disk chunk sizes: {'frequency': 50}

Define a load task for images

The load task receives load_params with:

"input_data_store" — the zarr store path
"data_selection" — {xds_name: {dim: slice}} at disk-chunk granularity

It must return {xds_name: xarray.Dataset} with the loaded disk chunk.

[7]:

def load_image_chunk(load_params):
    """Load one native zarr chunk of the image."""
    from xradio.image import load_image

    # data_selection["img"] contains the disk-chunk-level slice dict.
    return {
        "img": load_image(
            load_params["input_data_store"],
            block_des=load_params["data_selection"]["img"],
        )
    }

Map task with pre-loaded data

When the load layer is active, input_params["input_data"] is already populated with the sub-selected dataset — no disk read is needed inside the map task.

[8]:

def my_func(input_params):
    # input_params["input_data"] is set by the load layer.
    img = input_params["input_data"]["img"]
    return float(img["SKY"].values.sum())

input_params = {"input_data_store": "demo_simulated.im"}

node_task_data_mapping = interpolate_data_coords_onto_parallel_coords(
    parallel_coords, input_data
)

viper_graph = map(
    input_data=input_data,
    node_task_data_mapping=node_task_data_mapping,
    node_task=my_func,
    input_params=input_params,
    data_loading_task=load_image_chunk,
    disk_chunk_sizes=disk_chunk_sizes,
)

n_load = len(viper_graph["load"]["input_params"])
n_map  = len(viper_graph["map"]["input_params"])
print(f"Load nodes : {n_load}  (one per zarr chunk group)")
print(f"Map  nodes : {n_map}  (one per parallel_coords chunk)")
print(f"Load → map assignments: {viper_graph['map']['load_node_ids']}")

Load nodes : 1  (one per zarr chunk group)
Map  nodes : 10  (one per parallel_coords chunk)
Load → map assignments: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

Visualise the graph

The graph now has three layers: load → prepare → map. Load nodes appear as shared ancestors of the prepare nodes that feed into each map task.

[9]:

viper_graph_reduce = reduce(viper_graph, lambda inputs, _: sum(inputs), {}, mode="tree")
dask_graph = generate_dask_workflow(viper_graph_reduce)
dask.visualize(dask_graph, filename="image_map_graph_with_load_layer")

[9]:

_images/graph_building_tutorial_image_18_0.png

[10]:

result = dask.compute(dask_graph)[0]
print("Total SKY sum across all frequency chunks:", result)

Successful readonly open of default-locked table demo_simulated.im: 1 columns, 1 rows
Total SKY sum across all frequency chunks: 678252.171875

Scheduler plugin (distributed execution)

For multi-worker distributed runs, register ViperGraphPlugin before submitting the graph. It uses the viper_load_group and viper_map_pair annotations added automatically by generate_dask_workflow to:

Process one zarr chunk group at a time → only one decompressed chunk occupies worker memory at any moment.
Schedule reduction-adjacent task pairs together → the binary tree reduction progresses level-by-level with minimal intermediate results in flight.

from toolviper.dask.plugins.scheduler import ViperGraphPlugin
from toolviper.dask.client import local_client

client = local_client(cores=4, memory_limit="8GB")
client.register_scheduler_plugin(ViperGraphPlugin())
# submit graph normally ...