#!/usr/bin/env python # -*- coding: utf-8 -*- """Figure 2: (a) the global basin mask with the 30 largest basins labelled; (b) example rows of the Basin Attribute Table. Input: ../Data/bsn_global_060th.gpkg, Fig02_bat_rows_v11.csv, Fig02_top30_v11.csv. Output: ../Figures/Fig02_GlobalBSN_BAT_v11.png and .pdf. Run from this folder; ../Data holds the inputs and the figure is written to ../Figures. """ from pathlib import Path import time import numpy as np import pandas as pd import rasterio import geopandas as gpd import matplotlib as mpl import matplotlib.pyplot as plt import matplotlib.colors as mcolors import matplotlib.patheffects as pe from matplotlib.colors import ListedColormap # ============================================================================ # Paths # ============================================================================ HERE = Path(__file__).resolve().parent VECTOR = HERE.parent / "Data/bsn_global_060th.gpkg" ROWS = HERE.parent / "Data/Fig02_bat_rows_v11.csv" TOP30 = HERE.parent / "Data/Fig02_top30_v11.csv" OUT_DIR = HERE.parent / "Figures" # ============================================================================ # Helper — distinct soft colours (mirrors user's ipynb) # ============================================================================ def distinct_colors(n: int, s: float = 0.65, v: float = 0.85) -> np.ndarray: hues = np.linspace(0, 1, n, endpoint=False) hsv = np.stack([hues, np.full_like(hues, s), np.full_like(hues, v)], axis=1) return mcolors.hsv_to_rgb(hsv) # ============================================================================ # Panel (a) # ============================================================================ def draw_panel_a(ax, gdf, top30): t0 = time.time() MOD_COLORS = 100 cmap = ListedColormap(distinct_colors(MOD_COLORS, s=0.65, v=0.85)) gdf_fill = gdf.copy() gdf_fill["_mod"] = gdf_fill["global_basin_id"].astype(int) % MOD_COLORS # rasterized: the 238,076 polygons are drawn as an image in the PDF (55 MB as vectors); # text, axes and the table stay vector gdf_fill.plot(ax=ax, column="_mod", cmap=cmap, edgecolor="none", linewidth=0, rasterized=True) print(f" filled {len(gdf_fill):,} basins by mod-{MOD_COLORS}: {time.time()-t0:.1f}s") # Basin boundary overlay (44 MB polygonized geojson) t0 = time.time() gdf.boundary.plot(ax=ax, linewidth=0.2, color="k", rasterized=True) print(f" boundary overlay: {time.time()-t0:.1f}s") # Axis style — match user's ipynb ax.set_xlabel("Longitude", fontsize=20, fontweight="bold") ax.set_ylabel("Latitude", fontsize=20, fontweight="bold") ax.set_xlim(-180, 180) ax.set_ylim(-60, 85) ax.set_xticks(np.arange(-180, 181, 60)) ax.set_yticks(np.arange(-60, 86, 30)) ax.set_xticks(np.arange(-180, 181, 30), minor=True) ax.set_yticks(np.arange(-60, 86, 15), minor=True) ax.grid(which="major", linestyle="--", linewidth=0.4, alpha=0.6) ax.grid(which="minor", linestyle=":", linewidth=0.25, alpha=0.4) ax.tick_params(axis="both", which="major", labelsize=18, length=6, width=1) ax.tick_params(axis="both", which="minor", labelsize=18, length=3, width=0.8) # Top-30 labels at representative_point want = set(int(g) for g in top30["global_basin_id"]) gdf_lab = gdf[gdf["global_basin_id"].isin(want)].copy() for _, row in gdf_lab.iterrows(): rp = row.geometry.representative_point() ax.text( rp.x, rp.y, str(int(row["global_basin_id"])), fontsize=18, color="white", ha="center", va="center", fontweight="bold", path_effects=[pe.withStroke(linewidth=1.5, foreground="black")], ) # (a) title is drawn via fig.text() in main(), not ax.set_title # ============================================================================ # Panel (b) # ============================================================================ ELL = "…" def fetch_bat_picks() -> pd.DataFrame: """The eight rows, as written by Data_Fig02_Fig03_v11.py: basins 1 and 2, the first endorheic basin, the first basin of region 10001 and of island region 20101, the example basin 8595, and the last basins of 1 km2 and of 0.05 km2 or more.""" df = pd.read_csv(ROWS) assert len(df) == 8 return df.sort_values("global_basin_id").reset_index(drop=True) # 14 logical columns, mirror the original 20260226 draft DISPLAY_COLS = [ "bat_index", "region_id", "local_basin_id", "global_basin_id", "_ell1", "outlet_dir", "basin_cell_count", "basin_area_km2", "basin_nrow", "basin_ncol", "global_basin_col_min", "_ell2", "local_basin_col_min", "_ell3", ] # Relative column widths — chosen so the widest body strings # ("698,585,522", "5904430.00", "2,569,927", etc.) fit without bleeding # into neighbouring cells. Sum = 1.00. COL_WIDTHS = [ 0.07, # bat_index e.g. 2,569,926 0.06, # region_id e.g. 10,001 0.08, # local_basin_id e.g. 12,877 0.08, # global_basin_id e.g. 2,569,927 0.04, # _ell1 0.06, # outlet_dir 0 / 255 0.11, # basin_cell_count e.g. 698,585,522 0.08, # basin_area_km2 e.g. 5904430.00 0.06, # basin_nrow 0.06, # basin_ncol 0.11, # global_basin_col_min e.g. 270,183 / -5,202 (antimerid) 0.04, # _ell2 0.11, # local_basin_col_min e.g. 32,884 0.04, # _ell3 ] assert abs(sum(COL_WIDTHS) - 1.0) < 1e-9 assert len(COL_WIDTHS) == len(DISPLAY_COLS) def fmt_val(v, is_decimal: bool = False) -> str: if isinstance(v, str): return v try: return f"{float(v):.2f}" if is_decimal else f"{int(round(v)):,}" except Exception: return str(v) def wrap_at_first_underscore(name: str) -> str: """Mirrors user's ipynb: 'global_basin_row_min' → 'global\\nbasin_row_min'.""" if name.startswith("_ell"): return ELL if "_" not in name: return name a, b = name.split("_", 1) return f"{a}\n{b}" def draw_panel_b(ax, df_b): """Use plt.table() API (mirrors user's Figureb.ipynb).""" # --- Build display rows: 8 data records + 6 ellipsis rows --- body = [] for k in range(len(df_b)): if k > 0: prev_g = int(df_b.iloc[k-1]["global_basin_id"]) cur_g = int(df_b.iloc[k]["global_basin_id"]) if cur_g - prev_g > 1: body.append({c: ELL for c in DISPLAY_COLS}) rec = {} for c in DISPLAY_COLS: if c.startswith("_ell"): rec[c] = ELL elif c == "basin_area_km2": rec[c] = fmt_val(df_b.iloc[k][c], is_decimal=True) else: rec[c] = fmt_val(df_b.iloc[k][c]) body.append(rec) body.append({c: ELL for c in DISPLAY_COLS}) # trailing ellipsis row cell_text = [[r[c] for c in DISPLAY_COLS] for r in body] col_labels = [wrap_at_first_underscore(c) for c in DISPLAY_COLS] ax.axis("off") table = ax.table( cellText=cell_text, colLabels=col_labels, colWidths=COL_WIDTHS, # explicit per-column widths (sum=1) bbox=[0.0, 0.0, 1.0, 1.0], # fill the (b) axes exactly — no overflow cellLoc="center", colLoc="center", ) table.auto_set_font_size(False) # Style each cell — mirrors user's Figureb.ipynb for (r, c), cell in table._cells.items(): cell.set_edgecolor("black") cell.set_linewidth(0.6) if r == 0: # header cell.set_facecolor("#1F4E79") cell.get_text().set_color("white") cell.get_text().set_weight("bold") cell.get_text().set_fontsize(16) cell._text.set_multialignment("center") else: cell.get_text().set_fontsize(16) cell.set_facecolor("#f7f7f7" if r % 2 == 0 else "white") # Double-height header row (per user's ipynb). With bbox=[0,0,1,1] # cell heights sum to 1; redistribute so header takes 2 units and each # body row takes 1 unit out of (n_body + 2). n_rows = 1 + len(cell_text) # header + body rows body_h = 1.0 / (n_rows + 1) # body rows: 1 unit each head_h = 2.0 * body_h # header: 2 units for (r, c), cell in table._cells.items(): cell.set_height(head_h if r == 0 else body_h) # (b) title is drawn via fig.text() in main() # ============================================================================ # Main # ============================================================================ def main(): # --- data --- print("Loading inputs …") t0 = time.time() gdf = gpd.read_file(VECTOR) print(f" vector: {len(gdf):,} basins, {time.time()-t0:.1f}s") top30 = pd.read_csv(TOP30) df_b = fetch_bat_picks() # --- figure: (a) on top, (b) below --- # Vertical budget (inches, bottom-up): # 0.05 bottom margin # 4.60 (b) axes (table fills exactly via bbox) # 1.00 gap — "(b)" label sits in the middle of this gap # 5.80 (a) axes (world map + axis labels) # 0.55 top margin for "(a)" label fig_w_in = 18.0 bot_pad_in = 0.05 b_h_in = 4.60 gap_in = 1.00 a_h_in = 5.80 top_pad_in = 0.55 fig_h_in = bot_pad_in + b_h_in + gap_in + a_h_in + top_pad_in # 12.0 b_bot_in = bot_pad_in b_top_in = b_bot_in + b_h_in # 4.65 a_bot_in = b_top_in + gap_in # 5.65 a_top_in = a_bot_in + a_h_in # 11.45 fig = plt.figure(figsize=(fig_w_in, fig_h_in)) a_left, a_w = 0.06, 0.88 ax_a = fig.add_axes([ a_left, a_bot_in / fig_h_in, a_w, (a_top_in - a_bot_in) / fig_h_in, ]) b_left, b_w = 0.02, 0.96 ax_b = fig.add_axes([ b_left, b_bot_in / fig_h_in, b_w, (b_top_in - b_bot_in) / fig_h_in, ]) print("\nDraw panel (a) …") draw_panel_a(ax_a, gdf, top30) print("\nDraw panel (b) …") draw_panel_b(ax_b, df_b) # --- panel labels at fixed figure positions --- # "(a)" sits 0.10 in above the (a) axes top (still inside top_pad_in). fig.text(a_left, (a_top_in + 0.10) / fig_h_in, "(a)", fontsize=30, fontweight="bold", va="bottom") # "(b)" sits in the middle of the (a)-(b) gap, well clear of both. b_label_y_in = b_top_in + gap_in * 0.30 # 0.30 in above table top fig.text(b_left, b_label_y_in / fig_h_in, "(b)", fontsize=30, fontweight="bold", va="bottom") OUT_DIR.mkdir(parents=True, exist_ok=True) for ext in ("png", "pdf"): p = OUT_DIR / f"Fig02_GlobalBSN_BAT_v11.{ext}" fig.savefig(p, dpi=300, bbox_inches="tight") print(f"\n wrote {p}") print(f" fig {fig_w_in:.1f} × {fig_h_in:.2f} in") plt.close(fig) if __name__ == "__main__": main()