#!/usr/bin/env python3 """Plot drought-index forecast maps from Dr_Prob.nc using dr_ens. Main features ------------- * GrADS-like grid-cell filling (``gxout grfill`` style) via pcolormesh. * Three output figures, one for each drought timescale stored in ``timescale`` (expected values: 1, 3, 6 months). * Each output figure contains 6 columns for the 6 lead months. * Vietnam national outline and the 7 climate-region boundaries are drawn from Vietnam7.shp. * The sea label "Biển Đông" is added to every panel. * The title uses the full initialization date dd-mm-yyyy. Expected NetCDF structure ------------------------- dr_ens(lon, lat, timescale, lead) lon(lon), lat(lat), timescale(timescale), lead(lead) """ from __future__ import annotations import argparse import calendar import sys from dataclasses import dataclass from pathlib import Path from typing import Iterable import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib.colors import BoundaryNorm, ListedColormap from matplotlib.ticker import FuncFormatter, FixedLocator import numpy as np import shapefile # pyshp import xarray as xr try: import shapely from shapely.geometry import shape as shapely_shape from shapely.ops import unary_union except Exception as exc: # pragma: no cover raise SystemExit( "This script requires shapely (>=2 recommended). " f"Import failed: {exc}" ) # ----------------------------------------------------------------------------- # Drought color settings requested by the user. # Center interval [-0.5, 0.5] is white / no color. # ----------------------------------------------------------------------------- DR_BOUNDS = np.array([-9.0, -2.0, -1.5, -1.0, -0.5, 0.5, 1.0, 1.5, 2.0, 9.0], dtype=float) DR_TICKS = np.array([-2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0], dtype=float) DR_COLORS = [ "darkred", # <= -2 "red", # (-2, -1.5] "orange", # (-1.5, -1] "yellow", # (-1, -0.5] "white", # (-0.5, 0.5) "#39dbe0", # [0.5, 1) "#b5d1dc", # [1, 1.5) "blue", # [1.5, 2) "darkblue", # >= 2 ] CMAP_DR = ListedColormap(DR_COLORS, name="drought_discrete") NORM_DR = BoundaryNorm(DR_BOUNDS, ncolors=CMAP_DR.N, clip=False) @dataclass(frozen=True) class RegionGeometry: parts: list union: object bounds: tuple[float, float, float, float] def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser( description="Plot drought-index forecast maps from Dr_Prob.nc using dr_ens." ) p.add_argument("--year", type=int, required=True, help="Initial year, e.g. 2026") p.add_argument("--month", type=int, required=True, choices=range(1, 13), metavar="1..12", help="Initial month, e.g. 7") p.add_argument("--day", type=int, required=True, help="Initial day, e.g. 1") p.add_argument("--input", required=True, help="Input NetCDF file, e.g. Dr_Prob.nc") p.add_argument("--output-dir", required=True, help="Output directory, e.g. /home/.../PDF/20260701") p.add_argument("--shape", required=True, help="Climate-region shapefile, e.g. VN_SubRegs/Vietnam7.shp") p.add_argument("--nlead", type=int, default=6, help="Number of lead months to plot (default: 6)") p.add_argument("--lon-min", type=float, default=101.0) p.add_argument("--lon-max", type=float, default=111.0) p.add_argument("--lat-min", type=float, default=7.5) p.add_argument("--lat-max", type=float, default=24.0) p.add_argument("--dpi", type=int, default=200) p.add_argument("--mask-country", action="store_true", help=("Optionally mask grid-cell centres outside the union of Vietnam7 polygons. " "Default is OFF to preserve the original grid cells like GrADS grfill.")) return p.parse_args() def add_months(year: int, month: int, delta: int) -> tuple[int, int]: idx = year * 12 + (month - 1) + delta return idx // 12, idx % 12 + 1 def _repair_geom(g): if g.is_empty: return g if not g.is_valid: try: return shapely.make_valid(g) except Exception: return g.buffer(0) return g def load_regions(shp_path: str) -> RegionGeometry: shp = Path(shp_path) if shp.suffix.lower() != ".shp": shp = shp.with_suffix(".shp") if not shp.exists(): raise FileNotFoundError(f"Shapefile not found: {shp}") reader = shapefile.Reader(str(shp)) parts = [] for rec_shape in reader.shapes(): g = _repair_geom(shapely_shape(rec_shape.__geo_interface__)) if not g.is_empty: parts.append(g) if not parts: raise ValueError(f"No polygon geometry found in {shp}") country = _repair_geom(unary_union(parts)) bounds = tuple(float(v) for v in country.bounds) minx, miny, maxx, maxy = bounds if not (90 <= minx <= 120 and 90 <= maxx <= 130 and 0 <= miny <= 30 and 0 <= maxy <= 35): raise ValueError( "Vietnam7.shp does not appear to be in geographic lon/lat coordinates. " f"Detected bounds: {bounds}. Reproject it to lon/lat first." ) return RegionGeometry(parts=parts, union=country, bounds=bounds) def iter_polygons(geom) -> Iterable: if geom.is_empty: return gt = geom.geom_type if gt == "Polygon": yield geom elif gt in ("MultiPolygon", "GeometryCollection"): for child in geom.geoms: yield from iter_polygons(child) def draw_polygon_boundary(ax, geom, linewidth=0.55, zorder=7): for poly in iter_polygons(geom): x, y = poly.exterior.xy ax.plot(x, y, color="black", linewidth=linewidth, zorder=zorder) for ring in poly.interiors: x, y = ring.xy ax.plot(x, y, color="black", linewidth=linewidth, zorder=zorder) def draw_boundaries(ax, regions: RegionGeometry): for g in regions.parts: draw_polygon_boundary(ax, g, linewidth=0.45, zorder=7) draw_polygon_boundary(ax, regions.union, linewidth=0.9, zorder=8) def mask_to_country(z: np.ndarray, lon: np.ndarray, lat: np.ndarray, country_geom) -> np.ndarray: xx, yy = np.meshgrid(lon, lat) try: inside = shapely.intersects_xy(country_geom, xx, yy) except Exception: inside = shapely.contains_xy(country_geom, xx, yy) return np.where(inside, z, np.nan) def centers_to_edges(coord: np.ndarray) -> np.ndarray: c = np.asarray(coord, dtype=float) if c.ndim != 1 or c.size < 2: raise ValueError("Coordinate must be a 1-D array with at least two points") mid = 0.5 * (c[:-1] + c[1:]) edges = np.empty(c.size + 1, dtype=float) edges[1:-1] = mid edges[0] = c[0] - 0.5 * (c[1] - c[0]) edges[-1] = c[-1] + 0.5 * (c[-1] - c[-2]) return edges def get_2d(ds: xr.Dataset, varname: str, ts_index: int, lead_index: int) -> np.ndarray: if varname not in ds: raise KeyError(f"Variable '{varname}' not found. Available: {list(ds.data_vars)}") da = ds[varname] needed = {"lon", "lat", "timescale", "lead"} if not needed.issubset(da.dims): raise ValueError(f"Variable '{varname}' must contain dims {needed}; got {da.dims}") da = da.isel(timescale=ts_index, lead=lead_index) da = da.transpose("lat", "lon") arr = np.asarray(da.values, dtype=float) for attr in ("_FillValue", "missing_value"): val = da.attrs.get(attr) if val is not None and np.isfinite(val): arr[np.isclose(arr, float(val))] = np.nan arr[~np.isfinite(arr)] = np.nan return arr def lon_formatter(x, pos=None): return f"{abs(x):g}°{'E' if x >= 0 else 'W'}" def lat_formatter(y, pos=None): return f"{abs(y):g}°{'N' if y >= 0 else 'S'}" def setup_axis(ax, args, show_y: bool): ax.set_xlim(args.lon_min, args.lon_max) ax.set_ylim(args.lat_min, args.lat_max) ax.set_aspect("equal", adjustable="box") ax.xaxis.set_major_locator(FixedLocator([102, 104, 106, 108, 110])) ax.xaxis.set_major_formatter(FuncFormatter(lon_formatter)) ax.yaxis.set_major_locator(FixedLocator(np.arange(8, 25, 2))) ax.yaxis.set_major_formatter(FuncFormatter(lat_formatter)) if not show_y: ax.set_yticklabels([]) ax.tick_params(axis="both", labelsize=7, length=2.5, pad=1.5) for spine in ax.spines.values(): spine.set_linewidth(0.8) spine.set_color("black") def plot_panel(ax, lon, lat, z, args, regions, month_label, show_y): setup_axis(ax, args, show_y) lon_edge = centers_to_edges(lon) lat_edge = centers_to_edges(lat) pm = ax.pcolormesh( lon_edge, lat_edge, z, cmap=CMAP_DR, norm=NORM_DR, shading="flat", edgecolors="none", linewidth=0.0, antialiased=False, zorder=1, ) # Optional light contour hints around key thresholds, similar in spirit to the example image. try: ax.contour(lon, lat, z, levels=[-0.5, 0.5], colors=["#4a78ff", "#4a78ff"], linewidths=0.45, alpha=0.9, zorder=3) except Exception: pass draw_boundaries(ax, regions) ax.text( 109.15, 17.45, "Biển Đông", rotation=-55, rotation_mode="anchor", ha="center", va="center", fontsize=7.2, fontweight="bold", color="black", zorder=10, clip_on=False, ) ax.set_title(month_label, fontsize=8.3, pad=3) return pm def make_figure(ds: xr.Dataset, lon: np.ndarray, lat: np.ndarray, regions: RegionGeometry, args: argparse.Namespace, ts_index: int, ts_value: int): nlead = args.nlead fig = plt.figure(figsize=(12.8, 4.6)) gs = fig.add_gridspec( 1, nlead, left=0.045, right=0.985, bottom=0.16, top=0.85, wspace=0.02, hspace=0.0, ) axes = [fig.add_subplot(gs[0, c]) for c in range(nlead)] last_pm = None for i in range(nlead): ty, tm = add_months(args.year, args.month, i + 1) target_label = f"{tm}-{ty}" z = get_2d(ds, "dr_ens", ts_index, i) if args.mask_country: z = mask_to_country(z, lon, lat, regions.union) last_pm = plot_panel(axes[i], lon, lat, z, args, regions, target_label, show_y=(i == 0)) init_label = f"{args.day:02d}-{args.month:02d}-{args.year:04d}" fig.text(0.50, 0.91, f"Ensemble median drought index ({ts_value}-month timescale) initialized at {init_label}", ha="center", va="center", fontsize=10.5) cax = fig.add_axes([0.22, 0.075, 0.56, 0.038]) cb = fig.colorbar(last_pm, cax=cax, orientation="horizontal", extend="both") cb.set_ticks(DR_TICKS) cb.ax.tick_params(labelsize=7, length=2, pad=1.5) cb.outline.set_linewidth(0.7) cb.ax.set_xticklabels(["-2", "-1.5", "-1", "-0.5", "0", "0.5", "1", "1.5", "2"]) out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) out = out_dir / f"Dr_ens_{ts_value}month.png" fig.savefig(out, dpi=args.dpi, bbox_inches="tight", facecolor="white") plt.close(fig) print(f"Wrote: {out}") def validate_input(ds: xr.Dataset, args: argparse.Namespace): for coord in ("lon", "lat", "timescale", "lead"): if coord not in ds: raise KeyError(f"Coordinate '{coord}' not found in input NetCDF") if "dr_ens" not in ds: raise KeyError("Required variable 'dr_ens' not found in input NetCDF") da = ds["dr_ens"] needed = {"lon", "lat", "timescale", "lead"} if not needed.issubset(da.dims): raise ValueError(f"Variable 'dr_ens' must contain dims {needed}; got {da.dims}") nlead_file = int(ds.sizes.get("lead", 0)) nts_file = int(ds.sizes.get("timescale", 0)) if nlead_file < args.nlead: raise ValueError(f"Input contains only {nlead_file} lead(s), but --nlead={args.nlead}") if nts_file < 1: raise ValueError("No timescale found in input NetCDF") def main() -> int: args = parse_args() _, ndays = calendar.monthrange(args.year, args.month) if not 1 <= args.day <= ndays: raise ValueError(f"Invalid initial date: {args.year}-{args.month:02d}-{args.day:02d}") inp = Path(args.input) if not inp.exists(): raise FileNotFoundError(f"Input NetCDF not found: {inp}") regions = load_regions(args.shape) with xr.open_dataset(inp, mask_and_scale=True, decode_times=True) as ds: validate_input(ds, args) lon = np.asarray(ds["lon"].values, dtype=float) lat = np.asarray(ds["lat"].values, dtype=float) if np.any(np.diff(lon) < 0): ds = ds.sortby("lon") lon = np.asarray(ds["lon"].values, dtype=float) if np.any(np.diff(lat) < 0): ds = ds.sortby("lat") lat = np.asarray(ds["lat"].values, dtype=float) ts_values = np.asarray(ds["timescale"].values) ts_values = [int(round(float(v))) for v in ts_values] for its, ts in enumerate(ts_values): make_figure(ds, lon, lat, regions, args, its, ts) return 0 if __name__ == "__main__": try: sys.exit(main()) except Exception as exc: print(f"ERROR: {exc}", file=sys.stderr) sys.exit(2)