Multiple Flow Direction Methods in Python

pip install richdem>=0.3.4 rasterio>=1.3.0 numpy>=1.24.0 scipy>=1.10.0

conda install -c conda-forge richdem rasterio numpy scipy

f_ij = (tan β_ij)^p  /  Σ_k (tan β_ik)^p

import logging
import rasterio
import numpy as np

logger = logging.getLogger(__name__)

def load_dem(path: str) -> tuple[np.ndarray, np.ndarray, dict]:
    """Read a single-band elevation raster and return (dem, valid_mask, src_meta)."""
    with rasterio.open(path) as src:
        if src.count != 1:
            raise ValueError(f"Expected single-band DEM, got {src.count} bands: {path}")
        dem = src.read(1).astype(np.float64)
        nodata = src.nodata
        if nodata is not None:
            dem[dem == nodata] = np.nan
        mask = np.isfinite(dem)
        n_valid = int(mask.sum())
        if n_valid == 0:
            raise ValueError(f"DEM contains no finite values: {path}")
        logger.info("Loaded DEM %s — shape %s, %d valid cells", path, dem.shape, n_valid)
        return dem, mask, src.meta.copy()

import richdem as rd

def condition_dem(dem: np.ndarray) -> np.ndarray:
    """Apply priority-flood fill then flat resolution. Returns conditioned float64 array."""
    rd_dem = rd.rdarray(dem, no_data=np.nan)
    filled = rd.FillDepressions(rd_dem, epsilon=False, in_place=False)
    logger.info("Depression filling complete")
    resolved = rd.ResolveFlats(filled, in_place=False)
    logger.info("Flat resolution complete")
    return np.array(resolved, dtype=np.float64)

from scipy import sparse

def compute_mfd_proportions(
    dem: np.ndarray,
    cell_size: float = 1.0,
    p: float = 1.0
) -> sparse.csr_matrix:
    """
    Build the MFD proportion matrix P where P[i,j] is the fraction of flow
    sent from cell i to cell j.  All entries ≥ 0; each row sums to ≤ 1.0
    (rows for sink cells sum to 0).
    """
    rows, cols = dem.shape
    n_cells = rows * cols
    cell_ids = np.arange(n_cells).reshape(rows, cols)

    # 8-connected direction offsets and their Euclidean distances
    directions = [
        ((-1, -1), cell_size * np.sqrt(2)),
        ((-1,  0), cell_size),
        ((-1,  1), cell_size * np.sqrt(2)),
        (( 0, -1), cell_size),
        (( 0,  1), cell_size),
        (( 1, -1), cell_size * np.sqrt(2)),
        (( 1,  0), cell_size),
        (( 1,  1), cell_size * np.sqrt(2)),
    ]

    row_idx, col_idx, weights = [], [], []

    for (dr, dc), dist in directions:
        # Slice source and target windows to avoid boundary padding
        r_src = slice(max(0, -dr), min(rows, rows - dr))
        c_src = slice(max(0, -dc), min(cols, cols - dc))
        r_tgt = slice(max(0,  dr), min(rows, rows + dr))
        c_tgt = slice(max(0,  dc), min(cols, cols + dc))

        # Positive dz → neighbour is lower → valid downslope
        dz = dem[r_src, c_src] - dem[r_tgt, c_tgt]
        downslope = (dz > 0) & np.isfinite(dz)
        if not np.any(downslope):
            continue

        tan_beta = dz[downslope] / dist
        w = np.power(tan_beta, p)

        src_ids = cell_ids[r_src, c_src][downslope].ravel()
        tgt_ids = cell_ids[r_tgt, c_tgt][downslope].ravel()
        row_idx.extend(src_ids)
        col_idx.extend(tgt_ids)
        weights.extend(w)

    P = sparse.coo_matrix(
        (weights, (row_idx, col_idx)), shape=(n_cells, n_cells)
    ).tocsr()

    # Normalise each row so fractions sum to 1.0
    row_sums = np.array(P.sum(axis=1)).ravel()
    row_sums[row_sums == 0.0] = 1.0  # sink cells: avoid division by zero
    P = sparse.diags(1.0 / row_sums, format="csr") @ P

    nnz = P.nnz
    logger.info("Proportion matrix built — %d cells, %d nonzeros, p=%.2f", n_cells, nnz, p)
    return P

def compute_accumulation(P: sparse.csr_matrix, shape: tuple[int, int]) -> np.ndarray:
    """
    Solve (I - P) A = 1 for flow accumulation A.
    Raises RuntimeError if the linear system is singular (residual sinks present).
    """
    n = P.shape[0]
    M = sparse.eye(n, format="csr") - P
    ones = np.ones(n)
    try:
        A = sparse.linalg.spsolve(M, ones)
    except Exception as exc:
        raise RuntimeError(
            "Accumulation solver failed — check for residual sinks or "
            "disconnected DEM components. Consider re-running condition_dem()."
        ) from exc
    if not np.all(np.isfinite(A)):
        raise RuntimeError("Solver returned non-finite values — singular system.")
    logger.info(
        "Accumulation solved — min=%.1f, max=%.1f, mean=%.1f",
        float(A.min()), float(A.max()), float(A.mean())
    )
    return A.reshape(shape)

def export_accumulation(
    out_path: str,
    acc: np.ndarray,
    mask: np.ndarray,
    src_meta: dict,
) -> None:
    """Write accumulation GeoTIFF and run mass-conservation check."""
    out = acc.copy()
    out[~mask] = np.nan
    total = float(out[mask].sum())
    expected = float(mask.sum())
    ratio = total / expected if expected > 0 else 0.0
    if not np.isclose(ratio, 1.0, rtol=1e-3):
        logger.warning(
            "Mass conservation check failed: got %.3f, expected 1.000 (ratio=%.6f). "
            "Likely cause: unresolved sinks or NaN leakage.",
            total / expected, ratio
        )
    else:
        logger.info("Mass conservation OK (ratio=%.6f)", ratio)

    meta = src_meta.copy()
    meta.update(dtype="float64", count=1, nodata=np.nan)
    with rasterio.open(out_path, "w", **meta) as dst:
        dst.write(out, 1)
    logger.info("Accumulation raster written to %s", out_path)

import logging
import numpy as np
import rasterio
import richdem as rd
from scipy import sparse

logger = logging.getLogger(__name__)


def run_mfd_accumulation(
    dem_path: str,
    out_path: str,
    cell_size: float | None = None,
    p: float = 1.0,
) -> np.ndarray:
    """
    End-to-end MFD flow accumulation from a raw DEM.

    Parameters
    ----------
    dem_path  : Path to input single-band elevation GeoTIFF.
    out_path  : Destination path for the accumulation GeoTIFF.
    cell_size : Pixel size in metres. Inferred from raster metadata if None.
    p         : Freeman dispersion exponent. 1.0 = Quinn formulation (max dispersion).
                Increase toward 1.3 on convergent hillslopes.

    Returns
    -------
    acc : 2-D float64 array of flow accumulation values (NaN outside valid area).
    """
    logger.info("=== MFD accumulation starting: %s", dem_path)

    # --- Load ---
    with rasterio.open(dem_path) as src:
        if src.count != 1:
            raise ValueError(f"Expected single-band DEM, got {src.count} bands.")
        if cell_size is None:
            cell_size = abs(src.res[0])
            logger.info("Inferred cell size: %.4f m", cell_size)
        meta = src.meta.copy()
        dem_raw = src.read(1).astype(np.float64)
        if src.nodata is not None:
            dem_raw[dem_raw == src.nodata] = np.nan

    mask = np.isfinite(dem_raw)
    n_valid = int(mask.sum())
    if n_valid == 0:
        raise ValueError("DEM contains no finite values.")
    logger.info("Valid cells: %d of %d", n_valid, dem_raw.size)

    # --- Condition ---
    rd_dem = rd.rdarray(dem_raw, no_data=np.nan)
    filled = rd.FillDepressions(rd_dem, epsilon=False, in_place=False)
    conditioned = np.array(rd.ResolveFlats(filled, in_place=False), dtype=np.float64)
    logger.info("DEM conditioning complete")

    # --- Proportion matrix ---
    rows, cols = conditioned.shape
    n_cells = rows * cols
    cell_ids = np.arange(n_cells).reshape(rows, cols)
    directions = [
        ((-1, -1), cell_size * np.sqrt(2)),
        ((-1,  0), cell_size),
        ((-1,  1), cell_size * np.sqrt(2)),
        (( 0, -1), cell_size),
        (( 0,  1), cell_size),
        (( 1, -1), cell_size * np.sqrt(2)),
        (( 1,  0), cell_size),
        (( 1,  1), cell_size * np.sqrt(2)),
    ]
    row_idx, col_idx, wts = [], [], []
    for (dr, dc), dist in directions:
        r_s = slice(max(0, -dr), min(rows, rows - dr))
        c_s = slice(max(0, -dc), min(cols, cols - dc))
        r_t = slice(max(0,  dr), min(rows, rows + dr))
        c_t = slice(max(0,  dc), min(cols, cols + dc))
        dz = conditioned[r_s, c_s] - conditioned[r_t, c_t]
        down = (dz > 0) & np.isfinite(dz)
        if not np.any(down):
            continue
        row_idx.extend(cell_ids[r_s, c_s][down].ravel())
        col_idx.extend(cell_ids[r_t, c_t][down].ravel())
        wts.extend(np.power(dz[down] / dist, p))

    P = sparse.coo_matrix((wts, (row_idx, col_idx)), shape=(n_cells, n_cells)).tocsr()
    rs = np.array(P.sum(axis=1)).ravel()
    rs[rs == 0.0] = 1.0
    P = sparse.diags(1.0 / rs, format="csr") @ P
    logger.info("Proportion matrix: %d nonzeros (p=%.2f)", P.nnz, p)

    # --- Solve ---
    M = sparse.eye(n_cells, format="csr") - P
    try:
        A = sparse.linalg.spsolve(M, np.ones(n_cells))
    except Exception as exc:
        raise RuntimeError("Sparse solve failed — residual sinks suspected.") from exc
    if not np.all(np.isfinite(A)):
        raise RuntimeError("Solver returned non-finite values — singular system.")
    acc = A.reshape(rows, cols)
    logger.info("Accumulation: max=%.1f cells", float(acc[mask].max()))

    # --- Export ---
    out = acc.copy()
    out[~mask] = np.nan
    ratio = float(out[mask].sum()) / n_valid
    if not np.isclose(ratio, 1.0, rtol=1e-3):
        logger.warning("Mass conservation ratio=%.6f (expected 1.0)", ratio)
    meta.update(dtype="float64", count=1, nodata=np.nan)
    with rasterio.open(out_path, "w", **meta) as dst:
        dst.write(out, 1)
    logger.info("Output written: %s", out_path)
    return out

max_loc = np.unravel_index(np.nanargmax(acc), acc.shape)
logger.info("Peak accumulation at grid position %s", max_loc)

A, info = sparse.linalg.bicgstab(M, np.ones(n_cells), tol=1e-8, maxiter=500)
if info != 0:
    logger.warning("bicgstab did not converge (info=%d) — check for sinks", info)