perf: accelerate crop raster aggregation

CHANGELOG.md

@@ -86,6 +86,10 @@ introduce breaking changes to configuration and outputs.
 
 ### Changed
 
+- Crop-yield and harvested-area preparation now compute exact region and
+  resource-class cell coverage once per configuration and reuse it across
+  crops, substantially reducing build time and peak memory without changing
+  outputs.
 - 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

docs/crop_production.rst

@@ -115,12 +115,12 @@ The following figures show yield potential maps for three major crops, illustrat
 Yield Aggregation
 -----------------
 
-Yields are aggregated from the input resolution gridcells to (region, resource_class, water_supply) combinations by ``workflow/scripts/build_crop_yields.py``.
+Yields are aggregated from the input resolution gridcells to (region, resource_class, water_supply) combinations by ``workflow/scripts/build_crop_yields.py``. Exact region-to-cell coverage fractions are computed once per configuration by ``workflow/scripts/build_region_class_cell_mapping.py`` and reused for every crop-yield and harvested-area raster.
 
 Aggregation Process
 ~~~~~~~~~~~~~~~~~~~
 
-1. **Load resource classes**: Read the class assignment raster (see :doc:`land_use`)
+1. **Load the cell mapping**: Read the reusable region, resource-class, cell, and exact coverage arrays derived from the class raster (see :doc:`land_use`)
 
 2. **Load crop-specific rasters**:
 

docs/workflow.rst

@@ -71,13 +71,26 @@ Data Preparation Rules
   * **Script**: ``workflow/scripts/aggregate_class_areas.py``
   * **Purpose**: Compute available land area per (region, class, water, crop)
 
+**build_region_class_cell_mapping**
+  * **Input**: Resource classes, regions
+  * **Output**: ``processing/{name}/region_class_cell_mapping.npz``
+  * **Script**: ``workflow/scripts/build_region_class_cell_mapping.py``
+  * **Purpose**: Cache exact region and resource-class coverage by GAEZ grid cell for crop-yield and harvested-area aggregation
+
 **build_crop_yields**
   * **Wildcards**: ``{crop}`` (crop name), ``{water_supply}`` ("r" or "i")
-  * **Input**: Resource classes, GAEZ rasters (yield, suitability, water, growing season)
+  * **Input**: Reusable region/cell coverage mapping, GAEZ rasters (yield, suitability, water, growing season)
   * **Output**: ``processing/{name}/crop_yields/{crop}_{water_supply}.csv``
   * **Script**: ``workflow/scripts/build_crop_yields.py``
   * **Purpose**: Aggregate yields by (region, class) for each crop
 
+**build_harvested_area_gaez**
+  * **Wildcards**: ``{crop}`` (crop name), ``{water_supply}`` ("r" or "i")
+  * **Input**: Reusable region/cell coverage mapping, GAEZ harvested-area raster, crop shares
+  * **Output**: ``processing/{name}/harvested_area/gaez/{crop}_{water_supply}.csv``
+  * **Script**: ``workflow/scripts/build_harvested_area.py``
+  * **Purpose**: Aggregate harvested area by (region, class) and attribute pooled GAEZ crop modules
+
 **derive_mirca_multicropping**
   * **Input**: Annual harvested-area, footprint, and rice-subcrop grids from the MIRCA-OS release nearest ``baseline_year``; resource classes; regions; GAEZ RES01 multiple-cropping-zone rasters; the crop concordance; and the fixed combination catalog
   * **Output**: ``processing/{name}/multi_cropping/baseline_area.csv`` (observed physical link area), ``residual_multicrop.tif`` (unattributed extra-cycle area), and ``attribution_stats.csv`` (diagnostic totals)

tests/test_region_class_aggregation.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: 2026 Koen van Greevenbroek
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from pathlib import Path
+
+import geopandas as gpd
+import numpy as np
+from pyproj import CRS
+import pytest
+from rasterio.io import MemoryFile
+from rasterio.transform import from_origin
+from shapely.geometry import box
+import xarray as xr
+
+from workflow.scripts.build_harvested_area import _extract_harvested_area
+from workflow.scripts.build_region_class_cell_mapping import build_cell_mapping
+from workflow.scripts.region_class_aggregation import (
+    load_cell_mapping,
+    validate_raster_grid,
+    weighted_mean_by_group,
+    weighted_sum_by_group,
+)
+
+
+@pytest.fixture
+def cell_mapping(tmp_path: Path):
+    classes_path = tmp_path / "classes.nc"
+    regions_path = tmp_path / "regions.geojson"
+    mapping_path = tmp_path / "mapping.npz"
+
+    classes = xr.Dataset(
+        {"resource_class": (("y", "x"), np.array([[0, 1, -1], [0, 1, -1]]))},
+        attrs={
+            "transform": np.array([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]),
+            "crs_wkt": CRS.from_epsg(4326).to_wkt(),
+        },
+    )
+    classes.to_netcdf(classes_path)
+    regions = gpd.GeoDataFrame(
+        {"region": ["r0", "r1"]},
+        geometry=[box(0.0, 0.0, 1.3, 2.0), box(1.3, 0.0, 3.0, 2.0)],
+        crs="EPSG:4326",
+    )
+    regions.to_file(regions_path)
+
+    build_cell_mapping(str(classes_path), str(regions_path), str(mapping_path))
+    return load_cell_mapping(str(mapping_path))
+
+
+def test_cell_mapping_preserves_partial_region_coverage(cell_mapping):
+    values = np.array([[2.0, 10.0, 100.0], [4.0, 20.0, 200.0]])
+
+    means = weighted_mean_by_group(values, cell_mapping)
+    sums = weighted_sum_by_group(values, cell_mapping)
+
+    np.testing.assert_allclose(means, [3.0, 15.0, np.nan, 15.0], equal_nan=True)
+    np.testing.assert_allclose(sums, [6.0, 9.0, 0.0, 21.0])
+    assert cell_mapping.coverage.dtype == np.float64
+
+
+def test_group_aggregation_matches_exactextract_empty_group_semantics(cell_mapping):
+    values = np.full(cell_mapping.shape, np.nan)
+
+    means = weighted_mean_by_group(values, cell_mapping)
+    sums = weighted_sum_by_group(values, cell_mapping)
+
+    assert np.isnan(means).all()
+    np.testing.assert_array_equal(sums, np.zeros(cell_mapping.n_groups))
+
+
+def test_harvested_area_extraction_preserves_class_major_order(cell_mapping):
+    values = np.array([[2.0, 10.0, np.nan], [4.0, 20.0, np.nan]])
+
+    result = _extract_harvested_area(values, cell_mapping)
+
+    assert result[["region", "resource_class"]].to_records(index=False).tolist() == [
+        ("r0", 0),
+        ("r1", 0),
+        ("r0", 1),
+        ("r1", 1),
+    ]
+    np.testing.assert_allclose(result["value"], [6.0, 0.0, 9.0, 21.0])
+
+
+def test_raster_grid_validation_rejects_shifted_transform(cell_mapping):
+    values = np.zeros(cell_mapping.shape, dtype=np.float32)
+    with (
+        MemoryFile() as memory_file,
+        memory_file.open(
+            driver="GTiff",
+            height=2,
+            width=3,
+            count=1,
+            dtype="float32",
+            crs="EPSG:4326",
+            transform=from_origin(0.1, 2.0, 1.0, 1.0),
+        ) as source,
+        pytest.raises(ValueError, match="transform"),
+    ):
+        validate_raster_grid(values, source, cell_mapping)
+
+
+def test_cell_mapping_load_preserves_float64_coverage(tmp_path: Path):
+    mapping_path = tmp_path / "mapping.npz"
+    expected = np.array([0.30000000000000004], dtype=np.float64)
+    np.savez(
+        mapping_path,
+        cell_ids=np.array([0], dtype=np.int32),
+        coverage=expected,
+        group_ids=np.array([0], dtype=np.int32),
+        regions=np.array(["r0"]),
+        n_classes=np.array(1, dtype=np.int32),
+        height=np.array(1, dtype=np.int32),
+        width=np.array(1, dtype=np.int32),
+        transform=np.array([0.0, 1.0, 0.0, 1.0, 0.0, -1.0]),
+        crs_wkt=np.array(CRS.from_epsg(4326).to_wkt()),
+    )
+
+    mapping = load_cell_mapping(str(mapping_path))
+
+    np.testing.assert_array_equal(mapping.coverage, expected)
+    assert mapping.coverage.dtype == np.float64

workflow/rules/crops.smk

@@ -149,11 +149,29 @@ def yield_and_suitability_for_crop(w):
     return inputs
 
 
-rule build_crop_yields:
+rule build_region_class_cell_mapping:
     input:
-        unpack(yield_and_suitability_for_crop),
         classes="<processing>/{name}/resource_classes.nc",
         regions="<processing>/{name}/regions.geojson",
+    output:
+        mapping="<processing>/{name}/region_class_cell_mapping.npz",
+    group:
+        "prep"
+    resources:
+        runtime="1m",
+        mem_mb=600,
+    log:
+        "<logs>/{name}/build_region_class_cell_mapping.log",
+    benchmark:
+        "<benchmarks>/{name}/build_region_class_cell_mapping.tsv"
+    script:
+        "../scripts/build_region_class_cell_mapping.py"
+
+
+rule build_crop_yields:
+    input:
+        unpack(yield_and_suitability_for_crop),
+        cell_mapping="<processing>/{name}/region_class_cell_mapping.npz",
         yield_unit_conversions="data/curated/yield_unit_conversions.csv",
         moisture_content="data/curated/crop_moisture_content.csv",
     params:
@@ -168,7 +186,7 @@ rule build_crop_yields:
         "prep"
     resources:
         runtime="1m",
-        mem_mb=1300,
+        mem_mb=700,
     log:
         "<logs>/{name}/build_crop_yields_{crop}_{water_supply}.log",
     benchmark:
@@ -437,7 +455,7 @@ def _harvested_area_inputs(w):
     """Get inputs for build_harvested_area_gaez, including FDD shares when relevant."""
     inputs = {
         "harvested_area_raster": gaez_path("harvested_area", w.water_supply, w.crop),
-        "classes": f"<processing>/{w.name}/resource_classes.nc",
+        "cell_mapping": f"<processing>/{w.name}/region_class_cell_mapping.npz",
         "regions": f"<processing>/{w.name}/regions.geojson",
         "crop_mapping": "data/curated/gaez_crop_code_mapping.csv",
         "faostat_production": f"<processing>/{w.name}/faostat_crop_production.csv",
@@ -481,7 +499,7 @@ rule build_harvested_area_gaez:
         "prep"
     resources:
         runtime="1m",
-        mem_mb=700,
+        mem_mb=400,
     log:
         "<logs>/{name}/build_harvested_area_gaez_{crop}_{water_supply}.log",
     benchmark: