Using Slope-Area Thresholds for Channel Initiation

import logging

import numpy as np
import rasterio
from scipy import stats

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


def fit_slope_exponent(
    heads_area_cells: np.ndarray,
    heads_slope: np.ndarray,
) -> dict:
    """
    Fit the slope exponent from mapped channel heads.

    At a channel head A·S^k is at its critical value C, so log A = log C − k·log S.
    A linear regression of log A on log S returns −k as the gradient.
    """
    m = (heads_area_cells > 0) & (heads_slope > 1e-4) & np.isfinite(heads_slope)
    if m.sum() < 20:
        raise ValueError(f"only {int(m.sum())} usable channel heads — too few "
                         "to fit an exponent; use a published value instead")

    log_a = np.log10(heads_area_cells[m].astype(float))
    log_s = np.log10(heads_slope[m].astype(float))
    result = stats.linregress(log_s, log_a)

    k = -float(result.slope)
    C = float(10 ** result.intercept)
    log.info("Fitted from %d channel heads: k = %.3f, C = %.1f, r² = %.3f",
             int(m.sum()), k, C, result.rvalue ** 2)
    if result.rvalue ** 2 < 0.3:
        log.warning("r² of %.2f is weak — the heads may span several process "
                    "regimes, or the slope grid may be noisy", result.rvalue ** 2)
    if not 0.5 <= k <= 2.5:
        log.warning("Fitted exponent %.2f is outside the range usually reported "
                    "(1.0–2.0); check the slope units are a gradient, not degrees", k)
    return {"k": k, "C": C, "r_squared": float(result.rvalue ** 2),
            "n_heads": int(m.sum())}


def slope_area_streams(
    facc_path: str,
    slope_path: str,
    out_path: str,
    k: float = 1.7,
    C: float = 500.0,
    min_slope: float = 0.005,
) -> dict:
    """
    Extract a stream mask using a slope-area criterion.

    Parameters
    ----------
    facc_path  : Flow accumulation raster, in cells.
    slope_path : Slope raster as a GRADIENT (rise over run), not degrees.
    out_path    : Destination for the boolean stream mask.
    k, C       : Slope exponent and critical constant.
    min_slope  : Floor applied to slope, so near-zero slopes do not produce an
                 infinite threshold and exclude genuine valley-floor channels.
    """
    with rasterio.open(facc_path) as src:
        acc = src.read(1).astype("float64")
        profile = src.profile.copy()
        nodata = src.nodata
    with rasterio.open(slope_path) as src:
        slope = src.read(1).astype("float64")

    if slope.shape != acc.shape:
        raise ValueError("slope and accumulation grids are not on the same grid")
    if np.nanmax(slope) > 5.0:
        log.warning("Slope maximum is %.1f — this looks like degrees rather "
                    "than a gradient; converting", np.nanmax(slope))
        slope = np.tan(np.radians(slope))

    # --- Floor the slope. Without it a perfectly flat cell gives S^k = 0 and
    # an infinite required area, which deletes valley-floor channels — the
    # opposite of what the method is for. ---
    s = np.clip(slope, min_slope, None)
    metric = acc * np.power(s, k)
    mask = metric >= C

    if nodata is not None:
        mask &= acc != nodata

    n = int(mask.sum())
    constant_equivalent = float(np.median(acc[mask])) if n else float("nan")
    log.info("Slope-area extraction: %d stream cells (%.3f %% of the grid)",
             n, 100.0 * n / mask.size)
    log.info("Median accumulation at extracted cells: %.0f — a constant "
             "threshold of about this value would give a similar total length "
             "but distribute it differently", constant_equivalent)

    # Report how the threshold varies, which is the whole point of the method.
    for lo, hi, label in [(0.0, 0.03, "flat"), (0.03, 0.12, "moderate"),
                          (0.12, 10.0, "steep")]:
        band = (s >= lo) & (s < hi)
        if band.any():
            req = C / np.power(np.median(s[band]), k)
            log.info("  %-8s slope band: required area %.0f cells", label, req)

    profile.update(dtype="uint8", nodata=0, compress="LZW", tiled=True)
    with rasterio.open(out_path, "w", **profile) as dst:
        dst.write(mask.astype("uint8"), 1)
        dst.update_tags(THRESHOLD_METHOD="slope-area", SLOPE_EXPONENT=str(k),
                        CRITICAL_CONSTANT=str(C))
    log.info("Wrote %s", out_path)
    return {"stream_cells": n, "k": k, "C": C, "output": out_path}


# --- Example usage ---
# fit = fit_slope_exponent(head_areas, head_slopes)
# slope_area_streams("acc.tif", "slope.tif", "streams_sa.tif",
#                    k=fit["k"], C=fit["C"])