Computing NDVI with xarray on a Dask Cluster

Stack the Sentinel-2 scenes into one lazy (time, band, y, x) cube, express NDVI as ordinary xarray arithmetic, reduce over time to a composite, and trigger the whole thing with a single terminal call so the Dask cluster streams the work chunk-by-chunk and worker memory never holds the full stack:

import stackstac, xarray as xr, numpy as np
from dask.distributed import Client

client = Client()                                   # LocalCluster or distributed
cube = stackstac.stack(items, assets=["B04", "B08"], chunksize=2048)
red, nir = cube.sel(band="B04"), cube.sel(band="B08")
ndvi = xr.where(nir + red != 0, (nir - red) / (nir + red), np.nan)
composite = ndvi.median(dim="time")                 # still lazy
composite.rio.to_raster("ndvi_median.tif")          # one call runs the graph

This is the end-to-end companion to Tuning Dask Chunk Sizes for Raster Cubes: that page decides the chunk shape, this page runs a real multi-scene NDVI composite on top of it.

Why This Arises When Scaling Raster Work

Single-scene NDVI is a solved problem — divide two bands, guard the denominator, write a GeoTIFF. The difficulty appears the moment you want a cloud-free seasonal composite over dozens or hundreds of Sentinel-2 acquisitions covering the same tile. Loaded eagerly, that stack is tens of gigabytes; looped scene-by-scene, it wastes every core but one and leaves you stitching intermediate composites by hand. The distributed, lazy approach solves both: stackstac presents every scene as one array, xarray expresses NDVI once across the whole cube, and Dask’s scheduler streams chunks through the workers so peak memory tracks one chunk per thread rather than the whole series.

This differs from the sibling chunk-tuning page in intent. There, the question is what shape should a chunk be; here, the question is how does a full NDVI-to-composite job flow through the Dask cluster from STAC query to written raster. For cluster startup, the persist versus compute distinction, and dashboard monitoring, see the parent guide, Scaling Raster Processing with Dask, within Cloud Execution & Orchestration. The NDVI formula itself and its numerical guards come from Calculating NDVI directly from xarray DataArrays.

Dataflow: STAC query to written composite

The diagram traces the lazy pipeline. Every stage up to the terminal write only extends the graph; the single write is what dispatches reads to workers.

Lazy NDVI composite dataflow on a Dask cluster A STAC search yields item URLs; stackstac builds a lazy time-band-y-x cube; xarray computes NDVI lazily; a median over time yields a lazy composite; a single to_zarr or to_raster call triggers the Dask scheduler to stream chunk reads across workers and write the result. STAC search pystac-client → item list stackstac lazy cube (t,band,y,x) lazy NDVI (NIR−Red)/ (NIR+Red) median over time → composite write to_zarr / to_raster all lazy — graph only, no pixels read executes here The terminal write is the only line that dispatches chunk reads to the workers. Peak memory ≈ one chunk per worker thread, not the whole stack.

Environment & Setup

Package Minimum version Why required
dask[distributed] 2024.1 Scheduler, Client, worker processes
xarray 2024.1 Lazy cube and the NDVI arithmetic
stackstac 0.5 Turns a STAC item list into one lazy (time, band, y, x) cube
pystac-client 0.7 Runs the STAC search that yields the item list
rioxarray 0.15 .rio accessor for the georeferenced GeoTIFF write
zarr 2.16 Chunk-aligned out-of-core write for the full cube 
pip install "dask[distributed]>=2024.1" "xarray>=2024.1" "stackstac>=0.5" \
            "pystac-client>=0.7" "rioxarray>=0.15" "zarr>=2.16"

Complete Working Example

The script queries a STAC catalog for a season of Sentinel-2 L2A scenes over one tile, stacks them lazily, computes NDVI across the whole stack, reduces to a median composite, and writes it — triggering the graph exactly once. Item retrieval follows Using pystac-client to Filter Sentinel-2 Imagery by Date:

import numpy as np
import stackstac
import xarray as xr
from pystac_client import Client as StacClient
from dask.distributed import LocalCluster, Client


def ndvi_median_composite(
    bbox: tuple[float, float, float, float],
    date_range: str,
    out_path: str = "ndvi_median.tif",
    max_cloud: int = 20,
) -> str:
    """
    Build a cloud-filtered NDVI median composite over a Sentinel-2 stack,
    computed lazily on a Dask cluster and written with a single call.
    """
    # 1. Start a local cluster; swap this line for Client("tcp://…") to go distributed
    cluster = LocalCluster(n_workers=4, threads_per_worker=2, memory_limit="4GB")
    client = Client(cluster)
    print("dashboard:", client.dashboard_link)

    # 2. STAC search → item list (URLs only, no pixels yet)
    catalog = StacClient.open("https://earth-search.aws.element84.com/v1")
    search = catalog.search(
        collections=["sentinel-2-l2a"],
        bbox=bbox,
        datetime=date_range,
        query={"eo:cloud_cover": {"lt": max_cloud}},
    )
    items = list(search.items())
    if not items:
        raise ValueError("No scenes matched the search — widen the bbox or dates.")
    print(f"{len(items)} scenes matched")

    # 3. Stack every scene into ONE lazy (time, band, y, x) cube.
    #    chunksize sets the spatial chunk edge; band=-1 is implicit (assets kept whole).
    cube = stackstac.stack(
        items,
        assets=["red", "nir"],          # Sentinel-2 B04 and B08 at 10 m
        chunksize=2048,                 # ~ see the chunk-tuning guide for the target
        bounds_latlon=bbox,
        dtype="float32",
        fill_value=np.nan,              # missing pixels → NaN, propagate cleanly
    )

    # 4. Lazy NDVI across the whole stack — still no pixels read
    red = cube.sel(band="red")
    nir = cube.sel(band="nir")
    denom = nir + red
    ndvi = xr.where(denom != 0, (nir - red) / denom, np.nan).rename("ndvi")

    # 5. Median over time collapses clouds/gaps into one clean composite (lazy)
    composite = ndvi.median(dim="time", skipna=True)

    # 6. ONE terminal call runs the entire graph and streams chunks to disk.
    #    to_raster reduces to 2-D first, so the client never holds the full cube.
    composite.rio.write_crs(cube.rio.crs).rio.to_raster(
        out_path, driver="GTiff", tiled=True, compress="zstd",
    )
    client.close()
    cluster.close()
    return out_path


if __name__ == "__main__":
    path = ndvi_median_composite(
        bbox=(11.20, 46.05, 11.45, 46.25),   # a tile in the Italian Alps
        date_range="2024-06-01/2024-08-31",  # one growing season
    )
    print("wrote", path)

Every line before step 6 is lazy. The .to_raster call is the only point at which the scheduler dispatches chunk reads to the workers, so a hundred-scene stack is processed with peak memory of roughly one chunk per worker thread. Watch the Task Stream on the dashboard while it runs to confirm workers stay well under their memory_limit.

Variant Patterns

1. Writing the full time cube to Zarr instead of a composite

When you need every date’s NDVI preserved rather than a single composite — for phenology or temporal aggregation — write the lazy cube straight to Zarr, chunk-aligned, so no scene is ever held in the client:

ndvi.to_dataset(name="ndvi").chunk({"time": 1, "y": 2048, "x": 2048}).to_zarr(
    "ndvi_timeseries.zarr", mode="w", consolidated=True,
)

2. Persisting the stack when computing several outputs

If you want a median and a 90th-percentile composite from the same NDVI stack, persist it once so the expensive scene reads happen a single time, then branch:

from dask.distributed import wait

ndvi = ndvi.persist()          # pin the lazy stack in cluster RAM
wait(ndvi)                     # block until materialised on workers
median = ndvi.median(dim="time").compute()
p90 = ndvi.quantile(0.9, dim="time").compute()   # reuses persisted reads

3. Extending to other indices in one graph

Because the stack is already open, adding EVI or NDWI is just more lazy arithmetic on the same cube — the reads are shared across all indices in the final graph. The formulas match Spectral Index Calculation Pipelines:

green = cube.sel(band="green")
ndwi = xr.where(green + nir != 0, (green - nir) / (green + nir), np.nan)
both = xr.Dataset({"ndvi": ndvi.median("time"), "ndwi": ndwi.median("time")})
both.to_zarr("indices.zarr", mode="w")   # one graph reads each scene once

Common Errors

MemoryError / KilledWorker during .to_raster The composite reduction is pulling too much into one task, usually because the spatial chunks are too large. Reduce chunksize toward the 256–512 MB target from the chunk-tuning guide and rerun.

Composite is entirely NaN fill_value was left at 0 so cloud gaps divided to 0/0, or skipna was False in the median. Set fill_value=np.nan in stackstac.stack and skipna=True in .median so missing pixels are ignored rather than poisoning the result.

ValueError: assets ... not found in item The asset keys differ across STAC catalogs — Earth Search uses red/nir, other catalogs use B04/B08. Print items[0].assets.keys() and pass the exact keys your catalog exposes.