diff --git a/CHANGELOG.md b/CHANGELOG.md index b0a70692..bc834b69 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -86,6 +86,21 @@ introduce breaking changes to configuration and outputs. ### Changed +- Model regions are now built **basin-aware**: GADM provinces are first split + along AWARE hydrological basin boundaries, and each country is partitioned + into regions balancing geography against basin scarcity + (`aggregation.regions.basin_scarcity_weight`, default 2.0; 0 recovers the + previous purely geographic clustering). A province straddling an abundant and + a scarce basin is no longer pooled into one region, which used to average away + exactly the sub-provincial scarcity that constrains irrigation. Every region + is still either contained in one province or a union of whole provinces, so + regions remain comparable to political units. The + `aggregation.regions.allow_cross_border` option is removed; regions never + cross country borders. **This changes default region geometry for + every config**, so all `processing/` artefacts are rebuilt and the tracked + calibration sets must be regenerated. Requires the AWARE2.0 basin geopackage + (new automatic download). + - The **GDD-IA baseline-diet dataset is now retrieved automatically** from Zenodo ([10.5281/zenodo.20818140](https://doi.org/10.5281/zenodo.20818140), CC-BY-4.0) instead of being obtained on personal request and placed under diff --git a/config/default.yaml b/config/default.yaml index 17f4c432..387fa796 100644 --- a/config/default.yaml +++ b/config/default.yaml @@ -1364,8 +1364,17 @@ health: aggregation: regions: target_count: 750 - allow_cross_border: false method: "kmeans" + # Basin-aware clustering: split provinces that straddle hydrological basins so + # a scarce sub-basin is not pooled with an abundant one (which averages away + # the scarcity that drives irrigation constraints). Each model region is either + # contained in one GADM province or a union of whole provinces (never partial + # pieces across provinces). basin_scarcity_weight sets how strongly the AWARE + # basin scarcity (CF) separates regions relative to geography, in standardised + # units; 0 recovers plain geographic clustering. Separation saturates around 2 + # (area-weighted within-region CF spread 10.3 -> 7.6, against a floor of ~7.2), + # while larger values only inflate the biggest regions. + basin_scarcity_weight: 2.0 simplify_tolerance_km: 5 simplify_min_area_km: 25 resource_class_quantiles: [0.25, 0.5, 0.75] diff --git a/config/schemas/config.schema.yaml b/config/schemas/config.schema.yaml index 256d16bd..d600e6e1 100644 --- a/config/schemas/config.schema.yaml +++ b/config/schemas/config.schema.yaml @@ -1324,20 +1324,21 @@ properties: properties: regions: type: object - required: [target_count, allow_cross_border, method] + required: [target_count, method, basin_scarcity_weight] additionalProperties: false properties: target_count: type: integer minimum: 1 description: "Target number of regions to create through clustering" - allow_cross_border: - type: boolean - description: "Whether to allow regions to cross country borders (false = cluster within each country separately)" method: type: string enum: [kmeans, agglomerative] description: "Clustering method: 'kmeans' (faster, default) or 'agglomerative' (hierarchical clustering with ward linkage)" + basin_scarcity_weight: + type: number + minimum: 0 + description: "Weight of AWARE basin scarcity relative to geography when splitting provinces into regions (0 = plain geographic clustering)" simplify_tolerance_km: type: number minimum: 0 diff --git a/docs/data_sources.rst b/docs/data_sources.rst index 16ea32bd..4efae5bf 100644 --- a/docs/data_sources.rst +++ b/docs/data_sources.rst @@ -189,6 +189,24 @@ CROPGRIDS v1.08 **Usage**: Fallback source of harvested area and current cropland footprint for crops listed in ``config["cropgrids_crops"]`` (e.g. ``apple``), which are not covered by GAEZ. The CROPGRIDS ``harvarea`` raster drives both ``baseline_area_mha`` and ``suitable_area`` for these crops (see ``build_crop_yields_cropgrids.py``); per-country FAOSTAT QCL yields (item 515 for apple, element 5419 hg/ha) supply the dry-matter yield, broadcast uniformly to every (region, resource_class) cell within each country. +AWARE2.0 +~~~~~~~~ + +**Provider**: Seitfudem, Berger, Mueller Schmied & Boulay (2025) + +**Description**: The updated water-scarcity characterisation method for life-cycle assessment, built on WaterGAP 2.2e. The model uses the native-basin geopackage: hydrological basin polygons carrying annual and monthly characterisation factors. The annual agricultural CF is the basin scarcity signal used to split provinces along basin boundaries when building model regions (see :doc:`land_use`). + +**Coverage**: + * Spatial: Global, 11661 native hydrological basins + +**Access**: Zenodo record 15133241 (https://doi.org/10.5281/zenodo.15133241), file ``AWARE20_Native_CFs_geospatial.gpkg`` (5.3 MB). + +**License**: Creative Commons Attribution 4.0 International (CC BY 4.0) + +**Citation**: Seitfudem, G., Berger, M., Mueller Schmied, H., & Boulay, A.-M. (2025). *The updated and improved method for water scarcity impact assessment in LCA, AWARE2.0*. https://doi.org/10.5281/zenodo.15133241 + +**Retrieval**: ``download_aware2_basins`` fetches the geopackage from Zenodo. Basins whose annual agricultural CF is missing (about a fifth of polygons, mostly ice, desert and small islands) are filled with the global median. + MIRCA-OS v2 ~~~~~~~~~~~ diff --git a/docs/land_use.rst b/docs/land_use.rst index cedab86e..01996152 100644 --- a/docs/land_use.rst +++ b/docs/land_use.rst @@ -79,12 +79,19 @@ The model operates at sub-national regional resolution, balancing spatial detail Regional Clustering ~~~~~~~~~~~~~~~~~~~ -Optimization regions are created by clustering administrative units (GADM level 1) based on spatial proximity: +Optimization regions are created by clustering administrative units (GADM level 1), +basin-aware so that hydrological scarcity is not averaged away: 1. **Simplification**: Simplify GADM geometries to reduce complexity while preserving boundaries 2. **Country selection**: Filter to configured countries (``countries`` list in config) -3. **Clustering**: Aggregate administrative units using k-means clustering on centroids -4. **Output**: GeoJSON with region polygons (``processing/{name}/regions.geojson``) +3. **Basin split**: Overlay provinces with the AWARE hydrological basins, giving one piece per + (province, basin), so a province straddling an abundant and a scarce basin can be separated +4. **Clustering**: Per country, allocate a region budget by area and partition into exactly + ``target_count`` regions, balancing geography and basin scarcity. Every region is either + contained in one province (large provinces are split into sub-regions) or a union of whole + provinces (small provinces are merged) -- never a mix of partial pieces across provinces, so + regions stay comparable to political units +5. **Output**: GeoJSON with region polygons (``processing/{name}/regions.geojson``) .. figure:: https://github.com/Sustainable-Solutions-Lab/GLADE/releases/download/doc-figures/intro_global_coverage.png :width: 100% @@ -95,7 +102,7 @@ Optimization regions are created by clustering administrative units (GADM level Key configuration parameters: - ``aggregation.regions.target_count``: Number of regions to create (default 750) -- ``aggregation.regions.allow_cross_border``: Whether regions can span country boundaries (typically ``false``) +- ``aggregation.regions.basin_scarcity_weight``: Influence of AWARE basin scarcity relative to geography when partitioning a country (default 2.0; 0 recovers plain geographic clustering) - ``aggregation.simplify_tolerance_km``: Geometry simplification tolerance Resource Classes diff --git a/tests/test_build_regions.py b/tests/test_build_regions.py new file mode 100644 index 00000000..4aa711d6 --- /dev/null +++ b/tests/test_build_regions.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: 2026 Koen van Greevenbroek +# +# SPDX-License-Identifier: GPL-3.0-or-later + +"""Unit tests for the basin-aware region clustering in build_regions.py. + +Covers the exact-count guarantee, the nesting invariant (each region is either +contained in one GADM province or a union of whole provinces), and scarcity-based +splitting -- all on the pure ``cluster_country`` path, which works on a plain +DataFrame of province-basin pieces (no geometry needed). +""" + +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +from shapely.geometry import Polygon, box + +from workflow.scripts.build_regions import ( + _largest_remainder, + cluster_country, + cluster_regions, +) + + +def _pieces() -> pd.DataFrame: + """One country: a large province P1 straddling an abundant (CF 5) and a + scarce (CF 90) basin, plus two small provinces P2, P3 in a mid basin.""" + return pd.DataFrame( + { + "prov": ["P1", "P1", "P2", "P3"], + "px": [1.0, 3.0, 4.5, 4.5], + "py": [1.0, 1.0, 0.5, 1.5], + "area": [4.0, 4.0, 1.0, 1.0], + "cf": [5.0, 90.0, 40.0, 40.0], + } + ) + + +def _nesting_holds(pieces: pd.DataFrame, labels: np.ndarray) -> bool: + df = pieces.assign(region=labels) + for _, d in df.groupby("region"): + provs = d["prov"].unique() + if len(provs) == 1: + continue # contained in one province -> OK + # union of provinces: each must be wholly inside this region + region = d["region"].iloc[0] + if not all(df.loc[df["prov"] == p, "region"].eq(region).all() for p in provs): + return False + return True + + +def test_largest_remainder_sums_to_total(): + q = pd.Series({"a": 3.2, "b": 0.4, "c": 0.4}, dtype=float) + out = _largest_remainder(q, 4) + assert out.sum() == 4 + assert out["a"] == 3 # floor kept + assert (out >= 0).all() + + +@pytest.mark.parametrize("k", [2, 3, 4]) +def test_cluster_country_exact_count_and_nesting(k): + pieces = _pieces() + labels = cluster_country( + pieces, k, scarcity_weight=5.0, method="kmeans", random_state=0 + ) + assert len(set(labels)) == k # exact count + assert _nesting_holds(pieces, labels) + + +def test_cluster_country_splits_province_by_scarcity(): + pieces = _pieces() + labels = cluster_country( + pieces, 4, scarcity_weight=5.0, method="kmeans", random_state=0 + ) + p1 = pieces.assign(region=labels) + p1 = p1[p1["prov"] == "P1"] + # P1's abundant and scarce pieces land in different regions. + assert p1["region"].nunique() == 2 + by_region_cf = p1.groupby("region")["cf"].mean() + assert by_region_cf.max() - by_region_cf.min() > 50.0 + + +def test_cluster_country_reconciles_when_pieces_scarce(): + # P1's area earns it two sub-regions, but it is a single indivisible piece, + # so the split regime under-produces; reconciliation must recover the + # deficit by splitting the merge group (P2's pieces) instead. + pieces = pd.DataFrame( + { + "prov": ["P1", "P2", "P2", "P2"], + "px": [0.0, 10.0, 12.0, 14.0], + "py": [0.0, 0.0, 0.0, 0.0], + "area": [10.0, 1.0, 1.0, 1.0], + "cf": [5.0, 10.0, 50.0, 90.0], + } + ) + labels = cluster_country( + pieces, 3, scarcity_weight=2.0, method="kmeans", random_state=0 + ) + df = pieces.assign(region=labels) + assert df["region"].nunique() == 3 + assert _nesting_holds(pieces, labels) + # The indivisible P1 stays one region; P2's pieces supply the other two. + assert df.loc[df["prov"] == "P1", "region"].nunique() == 1 + assert df.loc[df["prov"] == "P2", "region"].nunique() == 2 + + +def test_country_budget_uses_full_province_area(): + """A country's region budget follows its land area, not its basin fragmentation. + + ``pieces`` holds province-basin intersections, so a province straddling many + basins contributes many rows. Summing the area of one row per province would + starve exactly those countries whose provinces are the most fragmented. + """ + # Two countries of equal total area and equal piece capacity, differing only + # in how that area is split between provinces and basins: A is one province + # cut into 4 basin pieces, B is two provinces cut into 2 pieces each. + pieces = gpd.GeoDataFrame( + { + "GID_0": ["A", "A", "A", "A", "B", "B", "B", "B"], + "prov": ["A1", "A1", "A1", "A1", "B1", "B1", "B2", "B2"], + "cf": [5.0, 20.0, 40.0, 90.0, 5.0, 20.0, 40.0, 90.0], + }, + geometry=[ + box(0, 0, 1, 1), + box(1, 0, 2, 1), + box(0, 1, 1, 2), + box(1, 1, 2, 2), + box(10, 0, 11, 1), + box(11, 0, 12, 1), + box(10, 1, 11, 2), + box(11, 1, 12, 2), + ], + crs="EPSG:4326", + ) + regions = cluster_regions(pieces, 4, scarcity_weight=3.0) + per_country = regions["country"].value_counts() + assert ( + per_country["A"] == 2 + ), f"basin-fragmented country starved of regions: {per_country.to_dict()}" + assert per_country["B"] == 2 + + +def test_dissolved_regions_are_valid_geometries(): + """Region geometries must be valid: exactextract segfaults on invalid ones. + + Downstream raster aggregation calls into native code that crashes rather + than raising on a self-intersecting polygon, so one bad region takes out an + unrelated rule with an empty log. Dissolving province-basin pieces can + produce them, hence the repair. + """ + # An asymmetric bowtie: self-intersecting like real dissolve artefacts, and + # with positive area so it survives the area-weighted clustering. + bad = Polygon([(0, 0), (4, 4), (4, 0), (0, 3)]) + assert not bad.is_valid, "fixture must actually be invalid" + pieces = gpd.GeoDataFrame( + { + "GID_0": ["A", "A"], + "prov": ["P1", "P2"], + "cf": [5.0, 50.0], + }, + geometry=[bad, box(10, 0, 11, 1)], + crs="EPSG:4326", + ) + regions = cluster_regions(pieces, 2, scarcity_weight=3.0) + assert regions.geometry.is_valid.all(), "invalid geometry escaped build_regions" diff --git a/workflow/rules/geography.smk b/workflow/rules/geography.smk index d020167d..bafcf8b1 100644 --- a/workflow/rules/geography.smk +++ b/workflow/rules/geography.smk @@ -57,18 +57,19 @@ rule simplify_gadm: rule build_regions: input: world="/shared/gadm-simplified.gpkg", + basins="data/downloads/aware2/AWARE20_Native_CFs_geospatial.gpkg", params: n_regions=config["aggregation"]["regions"]["target_count"], - allow_cross_border=config["aggregation"]["regions"]["allow_cross_border"], cluster_method=config["aggregation"]["regions"]["method"], + basin_scarcity_weight=config["aggregation"]["regions"]["basin_scarcity_weight"], countries=config["countries"], output: "/{name}/regions.geojson", group: "prep" resources: - runtime="1m", - mem_mb=400, + runtime="15m", + mem_mb=6000, log: "/{name}/build_regions.log", benchmark: diff --git a/workflow/rules/retrieve.smk b/workflow/rules/retrieve.smk index 5e3dff52..60226b39 100644 --- a/workflow/rules/retrieve.smk +++ b/workflow/rules/retrieve.smk @@ -947,6 +947,36 @@ rule extract_mirca_os_subcrop_monthly: "../scripts/extract_mirca_os_archive.py" +rule download_aware2_basins: + """Download the AWARE2.0 native-basin polygons from Zenodo. + + Source: Seitfudem, Berger, Mueller Schmied & Boulay (2025), "The updated + and improved method for water scarcity impact assessment in LCA, AWARE2.0", + https://doi.org/10.5281/zenodo.15133241 (WaterGAP2.2e). License: CC BY 4.0. + + The geopackage carries the native basin geometries and their + characterisation factors; the annual agricultural CF is used as the basin + scarcity signal for basin-aware region clustering. + """ + output: + basins="data/downloads/aware2/AWARE20_Native_CFs_geospatial.gpkg", + params: + base_url="https://zenodo.org/records/15133241/files", + resources: + runtime="30m", + mem_mb=500, + log: + "/shared/download_aware2_basins.log", + benchmark: + "/shared/download_aware2_basins.tsv" + shell: + r""" + mkdir -p "$(dirname {output.basins})" + curl -L --fail --progress-bar -o "{output.basins}" \ + "{params.base_url}/AWARE20_Native_CFs_geospatial.gpkg/content" > {log} 2>&1 + """ + + rule download_grassland_yield_data: """Retrieve historical managed-grassland yield from ISIMIP2a / LPJmL. diff --git a/workflow/scripts/build_regions.py b/workflow/scripts/build_regions.py index 45cdc0c6..042cb05b 100644 --- a/workflow/scripts/build_regions.py +++ b/workflow/scripts/build_regions.py @@ -2,16 +2,44 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +"""Basin-aware clustering of GADM level-1 provinces into model regions. + +Provinces are first split along hydrological basin boundaries (overlay with the +AWARE basins), so a province straddling an abundant and a scarce basin can be +separated -- otherwise pooling the two averages away the sub-provincial water +scarcity that drives groundwater use. The province-basin pieces are then +clustered per country into exactly ``target_count`` regions under a nesting +invariant: every model region is either contained in one GADM province (a large +province is split into sub-regions) or a union of whole provinces (small +provinces are merged), never a mix of partial pieces across provinces -- so +regions stay cleanly comparable to political units. + +Within each country the regions are balanced on geography and basin scarcity +(``basin_scarcity_weight`` sets their relative influence; 0 recovers plain +geographic clustering). A reconciliation step splits the largest splittable +region until the exact target is reached, since a province allocated more +sub-regions than it has basin pieces would otherwise under-produce. +""" + +import logging from pathlib import Path import geopandas as gpd import numpy as np import pandas as pd from pyproj import CRS, Geod +import shapely from sklearn.cluster import AgglomerativeClustering, KMeans +from workflow.scripts.logging_config import setup_script_logging + +logger = logging.getLogger(__name__) + GEOD = Geod(ellps="WGS84") +# Coordinate grid the GeoJSON writer rounds to (GDAL's default 7 decimals). +GEOJSON_PRECISION = 1e-7 + def _compute_country_geodesic_areas( gdf_wgs84: gpd.GeoDataFrame, country_col: str = "GID_0" @@ -51,7 +79,7 @@ def _allocate_per_country_targets_by_weight( if total_target < len(nonempty): raise ValueError( "target_count is smaller than the number of countries with regions; " - "cannot avoid cross-border clustering with this target." + "every country needs at least one region." ) # Respect global capacity (cannot exceed number of base units) @@ -104,11 +132,18 @@ def _allocate_per_country_targets_by_weight( def _cluster_coords( - coords: np.ndarray, k: int, method: str, random_state: int = 0 + coords: np.ndarray, + k: int, + method: str, + random_state: int = 0, + weights: np.ndarray | None = None, ) -> np.ndarray: """Cluster coordinate array into up to k clusters. Returns one label per row. If k <= 0 or k >= n, assigns unique labels. + ``weights`` (areas) weight the samples under k-means, so a sliver piece + does not pull a centroid as hard as a half-province piece; agglomerative + clustering has no sample weighting and ignores them. """ if coords.shape[0] <= k or k <= 0: return np.arange(coords.shape[0]) @@ -117,126 +152,350 @@ def _cluster_coords( if method == "kmeans": km = KMeans(n_clusters=k, n_init=10, random_state=random_state) - labels = km.fit_predict(coords) - return labels + return km.fit_predict(coords, sample_weight=weights) elif method == "agglomerative": # Ward linkage minimizes within-cluster variance (good heuristic) ac = AgglomerativeClustering(n_clusters=k, linkage="ward") - labels = ac.fit_predict(coords) - return labels + return ac.fit_predict(coords) else: raise ValueError(f"Unknown clustering method: {method}") +def _largest_remainder(quota: pd.Series, total: int) -> pd.Series: + """Integer apportionment summing exactly to ``total`` (floor 0). + + ``quota`` sums (approximately) to ``total``; floors are taken and the + remaining seats distributed to the largest fractional remainders. + """ + base = np.floor(quota).astype(int) + remaining = int(total - base.sum()) + if remaining > 0: + top = (quota - np.floor(quota)).sort_values(ascending=False).index[:remaining] + base.loc[top] += 1 + elif remaining < 0: + # floor already exceeds total (numerical): drop from smallest positive + drop = base[base > 0].sort_values().index[:-remaining] + base.loc[drop] -= 1 + return base + + +def _scarcity_features( + px: np.ndarray, py: np.ndarray, cf: np.ndarray, weight: float +) -> np.ndarray: + """Standardised ``[x, y, weight * scarcity]`` feature matrix for clustering.""" + + def z(a: np.ndarray) -> np.ndarray: + a = np.asarray(a, dtype=float) + s = a.std() + return (a - a.mean()) / (s if s > 0 else 1.0) + + return np.column_stack([z(px), z(py), weight * z(cf)]) + + +def _province_centroids(pieces: pd.DataFrame) -> pd.DataFrame: + """Area-weighted ``px``, ``py``, ``cf`` per province, plus total ``area``.""" + + def agg(g: pd.DataFrame) -> pd.Series: + w = g["area"] + return pd.Series( + { + "px": np.average(g["px"], weights=w), + "py": np.average(g["py"], weights=w), + "cf": np.average(g["cf"], weights=w), + "area": w.sum(), + } + ) + + return pieces.groupby("prov").apply(agg, include_groups=False) + + +def cluster_country( + pieces: pd.DataFrame, + k: int, + scarcity_weight: float, + method: str, + random_state: int, +) -> np.ndarray: + """Assign a local region id (0..k-1) to each province-basin piece of a country. + + ``pieces`` has one row per (province, basin) piece with columns ``prov``, + ``px``, ``py`` (equal-area centroid), ``area`` and ``cf`` (basin scarcity). + Provinces large enough for >=2 regions are *split*: their pieces are + clustered (on geography + weighted scarcity) into sub-regions that stay + within the province. All remaining provinces are grouped into whole-province + regions. Every region is therefore either contained in one province or a + union of whole provinces, and the total is exactly ``k``. + """ + prov_area = pieces.groupby("prov")["area"].sum() + quota = k * prov_area / prov_area.sum() + n_p = _largest_remainder(quota, k) + + split = n_p[n_p >= 2] + split_total = int(split.sum()) + remaining = [p for p in prov_area.index if n_p[p] < 2] + k_rem = k - split_total + # Ensure remaining provinces have at least one merge group to land in. + if remaining and k_rem < 1: + biggest = split.idxmax() + n_p[biggest] -= 1 + split = n_p[n_p >= 2] + split_total = int(split.sum()) + remaining = [p for p in prov_area.index if n_p[p] < 2] + k_rem = k - split_total + + labels = pd.Series(-1, index=pieces.index, dtype=int) + next_id = 0 + + # Split regime: partition each large province's own pieces. + for prov, n_sub in split.items(): + d = pieces[pieces["prov"] == prov] + n_sub = min(int(n_sub), len(d)) + feat = _scarcity_features( + d["px"].to_numpy(), d["py"].to_numpy(), d["cf"].to_numpy(), scarcity_weight + ) + sub = _cluster_coords( + feat, n_sub, method, random_state, weights=d["area"].to_numpy() + ) + labels.loc[d.index] = sub + next_id + next_id += int(sub.max()) + 1 + + # Merge regime: group remaining whole provinces. + if remaining: + rem = pieces[pieces["prov"].isin(remaining)] + pcen = _province_centroids(rem) + feat = _scarcity_features( + pcen["px"].to_numpy(), + pcen["py"].to_numpy(), + pcen["cf"].to_numpy(), + scarcity_weight, + ) + kk = min(int(k_rem), len(pcen)) + grp = _cluster_coords( + feat, kk, method, random_state, weights=pcen["area"].to_numpy() + ) + prov_to_group = dict(zip(pcen.index, grp + next_id)) + labels.loc[rem.index] = rem["prov"].map(prov_to_group).to_numpy() + + # A province allocated more sub-regions than it has basin pieces, or a merge + # group short of its target, under-produces. Reconcile to exactly k by + # repeatedly splitting the largest splittable region (a single-province + # region splits by piece; a whole-province group splits by province), which + # preserves nesting and also improves size balance. + labels = _reconcile_to_k(pieces, labels, k, scarcity_weight, method, random_state) + return labels.to_numpy() + + +def _reconcile_to_k( + pieces: pd.DataFrame, + labels: pd.Series, + k: int, + scarcity_weight: float, + method: str, + random_state: int, +) -> pd.Series: + """Split the largest splittable region until exactly ``k`` regions exist.""" + labels = labels.copy() + area = pieces["area"] + current = pd.Index(labels.unique()) + next_id = int(labels.max()) + 1 + while len(current) < k: + # Rank regions by area; split the largest one that can be split. + region_area = area.groupby(labels).sum().sort_values(ascending=False) + split_done = False + for region in region_area.index: + members = labels.index[labels == region] + sub = pieces.loc[members] + provs = sub["prov"].unique() + if len(provs) > 1: + # Whole-province group: split provinces into two, keep them whole. + pcen = _province_centroids(sub) + feat = _scarcity_features( + pcen["px"].to_numpy(), + pcen["py"].to_numpy(), + pcen["cf"].to_numpy(), + scarcity_weight, + ) + two = _cluster_coords( + feat, 2, method, random_state, weights=pcen["area"].to_numpy() + ) + move = pcen.index[two == 1] + labels.loc[sub.index[sub["prov"].isin(move)]] = next_id + elif len(sub) > 1: + # Single province with multiple pieces: split pieces into two. + feat = _scarcity_features( + sub["px"].to_numpy(), + sub["py"].to_numpy(), + sub["cf"].to_numpy(), + scarcity_weight, + ) + two = _cluster_coords( + feat, 2, method, random_state, weights=sub["area"].to_numpy() + ) + labels.loc[sub.index[two == 1]] = next_id + else: + continue # single indivisible piece + next_id += 1 + current = pd.Index(labels.unique()) + split_done = True + break + if not split_done: + break # no region can be split further (pieces exhausted) + return labels + + def cluster_regions( - gdf: gpd.GeoDataFrame, + pieces: gpd.GeoDataFrame, target_count: int, - allow_cross_border: bool, + scarcity_weight: float, method: str = "kmeans", random_state: int = 0, ) -> gpd.GeoDataFrame: - """Cluster level-1 administrative regions into target_count clusters. - - Clustering is based on centroids in a projected CRS (EPSG:3857) for a - reasonable Euclidean approximation. When cross-border clustering is not - allowed, clustering is performed per country and the per-country targets - are allocated proportionally to the number of base regions. + """Basin-aware clustering of province-basin pieces into ``target_count`` regions. + + ``pieces`` is the overlay of GADM level-1 provinces with AWARE basins, with + columns ``GID_0`` (country), ``prov`` (GADM province id), ``cf`` (basin + scarcity) and geometry. Per-country region counts are allocated + proportionally to area; each country is then partitioned (see + ``cluster_country``) so every region is either contained in one province or a + union of whole provinces. Basin scarcity separates regions within a province + so a scarce sub-basin is not pooled with an abundant one. """ if target_count <= 0: raise ValueError("target_count must be positive") + if "GID_0" not in pieces.columns: + raise ValueError("Expected GID_0 column for country codes") - if gdf.crs is None: - gdf = gdf.set_crs(4326, allow_override=True) - - # Project to equal-area for reasonable Euclidean approximation - gdf_proj = gdf.to_crs(6933) - cent = gdf_proj.geometry.centroid - coords = np.vstack([cent.x.values, cent.y.values]).T + proj = pieces.to_crs(6933) + cent = proj.geometry.centroid + pieces = pieces.assign( + px=cent.x.to_numpy(), py=cent.y.to_numpy(), area=proj.geometry.area.to_numpy() + ) - if allow_cross_border: - labels = _cluster_coords(coords, target_count, method, random_state) - # Create global cluster ids - cluster_ids = pd.Series(labels, index=gdf.index).astype(int) - gdf = gdf.assign(_cluster=cluster_ids) - else: - # Allocate targets per country and cluster within each - if "GID_0" not in gdf.columns: - raise ValueError( - "Expected GID_0 column for country codes in GADM level 1 data" - ) - - counts = gdf.groupby("GID_0").size() - # Compute geodesic area per country for proportional allocation - country_areas = _compute_country_geodesic_areas(gdf[["GID_0", "geometry"]]) - per_country = _allocate_per_country_targets_by_weight( - country_areas, counts, target_count - ) + # Per-country region budget, proportional to area (>=1 per country). Both + # inputs are taken over *all* pieces, which partition each province: one + # piece per province would undercount the area of every province straddling + # several basins, starving basin-fragmented countries of regions. The + # capacity cap is the piece count, since a province can now be split into + # as many regions as it has basin pieces. + counts = pieces.groupby("GID_0").size() + country_areas = _compute_country_geodesic_areas(pieces[["GID_0", "geometry"]]) + per_country = _allocate_per_country_targets_by_weight( + country_areas, counts, target_count + ) - cluster_ids = pd.Series(index=gdf.index, dtype=int) - next_cluster = 0 - for country, group in gdf.groupby("GID_0"): - k = int(per_country.get(country, 0)) - idx = group.index.values - if k <= 0: - # No clusters allocated (can happen if country had 0 units) - continue - sub_coords = coords[gdf.index.get_indexer(idx)] - labels = _cluster_coords(sub_coords, k, method, random_state) - # Offset labels to keep them globally unique - cluster_ids.loc[idx] = labels + next_cluster - next_cluster += int(labels.max()) + 1 - - gdf = gdf.assign(_cluster=cluster_ids.astype(int)) - - # Dissolve polygons by cluster id - # Keep only geometry and the minimal attributes needed to avoid duplicate columns. - keep_cols = [c for c in ["_cluster", "geometry"] if c in gdf.columns] - gdf_min = gdf[keep_cols].copy() - dissolved = gdf_min.dissolve(by="_cluster", as_index=False) - - # Assign unique region identifiers + cluster_ids = pd.Series(-1, index=pieces.index, dtype=int) + next_cluster = 0 + for country, group in pieces.groupby("GID_0"): + k = int(per_country.get(country, 0)) + if k <= 0: + continue + local = cluster_country(group, k, scarcity_weight, method, random_state) + cluster_ids.loc[group.index] = local + next_cluster + next_cluster += int(local.max()) + 1 + + pieces = pieces.assign(_cluster=cluster_ids.astype(int)) + dissolved = pieces[["_cluster", "geometry"]].dissolve(by="_cluster", as_index=False) dissolved["region"] = [f"region{int(i):04d}" for i in dissolved["_cluster"]] - - # Keep a representative country code if available, under a user-friendly name - if "GID_0" in gdf.columns: - # first country code within the cluster as representative - rep_country = gdf.groupby("_cluster")["GID_0"].first() - dissolved["country"] = dissolved["_cluster"].map(rep_country) - - dissolved = dissolved.set_index("region") - dissolved = dissolved.drop( - columns=[c for c in ["_cluster"] if c in dissolved.columns] + rep_country = pieces.groupby("_cluster")["GID_0"].first() + dissolved["country"] = dissolved["_cluster"].map(rep_country) + dissolved = dissolved.set_index("region").drop(columns="_cluster") + + # Repair, then snap to the GeoJSON writer's coordinate grid. Order matters. + # The raw dissolve leaves near-degenerate slivers that are valid in memory + # but tip into self-intersection once the writer truncates coordinates to 7 + # decimals, so an in-memory validity check passes while the file on disk is + # invalid. Snapping makes the in-memory geometry exactly what gets written, + # and set_precision's valid_output mode guarantees the snapped result is + # valid -- but it raises on genuinely invalid input, hence buffer(0) first. + # This matters because exactextract calls into native code that *segfaults* + # on invalid geometry rather than raising -- one bad region kills an + # unrelated downstream rule with an empty log and no traceback. + repaired = dissolved.geometry.buffer(0) + snapped = shapely.set_precision(repaired.values, GEOJSON_PRECISION) + dissolved["geometry"] = gpd.GeoSeries( + snapped, index=dissolved.index, crs=dissolved.crs ) + # Empty geometries are "valid", but a region collapsed to nothing by the + # repair or snap is just as fatal downstream as an invalid one. + still_bad = ~dissolved.geometry.is_valid | dissolved.geometry.is_empty + if still_bad.any(): + raise ValueError( + f"{int(still_bad.sum())} region geometries are invalid or empty " + "after normalisation: " + + ", ".join(map(str, dissolved.index[still_bad][:5])) + ) return dissolved if __name__ == "__main__": - # Use GADM level 1 (state/province boundaries). Geometries are pre-simplified. - gdf = gpd.read_file(snakemake.input.world) - - # Filter out invalid regions - gdf = gdf.rename({"GID_1": "region"}, axis=1) + logger = setup_script_logging(log_file=snakemake.log[0] if snakemake.log else None) + # GADM level 1 (state/province boundaries); geometries are pre-simplified. + provinces = gpd.read_file(snakemake.input.world).rename(columns={"GID_1": "prov"}) valid_mask = ( - (gdf["region"] != "?") - & gdf["region"].notna() - & (gdf["region"] != "") - & (gdf["region"] != "NA") + (provinces["prov"] != "?") + & provinces["prov"].notna() + & (provinces["prov"] != "") + & (provinces["prov"] != "NA") ) - gdf = gdf[valid_mask] + provinces = provinces[valid_mask] + if "GID_0" not in provinces.columns: + raise ValueError("Expected GID_0 column with ISO3 country codes in GADM data") + provinces = provinces[provinces["GID_0"].isin(list(snakemake.params.countries))] + if provinces.crs is None: + provinces = provinces.set_crs(4326, allow_override=True) - gdf = gdf.set_index("region", drop=True) + # AWARE basin scarcity (agricultural CF); fill missing with the median. + basins = gpd.read_file( + snakemake.input.basins, layer="AWARE20_Native_CFs_geospatial" + ) + basins["cf"] = pd.to_numeric(basins["CF_annual_agri"], errors="coerce") + basins["cf"] = basins["cf"].fillna(basins["cf"].median()) + basins = basins.to_crs(provinces.crs) + + # Split provinces along basin boundaries: one piece per (province, basin). + pieces = gpd.overlay( + provinces[["GID_0", "prov", "geometry"]], + basins[["cf", "geometry"]], + how="intersection", + keep_geom_type=True, + ) - # Narrowing by configured countries - if "GID_0" not in gdf.columns: - raise ValueError("Expected GID_0 column with ISO3 country codes in GADM data") - gdf = gdf[gdf["GID_0"].isin(list(snakemake.params.countries))] + # Coastal slivers and small islands can fall outside every AWARE basin and + # are dropped by the intersection. That is tolerable in aggregate but must + # never silently remove a whole country from the model. + missing = set(provinces["GID_0"]) - set(pieces["GID_0"]) + if missing: + raise ValueError( + "No AWARE basin overlaps any province of: " + f"{', '.join(sorted(missing))}. Check the basin geopackage." + ) + kept = _compute_country_geodesic_areas(pieces[["GID_0", "geometry"]]).sum() + total = _compute_country_geodesic_areas(provinces[["GID_0", "geometry"]]).sum() + logger.info( + "Basin overlay: %d provinces -> %d province-basin pieces, " + "%.2f%% of land area outside all basins", + provinces["prov"].nunique(), + len(pieces), + 100.0 * (1.0 - kept / total), + ) - gdf = cluster_regions( - gdf, + regions = cluster_regions( + pieces, snakemake.params.n_regions, - snakemake.params.allow_cross_border, + snakemake.params.basin_scarcity_weight, snakemake.params.cluster_method, ) Path(snakemake.output[0]).parent.mkdir(parents=True, exist_ok=True) - gdf.to_file(snakemake.output[0], driver="GeoJSON") + regions.to_file(snakemake.output[0], driver="GeoJSON") + + # The written file is what every downstream raster aggregation reads, so + # validate it rather than the in-memory frame. + written = gpd.read_file(snakemake.output[0]) + bad = ~written.geometry.is_valid | written.geometry.is_empty + if bad.any(): + raise ValueError( + f"{int(bad.sum())} region geometries are invalid or empty as written " + f"to {snakemake.output[0]}; exactextract would segfault on them." + ) + logger.info("Wrote %d regions, all geometries valid on disk", len(written))