#!/usr/bin/env python3 # -*- coding: utf-8 -*- """Figure 5: what the flow-tree index gives on the 8 x 8 example of Figure 4. No input data. Output: ../Figures/Fig05_DFS_Capabilities_v11.png and .pdf. Run from this folder; ../Data holds the inputs and the figure is written to ../Figures. """ import os from collections import deque import numpy as np import matplotlib as mpl import matplotlib.pyplot as plt import matplotlib.patches as mpatches # ============================================================================ # Style (ESSD / AGU – Helvetica, cm mathtext, 9 pt) # ============================================================================ mpl.rcParams.update({ "font.family": "Helvetica", "font.size": 9, "axes.labelsize": 9, "axes.titlesize": 10, "xtick.labelsize": 8, "ytick.labelsize": 8, "mathtext.fontset": "cm", "axes.linewidth": 0.6, }) # ============================================================================ # Paths # ============================================================================ HERE = os.path.dirname(os.path.abspath(__file__)) OUT_DIR = os.path.join(HERE, "..", "Figures") os.makedirs(OUT_DIR, exist_ok=True) # ============================================================================ # 8 × 8 schematic topology (same as Fig 4) # ============================================================================ NR, NC = 8, 8 size = NR * NC dir2d = np.array([ [ 1, 2, 4, 8, 8, 8, 16, 16], [128, 2, 4, 8, 8, 16, 32, 32], [ 2, 2, 4, 8, 8, 8, 8, 8], [ 4, 2, 4, 8, 8, 8, 8, 8], [ 2, 4, 8, 8, 16, 16, 8, 8], [ 2, 4, 8, 8, 4, 8, 8, 8], [ 0, 0, 0, 8, 2, 4, 8, 8], [247, 247, 0, 0, 0, 0, 0, 0], ], dtype=np.uint8) msk2d = np.array([ [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 1, 1, 1, 1], [0, 1, 0, 0, 1, 1, 1, 0], [0, 0, 0, 0, 0, 1, 0, 0], ], dtype=np.uint8) bsn2d = np.array([ [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 2], [1, 1, 1, 1, 1, 1, 2, 2], [1, 1, 1, 0, 2, 2, 2, 2], [0, 1, 0, 0, 2, 2, 2, 0], [0, 0, 0, 0, 0, 2, 0, 0], ], dtype=np.uint32) D8 = {1: (0, 1), 2: (1, 1), 4: (1, 0), 8: (1, -1), 16: (0, -1), 32: (-1, -1), 64: (-1, 0), 128: (-1, 1)} # ---- colours ---- BLUE = "#2563EB" GREEN = "#16A34A" GRAY = "#9CA3AF" RED = "#EF4444" ORANGE = "#F59E0B" TEAL = "#0D9488" BG_BLUE = "#DBEAFE" BG_GREEN = "#D1FAE5" # ============================================================================ # Build idxs_ds, children, UPA (Kahn), DFS (UPA-sorted), Cell Rank # ============================================================================ idxs_ds = np.full(size, -1, dtype=np.int32) for r in range(NR): for c in range(NC): i = r * NC + c if msk2d[r, c] == 0: continue d = int(dir2d[r, c]) if d == 0: idxs_ds[i] = i elif d in D8: dr, dc = D8[d] nr2, nc2 = r + dr, c + dc if 0 <= nr2 < NR and 0 <= nc2 < NC: idxs_ds[i] = nr2 * NC + nc2 ch_list = [[] for _ in range(size)] outs = [] for i in range(size): if idxs_ds[i] < 0: continue if idxs_ds[i] == i: outs.append(i) else: ch_list[idxs_ds[i]].append(i) upg = np.zeros(size, dtype=np.int32) for r in range(NR): for c in range(NC): if msk2d[r, c] == 1: upg[r * NC + c] = 1 in_deg = np.zeros(size, dtype=np.int32) for i in range(size): if idxs_ds[i] < 0 or idxs_ds[i] == i: continue in_deg[idxs_ds[i]] += 1 q = deque(i for i in range(size) if msk2d[i // NC, i % NC] == 1 and in_deg[i] == 0) while q: n = q.popleft() ds = idxs_ds[n] if ds < 0 or ds == n: continue upg[ds] += upg[n] in_deg[ds] -= 1 if in_deg[ds] == 0: q.append(ds) for i in range(size): if len(ch_list[i]) > 1: ch_list[i].sort(key=lambda k: -int(upg[k])) # one number per pixel: the counter advances only when a pixel is discovered; # basins by the pixel index of their outlet, larger UPA first (as Fig. 4 v11) din = np.full(size, -1, dtype=np.int32) t = 0 for ol in outs: stk = [ol] while stk: nd = stk.pop() din[nd] = t; t += 1 for c2 in reversed(ch_list[nd]): stk.append(c2) def _subtree(i): out = [i] for k in ch_list[i]: out += _subtree(k) return out # dfs_out = largest dfs_in in the upstream catchment (Haag and Shokoufandeh 2019) = dfs_in + upg - 1 dout = np.full(size, -1, dtype=np.int32) for i in range(size): if din[i] >= 0: sub = _subtree(i) assert sorted(din[sub]) == list(range(din[i], din[i] + upg[i])), i dout[i] = max(din[sub]) assert dout[i] == din[i] + upg[i] - 1, i cell_rank = np.full(size, -1, dtype=np.int32) for i in range(size): r, c = divmod(i, NC) if msk2d[r, c] == 0: continue cur, d = i, 0 while idxs_ds[cur] != cur: cur = idxs_ds[cur]; d += 1 cell_rank[i] = d # ============================================================================ # Drawing helpers # ============================================================================ FS_CELL = 8.5 FS_SMALL = 6.5 FS_ANNOT = 6.5 def rc(i): return divmod(i, NC) def setup_grid(ax): ax.set_xlim(-0.5, NC - 0.5); ax.set_ylim(NR - 0.5, -0.5) ax.set_aspect("equal") ax.set_xticks(np.arange(-0.5, NC, 1)) ax.set_yticks(np.arange(-0.5, NR, 1)) ax.set_xticklabels([]); ax.set_yticklabels([]) ax.grid(True, color="0.4", lw=0.35, alpha=0.45, zorder=1) ax.tick_params(length=0) def fill_cell(ax, px, color, alpha=0.3, zorder=0): r, c = rc(px) ax.add_patch(mpatches.Rectangle((c - 0.5, r - 0.5), 1, 1, fc=color, ec="none", alpha=alpha, zorder=zorder)) def label_cell(ax, px, text, color=BLUE, fs=FS_CELL, bold=True): r, c = rc(px) ax.text(c, r, text, ha="center", va="center", fontsize=fs, fontweight=("bold" if bold else "normal"), color=color, zorder=7) def annot_below(ax, text): ax.text(0.5, -0.04, text, fontsize=FS_ANNOT, va="top", ha="center", transform=ax.transAxes, bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="0.6", lw=0.4, alpha=0.95), zorder=10, clip_on=False) def draw_flow_path(ax, path, color, lw=2.0, alpha=0.85): xs = [p % NC for p in path] ys = [p // NC for p in path] ax.plot(xs, ys, "-", color=color, lw=lw, alpha=alpha, zorder=5, solid_capstyle="round") def get_upstream(target): lo, hi = din[target], dout[target] return {i for i in range(size) if 0 <= din[i] and lo <= din[i] <= hi} def trace_main_stem(start): path = [start]; cur = start while ch_list[cur]: cur = ch_list[cur][0] path.append(cur) return path def edge_is_boundary(us_set, r, c, dr, dc): nr, nc = r + dr, c + dc if nr < 0 or nr >= NR or nc < 0 or nc >= NC: return True if msk2d[nr, nc] == 0: return True return (nr * NC + nc) not in us_set def draw_4conn_boundary(ax, us_set, color="black", lw=1.5, zorder=6): for px in us_set: r, c = rc(px) if edge_is_boundary(us_set, r, c, -1, 0): ax.plot([c - 0.5, c + 0.5], [r - 0.5, r - 0.5], color=color, lw=lw, zorder=zorder, solid_capstyle="round") if edge_is_boundary(us_set, r, c, 0, 1): ax.plot([c + 0.5, c + 0.5], [r - 0.5, r + 0.5], color=color, lw=lw, zorder=zorder, solid_capstyle="round") if edge_is_boundary(us_set, r, c, 1, 0): ax.plot([c - 0.5, c + 0.5], [r + 0.5, r + 0.5], color=color, lw=lw, zorder=zorder, solid_capstyle="round") if edge_is_boundary(us_set, r, c, 0, -1): ax.plot([c - 0.5, c - 0.5], [r - 0.5, r + 0.5], color=color, lw=lw, zorder=zorder, solid_capstyle="round") def boundary_pixels_8conn(us_set): bnd = set() nbrs_8 = [(-1, -1), (-1, 0), (-1, 1), ( 0, -1), ( 0, 1), ( 1, -1), ( 1, 0), ( 1, 1)] for px in us_set: r, c = rc(px) for dr, dc in nbrs_8: nr, nc = r + dr, c + dc if (nr < 0 or nr >= NR or nc < 0 or nc >= NC or msk2d[nr, nc] == 0 or (nr * NC + nc) not in us_set): bnd.add(px); break return bnd # ============================================================================ # Panels # ============================================================================ def panel_a(ax): """(a) D8 flow direction: the input; the index is checked against it at every pixel.""" for i in range(size): r, c = rc(i) if msk2d[r, c] == 0: continue col = BLUE if bsn2d[r, c] == 1 else GREEN label_cell(ax, i, str(int(dir2d[r, c])), col, fs=FS_CELL) # the check: each pixel's D8 neighbour is the neighbour whose interval is # the tightest one enclosing the pixel (outlets: none) for i in range(size): r, c = rc(i) if msk2d[r, c] == 0: continue enclosing = [k for k in range(size) if k != i and din[k] >= 0 and din[k] <= din[i] <= dout[k] and max(abs(rc(k)[0] - r), abs(rc(k)[1] - c)) == 1] parent = min(enclosing, key=lambda k: dout[k] - din[k]) if enclosing else i assert parent == idxs_ds[i], i setup_grid(ax) annot_below(ax, "input; the index is checked against it on this grid") ax.set_title("(a) D8 flow direction", loc="left", fontweight="bold", pad=3) def panel_b(ax): """(b) Distance to outlet in pixels.""" for i in range(size): r, c = rc(i) if msk2d[r, c] == 0: continue col = BLUE if bsn2d[r, c] == 1 else GREEN n_enc = sum(1 for k in range(size) if k != i and din[k] >= 0 and din[k] <= din[i] <= dout[k]) assert n_enc == cell_rank[i], i label_cell(ax, i, str(cell_rank[i]), col, fs=FS_CELL) setup_grid(ax) annot_below(ax, r"count of other pixels' [$\mathit{dfs\_in}$, $\mathit{dfs\_out}$] enclosing this pixel") ax.set_title("(b) Distance to outlet in pixels", loc="left", fontweight="bold", pad=3) def panel_c(ax): """(c) Upstream pixel count.""" for i in range(size): r, c = rc(i) if msk2d[r, c] == 0: continue val = dout[i] - din[i] + 1 assert val == upg[i], i col = BLUE if bsn2d[r, c] == 1 else GREEN label_cell(ax, i, str(val), col, fs=FS_CELL) setup_grid(ax) annot_below(ax, r"$\mathit{dfs\_out} - \mathit{dfs\_in} + 1$") ax.set_title("(c) Upstream pixel count", loc="left", fontweight="bold", pad=3) def panel_d(ax): """(d) Upstream catchment via interval containment + 4-conn boundary.""" target = 26 upstream = get_upstream(target) for px in upstream: fill_cell(ax, px, BG_BLUE, alpha=0.55) draw_4conn_boundary(ax, upstream, color="black", lw=1.4, zorder=5) for i in range(size): r, c = rc(i) if msk2d[r, c] == 0: continue if i == target: label_cell(ax, i, r"$T_1$", RED, fs=FS_CELL + 1.5) elif i in upstream: label_cell(ax, i, str(din[i]), BLUE, fs=FS_CELL) else: label_cell(ax, i, str(din[i]), GRAY, fs=FS_SMALL - 0.5, bold=False) rt, ct = rc(target) ax.add_patch(mpatches.Rectangle((ct - 0.5, rt - 0.5), 1, 1, fc="none", ec=RED, lw=1.5, zorder=8)) setup_grid(ax) annot_below(ax, rf"pixel with $\mathit{{dfs\_in}} \in [\mathit{{dfs\_in}}(T_1)={din[target]}, \mathit{{dfs\_out}}(T_1)={dout[target]}]$") ax.set_title("(d) Upstream catchment extraction", loc="left", fontweight="bold", pad=3) def panel_e(ax): """(e) Boundary identification (8-conn neighbour outside interval).""" target = 26 upstream = get_upstream(target) bnd = boundary_pixels_8conn(upstream) interior = upstream - bnd for px in interior: fill_cell(ax, px, BG_BLUE, alpha=0.55) for px in bnd: fill_cell(ax, px, BG_BLUE, alpha=0.55) fill_cell(ax, px, BLUE, alpha=0.45, zorder=1) draw_4conn_boundary(ax, upstream, color="black", lw=1.4, zorder=5) for i in range(size): r, c = rc(i) if msk2d[r, c] == 0: continue if i == target: label_cell(ax, i, r"$T_1$", RED, fs=FS_CELL + 1.5) elif i in upstream: # boundary pixels: blue bold; interior catchment: blue not bold label_cell(ax, i, str(din[i]), BLUE, fs=FS_CELL, bold=(i in bnd)) else: # outside catchment label_cell(ax, i, str(din[i]), GRAY, fs=FS_SMALL - 0.5, bold=False) rt, ct = rc(target) ax.add_patch(mpatches.Rectangle((ct - 0.5, rt - 0.5), 1, 1, fc="none", ec=RED, lw=1.5, zorder=8)) setup_grid(ax) annot_below(ax, rf"catchment pixel + neighbour's $\mathit{{dfs\_in}} \notin [{din[target]},{dout[target]}]$") ax.set_title("(e) Boundary identification", loc="left", fontweight="bold", pad=3) def panel_f(ax): """(f) Largest-tributary chains — runs of consecutive dfs_in.""" stem0 = trace_main_stem(49) # Basin A main stem (dfs_in 0..) stem_a1 = trace_main_stem(42) # Basin A tributary 1 stem_a2 = trace_main_stem(19) # Basin A tributary 2 bsn_b_outs = [o for o in outs if bsn2d[rc(o)[0], rc(o)[1]] == 2] stem_b = trace_main_stem(bsn_b_outs[0]) for i in range(size): r, c = rc(i) if msk2d[r, c] == 0: continue bg = BG_BLUE if bsn2d[r, c] == 1 else BG_GREEN fill_cell(ax, i, bg, alpha=0.25) draw_flow_path(ax, stem0, BLUE, lw=3.5, alpha=0.75) draw_flow_path(ax, stem_a1, TEAL, lw=2.0, alpha=0.65) draw_flow_path(ax, stem_a2, ORANGE, lw=2.0, alpha=0.65) draw_flow_path(ax, stem_b, GREEN, lw=2.8, alpha=0.75) for px in stem0: label_cell(ax, px, str(din[px]), BLUE, fs=FS_CELL - 0.5) for px in stem_a1: label_cell(ax, px, str(din[px]), TEAL, fs=FS_SMALL) for px in stem_a2: label_cell(ax, px, str(din[px]), ORANGE, fs=FS_SMALL) for px in stem_b: label_cell(ax, px, str(din[px]), GREEN, fs=FS_SMALL) labelled = set(stem0) | set(stem_a1) | set(stem_a2) | set(stem_b) for i in range(size): r, c = rc(i) if msk2d[r, c] == 0 or i in labelled: continue label_cell(ax, i, str(din[i]), GRAY, fs=FS_SMALL - 1, bold=False) setup_grid(ax) annot_below(ax, r"each chain: consecutive $\mathit{dfs\_in}$") ax.set_title("(f) Largest-tributary chains", loc="left", fontweight="bold", pad=3) # ============================================================================ # Main # ============================================================================ def main(): fig, axes = plt.subplots(2, 3, figsize=(7.4, 5.7)) panel_a(axes[0, 0]); panel_b(axes[0, 1]); panel_c(axes[0, 2]) panel_d(axes[1, 0]); panel_e(axes[1, 1]); panel_f(axes[1, 2]) fig.subplots_adjust(left=0.025, right=0.975, top=0.95, bottom=0.06, wspace=0.12, hspace=0.42) for ext in ("png", "pdf"): out = os.path.join(OUT_DIR, f"Fig05_DFS_Capabilities_v11.{ext}") fig.savefig(out, dpi=300, bbox_inches="tight", facecolor="white") print(f"Saved: {out}") target = 26 n_up = len(get_upstream(target)) print(f" T = i{target} dfs_in={din[target]} dfs_out={dout[target]} " f"upstream={n_up} px") print(f" outlets = {outs} (Basin A first, then B)") plt.close(fig) if __name__ == "__main__": main()