Processing Large DEMs with Dask and rioxarray

rioxarray opens a raster as a lazily-chunked xarray object, and Dask executes operations on those chunks in parallel. For the local and neighbourhood stages of a hydrologic pipeline that combination does real work with very little code. For the accumulation stage it produces confidently wrong answers. This guide covers both halves — as part of the tiled and parallel watershed processing topic within watershed delineation and catchment synchronization.

Prerequisites

• A Cloud-Optimized GeoTIFF, or at least an internally tiled GeoTIFF. A stripe-oriented file defeats chunked reading entirely.
• rioxarray, dask, distributed, and richdem or whitebox for the hydrology.
• A stated per-worker memory budget. Everything below depends on it.

Core Technique: Chunks That Match the File

The single decision that determines whether a chunked pipeline is fast is whether the chunks align with the file’s internal blocks.

chunks="auto" in rioxarray.open_rasterio reads the internal blocking from the file and aligns to it, which is almost always what you want. An explicit chunk size chosen for tidiness — 1000, say — is the most common cause of a chunked pipeline that runs slower than a serial one.

Annotated Code Example

import logging

import dask.array as da
import numpy as np
import richdem as rd
import rioxarray
from dask.distributed import Client, LocalCluster

log = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")


def condition_dem_chunked(
    dem_path: str,
    out_path: str,
    halo: int = 128,
    n_workers: int = 4,
    memory_limit: str = "4GB",
) -> dict:
    """
    Fill depressions and compute slope across a larger-than-memory DEM.

    Both operations have a bounded neighbourhood, so map_overlap expresses
    them correctly. Flow accumulation deliberately does NOT appear here.
    """
    cluster = LocalCluster(n_workers=n_workers, threads_per_worker=1,
                           memory_limit=memory_limit)
    client = Client(cluster)
    log.info("Dask cluster: %d workers × %s", n_workers, memory_limit)

    try:
        # --- chunks="auto" aligns to the file's internal blocking. An explicit
        # size that straddles blocks makes every read decompress data it
        # immediately discards. ---
        dem = rioxarray.open_rasterio(dem_path, chunks="auto", masked=True).squeeze()
        arr = dem.data
        log.info("Opened %s: shape %s, chunks %s, dtype %s",
                 dem_path, arr.shape, arr.chunksize, arr.dtype)

        nbytes_per_chunk = np.prod(arr.chunksize) * arr.dtype.itemsize
        log.info("Chunk footprint %.1f MB before intermediates; peak per worker "
                 "will be roughly 5× that", nbytes_per_chunk / 1e6)

        nodata = float(dem.rio.nodata) if dem.rio.nodata is not None else -9999.0

        def _fill(block: np.ndarray) -> np.ndarray:
            """Fill depressions on one chunk plus its halo."""
            a = np.nan_to_num(block, nan=nodata).astype(np.float32)
            grid = rd.rdarray(a, no_data=nodata)
            rd.FillDepressions(grid, epsilon=True, in_place=True)
            return np.asarray(grid, dtype=np.float32)

        # --- The halo must exceed the widest depression that straddles a chunk
        # boundary. Measure it on a sample tile rather than guessing: too narrow
        # and the two sides of a boundary depression fill to different levels. ---
        filled = da.map_overlap(
            _fill, arr.astype(np.float32),
            depth=halo, boundary="nearest", dtype=np.float32,
        )

        # --- Slope is a 3×3 operation, so a halo of 1 would do; reusing the
        # same halo keeps the graph simple and costs one extra ring of cells. ---
        def _slope(block: np.ndarray) -> np.ndarray:
            grid = rd.rdarray(block.astype(np.float32), no_data=nodata)
            return np.asarray(rd.TerrainAttribute(grid, attrib="slope_degrees"),
                              dtype=np.float32)

        slope = da.map_overlap(
            _slope, filled, depth=halo, boundary="nearest", dtype=np.float32,
        )

        out = dem.copy(data=filled)
        out.rio.write_nodata(nodata, inplace=True)
        out.rio.to_raster(out_path, tiled=True, blockxsize=512, blockysize=512,
                          compress="LZW", BIGTIFF="IF_SAFER")
        log.info("Wrote conditioned DEM: %s", out_path)

        # --- Reductions are safe on the Dask graph: they need no neighbourhood. ---
        mean_slope = float(slope.mean().compute())
        log.info("Mean slope over the domain: %.2f degrees", mean_slope)

        return {"output": out_path, "chunks": str(arr.chunksize),
                "mean_slope_deg": mean_slope}
    finally:
        client.close()
        cluster.close()


# --- Example usage ---
# condition_dem_chunked("state_1m_cog.tif", "state_1m_filled.tif",
#                       halo=128, n_workers=6, memory_limit="6GB")

Parameter Reference

Worked Example: What Belongs in the Graph

Building the pipeline as one Dask graph is tempting and partly wrong. The table below is the division that works.

The two “no” rows are the whole reason this page has a companion topic page. They need the ordered edge-exchange pass described in tiled and parallel watershed processing, which cannot be expressed as a chunk-local function.

Peak memory is the number that decides how many workers fit, and it is several times the chunk size. Counting the arrays that exist simultaneously is the only reliable way to size a cluster.

Watching the dashboard rather than the log

Dask’s diagnostic dashboard is the fastest way to see what a chunked hydrology pipeline is actually doing, and the two panels worth watching are the task stream and the memory plot. A healthy conditioning run shows solid blocks of compute with brief transfer gaps; a run whose chunks are misaligned with the file shows long, thin read tasks dominating the stream, which is the visual signature of the block-alignment problem.

The memory plot answers a different question. Workers that climb steadily toward their limit and then spill to disk are being asked to hold too much, and the fix is a smaller chunk rather than a bigger machine — spilling turns a compute-bound stage into a disk-bound one and can make a run several times slower than the serial equivalent.

Keeping the graph small enough to schedule

Dask’s scheduler holds the whole task graph in memory on the client, and a graph with hundreds of thousands of tasks costs real time to build and traverse before any compute starts. A pipeline that chunks a large raster into very small pieces and then applies a dozen operations to each can spend minutes constructing a graph for a job that runs in seconds.

The practical rule is to keep the task count in the thousands rather than the hundreds of thousands, which usually means larger chunks rather than fewer operations. Where the chunk size is fixed by memory, fusing operations — doing the fill and the direction computation inside one map_overlap callable rather than two — halves the graph without changing the result, and it also avoids materialising the intermediate array.

Gotchas and Edge Cases

• map_overlap for accumulation. Every chunk’s accumulation restarts near zero at its upstream edge. The numbers are plausible chunk by chunk, which is what makes this the most dangerous mistake in the whole area.
• Chunk size unaligned to blocks. Halves or thirds the read throughput on a job that is usually I/O bound to begin with.
• boundary="constant" with nodata. Surrounds each chunk with a nodata ring that the fill algorithm treats as terrain, producing an artificial rim. "nearest" avoids it.
• Memory sized from the file. The peak is roughly five times the chunk size per worker, not one. A cluster sized from the input size will spill to disk and then die.
• threads_per_worker above 1. The compiled hydrology libraries do not consistently release the GIL, so threads contend rather than parallelise.
• Writing untiled output. The next stage in the pipeline then cannot read it chunk-wise, and the problem compounds along the chain. Always write tiled, compressed output.
• Computing inside a loop. Calling .compute() per chunk in Python defeats the scheduler entirely. Build the graph, then compute once.

Related Topics

• Tiled and Parallel Watershed Processing — the parent topic, including the ordered exchange that accumulation needs
• Writing Cloud-Optimized GeoTIFFs for Hydrology Pipelines — producing the internally tiled files this depends on
• Choosing Tile Size and Halo Width for Flow Accumulation — measuring the halo instead of guessing it
• Pipeline Orchestration: DEM to Watershed — the scheduler layer above this one