Xarray Linked Data

Xarray Linked Data

Xarray is the natural home for highly linked data: measurements, derived estimates, features, and model outputs sharing coordinates and staying synchronized through selection, computation, and storage.

When to use this skill

• Designing xarray data structures for multi-dimensional, multi-assay data
• Linking data from different experiments through shared coordinates
• Writing custom xarray indexes (when default PandasIndex is insufficient)
• Reading/manipulating xarray to accomplish analysis goals
• Converting scattered-file workflows into unified coordinate-aligned datasets

Core mental model

Experiment Design          xarray Structure
─────────────────          ────────────────
Samples/subjects    →      dimension coordinates (e.g. molecule, sample)
Experimental factors →     dimension coordinates (e.g. target, cell_line, treatment)
Raw measurements    →      data variables (e.g. od450, intensity, uptake_pct_positive)
Derived estimates   →      data variables (e.g. ec50_nm, kd_nm)
Uncertainty         →      data variables along a `draw` dim (e.g. ec50_nm_posterior)
Features            →      data variables along a feature dimension
Splits/assignments  →      boolean mask variables along a split_type dimension
Physical units      →      coordinate attrs or CoordinateTransform indexes

Key principle: every piece of data knows its own coordinates. When you slice, everything stays aligned. No manual index-matching across files.

Quick reference

Design data structures

Read references/data-structure-design.md for full guidance.

# Progressive accumulation: start with raw data, add layers
ds = xr.Dataset(
    {"expression": (["molecule", "treatment", "replicate"], raw_data)},
    coords={
        "molecule": mol_ids,
        "treatment": ["control", "low", "high"],
        "replicate": ["r1", "r2", "r3"],
    },
)
# Add derived estimates — auto-aligned by coordinates
ds = ds.assign({"ec50": (["molecule"], ec50_estimates)})
# Add features along a new dimension
ds = ds.assign({"features": (["molecule", "feature"], feat_matrix)})

When to use Dataset vs DataTree:

• Dataset: all variables share compatible dimensions/coordinates
• DataTree: variables live on different grids (different resolutions, different assay types with incompatible shapes); use tree nodes with coordinate inheritance

Select and manipulate

Read references/selection-and-manipulation.md for full guidance.

# Label-based selection — everything stays aligned
# (.drop_vars() strips the leftover coordinate so the mask aligns cleanly)
train = ds.where(
    ds.train_mask.sel(split_type="random_80_20").drop_vars("split_type"), drop=True
)

# Pointwise indexing with shared dimension
targets = ds.sel(
    molecule=xr.DataArray(["LNP_003", "LNP_007"], dims="query"),
    method="nearest",
)

# Alignment with explicit join type
aligned_a, aligned_b = xr.align(ds_a, ds_b, join="inner")

Build custom indexes

Read references/custom-indexes.md for full guidance, including complete reference implementations.

When default PandasIndex (1-D monotonic labels) is insufficient, build a custom index. Three patterns, escalating complexity:

1. Index subclass + hand-written sel — for custom selection logic (periodic/wrapping coordinates, domain-specific nearest-neighbor)
2. CoordinateTransform subclass — for coordinate mappings with forward/reverse (pixel to physical, lazy z-stacks, nonlinear warps)
3. Meta-index — wraps other indexes for linked multi-coordinate selection (interval cross-slicing, N-D derived coordinates)

# Register a custom index
ds = ds.drop_indexes("angle").set_xindex("angle", PeriodicIndex, period=360.0)

Link across experiments

Read references/cross-experiment-linking.md for full guidance.

# DataTree for multi-assay hierarchy: root defines the shared `molecule`
# coordinate, child nodes inherit it (a "deep" tree -- see nb5)
tree = xr.DataTree.from_dict({
    "/": xr.Dataset(coords={"molecule": mol_ids}),
    "/characterization/dls": dls_ds,
    "/characterization/hplc": hplc_ds,
    "/bioassays/binding": binding_ds,
    "/bioassays/flow_cytometry": flow_ds,
})
# One selection propagates to every node
tree.sel(molecule="LNP_007")

Reference notebooks

Self-verifying marimo notebooks with synthetic data covering diverse assay domains (cargo delivery, binding assays, cell-based assays, genomics, functional readouts, analytical chemistry). Run each with uv run:

API instability note

Custom indexes are an experimental xarray feature (since v2022.06.0). The API may change without deprecation notice. Pin xarray versions in notebooks and check the official docs for current API.

Key references

• xarray custom index how-to
• Ian Hunt-Isaak's linked-indexes -- DimensionInterval, NDIndex
• xarray-tutorial: why-an-index -- PeriodicIndex, CoordinateTransform examples
• xarray-indexes gallery -- NDPointIndex (KD-tree/Ball-tree)