Distributed (dask) out-of-core backend¶

Setting workers > 1 (or passing a distributed.Client) runs the embarrassingly-parallel per-tile passes of the tiled fill and accumulation through dask.delayed, on top of pyramids' dask plumbing (pyramids.configure(client=...)). The edge reduction stays serial and the producer writes the finished tiles — so the result is identical to the single-process path.

The dask backend reopens the source by path, so the input must be file-backed (a GeoTIFF), not an in-memory dataset.

In [ ]:

import matplotlib.pyplot as plt
# keep the executed notebook small
plt.rcParams["figure.dpi"] = 80
plt.rcParams["savefig.dpi"] = 80

import matplotlib.pyplot as plt # keep the executed notebook small plt.rcParams["figure.dpi"] = 80 plt.rcParams["savefig.dpi"] = 80

Setup (file-backed DEM)¶

In [ ]:

import gzip, os, shutil, tempfile, urllib.request
from pathlib import Path

import numpy as np
from pyramids.dataset import Dataset
from digitalrivers import DEM


def _data_dir() -> Path:
    """Cache downloads under <project>/examples/data/ (project root marked by pyproject.toml)."""
    cwd = Path.cwd().resolve()
    for parent in (cwd, *cwd.parents):
        if (parent / "pyproject.toml").exists():
            d = parent / "examples" / "data"
            d.mkdir(parents=True, exist_ok=True)
            return d
    d = cwd / "data"
    d.mkdir(exist_ok=True)
    return d


# Download a real 1-arc-second SRTM tile (AWS open data, no auth), cached locally.
DATA = _data_dir()
hgt, gz = DATA / "N45W122.hgt", DATA / "N45W122.hgt.gz"
URL = "https://elevation-tiles-prod.s3.amazonaws.com/skadi/N45/N45W122.hgt.gz"
if not hgt.exists():
    if not gz.exists():
        print(f"Downloading {URL} ...")
        urllib.request.urlretrieve(URL, gz)
    with gzip.open(gz, "rb") as fi, open(hgt, "wb") as fo:
        shutil.copyfileobj(fi, fo)

# Crop a ~0.1 degree (360x360) window around Mount Hood and save it as a file-backed GeoTIFF DEM.
raw = np.fromfile(hgt, dtype=">i2").reshape(3601, 3601).astype(np.float32)
crop = raw[940:1300, 1080:1440].copy()
WORK = Path(tempfile.mkdtemp(prefix="dr_ooc_"))
DEM_PATH = WORK / "mt_hood_dem.tif"
Dataset.create_from_array(
    crop, top_left_corner=(-121.9, 45.5), cell_size=1 / 3600, epsg=4326,
    no_data_value=-32768.0, driver_type="GTiff", path=str(DEM_PATH),
).close()
dem = DEM(Dataset.read_file(str(DEM_PATH)).raster)
print(f"DEM: {dem.rows} x {dem.columns} cells, elevation {crop.min():.0f}-{crop.max():.0f} m")

import gzip, os, shutil, tempfile, urllib.request from pathlib import Path import numpy as np from pyramids.dataset import Dataset from digitalrivers import DEM def _data_dir() -> Path: """Cache downloads under <project>/examples/data/ (project root marked by pyproject.toml).""" cwd = Path.cwd().resolve() for parent in (cwd, *cwd.parents): if (parent / "pyproject.toml").exists(): d = parent / "examples" / "data" d.mkdir(parents=True, exist_ok=True) return d d = cwd / "data" d.mkdir(exist_ok=True) return d # Download a real 1-arc-second SRTM tile (AWS open data, no auth), cached locally. DATA = _data_dir() hgt, gz = DATA / "N45W122.hgt", DATA / "N45W122.hgt.gz" URL = "https://elevation-tiles-prod.s3.amazonaws.com/skadi/N45/N45W122.hgt.gz" if not hgt.exists(): if not gz.exists(): print(f"Downloading {URL} ...") urllib.request.urlretrieve(URL, gz) with gzip.open(gz, "rb") as fi, open(hgt, "wb") as fo: shutil.copyfileobj(fi, fo) # Crop a ~0.1 degree (360x360) window around Mount Hood and save it as a file-backed GeoTIFF DEM. raw = np.fromfile(hgt, dtype=">i2").reshape(3601, 3601).astype(np.float32) crop = raw[940:1300, 1080:1440].copy() WORK = Path(tempfile.mkdtemp(prefix="dr_ooc_")) DEM_PATH = WORK / "mt_hood_dem.tif" Dataset.create_from_array( crop, top_left_corner=(-121.9, 45.5), cell_size=1 / 3600, epsg=4326, no_data_value=-32768.0, driver_type="GTiff", path=str(DEM_PATH), ).close() dem = DEM(Dataset.read_file(str(DEM_PATH)).raster) print(f"DEM: {dem.rows} x {dem.columns} cells, elevation {crop.min():.0f}-{crop.max():.0f} m")

Dask tiled fill == serial tiled fill¶

workers=4 uses dask's threaded scheduler by default; pass scheduler="synchronous" or a real distributed.Client via client= for a cluster.

In [ ]:

serial = dem.fill_depressions(
    engine="tiled", out_path=str(WORK / "fill_serial.tif"), tile_size=96
)
dask_filled = dem.fill_depressions(
    engine="tiled", out_path=str(WORK / "fill_dask.tif"), tile_size=96, workers=4
)
print("dask fill == serial fill:", np.array_equal(
    np.asarray(serial.read_array()), np.asarray(dask_filled.read_array())
))

serial = dem.fill_depressions( engine="tiled", out_path=str(WORK / "fill_serial.tif"), tile_size=96 ) dask_filled = dem.fill_depressions( engine="tiled", out_path=str(WORK / "fill_dask.tif"), tile_size=96, workers=4 ) print("dask fill == serial fill:", np.array_equal( np.asarray(serial.read_array()), np.asarray(dask_filled.read_array()) ))

Dask tiled accumulation¶

Flow direction must also be file-backed for the workers, so we write it to a GeoTIFF and reopen it as a typed FlowDirection.

In [ ]:

from digitalrivers import FlowDirection

filled = dem.fill_depressions(engine="in_memory")
fd_mem = filled.flow_direction(method="d8")
fd_path = WORK / "fdir.tif"
Dataset.create_from_array(
    np.asarray(fd_mem.read_array()), top_left_corner=(-121.9, 45.5), cell_size=1 / 3600,
    epsg=4326, no_data_value=fd_mem.no_data_value[0], driver_type="GTiff", path=str(fd_path),
).close()
fdir = FlowDirection.open(str(fd_path), routing="d8")

acc_serial = fdir.accumulate(engine="tiled", out_path=str(WORK / "acc_serial.tif"), tile_size=96)
acc_dask = fdir.accumulate(
    engine="tiled", out_path=str(WORK / "acc_dask.tif"), tile_size=96, workers=4
)
print("dask accumulation == serial:", np.allclose(
    np.asarray(acc_serial.read_array()).astype(float),
    np.asarray(acc_dask.read_array()).astype(float),
))

from digitalrivers import FlowDirection filled = dem.fill_depressions(engine="in_memory") fd_mem = filled.flow_direction(method="d8") fd_path = WORK / "fdir.tif" Dataset.create_from_array( np.asarray(fd_mem.read_array()), top_left_corner=(-121.9, 45.5), cell_size=1 / 3600, epsg=4326, no_data_value=fd_mem.no_data_value[0], driver_type="GTiff", path=str(fd_path), ).close() fdir = FlowDirection.open(str(fd_path), routing="d8") acc_serial = fdir.accumulate(engine="tiled", out_path=str(WORK / "acc_serial.tif"), tile_size=96) acc_dask = fdir.accumulate( engine="tiled", out_path=str(WORK / "acc_dask.tif"), tile_size=96, workers=4 ) print("dask accumulation == serial:", np.allclose( np.asarray(acc_serial.read_array()).astype(float), np.asarray(acc_dask.read_array()).astype(float), ))

A MEM source is rejected up front¶

Because workers reopen by path, an in-memory dataset raises a clear error rather than failing on a worker.

In [ ]:

mem_dem = DEM(Dataset.create_from_array(
    crop, top_left_corner=(-121.9, 45.5), cell_size=1 / 3600, epsg=4326, no_data_value=-32768.0
).raster)
try:
    mem_dem.fill_depressions(engine="tiled", out_path=str(WORK / "x.tif"), workers=4)
except ValueError as exc:
    print("rejected:", exc)

mem_dem = DEM(Dataset.create_from_array( crop, top_left_corner=(-121.9, 45.5), cell_size=1 / 3600, epsg=4326, no_data_value=-32768.0 ).raster) try: mem_dem.fill_depressions(engine="tiled", out_path=str(WORK / "x.tif"), workers=4) except ValueError as exc: print("rejected:", exc)

In [ ]:

for d in (serial, dask_filled, filled, acc_serial, acc_dask):
    d.close()

for d in (serial, dask_filled, filled, acc_serial, acc_dask): d.close()