perf: accelerate model build workflow

CHANGELOG.md

@@ -86,6 +86,10 @@ introduce breaking changes to configuration and outputs.
 
 ### Changed
 
+- Spatial preprocessing and model construction now reuse region/class cell
+  mappings, bound raster cache memory, and vectorize repeated aggregation and
+  crop-link operations. This substantially reduces the time and peak memory
+  needed to build a default model without changing its contents.
 - 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

tests/test_compute_resource_classes.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: 2026 Koen van Greevenbroek
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import numpy as np
+
+from workflow.scripts.compute_resource_classes import classify_by_region
+
+
+def test_classify_by_region_handles_ties_and_invalid_cells():
+    score = np.array(
+        [
+            [1.0, 2.0, 3.0, 4.0, 1.0, 1.0],
+            [np.nan, 0.0, -1.0, 1.0, 1.0, 99.0],
+        ]
+    )
+    regions = np.array(
+        [
+            [0, 0, 0, 0, 1, 1],
+            [0, 0, 0, 1, 1, -1],
+        ]
+    )
+
+    result = classify_by_region(score, regions, [0.0, 0.25, 0.5, 0.75, 1.0])
+
+    np.testing.assert_array_equal(
+        result,
+        np.array(
+            [
+                [0, 1, 2, 3, 3, 3],
+                [-1, -1, -1, 3, 3, -1],
+            ],
+            dtype=np.int8,
+        ),
+    )
+
+
+def test_classify_by_region_returns_empty_grid_when_no_score_is_positive():
+    result = classify_by_region(
+        np.array([[np.nan, 0.0]]),
+        np.array([[0, 0]]),
+        [0.0, 0.5, 1.0],
+    )
+
+    np.testing.assert_array_equal(result, np.array([[-1, -1]], dtype=np.int8))

tests/test_luc_carbon.py

@@ -7,22 +7,22 @@
 import textwrap
 
 from affine import Affine
-import geopandas as gpd
 import numpy as np
 from osgeo import gdal, osr
 import pandas as pd
 import pytest
-from shapely.geometry import box
 import xarray as xr
 
 from workflow.scripts.build_luc_carbon_coefficients import (
     CO2_PER_C,
     _correct_subpixel_soc,
     _decompose_agb,
     _ensure_mode_zero,
+    _weighted_mean_by_group,
     _zone_index,
     _zone_parameters,
 )
+from workflow.scripts.region_class_aggregation import CellMapping
 
 gdal.UseExceptions()
 osr.UseExceptions()
@@ -45,6 +45,32 @@ def test_value(self):
         assert pytest.approx(3.66667, rel=1e-4) == CO2_PER_C
 
 
+# ---------------------------------------------------------------------------
+# Tests: _weighted_mean_by_group
+# ---------------------------------------------------------------------------
+
+
+def test_weighted_mean_by_group_uses_cell_coverage_and_weights():
+    """Group means combine explicit weights with fractional cell coverage."""
+    mapping = CellMapping(
+        cell_ids=np.array([0, 1, 2, 3], dtype=np.int32),
+        coverage=np.array([1.0, 0.5, 0.25, 1.0]),
+        group_ids=np.array([0, 0, 1, 1], dtype=np.int32),
+        regions=np.array(["north", "south"]),
+        n_classes=1,
+        shape=(2, 2),
+        transform=(0.0, 1.0, 0.0, 2.0, 0.0, -1.0),
+        crs_wkt="",
+    )
+    values = np.array([[10.0, 20.0], [30.0, np.nan]])
+    weights = np.array([[1.0, 2.0], [4.0, 1.0]])
+
+    result = _weighted_mean_by_group(values, weights, mapping)
+
+    assert result[0] == pytest.approx(15.0)
+    assert result[1] == pytest.approx(30.0)
+
+
 # ---------------------------------------------------------------------------
 # Tests: _zone_index
 # ---------------------------------------------------------------------------
@@ -756,18 +782,24 @@ def _make_synthetic_inputs(tmp_path, *, height=2, width=2):
     zone_path = tmp_path / "zone_params.csv"
     zone_path.write_text(zone_csv)
 
-    # Region GeoDataFrame: one region covering the full grid
-    region_gdf = gpd.GeoDataFrame(
-        {"region": ["test_region"]},
-        geometry=[box(0, 0, 2, 2)],
-        crs="EPSG:4326",
+    mapping_path = tmp_path / "cell_mapping.npz"
+    cell_ids = np.arange(height * width, dtype=np.int32)
+    np.savez(
+        mapping_path,
+        cell_ids=cell_ids,
+        coverage=np.ones(height * width, dtype=np.float64),
+        group_ids=np.zeros(height * width, dtype=np.int32),
+        regions=np.array(["test_region"]),
+        n_classes=np.array(1, dtype=np.int32),
+        height=np.array(height, dtype=np.int32),
+        width=np.array(width, dtype=np.int32),
+        transform=np.asarray(transform.to_gdal()),
+        crs_wkt=np.array(WGS84_WKT),
     )
-    regions_path = tmp_path / "regions.geojson"
-    region_gdf.to_file(str(regions_path), driver="GeoJSON")
 
     return {
         "classes": str(classes_path),
-        "regions": str(regions_path),
+        "cell_mapping": str(mapping_path),
         "agb": str(agb_path),
         "soc": str(soc_path),
         "regrowth": str(regrowth_path),

tests/test_luicube_grassland_yields.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: 2026 Koen van Greevenbroek
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import numpy as np
+import pytest
+
+from workflow.scripts.build_luicube_grassland_yields import (
+    aggregate_grassland_yields,
+)
+from workflow.scripts.region_class_aggregation import CellMapping
+
+
+def _mapping() -> CellMapping:
+    return CellMapping(
+        cell_ids=np.arange(4, dtype=np.int32),
+        coverage=np.array([1.0, 0.5, 0.25, 1.0]),
+        group_ids=np.array([0, 0, 1, 1], dtype=np.int32),
+        regions=np.array(["region0"]),
+        n_classes=2,
+        shape=(2, 2),
+        transform=(0.0, 1.0, 0.0, 2.0, 0.0, -1.0),
+        crs_wkt="",
+    )
+
+
+def test_aggregate_grassland_yields_uses_exact_cell_coverage():
+    result = aggregate_grassland_yields(
+        area_km2=np.array([[1.0, 2.0], [3.0, 4.0]]),
+        npp_act=np.array([[10.0, 20.0], [0.0, 40.0]]),
+        hanpp_harv=np.array([[5.0, 10.0], [1.0, 80.0]]),
+        mapping=_mapping(),
+    )
+
+    assert result.loc[("region0", 0), "yield"] == pytest.approx(10.0 / 100.0 / 0.45)
+    assert result.loc[("region0", 0), "suitable_area"] == pytest.approx(200.0)
+    assert result.loc[("region0", 0), "grazing_intensity"] == pytest.approx(0.5)
+    assert result.loc[("region0", 1), "yield"] == pytest.approx(80.25 / 400.0 / 0.45)
+    assert result.loc[("region0", 1), "suitable_area"] == pytest.approx(475.0)
+    assert result.loc[("region0", 1), "grazing_intensity"] == pytest.approx(1.0)
+
+
+def test_aggregate_grassland_yields_rejects_grid_mismatch():
+    with pytest.raises(ValueError, match="does not match"):
+        aggregate_grassland_yields(
+            area_km2=np.ones((1, 1)),
+            npp_act=np.ones((1, 1)),
+            hanpp_harv=np.ones((1, 1)),
+            mapping=_mapping(),
+        )

tests/test_process_huang_irrigation_water.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: 2026 Koen van Greevenbroek
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from pathlib import Path
+
+import pandas as pd
+
+from workflow.scripts.process_huang_irrigation_water import (
+    load_crop_growing_seasons,
+)
+
+
+def _write_crop_yield(path: Path, *, include_season: bool = True) -> None:
+    rows = []
+    values = {
+        "suitable_area": [1.0, 3.0],
+        "growing_season_start_day": [10.0, 30.0],
+        "growing_season_length_days": [100.0, 200.0],
+        "yield": [2.0, 4.0],
+    }
+    if not include_season:
+        values = {"suitable_area": [1.0, 3.0], "yield": [2.0, 4.0]}
+    for variable, variable_values in values.items():
+        for resource_class, value in enumerate(variable_values):
+            rows.append(
+                {
+                    "region": "r0",
+                    "resource_class": resource_class,
+                    "variable": variable,
+                    "unit": "test",
+                    "value": value,
+                }
+            )
+    pd.DataFrame(rows).to_csv(path, index=False)
+
+
+def test_crop_growing_seasons_use_irrigated_files_when_available(tmp_path: Path):
+    rainfed = tmp_path / "wheat_r.csv"
+    irrigated = tmp_path / "wheat_i.csv"
+    _write_crop_yield(rainfed)
+    _write_crop_yield(irrigated)
+
+    result = load_crop_growing_seasons([rainfed, irrigated])
+
+    assert result.to_dict("records") == [
+        {
+            "region": "r0",
+            "crop": "wheat",
+            "water_supply": "i",
+            "total_area": 4.0,
+            "growing_season_start_day": 25.0,
+            "growing_season_length_days": 175.0,
+        }
+    ]
+
+
+def test_crop_growing_seasons_fall_back_to_rainfed_files(tmp_path: Path):
+    rainfed = tmp_path / "wheat_r.csv"
+    irrigated = tmp_path / "wheat_i.csv"
+    _write_crop_yield(rainfed)
+    _write_crop_yield(irrigated, include_season=False)
+
+    result = load_crop_growing_seasons([rainfed, irrigated])
+
+    assert result["water_supply"].tolist() == ["r"]
+    assert result["growing_season_start_day"].tolist() == [25.0]

workflow/rules/crops.smk

@@ -636,8 +636,8 @@ rule derive_mirca_multicropping:
     params:
         source_year=MIRCA_MULTICROPPING_YEAR,
     resources:
-        runtime="15m",
-        mem_mb=8000,
+        runtime="2m",
+        mem_mb=3000,
     log:
         "<logs>/{name}/derive_mirca_multicropping.log",
     benchmark:
@@ -671,8 +671,7 @@ def multi_cropping_inputs(_wildcards):
         "actual_yield" if config["validation"]["use_actual_yields"] else "yield"
     )
     inputs = {
-        "classes": "<processing>/{name}/resource_classes.nc",
-        "regions": "<processing>/{name}/regions.geojson",
+        "cell_mapping": "<processing>/{name}/region_class_cell_mapping.npz",
         "yield_unit_conversions": "data/curated/yield_unit_conversions.csv",
         "combinations": multicropping_combinations_yaml(),
     }
@@ -705,7 +704,7 @@ rule build_multi_cropping:
         "prep"
     resources:
         runtime="2m",
-        mem_mb=5500,
+        mem_mb=2500,
     log:
         "<logs>/{name}/build_multi_cropping.log",
     benchmark:
@@ -737,15 +736,14 @@ rule build_grassland_yields:
 rule build_luicube_grassland_yields:
     input:
         luicube="<processing>/shared/luc/luicube_grassland.nc",
-        classes="<processing>/{name}/resource_classes.nc",
-        regions="<processing>/{name}/regions.geojson",
+        cell_mapping="<processing>/{name}/region_class_cell_mapping.npz",
     output:
         "<processing>/{name}/luicube_grassland_yields.csv",
     group:
         "prep"
     resources:
         runtime="1m",
-        mem_mb=2600,
+        mem_mb=1000,
     log:
         "<logs>/{name}/build_luicube_grassland_yields.log",
     benchmark:

workflow/rules/geography.smk

@@ -132,7 +132,7 @@ rule compute_resource_classes:
         "prep"
     resources:
         runtime="1m",
-        mem_mb=1900,
+        mem_mb=900,
     log:
         "<logs>/{name}/compute_resource_classes.log",
     benchmark:
@@ -143,11 +143,10 @@ rule compute_resource_classes:
 
 rule aggregate_class_areas:
     input:
-        classes="<processing>/{name}/resource_classes.nc",
+        cell_mapping="<processing>/{name}/region_class_cell_mapping.npz",
         sr=[gaez_path("suitability", "r", crop) for crop in gaez_crops()],
         si=[gaez_path("suitability", "i", crop) for crop in gaez_crops()],
         irrigated_share="data/downloads/gaez_land_equipped_for_irrigation_share.tif",
-        regions="<processing>/{name}/regions.geojson",
     params:
         irrigated_area_source=config["aggregation"]["irrigated_area_source"],
     output:
@@ -156,7 +155,7 @@ rule aggregate_class_areas:
         "prep"
     resources:
         runtime="1m",
-        mem_mb=3000,
+        mem_mb=800,
     log:
         "<logs>/{name}/aggregate_class_areas.log",
     benchmark:

workflow/rules/luc.smk

@@ -304,7 +304,7 @@ rule build_grazing_only_land:
 rule build_luc_carbon_coefficients:
     input:
         classes="<processing>/{name}/resource_classes.nc",
-        regions="<processing>/{name}/regions.geojson",
+        cell_mapping=rules.build_region_class_cell_mapping.output.mapping,
         agb=rules.prepare_luc_inputs.output.agb,
         soc=rules.prepare_luc_inputs.output.soc,
         regrowth=rules.prepare_luc_inputs.output.regrowth,

workflow/rules/model.smk

@@ -230,7 +230,7 @@ rule build_model:
         "build_model"
     resources:
         runtime="1m",
-        mem_mb=900,
+        mem_mb=1500,
     log:
         "<logs>/{name}/build_model.log",
     benchmark: