Chunked Raster Processing with xarray and dask for Forest Time Series

A lazy stack is only lazy in the useful sense if its chunks line up with how the data is stored and how you intend to reduce it. Get the chunk shape wrong and every operation triggers a rechunk — a full shuffle of the array through memory — which is reliably slower than downloading the scenes would have been. This guide is about choosing that shape and keeping memory bounded, as the execution half of Cloud-Native Raster Workflows for Forestry, inside Ecological GIS Data Foundations in Python.

When to use this approach

The honest threshold is memory. If the clipped, masked stack fits in a third of the machine’s RAM, .compute() once and work eagerly — the code is simpler and every subsequent operation is faster.

Chunk shape follows the reduction

A chunk is the unit of work. Two rules decide the shape, and they can conflict:

1. Spatial chunks should be a multiple of the file’s internal tile size — 512 or 1024 for most COGs — so a chunk maps onto whole tiles and no tile is fetched twice.
2. The reduction dimension should not be split. A median over time wants every date for a given pixel block in one chunk; splitting time means dask must combine partial results across chunks, which for a median means holding them all anyway.

For a forest time series the resolution is usually: chunk the whole time dimension together, and chunk space at 1024 × 1024.

$$\text{chunk bytes}=n_t\times n_b\times h\times w\times \text{itemsize}$$

With 60 dates, 3 bands, 1024 × 1024 pixels and float32, that is 60 × 3 × 1024 × 1024 × 4 ≈ 750 MB per chunk — far too large. The arithmetic is the point: chunk sizing is a budget, and something has to give. Either drop to 512 × 512 spatial chunks (≈ 190 MB) or split time into groups and reduce in two stages.

Minimal reproducible example

Build the stack with an explicit chunk shape, clip immediately, and only then reduce.

import pystac_client
import stackstac
import geopandas as gpd
from shapely.ops import unary_union

CATALOG = "https://earth-search.aws.element84.com/v1"


def seasonal_composite(stands_path: str, start: str, end: str,
                       epsg: int = 32633, res: int = 20, chunk: int = 512):
    """Median NDVI over a growing season, computed in bounded memory."""
    stands = gpd.read_file(stands_path).to_crs(epsg)
    aoi_ll = gpd.GeoSeries([unary_union(stands.geometry.values)], crs=epsg).to_crs(4326)
    bbox = tuple(aoi_ll.total_bounds)

    items = list(pystac_client.Client.open(CATALOG).search(
        collections=["sentinel-2-l2a"], bbox=bbox, datetime=f"{start}/{end}",
        query={"eo:cloud_cover": {"lt": 60}}).items())

    stack = stackstac.stack(
        items,
        assets=["red", "nir", "scl"],
        epsg=epsg,
        resolution=res,
        bounds=tuple(stands.total_bounds),   # clip at construction, not later
        chunksize=(-1, 1, chunk, chunk),     # -1 = the whole time axis in one chunk
        dtype="float32",
        fill_value=float("nan"),
    )

    clear = stack.sel(band="scl").isin([4, 5, 6, 7])
    red = stack.sel(band="red").where(clear)
    nir = stack.sel(band="nir").where(clear)
    ndvi = (nir - red) / (nir + red)
    return ndvi.median(dim="time", skipna=True)

bounds= at construction is the highest-leverage line in that function. Clipping after the stack is built means dask plans work over the full scene footprint and discards most of it — the graph is enormous, the scheduler spends longer planning than computing, and memory spikes on chunks you never wanted.

Parameter reference

Keep memory bounded

Two habits keep a long time series inside a laptop’s memory. First, reduce before you widen: apply the cloud mask and compute the index before any operation that increases the array’s footprint. Second, persist the small intermediate, not the large one.

import dask


def composite_by_year(ndvi, years: list[int]):
    """Reduce year by year so peak memory is one year, not the whole archive."""
    out = {}
    for year in years:
        sel = ndvi.sel(time=str(year))
        # Compute this year fully, then release its graph before the next.
        out[year] = sel.median(dim="time", skipna=True).compute()
    return out


with dask.config.set(scheduler="threads", num_workers=8):
    yearly = composite_by_year(ndvi, list(range(2018, 2027)))

Looping deliberately, rather than expressing the whole decade as one graph, trades a little parallelism for a memory ceiling you can predict. On a machine that is not a cluster, that trade is nearly always right.

For an operation dask has no built-in for — a per-pixel regression slope, say — map_blocks applies a function to each chunk with an explicit output template:

import numpy as np
import xarray as xr


def _slope_block(block: np.ndarray) -> np.ndarray:
    """Per-pixel linear trend over the leading (time) axis of one chunk."""
    t = np.arange(block.shape[0], dtype="float32")
    t = t - t.mean()
    y = block - np.nanmean(block, axis=0, keepdims=True)
    denom = float((t ** 2).sum())
    return np.nansum(t[:, None, None] * y, axis=0) / denom


def ndvi_trend(ndvi: xr.DataArray) -> xr.DataArray:
    """NDVI change per acquisition, computed chunk by chunk."""
    template = ndvi.isel(time=0, drop=True)
    return xr.map_blocks(
        lambda b: xr.DataArray(_slope_block(b.values), coords=template.coords,
                               dims=template.dims),
        ndvi,
        template=template,
    )

The template argument is what keeps this lazy. Without it, dask computes one chunk eagerly just to learn the output shape and dtype, which on a remote stack means a round trip before the graph is even built.

Expected output and verification

Inspect the graph before running it. Three numbers tell you whether the plan is sane:

def stack_report(arr) -> dict:
    """Chunk geometry and graph size — check these before calling compute()."""
    chunk_elems = 1
    for dim_chunks in arr.chunks:
        chunk_elems *= max(dim_chunks)
    return {
        "shape": tuple(arr.shape),
        "chunks": tuple(tuple(sorted(set(c))) for c in arr.chunks),
        "n_tasks": len(arr.data.dask),
        "max_chunk_mb": round(chunk_elems * arr.dtype.itemsize / 1e6, 1),
    }


report = stack_report(ndvi)
assert report["max_chunk_mb"] < 300, f"chunks too large: {report}"
assert report["n_tasks"] < 500_000, f"graph too big — clip or group first: {report}"
print(report)

A task count in the millions means the graph itself will exhaust memory before any data is read; the cause is nearly always a stack built over the full scene footprint instead of the study area. A max_chunk_mb above a few hundred means one worker will hold that much at a time, and with eight threads that is the whole machine.

Common pitfalls

• Clipping after building the stack. The graph is planned over the full footprint, so planning alone can take longer than the computation.
• Chunking finely in time and then taking a median. Every temporal reduction becomes a rechunk, and the shuffle moves the whole array through memory.
• Spatial chunks that are not a multiple of the internal tile size. Chunk boundaries then cut tiles in half and each tile is fetched by two chunks.
• Using the process scheduler for I/O-bound work. Serialising array chunks between processes costs more than the concurrency buys when the bottleneck is the network.
• Calling .compute() inside a loop over pixels or dates without releasing. Keep the loop, but let each iteration’s result be small and let its graph go out of scope.

Frequently Asked Questions

Related

• Cloud-Native Raster Workflows for Forestry — the wider pattern this executes
• Searching STAC Catalogs for Sentinel-2 with pystac-client — produces the item list a stack is built from
• Reading Cloud Optimized GeoTIFFs with rasterio Windows — the per-chunk read underneath
• Masking Clouds with the Sentinel-2 Scene Classification Layer — the per-pixel masking step these stacks apply

Up: Cloud-Native Raster Workflows · Ecological GIS Data Foundations in Python