From 8bb20e10db86837240fb634c2ac6fbda77d4bb34 Mon Sep 17 00:00:00 2001
From: Koen van Greevenbroek <koenvg@stanford.edu>
Date: Mon, 20 Jul 2026 17:41:57 -0700
Subject: [PATCH 1/6] feat: basin-aware region clustering

Split GADM provinces along AWARE hydrological basin boundaries before
clustering, so a province straddling an abundant and a scarce basin is
no longer pooled into one region -- which averaged away exactly the
sub-provincial scarcity that constrains irrigation. Each country is then
partitioned into exactly target_count regions balancing geography
against basin scarcity (aggregation.regions.basin_scarcity_weight; 0
recovers plain geographic clustering).

Every region remains either contained in one province (large provinces
split into sub-regions) or a union of whole provinces (small provinces
merged), never a mix of partial pieces across provinces, so regions stay
comparable to political units. allow_cross_border: true has no
basin-aware path and now raises rather than silently ignoring basins.

Two corrections relative to the water-aware original:

- The per-country region budget summed the geodesic area of one
  arbitrary basin fragment per province, so countries whose provinces
  are cut into many basins were systematically starved of regions
  (Russia -27, China -12, Indonesia -10; DRC +12, Nigeria +11). Area is
  now summed over all pieces, and the capacity cap is the piece count,
  since a province can now yield more than one region.
- basin_scarcity_weight is 2.0 rather than 3.0. Swept on the full
  750-region set, the area-weighted within-region scarcity spread falls
  10.34 (weight 0) -> 8.27 -> 7.60 -> 7.40 -> 7.18 (weight 5) while the
  largest region grows 92 -> 91 -> 91 -> 110 -> 133 Mha: separation
  saturates by 2, beyond which only the size tail degrades.

Also guards the overlay: a configured country with no basin overlap is
an error rather than a silent disappearance, and the dropped sliver area
(0.05% of land on the default country set) is logged.

This changes default region geometry for every config.
---
 CHANGELOG.md                      |  14 ++
 config/default.yaml               |  10 +
 config/schemas/config.schema.yaml |   6 +-
 docs/data_sources.rst             |  18 ++
 docs/land_use.rst                 |  16 +-
 tests/test_build_regions.py       | 146 ++++++++++++
 workflow/rules/geography.smk      |   6 +-
 workflow/rules/retrieve.smk       |  30 +++
 workflow/scripts/build_regions.py | 376 +++++++++++++++++++++++-------
 9 files changed, 532 insertions(+), 90 deletions(-)
 create mode 100644 tests/test_build_regions.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b0a70692..be616f41 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -86,6 +86,20 @@ 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. `allow_cross_border: true` is no
+  longer supported and now raises. **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..f8f39ce5 100644
--- a/config/default.yaml
+++ b/config/default.yaml
@@ -1366,6 +1366,16 @@ aggregation:
     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..b6bea24a 100644
--- a/config/schemas/config.schema.yaml
+++ b/config/schemas/config.schema.yaml
@@ -1324,7 +1324,7 @@ properties:
     properties:
       regions:
         type: object
-        required: [target_count, allow_cross_border, method]
+        required: [target_count, allow_cross_border, method, basin_scarcity_weight]
         additionalProperties: false
         properties:
           target_count:
@@ -1338,6 +1338,10 @@ properties:
             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..34fe6bec 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,8 @@ 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.allow_cross_border``: Whether regions can span country boundaries (must be ``false``; basin-aware clustering has no cross-border path and raises otherwise)
+- ``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..53a41051
--- /dev/null
+++ b/tests/test_build_regions.py
@@ -0,0 +1,146 @@
+# 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 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():
+    # A single one-piece province asked for more regions than it has pieces:
+    # reconciliation cannot exceed piece count, so it produces what it can.
+    pieces = pd.DataFrame(
+        {
+            "prov": ["P1", "P2"],
+            "px": [0.0, 1.0],
+            "py": [0.0, 0.0],
+            "area": [1.0, 1.0],
+            "cf": [5.0, 5.0],
+        }
+    )
+    labels = cluster_country(
+        pieces, 2, scarcity_weight=3.0, method="kmeans", random_state=0
+    )
+    assert len(set(labels)) == 2  # two one-piece provinces -> two regions
+    assert _nesting_holds(pieces, labels)
+
+
+def test_allow_cross_border_rejected():
+    pieces = gpd.GeoDataFrame(
+        {"GID_0": ["A"], "prov": ["P1"], "cf": [5.0]},
+        geometry=[box(0, 0, 1, 1)],
+        crs="EPSG:4326",
+    )
+    with pytest.raises(ValueError, match="per-country"):
+        cluster_regions(pieces, 2, allow_cross_border=True, scarcity_weight=3.0)
+
+
+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, allow_cross_border=False, 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
diff --git a/workflow/rules/geography.smk b/workflow/rules/geography.smk
index d020167d..5b910fb7 100644
--- a/workflow/rules/geography.smk
+++ b/workflow/rules/geography.smk
@@ -57,18 +57,20 @@ 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:
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:
+        "<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.
 
diff --git a/workflow/scripts/build_regions.py b/workflow/scripts/build_regions.py
index 45cdc0c6..7da6e39d 100644
--- a/workflow/scripts/build_regions.py
+++ b/workflow/scripts/build_regions.py
@@ -2,6 +2,26 @@
 #
 # 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
@@ -10,6 +30,10 @@
 from pyproj import CRS, Geod
 from sklearn.cluster import AgglomerativeClustering, KMeans
 
+from workflow.scripts.logging_config import setup_script_logging
+
+logger = logging.getLogger(__name__)
+
 GEOD = Geod(ellps="WGS84")
 
 
@@ -128,115 +152,301 @@ def _cluster_coords(
         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 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)
+        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)]
+        prov_px = rem.groupby("prov").apply(
+            lambda g: np.average(g["px"], weights=g["area"]), include_groups=False
+        )
+        prov_py = rem.groupby("prov").apply(
+            lambda g: np.average(g["py"], weights=g["area"]), include_groups=False
+        )
+        prov_cf = rem.groupby("prov").apply(
+            lambda g: np.average(g["cf"], weights=g["area"]), include_groups=False
+        )
+        feat = _scarcity_features(
+            prov_px.to_numpy(), prov_py.to_numpy(), prov_cf.to_numpy(), scarcity_weight
+        )
+        kk = min(int(k_rem), len(prov_px))
+        grp = _cluster_coords(feat, kk, method, random_state)
+        prov_to_group = dict(zip(prov_px.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 = sub.groupby("prov").apply(
+                    lambda g: pd.Series(
+                        {
+                            "px": np.average(g["px"], weights=g["area"]),
+                            "py": np.average(g["py"], weights=g["area"]),
+                            "cf": np.average(g["cf"], weights=g["area"]),
+                        }
+                    ),
+                    include_groups=False,
+                )
+                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)
+                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)
+                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 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
-
     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
+        raise ValueError(
+            "Basin-aware clustering requires per-country partitioning; set "
+            "aggregation.regions.allow_cross_border to false."
         )
+    if "GID_0" not in pieces.columns:
+        raise ValueError("Expected GID_0 column for country codes")
 
-        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
-    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)
+    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()
+    )
 
-    dissolved = dissolved.set_index("region")
-    dissolved = dissolved.drop(
-        columns=[c for c in ["_cluster"] if c in dissolved.columns]
+    # 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(-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"]]
+    rep_country = pieces.groupby("_cluster")["GID_0"].first()
+    dissolved["country"] = dissolved["_cluster"].map(rep_country)
+    dissolved = dissolved.set_index("region").drop(columns="_cluster")
     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")

From 74bf11124b6a6ae98046cd152bdc64097f9b2d86 Mon Sep 17 00:00:00 2001
From: Koen van Greevenbroek <koenvg@stanford.edu>
Date: Mon, 20 Jul 2026 19:03:53 -0700
Subject: [PATCH 2/6] fix: repair invalid region geometries before they leave
 build_regions

Dissolving province-basin pieces leaves a handful of self-intersecting
polygons -- 3 of 750 regions at the default resolution, 1 of 200 in the
test config. That is not a cosmetic problem: exactextract calls into
native code that SEGFAULTS on invalid geometry rather than raising, so a
single bad region kills an unrelated downstream rule with an empty log
and no traceback. It surfaced as a crash in the Huang water aggregation
during a calibration run, several rules removed from its cause.

Repairing at the source rather than in each consumer, since every rule
that aggregates a raster over regions is exposed. Regions that cannot be
repaired now raise, so a real geometry problem still fails loudly.
---
 tests/test_build_regions.py       | 27 ++++++++++++++++++++-
 workflow/scripts/build_regions.py | 39 +++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/tests/test_build_regions.py b/tests/test_build_regions.py
index 53a41051..7d8537de 100644
--- a/tests/test_build_regions.py
+++ b/tests/test_build_regions.py
@@ -14,7 +14,7 @@
 import numpy as np
 import pandas as pd
 import pytest
-from shapely.geometry import box
+from shapely.geometry import Polygon, box
 
 from workflow.scripts.build_regions import (
     _largest_remainder,
@@ -144,3 +144,28 @@ def test_country_budget_uses_full_province_area():
         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, allow_cross_border=False, scarcity_weight=3.0)
+    assert regions.geometry.is_valid.all(), "invalid geometry escaped build_regions"
diff --git a/workflow/scripts/build_regions.py b/workflow/scripts/build_regions.py
index 7da6e39d..8ecb50bf 100644
--- a/workflow/scripts/build_regions.py
+++ b/workflow/scripts/build_regions.py
@@ -28,6 +28,7 @@
 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
@@ -36,6 +37,9 @@
 
 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"
@@ -385,6 +389,30 @@ def cluster_regions(
     rep_country = pieces.groupby("_cluster")["GID_0"].first()
     dissolved["country"] = dissolved["_cluster"].map(rep_country)
     dissolved = dissolved.set_index("region").drop(columns="_cluster")
+
+    # Snap to the GeoJSON writer's coordinate grid, then repair. 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 first makes the in-memory geometry exactly what gets
+    # written, so repairing it afterwards actually holds. 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.
+    # buffer(0) first: set_precision raises on genuinely invalid input rather
+    # than repairing it. Then snap, then repair again, since snapping can itself
+    # introduce self-intersections.
+    repaired = dissolved.geometry.buffer(0)
+    snapped = shapely.set_precision(repaired.values, GEOJSON_PRECISION)
+    dissolved["geometry"] = gpd.GeoSeries(
+        snapped, index=dissolved.index, crs=dissolved.crs
+    ).buffer(0)
+    still_bad = ~dissolved.geometry.is_valid
+    if still_bad.any():
+        raise ValueError(
+            f"{int(still_bad.sum())} region geometries are invalid after "
+            "normalisation: " + ", ".join(map(str, dissolved.index[still_bad][:5]))
+        )
     return dissolved
 
 
@@ -450,3 +478,14 @@ def cluster_regions(
 
     Path(snakemake.output[0]).parent.mkdir(parents=True, exist_ok=True)
     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
+    if bad.any():
+        raise ValueError(
+            f"{int(bad.sum())} region geometries are invalid as written to "
+            f"{snakemake.output[0]}; exactextract would segfault on them."
+        )
+    logger.info("Wrote %d regions, all geometries valid on disk", len(written))

From c8396ff89a3d474315e98e2c2a42c383018b1371 Mon Sep 17 00:00:00 2001
From: Koen van Greevenbroek <koenvg@stanford.edu>
Date: Wed, 22 Jul 2026 10:54:03 -0700
Subject: [PATCH 3/6] refactor: remove the allow_cross_border config option

Basin-aware clustering has no cross-border path, so the option's only
valid value was false; drop the key from config, schema, rule and script
instead of rejecting it at runtime.
---
 CHANGELOG.md                      |  5 +++--
 config/default.yaml               |  1 -
 config/schemas/config.schema.yaml |  5 +----
 docs/land_use.rst                 |  1 -
 tests/test_build_regions.py       | 14 ++------------
 workflow/rules/geography.smk      |  1 -
 workflow/scripts/build_regions.py |  9 +--------
 7 files changed, 7 insertions(+), 29 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index be616f41..bc834b69 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -94,8 +94,9 @@ introduce breaking changes to configuration and outputs.
   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. `allow_cross_border: true` is no
-  longer supported and now raises. **This changes default region geometry for
+  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).
diff --git a/config/default.yaml b/config/default.yaml
index f8f39ce5..387fa796 100644
--- a/config/default.yaml
+++ b/config/default.yaml
@@ -1364,7 +1364,6 @@ 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
diff --git a/config/schemas/config.schema.yaml b/config/schemas/config.schema.yaml
index b6bea24a..d600e6e1 100644
--- a/config/schemas/config.schema.yaml
+++ b/config/schemas/config.schema.yaml
@@ -1324,16 +1324,13 @@ properties:
     properties:
       regions:
         type: object
-        required: [target_count, allow_cross_border, method, basin_scarcity_weight]
+        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]
diff --git a/docs/land_use.rst b/docs/land_use.rst
index 34fe6bec..01996152 100644
--- a/docs/land_use.rst
+++ b/docs/land_use.rst
@@ -102,7 +102,6 @@ basin-aware so that hydrological scarcity is not averaged away:
 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 (must be ``false``; basin-aware clustering has no cross-border path and raises otherwise)
 - ``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
 
diff --git a/tests/test_build_regions.py b/tests/test_build_regions.py
index 7d8537de..a1fa3cb6 100644
--- a/tests/test_build_regions.py
+++ b/tests/test_build_regions.py
@@ -100,16 +100,6 @@ def test_cluster_country_reconciles_when_pieces_scarce():
     assert _nesting_holds(pieces, labels)
 
 
-def test_allow_cross_border_rejected():
-    pieces = gpd.GeoDataFrame(
-        {"GID_0": ["A"], "prov": ["P1"], "cf": [5.0]},
-        geometry=[box(0, 0, 1, 1)],
-        crs="EPSG:4326",
-    )
-    with pytest.raises(ValueError, match="per-country"):
-        cluster_regions(pieces, 2, allow_cross_border=True, scarcity_weight=3.0)
-
-
 def test_country_budget_uses_full_province_area():
     """A country's region budget follows its land area, not its basin fragmentation.
 
@@ -138,7 +128,7 @@ def test_country_budget_uses_full_province_area():
         ],
         crs="EPSG:4326",
     )
-    regions = cluster_regions(pieces, 4, allow_cross_border=False, scarcity_weight=3.0)
+    regions = cluster_regions(pieces, 4, scarcity_weight=3.0)
     per_country = regions["country"].value_counts()
     assert (
         per_country["A"] == 2
@@ -167,5 +157,5 @@ def test_dissolved_regions_are_valid_geometries():
         geometry=[bad, box(10, 0, 11, 1)],
         crs="EPSG:4326",
     )
-    regions = cluster_regions(pieces, 2, allow_cross_border=False, scarcity_weight=3.0)
+    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 5b910fb7..bafcf8b1 100644
--- a/workflow/rules/geography.smk
+++ b/workflow/rules/geography.smk
@@ -60,7 +60,6 @@ rule build_regions:
         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"],
diff --git a/workflow/scripts/build_regions.py b/workflow/scripts/build_regions.py
index 8ecb50bf..0026724c 100644
--- a/workflow/scripts/build_regions.py
+++ b/workflow/scripts/build_regions.py
@@ -79,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)
@@ -330,7 +330,6 @@ def _reconcile_to_k(
 def cluster_regions(
     pieces: gpd.GeoDataFrame,
     target_count: int,
-    allow_cross_border: bool,
     scarcity_weight: float,
     method: str = "kmeans",
     random_state: int = 0,
@@ -347,11 +346,6 @@ def cluster_regions(
     """
     if target_count <= 0:
         raise ValueError("target_count must be positive")
-    if allow_cross_border:
-        raise ValueError(
-            "Basin-aware clustering requires per-country partitioning; set "
-            "aggregation.regions.allow_cross_border to false."
-        )
     if "GID_0" not in pieces.columns:
         raise ValueError("Expected GID_0 column for country codes")
 
@@ -471,7 +465,6 @@ def cluster_regions(
     regions = cluster_regions(
         pieces,
         snakemake.params.n_regions,
-        snakemake.params.allow_cross_border,
         snakemake.params.basin_scarcity_weight,
         snakemake.params.cluster_method,
     )

From c7c9bb511fd41ed01146a4579ef8b00b4c5fb440 Mon Sep 17 00:00:00 2001
From: Koen van Greevenbroek <koenvg@stanford.edu>
Date: Wed, 22 Jul 2026 10:54:33 -0700
Subject: [PATCH 4/6] fix: flag empty region geometries and drop a redundant
 repair pass

set_precision's valid_output mode already guarantees valid snapped
geometry, so the trailing buffer(0) did nothing; conversely a region
collapsed to POLYGON EMPTY passes is_valid, so check emptiness
explicitly both in memory and on the written file.
---
 workflow/scripts/build_regions.py | 33 ++++++++++++++++---------------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/workflow/scripts/build_regions.py b/workflow/scripts/build_regions.py
index 0026724c..b1051951 100644
--- a/workflow/scripts/build_regions.py
+++ b/workflow/scripts/build_regions.py
@@ -384,28 +384,29 @@ def cluster_regions(
     dissolved["country"] = dissolved["_cluster"].map(rep_country)
     dissolved = dissolved.set_index("region").drop(columns="_cluster")
 
-    # Snap to the GeoJSON writer's coordinate grid, then repair. Order matters.
+    # 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 first makes the in-memory geometry exactly what gets
-    # written, so repairing it afterwards actually holds. 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.
-    # buffer(0) first: set_precision raises on genuinely invalid input rather
-    # than repairing it. Then snap, then repair again, since snapping can itself
-    # introduce self-intersections.
+    # 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
-    ).buffer(0)
-    still_bad = ~dissolved.geometry.is_valid
+    )
+    # 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 after "
-            "normalisation: " + ", ".join(map(str, dissolved.index[still_bad][:5]))
+            f"{int(still_bad.sum())} region geometries are invalid or empty "
+            "after normalisation: "
+            + ", ".join(map(str, dissolved.index[still_bad][:5]))
         )
     return dissolved
 
@@ -475,10 +476,10 @@ def cluster_regions(
     # 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
+    bad = ~written.geometry.is_valid | written.geometry.is_empty
     if bad.any():
         raise ValueError(
-            f"{int(bad.sum())} region geometries are invalid as written to "
-            f"{snakemake.output[0]}; exactextract would segfault on them."
+            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))

From 08f8e0c628b4f16af77b1dedc90c937a6f04d8b1 Mon Sep 17 00:00:00 2001
From: Koen van Greevenbroek <koenvg@stanford.edu>
Date: Wed, 22 Jul 2026 10:56:31 -0700
Subject: [PATCH 5/6] refactor: share province centroid maths and weight
 clustering by area

Area-weighted province centroids were computed twice in different
idioms; factor them into _province_centroids. Pass piece/province areas
as k-means sample weights so a coastal sliver does not pull a centroid
as hard as a half-province piece.
---
 workflow/scripts/build_regions.py | 78 +++++++++++++++++++------------
 1 file changed, 47 insertions(+), 31 deletions(-)

diff --git a/workflow/scripts/build_regions.py b/workflow/scripts/build_regions.py
index b1051951..042cb05b 100644
--- a/workflow/scripts/build_regions.py
+++ b/workflow/scripts/build_regions.py
@@ -132,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])
@@ -145,13 +152,11 @@ 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}")
 
@@ -187,6 +192,23 @@ def z(a: np.ndarray) -> np.ndarray:
     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,
@@ -231,28 +253,27 @@ def cluster_country(
         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)
+        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)]
-        prov_px = rem.groupby("prov").apply(
-            lambda g: np.average(g["px"], weights=g["area"]), include_groups=False
-        )
-        prov_py = rem.groupby("prov").apply(
-            lambda g: np.average(g["py"], weights=g["area"]), include_groups=False
-        )
-        prov_cf = rem.groupby("prov").apply(
-            lambda g: np.average(g["cf"], weights=g["area"]), include_groups=False
-        )
+        pcen = _province_centroids(rem)
         feat = _scarcity_features(
-            prov_px.to_numpy(), prov_py.to_numpy(), prov_cf.to_numpy(), scarcity_weight
+            pcen["px"].to_numpy(),
+            pcen["py"].to_numpy(),
+            pcen["cf"].to_numpy(),
+            scarcity_weight,
         )
-        kk = min(int(k_rem), len(prov_px))
-        grp = _cluster_coords(feat, kk, method, random_state)
-        prov_to_group = dict(zip(prov_px.index, grp + next_id))
+        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
@@ -287,23 +308,16 @@ def _reconcile_to_k(
             provs = sub["prov"].unique()
             if len(provs) > 1:
                 # Whole-province group: split provinces into two, keep them whole.
-                pcen = sub.groupby("prov").apply(
-                    lambda g: pd.Series(
-                        {
-                            "px": np.average(g["px"], weights=g["area"]),
-                            "py": np.average(g["py"], weights=g["area"]),
-                            "cf": np.average(g["cf"], weights=g["area"]),
-                        }
-                    ),
-                    include_groups=False,
-                )
+                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)
+                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:
@@ -314,7 +328,9 @@ def _reconcile_to_k(
                     sub["cf"].to_numpy(),
                     scarcity_weight,
                 )
-                two = _cluster_coords(feat, 2, method, random_state)
+                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

From d4d921b848f08881efe59eb4515d80c0e7331628 Mon Sep 17 00:00:00 2001
From: Koen van Greevenbroek <koenvg@stanford.edu>
Date: Wed, 22 Jul 2026 10:57:11 -0700
Subject: [PATCH 6/6] tests: exercise the reconciliation split path

The pieces-scarce fixture resolved entirely in the merge regime, so
_reconcile_to_k was never hit; give the indivisible province a quota of
two so reconciliation must recover the deficit from the merge group.
---
 tests/test_build_regions.py | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/tests/test_build_regions.py b/tests/test_build_regions.py
index a1fa3cb6..4aa711d6 100644
--- a/tests/test_build_regions.py
+++ b/tests/test_build_regions.py
@@ -82,22 +82,27 @@ def test_cluster_country_splits_province_by_scarcity():
 
 
 def test_cluster_country_reconciles_when_pieces_scarce():
-    # A single one-piece province asked for more regions than it has pieces:
-    # reconciliation cannot exceed piece count, so it produces what it can.
+    # 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"],
-            "px": [0.0, 1.0],
-            "py": [0.0, 0.0],
-            "area": [1.0, 1.0],
-            "cf": [5.0, 5.0],
+            "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, 2, scarcity_weight=3.0, method="kmeans", random_state=0
+        pieces, 3, scarcity_weight=2.0, method="kmeans", random_state=0
     )
-    assert len(set(labels)) == 2  # two one-piece provinces -> two regions
+    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():