#!/usr/bin/env python3 # -*- coding: utf-8 -*- """Figure 4: the flow-tree index (dfs_in, dfs_out) on an 8 x 8 D8 example with two basins. The example grid is written in this script; no input data. Output: ../Figures/Fig04_DFS_Interval_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 # ============================================================================ # 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 (carry-over from v6) # ============================================================================ 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" # Basin A GREEN = "#16A34A" # Basin B GRAY = "#9CA3AF" # invalid RED = "#EF4444" # arrows # ============================================================================ # Build idxs_ds, children list, UPA (Kahn topo), then sort children desc by UPA # ============================================================================ 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) upa = np.zeros(size, dtype=np.int32) for r in range(NR): for c in range(NC): if msk2d[r, c] == 1: upa[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 upa[ds] += upa[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(upa[k])) # ============================================================================ # DFS interval (UPA-sorted children) → dfs_in / dfs_out + event log # ============================================================================ events = [] din = np.full(size, -1, dtype=np.int32) # discovery order (one number per pixel) seq = np.full(size, -1, dtype=np.int32) # seq_dfs t = 0 for ol in outs: # basins by the pixel index of their outlet, as FlowTopo first = t stk = [ol] while stk: nd = stk.pop() din[nd] = t events.append((t, nd, "in")) t += 1 for c2 in reversed(ch_list[nd]): # larger UPA first; ties: smaller pixel index stk.append(c2) n_basin = t - first members = [i for i in range(size) if din[i] >= first and din[i] < t] for i in members: seq[i] = first + n_basin - 1 - (din[i] - first) # every pixel's upstream pixels are the numbers [seq - upg + 1, seq] def _upstream(i): out = [i] for k in ch_list[i]: out += _upstream(k) return out for i in range(size): if seq[i] >= 0: assert sorted(seq[_upstream(i)]) == list(range(seq[i] - upa[i] + 1, seq[i] + 1)), i T1 = 26 # the example target T1 BLOCK = set(_upstream(T1)) for i in range(size): # the block [dfs_in, dfs_in + upg - 1] is the same set if din[i] >= 0: assert sorted(din[_upstream(i)]) == list(range(din[i], din[i] + upa[i])), i # dfs_out by Haag's definition: the largest dfs_in in the catchment upstream of the pixel (closed interval; # a source pixel has dfs_out = dfs_in); it equals dfs_in + upg - 1 at every pixel dout = np.full(size, -1, dtype=np.int32) for i in range(size): if din[i] >= 0: dout[i] = max(din[_upstream(i)]) assert dout[i] == din[i] + upa[i] - 1, i DIN = din.reshape(NR, NC) DOUT = dout.reshape(NR, NC) # ============================================================================ # Drawing helpers # ============================================================================ FS_CELL = 8.5 # D8 / index in (a)(b) FS_TS = 7.5 # timestamps in (c)(d)(e)(f) FS_OUT = 6.0 # boxed exit timestamps (smaller to fit) def bsn_color(r, c): if msk2d[r, c] == 0: return GRAY return BLUE if bsn2d[r, c] == 1 else GREEN 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 draw_arrows(ax, alpha=0.7, pixel_set=None): for r in range(NR): for c in range(NC): if msk2d[r, c] == 0: continue i = r * NC + c if pixel_set is not None and i not in pixel_set: continue d = int(dir2d[r, c]) if d == 0 or d not in D8: continue dr, dc = D8[d] ax.arrow(c + 0.15 * dc, r + 0.15 * dr, 0.38 * dc, 0.38 * dr, head_width=0.18, head_length=0.18, fc=RED, ec="black", lw=0.35, length_includes_head=True, zorder=3, alpha=alpha) def panel_a(ax): draw_arrows(ax) for r in range(NR): for c in range(NC): v = int(dir2d[r, c]); col = bsn_color(r, c) if msk2d[r, c] == 0: ax.text(c, r, str(v), ha="center", va="center", fontsize=FS_CELL - 1.5, color=GRAY, alpha=0.5, zorder=2) else: ax.text(c, r, str(v), ha="center", va="center", fontsize=FS_CELL, color=col, zorder=2) setup_grid(ax) ax.set_title("(a) D8 flow direction", loc="left", fontweight="bold", pad=3) def panel_b(ax): draw_arrows(ax) for r in range(NR): for c in range(NC): i = r * NC + c; col = bsn_color(r, c) label = f"i{i}" if msk2d[r, c] == 0: ax.text(c, r, label, ha="center", va="center", fontsize=FS_CELL - 2, color=GRAY, alpha=0.5, zorder=2) else: ax.text(c, r, label, ha="center", va="center", fontsize=FS_CELL - 1, color=col, zorder=2) setup_grid(ax) ax.set_title("(b) Linear pixel index", loc="left", fontweight="bold", pad=3) def draw_snapshot(ax, after_idx, title): entered, exited = {}, {} for j in range(after_idx + 1): s, nd, typ = events[j] (entered if typ == "in" else exited)[nd] = s processed = set(entered) draw_arrows(ax, alpha=0.4, pixel_set=processed) for r in range(NR): for c in range(NC): i = r * NC + c; col = bsn_color(r, c) if msk2d[r, c] == 0: continue if i in entered: ax.text(c, r, str(entered[i]), ha="center", va="center", fontsize=FS_TS, fontweight="bold", color=col, zorder=5) setup_grid(ax) ax.set_title(title, loc="left", fontweight="bold", pad=3) def draw_final(ax, vals, title, framed=False, block=None): for r in range(NR): for c in range(NC): if msk2d[r, c] == 0: continue col = bsn_color(r, c); v = vals[r, c] if block is not None and r * NC + c in block: ax.add_patch(plt.Rectangle((c - 0.5, r - 0.5), 1, 1, facecolor="#fff1b8", edgecolor="none", zorder=0.5)) if r * NC + c == T1: ax.add_patch(plt.Rectangle((c - 0.46, r - 0.46), 0.92, 0.92, fill=False, edgecolor="black", lw=0.9, zorder=6)) ax.text(c - 0.5, r - 0.5, "T$_1$", ha="center", va="center", fontsize=FS_OUT - 0.5, color="black", zorder=7, bbox=dict(boxstyle="round,pad=0.12", fc="white", ec="black", lw=0.5)) if framed: ax.text(c, r, str(v), ha="center", va="center", fontsize=FS_CELL - 1, fontweight="bold", color=col, bbox=dict(boxstyle="round,pad=0.04", fc="white", ec="black", lw=0.35), zorder=5) else: ax.text(c, r, str(v), ha="center", va="center", fontsize=FS_CELL, fontweight="bold", color=col, zorder=5) setup_grid(ax) ax.set_title(title, loc="left", fontweight="bold", pad=3) # ============================================================================ # Main # ============================================================================ def main(): fig, axes = plt.subplots(2, 3, figsize=(7.2, 4.7)) panel_a(axes[0, 0]) panel_b(axes[0, 1]) draw_snapshot(axes[0, 2], 11, "(c) DFS steps 0–11") draw_snapshot(axes[1, 0], len(events) - 1, "(d) DFS complete") draw_final(axes[1, 1], DIN, r"(e) $\mathit{dfs\_in}$ raster", framed=False, block=BLOCK) draw_final(axes[1, 2], DOUT, r"(f) $\mathit{dfs\_out}$ raster", framed=True, block=BLOCK) fig.subplots_adjust(left=0.02, right=0.98, top=0.94, bottom=0.03, wspace=0.12, hspace=0.20) for ext in ("png", "pdf"): out = os.path.join(OUT_DIR, f"Fig04_DFS_Interval_v11.{ext}") fig.savefig(out, dpi=300, bbox_inches="tight", facecolor="white") print(f"Saved: {out}") n_valid = int(msk2d.sum()) n_a = int((bsn2d == 1).sum() & msk2d.astype(bool).sum()) # info only n_b = int(((bsn2d == 2) & msk2d.astype(bool)).sum()) n_a = int(((bsn2d == 1) & msk2d.astype(bool)).sum()) print(f" valid pixels = {n_valid} (Basin A = {n_a}, Basin B = {n_b}, " f"invalid = {size - n_valid})") print(f" numbers 0..{len(events) - 1}; T1 = i{T1}: dfs_in {din[T1]}, upg {upa[T1]}, dfs_out {din[T1] + upa[T1] - 1}; " f"seq_dfs {seq[T1]}; outlets (dfs_in) {[(o, int(din[o])) for o in outs]}") plt.close(fig) if __name__ == "__main__": main()