Hydrologic indices: TWI, SPI, STI¶

Three classic catchment-scale indices computed from accumulation × slope:

Index	Formula	Use
TWI	ln(SCA / tan(slope))	Wetness / saturation likelihood (Beven & Kirkby 1979)
SPI	SCA × tan(slope)	Erosion / stream-power proxy
STI	(SCA/22.13)^0.6 × (sin(slope)/0.0896)^1.3	Sediment transport (Moore & Burch 1986)

All three share the _area_slope_index kernel inside DEM — caller supplies a precomputed slope raster (in degrees) or lets the method compute it from DEM.slope().

In [1]:

import matplotlib.pyplot as plt
plt.rcParams["figure.dpi"] = 80
plt.rcParams["savefig.dpi"] = 80

import matplotlib.pyplot as plt plt.rcParams["figure.dpi"] = 80 plt.rcParams["savefig.dpi"] = 80

In [2]:

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

# Catchment with a clear main-stem chain and tributaries.
rows, cols = 15, 15
z = np.full((rows, cols), 100.0, dtype=np.float32)
for r in range(rows):
    z[r, 7] = float(40 - 2.5 * r)
for c in range(7):
    z[5, c] = float(30 - c)
for c in range(8, cols):
    z[10, c] = float(20 - (c - 8))

ds = Dataset.create_from_array(
    z, top_left_corner=(0.0, 0.0), cell_size=30.0,
    epsg=32618, no_data_value=-9999.0,
)
dem = DEM(ds.raster)
filled = dem.fill_depressions(method="priority_flood")
fd = filled.flow_direction(method="d8")
acc = fd.accumulate()
print(f"Accumulation max: {int(acc.read_array().max())} cells")

import numpy as np from pyramids.dataset import Dataset from digitalrivers import DEM # Catchment with a clear main-stem chain and tributaries. rows, cols = 15, 15 z = np.full((rows, cols), 100.0, dtype=np.float32) for r in range(rows): z[r, 7] = float(40 - 2.5 * r) for c in range(7): z[5, c] = float(30 - c) for c in range(8, cols): z[10, c] = float(20 - (c - 8)) ds = Dataset.create_from_array( z, top_left_corner=(0.0, 0.0), cell_size=30.0, epsg=32618, no_data_value=-9999.0, ) dem = DEM(ds.raster) filled = dem.fill_depressions(method="priority_flood") fd = filled.flow_direction(method="d8") acc = fd.accumulate() print(f"Accumulation max: {int(acc.read_array().max())} cells")

2026-06-09 23:45:22 | INFO | pyramids.base.config | Logging is configured.

Accumulation max: 27 cells

Visualise: DEM + accumulation¶

In [3]:

import matplotlib.pyplot as plt
import numpy as np
from pyramids.dataset import Dataset

fig, axes = plt.subplots(1, 2, figsize=(10, 4))

filled.plot(band=0, ax=axes[0], title="Filled DEM", cmap="terrain")
axes[0].set_xticks([])
axes[0].set_yticks([])

acc_arr = acc.read_array()
acc_masked = np.where(acc_arr < 1, np.nan, acc_arr).astype(np.float32)
Dataset.dataset_like(acc, acc_masked).plot(
    band=0, ax=axes[1], title="Flow accumulation (log)",
    cmap="viridis", color_scale="power",
)
axes[1].set_xticks([])
axes[1].set_yticks([])
fig.tight_layout()
plt.show()

import matplotlib.pyplot as plt import numpy as np from pyramids.dataset import Dataset fig, axes = plt.subplots(1, 2, figsize=(10, 4)) filled.plot(band=0, ax=axes[0], title="Filled DEM", cmap="terrain") axes[0].set_xticks([]) axes[0].set_yticks([]) acc_arr = acc.read_array() acc_masked = np.where(acc_arr < 1, np.nan, acc_arr).astype(np.float32) Dataset.dataset_like(acc, acc_masked).plot( band=0, ax=axes[1], title="Flow accumulation (log)", cmap="viridis", color_scale="power", ) axes[1].set_xticks([]) axes[1].set_yticks([]) fig.tight_layout() plt.show()

C:\Users\main\AppData\Local\Temp\ipykernel_2724\426115222.py:19: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
  fig.tight_layout()

Topographic Wetness Index (TWI)¶

TWI = ln(SCA / tan(slope)). High TWI marks cells likely to saturate first — large upstream area AND low slope. The classic stream-bottom cells score highest.

In [4]:

twi = filled.twi(acc).read_array()
valid_twi = twi[twi != -9999.0]
valid_twi = valid_twi[np.isfinite(valid_twi)]
print(f"TWI range:     [{valid_twi.min():.3f}, {valid_twi.max():.3f}]")
print(f"Mean TWI:      {valid_twi.mean():.3f}")

twi = filled.twi(acc).read_array() valid_twi = twi[twi != -9999.0] valid_twi = valid_twi[np.isfinite(valid_twi)] print(f"TWI range: [{valid_twi.min():.3f}, {valid_twi.max():.3f}]") print(f"Mean TWI: {valid_twi.mean():.3f}")

TWI range:     [5.886, 13.605]
Mean TWI:      8.949

Stream Power Index (SPI)¶

SPI = SCA × tan(slope). Proportional to the rate at which overland flow does work — a proxy for erosion risk. Steep cells with large upstream area score highest.

In [5]:

spi = filled.spi(acc).read_array()
valid_spi = spi[(spi != -9999.0) & np.isfinite(spi)]
print(f"SPI range:     [{valid_spi.min():.3f}, {valid_spi.max():.3f}]")
print(f"Mean SPI:      {valid_spi.mean():.3f}")

# SPI must be ≥ 0 (SCA and tan(slope) both non-negative on valid cells)
assert (valid_spi >= 0).all()

spi = filled.spi(acc).read_array() valid_spi = spi[(spi != -9999.0) & np.isfinite(spi)] print(f"SPI range: [{valid_spi.min():.3f}, {valid_spi.max():.3f}]") print(f"Mean SPI: {valid_spi.mean():.3f}") # SPI must be ≥ 0 (SCA and tan(slope) both non-negative on valid cells) assert (valid_spi >= 0).all()

SPI range:     [0.000, 60.000]
Mean SPI:      2.059

Sediment Transport Index (STI)¶

STI = (SCA / 22.13)^0.6 × (sin(slope) / 0.0896)^1.3 — the per-unit-area form of the RUSLE LS factor. Predicts where flow will transport sediment rather than deposit it.

In [6]:

sti = filled.sti(acc).read_array()
valid_sti = sti[(sti != -9999.0) & np.isfinite(sti)]
print(f"STI range:     [{valid_sti.min():.3f}, {valid_sti.max():.3f}]")
print(f"Mean STI:      {valid_sti.mean():.3f}")

assert (valid_sti >= 0).all()

sti = filled.sti(acc).read_array() valid_sti = sti[(sti != -9999.0) & np.isfinite(sti)] print(f"STI range: [{valid_sti.min():.3f}, {valid_sti.max():.3f}]") print(f"Mean STI: {valid_sti.mean():.3f}") assert (valid_sti >= 0).all()

STI range:     [0.000, 7.916]
Mean STI:      0.323

Visualise: TWI / SPI / STI¶

In [7]:

import matplotlib.pyplot as plt
import numpy as np
from pyramids.dataset import Dataset

def _mask(a):
    return np.where((a == -9999.0) | (~np.isfinite(a)), np.nan, a).astype(np.float32)

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

twi_show = _mask(twi)
Dataset.dataset_like(filled, twi_show).plot(
    band=0, ax=axes[0], title="TWI", cmap="YlGnBu",
)
axes[0].set_xticks([])
axes[0].set_yticks([])

spi_show = _mask(spi)
spi_show = np.where(spi_show > 0, spi_show, np.nan).astype(np.float32)
Dataset.dataset_like(filled, spi_show).plot(
    band=0, ax=axes[1], title="SPI (log)", cmap="hot_r", color_scale="power",
)
axes[1].set_xticks([])
axes[1].set_yticks([])

sti_show = _mask(sti)
sti_show = np.where(sti_show > 0, sti_show, np.nan).astype(np.float32)
Dataset.dataset_like(filled, sti_show).plot(
    band=0, ax=axes[2], title="STI (log)", cmap="hot_r", color_scale="power",
)
axes[2].set_xticks([])
axes[2].set_yticks([])
fig.tight_layout()
plt.show()

import matplotlib.pyplot as plt import numpy as np from pyramids.dataset import Dataset def _mask(a): return np.where((a == -9999.0) | (~np.isfinite(a)), np.nan, a).astype(np.float32) fig, axes = plt.subplots(1, 3, figsize=(15, 4)) twi_show = _mask(twi) Dataset.dataset_like(filled, twi_show).plot( band=0, ax=axes[0], title="TWI", cmap="YlGnBu", ) axes[0].set_xticks([]) axes[0].set_yticks([]) spi_show = _mask(spi) spi_show = np.where(spi_show > 0, spi_show, np.nan).astype(np.float32) Dataset.dataset_like(filled, spi_show).plot( band=0, ax=axes[1], title="SPI (log)", cmap="hot_r", color_scale="power", ) axes[1].set_xticks([]) axes[1].set_yticks([]) sti_show = _mask(sti) sti_show = np.where(sti_show > 0, sti_show, np.nan).astype(np.float32) Dataset.dataset_like(filled, sti_show).plot( band=0, ax=axes[2], title="STI (log)", cmap="hot_r", color_scale="power", ) axes[2].set_xticks([]) axes[2].set_yticks([]) fig.tight_layout() plt.show()

C:\Users\main\AppData\Local\Temp\ipykernel_2724\776328841.py:32: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
  fig.tight_layout()

Reusing a precomputed slope raster¶

All three methods accept slope_deg= as a kwarg to skip recomputation when you already have a slope raster handy.

In [8]:

slope_ratio = filled.slope().read_array()  # m/m
slope_deg_arr = np.degrees(np.arctan(np.where(
    np.isfinite(slope_ratio), slope_ratio, 0.0,
)))
slope_deg = Dataset.create_from_array(
    slope_deg_arr.astype(np.float32),
    geo=filled.geotransform, epsg=filled.epsg, no_data_value=-9999.0,
)
twi_reused = filled.twi(acc, slope_deg=slope_deg).read_array()
np.testing.assert_allclose(
    twi_reused[twi_reused != -9999.0],
    twi[twi != -9999.0],
    atol=1e-3,
)
print("TWI with reused slope matches default-computed TWI.")

slope_ratio = filled.slope().read_array() # m/m slope_deg_arr = np.degrees(np.arctan(np.where( np.isfinite(slope_ratio), slope_ratio, 0.0, ))) slope_deg = Dataset.create_from_array( slope_deg_arr.astype(np.float32), geo=filled.geotransform, epsg=filled.epsg, no_data_value=-9999.0, ) twi_reused = filled.twi(acc, slope_deg=slope_deg).read_array() np.testing.assert_allclose( twi_reused[twi_reused != -9999.0], twi[twi != -9999.0], atol=1e-3, ) print("TWI with reused slope matches default-computed TWI.")

TWI with reused slope matches default-computed TWI.

Summary¶

Three classic indices in three one-line calls. The shared _area_slope_index kernel handles no-data masking, slope-floor clamping (to avoid log(0) on flats), and float32 output wrapping.