feat: basin-aware region clustering

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

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]

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

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
 ~~~~~~~~~~~
 

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

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"

workflow/rules/geography.smk

@@ -57,18 +57,19 @@ rule simplify_gadm:
 rule build_regions:
     input:
         world="<processing>/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:
         "<processing>/{name}/regions.geojson",
     group:
         "prep"
     resources:
-        runtime="1m",
-        mem_mb=400,
+        runtime="15m",
+        mem_mb=6000,
     log:
         "<logs>/{name}/build_regions.log",
     benchmark:

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:
+        "<logs>/shared/download_aware2_basins.log",
+    benchmark:
+        "<benchmarks>/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.