Dask

Every xrtoolz operator accepts dask-backed (chunked) xarray inputs, not just eager numpy ones. Where an operator's kernel is naturally per-chunk it stays lazy — the work is deferred into the dask graph and only runs on .compute(). Where the algorithm needs the whole array (a global neighbour search, a global fit), the operator still accepts dask input but materialises internally, with an identical result.

The one rule: keep core dims in a single chunk

An operator transforms along its core dimensions — the spatial (lat, lon) for a 2-D gap-filler, the transform dim for an FFT or wavelet, the smoothing dim for a filter. Chunk the batch dims (e.g. time) and keep each core dim whole:

ds = ds.chunk({"time": 1, "lat": -1, "lon": -1})   # -1 = one chunk

Operators that need a whole core slice are built with allow_rechunk=False, so a split core dim raises a clear error rather than silently rechunking (which would corrupt a windowed transform).

A lazy pipeline

The composition layer is dask-transparent — chunk once at the source and the whole Sequential / Graph stays lazy until you ask for a result:

import xarray as xr
from pipekit import Sequential
from xrops.anomaly import RemoveMean
from xrops.ocn.operators import Streamfunction, GeostrophicVelocities

ds = xr.open_mfdataset("ssh_*.nc").chunk({"time": 1, "lat": -1, "lon": -1})

pipeline = Sequential(
    RemoveMean("time"),
    Streamfunction(),
    GeostrophicVelocities(),
)

lazy = pipeline(ds)        # nothing computed yet — a dask graph
result = lazy.compute()    # run it (optionally under a dask scheduler/cluster)

Declared backends

Every xrops operator declares backends from pipekit's vocabulary (OP-O6); today all declare ("numpy", "dask"). A DataArray counts by the backend of its data and a Dataset by its variables', so a JAX-backed input raises a BackendError naming the conversion, while a NumPy operand beside a dask one (a static mask, an in-memory reference) joins the lazy computation. geoops geometry operators, which need a GeoArray per item, run over a cube through xrops.lift.

Per-operator tiers

Every operator falls into one of three tiers. The tests/test_dask_compat.py sweep asserts numpy/dask result parity for all of them, and asserts laziness for the lazy tiers.

Tier	Behaviour	Operators
Lazy	stays a dask graph	xrops.grad (partial, gradient, divergence, curl, laplacian, Gradient, Difference; stencils run per chunk with a halo, so the result equals the in-memory one) and the ocn diagnostics on it (KE, streamfunction, velocity magnitude, RelativeVorticity, Divergence, GeostrophicVelocities, OkuboWeiss); geo (RemoveMean, climatology, validators, masks); transforms.PowerSpectrum; Coarsen; RegridLike; all pixel/spectral metrics; the gap-fillers (FillNaN{Laplacian,Spatial,RBF,Biharmonic,IDW}); the smoothers (GaussianSmooth, MovingAverage, fir_filter, lowpass_filter); morphology
Eager (dask-safe)	accepts dask, computes internally, identical result	binning (Bin2D, Histogram2D); the sklearn bridge (XarrayEstimator, decompositions, SklearnModelOp)
N/A	numpy-array primitives — no xarray, so no dask	points_to_grid, idw_to_grid, kde_to_grid (use the *ToGrid operators for the xarray surface)

Why some operators can't be fully lazy

Bin2D, the IDW/KNN regridders, RBF infilling, and the scikit-learn estimators rest on a global computation — a KD-tree over every point, or a single fit across the whole sample. Those have no correct per-chunk decomposition, so the operator gathers what it needs and computes it. They remain perfectly usable on dask-backed data; they just aren't a deferred graph. Making them distributed would mean reimplementing the underlying KD-tree / least-squares solvers, which is out of scope.

Testing your own pipelines

The array_backend / maybe_chunk fixtures (in tests/conftest.py) let a single test body run against both backends:

def test_my_pipeline(array_backend, maybe_chunk):
    ds = maybe_chunk(build_dataset(), array_backend, {"time": 1})
    out = my_pipeline(ds)
    if array_backend == "dask":
        assert out["u"].chunks is not None     # stayed lazy
        out = out.compute()
    xr.testing.assert_allclose(out, expected)