From 0c1eb6d44d88a98b2e9020fac487444fbcae3ec3 Mon Sep 17 00:00:00 2001 From: Koen van Greevenbroek Date: Wed, 22 Jul 2026 12:06:53 -0700 Subject: [PATCH 1/4] perf: fuse crop yield zonal extraction --- workflow/scripts/build_crop_yields.py | 179 ++++++++++++-------------- 1 file changed, 83 insertions(+), 96 deletions(-) diff --git a/workflow/scripts/build_crop_yields.py b/workflow/scripts/build_crop_yields.py index 89db4884..9694d4e8 100644 --- a/workflow/scripts/build_crop_yields.py +++ b/workflow/scripts/build_crop_yields.py @@ -11,7 +11,7 @@ gdal.UseExceptions() osr.UseExceptions() -from exactextract import exact_extract # noqa: E402 +from exactextract import Operation, exact_extract # noqa: E402 from exactextract.raster import NumPyRasterSource # noqa: E402 import geopandas as gpd # noqa: E402 import numpy as np # noqa: E402 @@ -47,11 +47,11 @@ KG_TO_TONNE = 0.001 # Load classes - ds = xr.load_dataset(classes_nc) - class_labels = ds["resource_class"].values.astype(np.int16) + with xr.open_dataset(classes_nc) as ds: + class_labels = ds["resource_class"].load().values # Load rasters - y_raw, y_src = read_raster_float(yield_path) + y_tpha, y_src = read_raster_float(yield_path) conversion_overrides: dict[str, float] = {} if conv_csv: conversion_overrides = ( @@ -72,7 +72,7 @@ def _yield_multiplier(crop: str) -> float: # GAEZ publishes RES05 potential yields in kg/ha but the historical - # “actual yield” variant in t/ha. Validation runs toggle + # "actual yield" variant in t/ha. Validation runs toggle # ``use_actual_yields`` so we keep the raster units untouched in that # mode while the standard pathway still divides by 1_000. base_scale = 1.0 if use_actual_yields else KG_TO_TONNE @@ -86,17 +86,17 @@ def _yield_multiplier(crop: str) -> float: # multiplier so the same table works for both actual and potential runs. return base_scale * (override / KG_TO_TONNE) - y_tpha = y_raw * _yield_multiplier(crop_code) + y_tpha *= _yield_multiplier(crop_code) if use_actual_yields: moisture_fraction = float(moisture_lookup[crop_code]) - y_tpha = y_tpha * (1.0 - moisture_fraction) + y_tpha *= 1.0 - moisture_fraction s_raw, _ = read_raster_float(suit_path) s_frac = scale_fraction(s_raw) if water_path: - water_raw_mm, _ = read_raster_float(water_path) - water_m3_per_ha = water_raw_mm * 10.0 # 1 mm over 1 ha equals 10 m³ + water_m3_per_ha, _ = read_raster_float(water_path) + water_m3_per_ha *= 10.0 # 1 mm over 1 ha equals 10 m3 else: - water_m3_per_ha = np.zeros_like(y_raw) + water_m3_per_ha = np.zeros_like(y_tpha) gs_start_raw, _ = read_raster_float(gs_start_path) gs_length_raw, _ = read_raster_float(gs_length_path) @@ -108,7 +108,8 @@ def _yield_multiplier(crop: str) -> float: # Use 1D cell areas and broadcast to save memory cell_area_ha_1d = calculate_all_cell_areas(y_src, repeat=False) - area_ha = s_frac * cell_area_ha_1d[:, np.newaxis] + s_frac *= cell_area_ha_1d[:, np.newaxis] + area_ha = s_frac # Regions regions_gdf = gpd.read_file(regions_path) @@ -116,7 +117,8 @@ def _yield_multiplier(crop: str) -> float: regions_gdf = regions_gdf.to_crs(crs) regions_for_extract = regions_gdf.reset_index() - # Create raster sources once (before the loop) to avoid repeated allocations + # Build every class-specific operation up front so exactextract traverses each + # region geometry only once for all variables and resource classes. raster_kwargs = { "xmin": xmin, "ymin": ymin, @@ -125,108 +127,93 @@ def _yield_multiplier(crop: str) -> float: "nodata": np.nan, "srs_wkt": crs_wkt, } - y_src_np = NumPyRasterSource(y_tpha, **raster_kwargs) - a_src_np = NumPyRasterSource(area_ha.astype(np.float32), **raster_kwargs) - water_src_np = NumPyRasterSource(water_m3_per_ha, **raster_kwargs) - gs_start_src_np = NumPyRasterSource(gs_start_raw, **raster_kwargs) - gs_length_src_np = NumPyRasterSource(gs_length_raw, **raster_kwargs) + y_src_np = NumPyRasterSource(y_tpha, name="yield", **raster_kwargs) + a_src_np = NumPyRasterSource(area_ha, name="suitable_area", **raster_kwargs) + water_src_np = NumPyRasterSource( + water_m3_per_ha, name="water_requirement", **raster_kwargs + ) + gs_start_src_np = NumPyRasterSource( + gs_start_raw, name="growing_season_start", **raster_kwargs + ) + gs_length_src_np = NumPyRasterSource( + gs_length_raw, name="growing_season_length", **raster_kwargs + ) + value_sources = [ + y_src_np, + a_src_np, + water_src_np, + gs_start_src_np, + gs_length_src_np, + ] - # Aggregate mean yield and sum area per class using weighted extraction - out = [] n_classes = ( int(np.nanmax(class_labels)) + 1 if np.isfinite(class_labels).any() else 0 ) + # Operation borrows its weight RasterSource, so keep the owners alive. + class_sources = [] + operations = [] + valid_classes = [] for cls in range(n_classes): - # Create binary mask as weight (only allocation per iteration) - mask_float = (class_labels == cls).astype(np.float32) - if not np.any(mask_float > 0): + class_mask = class_labels == cls + if not np.any(class_mask): continue - - # Use the mask as weights (0=not this class, 1=this class) - # Don't set nodata so 0s are treated as zero weight, not missing data mask_src = NumPyRasterSource( - mask_float, + class_mask, xmin=xmin, ymin=ymin, xmax=xmax, ymax=ymax, + name=f"resource_class_{cls}", srs_wkt=crs_wkt, ) - - # Use weighted operations with class mask as weight - y_stats = exact_extract( - y_src_np, - regions_for_extract, - ["weighted_mean"], - weights=mask_src, - include_cols=["region"], - output="pandas", - ) - a_stats = exact_extract( - a_src_np, - regions_for_extract, - ["weighted_sum"], - weights=mask_src, - include_cols=["region"], - output="pandas", - ) - water_stats = exact_extract( - water_src_np, - regions_for_extract, - ["weighted_mean"], - weights=mask_src, - include_cols=["region"], - output="pandas", - ) - gs_start_stats = exact_extract( - gs_start_src_np, - regions_for_extract, - ["weighted_mean"], - weights=mask_src, - include_cols=["region"], - output="pandas", + class_sources.append(mask_src) + valid_classes.append(cls) + operations.extend( + [ + Operation("weighted_mean", f"yield_{cls}", y_src_np, mask_src), + Operation("weighted_sum", f"suitable_area_{cls}", a_src_np, mask_src), + Operation( + "weighted_mean", + f"water_requirement_m3_per_ha_{cls}", + water_src_np, + mask_src, + ), + Operation( + "weighted_mean", + f"growing_season_start_day_{cls}", + gs_start_src_np, + mask_src, + ), + Operation( + "weighted_mean", + f"growing_season_length_days_{cls}", + gs_length_src_np, + mask_src, + ), + ] ) - gs_length_stats = exact_extract( - gs_length_src_np, + + out = [] + if operations: + stats = exact_extract( + value_sources, regions_for_extract, - ["weighted_mean"], - weights=mask_src, + operations, include_cols=["region"], output="pandas", ) - if y_stats.empty or a_stats.empty: - continue - merged = ( - y_stats.rename(columns={"weighted_mean": "yield"}) - .merge( - a_stats.rename(columns={"weighted_sum": "suitable_area"}), - on="region", - how="inner", - ) - .merge( - water_stats.rename( - columns={"weighted_mean": "water_requirement_m3_per_ha"} - ), - on="region", - how="left", - ) - .merge( - gs_start_stats.rename( - columns={"weighted_mean": "growing_season_start_day"} - ), - on="region", - how="left", - ) - .merge( - gs_length_stats.rename( - columns={"weighted_mean": "growing_season_length_days"} - ), - on="region", - how="left", - ) - ) - merged["resource_class"] = cls - out.append(merged) + variables = [ + "yield", + "suitable_area", + "water_requirement_m3_per_ha", + "growing_season_start_day", + "growing_season_length_days", + ] + for cls in valid_classes: + columns = {f"{variable}_{cls}": variable for variable in variables} + class_stats = stats[["region", *columns]].rename(columns=columns) + class_stats["resource_class"] = cls + out.append(class_stats) if out: df = ( From d8be1674fa858d62dc6286e8886fe06a06b1fb87 Mon Sep 17 00:00:00 2001 From: Koen van Greevenbroek Date: Wed, 22 Jul 2026 12:10:52 -0700 Subject: [PATCH 2/4] perf: cache crop yield cell coverage --- docs/crop_production.rst | 4 +- docs/workflow.rst | 8 +- tests/test_crop_yield_aggregation.py | 64 +++++ workflow/rules/crops.smk | 24 +- .../scripts/build_crop_yield_cell_mapping.py | 108 ++++++++ workflow/scripts/build_crop_yields.py | 247 +++++------------- workflow/scripts/crop_yield_aggregation.py | 79 ++++++ 7 files changed, 346 insertions(+), 188 deletions(-) create mode 100644 tests/test_crop_yield_aggregation.py create mode 100644 workflow/scripts/build_crop_yield_cell_mapping.py create mode 100644 workflow/scripts/crop_yield_aggregation.py diff --git a/docs/crop_production.rst b/docs/crop_production.rst index 014d617c..d2901295 100644 --- a/docs/crop_production.rst +++ b/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_crop_yield_cell_mapping.py`` and reused for every crop and water supply. 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**: diff --git a/docs/workflow.rst b/docs/workflow.rst index 14dc896d..34c947ee 100644 --- a/docs/workflow.rst +++ b/docs/workflow.rst @@ -71,9 +71,15 @@ Data Preparation Rules * **Script**: ``workflow/scripts/aggregate_class_areas.py`` * **Purpose**: Compute available land area per (region, class, water, crop) +**build_crop_yield_cell_mapping** + * **Input**: Resource classes, regions + * **Output**: ``processing/{name}/crop_yield_cell_mapping.npz`` + * **Script**: ``workflow/scripts/build_crop_yield_cell_mapping.py`` + * **Purpose**: Cache exact region and resource-class coverage by GAEZ grid cell + **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 diff --git a/tests/test_crop_yield_aggregation.py b/tests/test_crop_yield_aggregation.py new file mode 100644 index 00000000..11e7520f --- /dev/null +++ b/tests/test_crop_yield_aggregation.py @@ -0,0 +1,64 @@ +# 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 shapely.geometry import box +import xarray as xr + +from workflow.scripts.build_crop_yield_cell_mapping import build_cell_mapping +from workflow.scripts.crop_yield_aggregation import ( + load_cell_mapping, + 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.5, 2.0), box(1.5, 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, 15.0, 0.0, 15.0]) + + +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)) diff --git a/workflow/rules/crops.smk b/workflow/rules/crops.smk index 82b4df4d..80a3803c 100644 --- a/workflow/rules/crops.smk +++ b/workflow/rules/crops.smk @@ -149,11 +149,29 @@ def yield_and_suitability_for_crop(w): return inputs -rule build_crop_yields: +rule build_crop_yield_cell_mapping: input: - unpack(yield_and_suitability_for_crop), classes="/{name}/resource_classes.nc", regions="/{name}/regions.geojson", + output: + mapping="/{name}/crop_yield_cell_mapping.npz", + group: + "prep" + resources: + runtime="1m", + mem_mb=600, + log: + "/{name}/build_crop_yield_cell_mapping.log", + benchmark: + "/{name}/build_crop_yield_cell_mapping.tsv" + script: + "../scripts/build_crop_yield_cell_mapping.py" + + +rule build_crop_yields: + input: + unpack(yield_and_suitability_for_crop), + cell_mapping="/{name}/crop_yield_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: "/{name}/build_crop_yields_{crop}_{water_supply}.log", benchmark: diff --git a/workflow/scripts/build_crop_yield_cell_mapping.py b/workflow/scripts/build_crop_yield_cell_mapping.py new file mode 100644 index 00000000..c75081d4 --- /dev/null +++ b/workflow/scripts/build_crop_yield_cell_mapping.py @@ -0,0 +1,108 @@ +""" +SPDX-FileCopyrightText: 2026 Koen van Greevenbroek + +SPDX-License-Identifier: GPL-3.0-or-later + +Build exact region/resource-class coverage arrays for the common GAEZ grid. + +Inputs +------ +``classes`` + NetCDF resource-class raster, including the grid transform and CRS. +``regions`` + GeoJSON optimization regions. + +Output +------ +``mapping`` + NPZ arrays used by every crop-yield aggregation for this configuration. +""" + +from pathlib import Path + +from osgeo import gdal, osr + +gdal.UseExceptions() +osr.UseExceptions() + +from exactextract import exact_extract # noqa: E402 +from exactextract.raster import NumPyRasterSource # noqa: E402 +import geopandas as gpd # noqa: E402 +import numpy as np # noqa: E402 +from pyproj import CRS # noqa: E402 +import xarray as xr # noqa: E402 + + +def build_cell_mapping(classes_path: str, regions_path: str, output_path: str) -> None: + """Write exact region/class coverage for each relevant raster cell.""" + classes_ds = xr.load_dataset(classes_path) + class_labels = classes_ds["resource_class"].values + height, width = class_labels.shape + transform = np.asarray(classes_ds.attrs["transform"], dtype=float) + if transform[2] != 0 or transform[4] != 0: + raise ValueError("Rotated resource-class grids are not supported") + + xmin = transform[0] + xmax = xmin + width * transform[1] + ymax = transform[3] + ymin = ymax + height * transform[5] + crs_wkt = str(classes_ds.attrs["crs_wkt"]) + + regions = gpd.read_file(regions_path) + grid_crs = CRS.from_wkt(crs_wkt) + if regions.crs and regions.crs != grid_crs: + regions = regions.to_crs(grid_crs) + regions = regions.reset_index() + + grid = NumPyRasterSource( + class_labels, + xmin=xmin, + ymin=ymin, + xmax=xmax, + ymax=ymax, + srs_wkt=crs_wkt, + ) + extracted = exact_extract( + grid, + regions, + ["cell_id", "coverage"], + include_cols=["region"], + output="pandas", + ) + + lengths = np.fromiter( + (len(cell_ids) for cell_ids in extracted["cell_id"]), + dtype=np.int64, + ) + cell_ids = np.concatenate(extracted["cell_id"].to_numpy()).astype(np.int32) + coverage = np.concatenate(extracted["coverage"].to_numpy()).astype(np.float32) + region_ids = np.repeat(np.arange(len(extracted), dtype=np.int32), lengths) + class_ids = class_labels.ravel()[cell_ids] + valid = class_ids >= 0 + cell_ids = cell_ids[valid] + coverage = coverage[valid] + class_ids = class_ids[valid] + + n_classes = int(class_labels.max()) + 1 + group_ids = (region_ids[valid] * n_classes + class_ids).astype(np.int32) + region_names = extracted["region"].astype(str).to_numpy(dtype=str) + + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + np.savez( + output_path, + cell_ids=cell_ids, + coverage=coverage, + group_ids=group_ids, + regions=region_names, + n_classes=np.array(n_classes, dtype=np.int32), + height=np.array(height, dtype=np.int32), + width=np.array(width, dtype=np.int32), + ) + + +if __name__ == "__main__": + build_cell_mapping( + snakemake.input.classes, # type: ignore[name-defined] + snakemake.input.regions, # type: ignore[name-defined] + snakemake.output.mapping, # type: ignore[name-defined] + ) diff --git a/workflow/scripts/build_crop_yields.py b/workflow/scripts/build_crop_yields.py index 9694d4e8..820264e5 100644 --- a/workflow/scripts/build_crop_yields.py +++ b/workflow/scripts/build_crop_yields.py @@ -6,34 +6,28 @@ from pathlib import Path -from osgeo import gdal, osr +import numpy as np +import pandas as pd -gdal.UseExceptions() -osr.UseExceptions() - -from exactextract import Operation, exact_extract # noqa: E402 -from exactextract.raster import NumPyRasterSource # noqa: E402 -import geopandas as gpd # noqa: E402 -import numpy as np # noqa: E402 -import pandas as pd # noqa: E402 -import xarray as xr # noqa: E402 - -from workflow.scripts.raster_utils import ( # noqa: E402 +from workflow.scripts.crop_yield_aggregation import ( + load_cell_mapping, + weighted_mean_by_group, + weighted_sum_by_group, +) +from workflow.scripts.raster_utils import ( calculate_all_cell_areas, - raster_bounds, read_raster_float, scale_fraction, ) if __name__ == "__main__": # Inputs - classes_nc: str = snakemake.input.classes # type: ignore[name-defined] + mapping_path: str = snakemake.input.cell_mapping # type: ignore[name-defined] yield_path: str = snakemake.input.yield_raster # type: ignore[name-defined] suit_path: str = snakemake.input.suitability_raster # type: ignore[name-defined] water_path: str | None = getattr( # type: ignore[attr-defined] snakemake.input, "water_requirement_raster", None ) - regions_path: str = snakemake.input.regions # type: ignore[name-defined] gs_start_path: str = snakemake.input.growing_season_start_raster # type: ignore[name-defined] gs_length_path: str = snakemake.input.growing_season_length_raster # type: ignore[name-defined] crop_code: str = snakemake.wildcards.crop # type: ignore[name-defined] @@ -46,12 +40,8 @@ KG_TO_TONNE = 0.001 - # Load classes - with xr.open_dataset(classes_nc) as ds: - class_labels = ds["resource_class"].load().values + mapping = load_cell_mapping(mapping_path) - # Load rasters - y_tpha, y_src = read_raster_float(yield_path) conversion_overrides: dict[str, float] = {} if conv_csv: conversion_overrides = ( @@ -86,179 +76,72 @@ def _yield_multiplier(crop: str) -> float: # multiplier so the same table works for both actual and potential runs. return base_scale * (override / KG_TO_TONNE) - y_tpha *= _yield_multiplier(crop_code) + y_raw, y_src = read_raster_float(yield_path) + y_tpha = y_raw * _yield_multiplier(crop_code) if use_actual_yields: moisture_fraction = float(moisture_lookup[crop_code]) - y_tpha *= 1.0 - moisture_fraction - s_raw, _ = read_raster_float(suit_path) - s_frac = scale_fraction(s_raw) - if water_path: - water_m3_per_ha, _ = read_raster_float(water_path) - water_m3_per_ha *= 10.0 # 1 mm over 1 ha equals 10 m3 - else: - water_m3_per_ha = np.zeros_like(y_tpha) - gs_start_raw, _ = read_raster_float(gs_start_path) - gs_length_raw, _ = read_raster_float(gs_length_path) - - height, width = y_tpha.shape - transform = y_src.transform - crs = y_src.crs - crs_wkt = crs.to_wkt() if crs else None - xmin, ymin, xmax, ymax = raster_bounds(transform, width, height) - # Use 1D cell areas and broadcast to save memory + y_tpha = y_tpha * (1.0 - moisture_fraction) + yield_by_group = weighted_mean_by_group(y_tpha, mapping) cell_area_ha_1d = calculate_all_cell_areas(y_src, repeat=False) + y_src.close() + del y_raw, y_tpha - s_frac *= cell_area_ha_1d[:, np.newaxis] - area_ha = s_frac - - # Regions - regions_gdf = gpd.read_file(regions_path) - if regions_gdf.crs and crs and regions_gdf.crs != crs: - regions_gdf = regions_gdf.to_crs(crs) - regions_for_extract = regions_gdf.reset_index() - - # Build every class-specific operation up front so exactextract traverses each - # region geometry only once for all variables and resource classes. - raster_kwargs = { - "xmin": xmin, - "ymin": ymin, - "xmax": xmax, - "ymax": ymax, - "nodata": np.nan, - "srs_wkt": crs_wkt, - } - y_src_np = NumPyRasterSource(y_tpha, name="yield", **raster_kwargs) - a_src_np = NumPyRasterSource(area_ha, name="suitable_area", **raster_kwargs) - water_src_np = NumPyRasterSource( - water_m3_per_ha, name="water_requirement", **raster_kwargs - ) - gs_start_src_np = NumPyRasterSource( - gs_start_raw, name="growing_season_start", **raster_kwargs - ) - gs_length_src_np = NumPyRasterSource( - gs_length_raw, name="growing_season_length", **raster_kwargs - ) - value_sources = [ - y_src_np, - a_src_np, - water_src_np, - gs_start_src_np, - gs_length_src_np, - ] - - n_classes = ( - int(np.nanmax(class_labels)) + 1 if np.isfinite(class_labels).any() else 0 - ) - # Operation borrows its weight RasterSource, so keep the owners alive. - class_sources = [] - operations = [] - valid_classes = [] - for cls in range(n_classes): - class_mask = class_labels == cls - if not np.any(class_mask): - continue - mask_src = NumPyRasterSource( - class_mask, - xmin=xmin, - ymin=ymin, - xmax=xmax, - ymax=ymax, - name=f"resource_class_{cls}", - srs_wkt=crs_wkt, - ) - class_sources.append(mask_src) - valid_classes.append(cls) - operations.extend( - [ - Operation("weighted_mean", f"yield_{cls}", y_src_np, mask_src), - Operation("weighted_sum", f"suitable_area_{cls}", a_src_np, mask_src), - Operation( - "weighted_mean", - f"water_requirement_m3_per_ha_{cls}", - water_src_np, - mask_src, - ), - Operation( - "weighted_mean", - f"growing_season_start_day_{cls}", - gs_start_src_np, - mask_src, - ), - Operation( - "weighted_mean", - f"growing_season_length_days_{cls}", - gs_length_src_np, - mask_src, - ), - ] - ) - - out = [] - if operations: - stats = exact_extract( - value_sources, - regions_for_extract, - operations, - include_cols=["region"], - output="pandas", - ) - variables = [ - "yield", - "suitable_area", - "water_requirement_m3_per_ha", - "growing_season_start_day", - "growing_season_length_days", - ] - for cls in valid_classes: - columns = {f"{variable}_{cls}": variable for variable in variables} - class_stats = stats[["region", *columns]].rename(columns=columns) - class_stats["resource_class"] = cls - out.append(class_stats) + s_raw, s_src = read_raster_float(suit_path) + s_src.close() + s_frac = scale_fraction(s_raw) + area_ha = s_frac * cell_area_ha_1d[:, np.newaxis] + area_by_group = weighted_sum_by_group(area_ha, mapping) + del s_raw, s_frac, area_ha - if out: - df = ( - pd.concat(out, ignore_index=True) - .set_index(["region", "resource_class"]) - .sort_index() - ) + if water_path: + water_raw_mm, water_src = read_raster_float(water_path) + water_src.close() + water_m3_per_ha = water_raw_mm * 10.0 # 1 mm over 1 ha equals 10 m3 + water_by_group = weighted_mean_by_group(water_m3_per_ha, mapping) + del water_raw_mm, water_m3_per_ha else: - df = pd.DataFrame( - columns=[ - "region", - "resource_class", - "yield", - "suitable_area", - "water_requirement_m3_per_ha", - "growing_season_start_day", - "growing_season_length_days", - ] - ).set_index(["region", "resource_class"]) # type: ignore[name-defined] - - df_reset = df.reset_index() - df_reset["resource_class"] = df_reset["resource_class"].astype(int) - - variable_units = { - "yield": "t/ha (DM)", - "suitable_area": "ha", - "water_requirement_m3_per_ha": "m^3/ha", - "growing_season_start_day": "day-of-year", - "growing_season_length_days": "days", + weight = np.bincount( + mapping.group_ids, + weights=mapping.coverage, + minlength=mapping.n_groups, + ) + water_by_group = np.full(mapping.n_groups, np.nan) + water_by_group[weight != 0] = 0.0 + + gs_start_raw, gs_start_src = read_raster_float(gs_start_path) + gs_start_src.close() + gs_start_by_group = weighted_mean_by_group(gs_start_raw, mapping) + del gs_start_raw + gs_length_raw, gs_length_src = read_raster_float(gs_length_path) + gs_length_src.close() + gs_length_by_group = weighted_mean_by_group(gs_length_raw, mapping) + del gs_length_raw + + variable_values_and_units = { + "yield": (yield_by_group, "t/ha (DM)"), + "suitable_area": (area_by_group, "ha"), + "water_requirement_m3_per_ha": (water_by_group, "m^3/ha"), + "growing_season_start_day": (gs_start_by_group, "day-of-year"), + "growing_season_length_days": (gs_length_by_group, "days"), } + region_index = np.repeat(mapping.regions, mapping.n_classes) + class_index = np.tile(np.arange(mapping.n_classes), len(mapping.regions)) tidy_frames = [] - for variable, unit in variable_units.items(): - if variable not in df_reset.columns: - continue - subset = df_reset[["region", "resource_class", variable]].dropna( - subset=[variable] - ) - if subset.empty: + for variable, (values, unit) in variable_values_and_units.items(): + valid = ~np.isnan(values) + if not np.any(valid): continue - subset = subset.rename(columns={variable: "value"}) - subset["variable"] = variable - subset["unit"] = unit tidy_frames.append( - subset[["region", "resource_class", "variable", "unit", "value"]] + pd.DataFrame( + { + "region": region_index[valid], + "resource_class": class_index[valid], + "variable": variable, + "unit": unit, + "value": values[valid], + } + ) ) if tidy_frames: diff --git a/workflow/scripts/crop_yield_aggregation.py b/workflow/scripts/crop_yield_aggregation.py new file mode 100644 index 00000000..c358f8ba --- /dev/null +++ b/workflow/scripts/crop_yield_aggregation.py @@ -0,0 +1,79 @@ +""" +SPDX-FileCopyrightText: 2026 Koen van Greevenbroek + +SPDX-License-Identifier: GPL-3.0-or-later +""" + +from typing import NamedTuple + +import numpy as np + + +class CellMapping(NamedTuple): + """Exact region/class coverage for cells on the common GAEZ grid.""" + + cell_ids: np.ndarray + coverage: np.ndarray + group_ids: np.ndarray + regions: np.ndarray + n_classes: int + shape: tuple[int, int] + + @property + def n_groups(self) -> int: + return len(self.regions) * self.n_classes + + +def load_cell_mapping(path: str) -> CellMapping: + """Load the config-specific region/class cell mapping from an NPZ file.""" + with np.load(path, allow_pickle=False) as data: + return CellMapping( + cell_ids=data["cell_ids"], + coverage=data["coverage"].astype(np.float64), + group_ids=data["group_ids"], + regions=data["regions"], + n_classes=int(data["n_classes"]), + shape=(int(data["height"]), int(data["width"])), + ) + + +def _mapped_values(values: np.ndarray, mapping: CellMapping) -> np.ndarray: + if values.shape != mapping.shape: + raise ValueError( + f"Raster shape {values.shape} does not match cell mapping " + f"shape {mapping.shape}" + ) + return values.ravel()[mapping.cell_ids] + + +def weighted_mean_by_group(values: np.ndarray, mapping: CellMapping) -> np.ndarray: + """Return coverage-weighted means for every region/resource-class group.""" + mapped = _mapped_values(values, mapping) + valid = ~np.isnan(mapped) + numerator = np.bincount( + mapping.group_ids[valid], + weights=mapped[valid] * mapping.coverage[valid], + minlength=mapping.n_groups, + ) + denominator = np.bincount( + mapping.group_ids[valid], + weights=mapping.coverage[valid], + minlength=mapping.n_groups, + ) + return np.divide( + numerator, + denominator, + out=np.full(mapping.n_groups, np.nan), + where=denominator != 0, + ) + + +def weighted_sum_by_group(values: np.ndarray, mapping: CellMapping) -> np.ndarray: + """Return coverage-weighted sums for every region/resource-class group.""" + mapped = _mapped_values(values, mapping) + valid = ~np.isnan(mapped) + return np.bincount( + mapping.group_ids[valid], + weights=mapped[valid] * mapping.coverage[valid], + minlength=mapping.n_groups, + ) From 920b932e1d7997afa81471d7abbcc92e8a1bf68e Mon Sep 17 00:00:00 2001 From: Koen van Greevenbroek Date: Wed, 22 Jul 2026 12:16:14 -0700 Subject: [PATCH 3/4] perf: tighten crop yield cache handling --- CHANGELOG.md | 3 ++ .../scripts/build_crop_yield_cell_mapping.py | 15 ++++---- workflow/scripts/build_crop_yields.py | 34 +++++++------------ workflow/scripts/crop_yield_aggregation.py | 2 +- 4 files changed, 25 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bc834b69..837b783f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -86,6 +86,9 @@ introduce breaking changes to configuration and outputs. ### Changed +- Crop-yield preparation now computes exact region and resource-class cell + coverage once per configuration and reuses it across crops, substantially + reducing build time and peak memory without changing crop-yield 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 diff --git a/workflow/scripts/build_crop_yield_cell_mapping.py b/workflow/scripts/build_crop_yield_cell_mapping.py index c75081d4..6bf22a13 100644 --- a/workflow/scripts/build_crop_yield_cell_mapping.py +++ b/workflow/scripts/build_crop_yield_cell_mapping.py @@ -35,10 +35,11 @@ def build_cell_mapping(classes_path: str, regions_path: str, output_path: str) -> None: """Write exact region/class coverage for each relevant raster cell.""" - classes_ds = xr.load_dataset(classes_path) - class_labels = classes_ds["resource_class"].values + with xr.open_dataset(classes_path) as classes_ds: + class_labels = classes_ds["resource_class"].load().values + transform = np.asarray(classes_ds.attrs["transform"], dtype=float) + crs_wkt = str(classes_ds.attrs["crs_wkt"]) height, width = class_labels.shape - transform = np.asarray(classes_ds.attrs["transform"], dtype=float) if transform[2] != 0 or transform[4] != 0: raise ValueError("Rotated resource-class grids are not supported") @@ -46,8 +47,6 @@ def build_cell_mapping(classes_path: str, regions_path: str, output_path: str) - xmax = xmin + width * transform[1] ymax = transform[3] ymin = ymax + height * transform[5] - crs_wkt = str(classes_ds.attrs["crs_wkt"]) - regions = gpd.read_file(regions_path) grid_crs = CRS.from_wkt(crs_wkt) if regions.crs and regions.crs != grid_crs: @@ -75,7 +74,7 @@ def build_cell_mapping(classes_path: str, regions_path: str, output_path: str) - dtype=np.int64, ) cell_ids = np.concatenate(extracted["cell_id"].to_numpy()).astype(np.int32) - coverage = np.concatenate(extracted["coverage"].to_numpy()).astype(np.float32) + coverage = np.concatenate(extracted["coverage"].to_numpy()).astype(np.float64) region_ids = np.repeat(np.arange(len(extracted), dtype=np.int32), lengths) class_ids = class_labels.ravel()[cell_ids] valid = class_ids >= 0 @@ -83,7 +82,9 @@ def build_cell_mapping(classes_path: str, regions_path: str, output_path: str) - coverage = coverage[valid] class_ids = class_ids[valid] - n_classes = int(class_labels.max()) + 1 + if not np.any(valid): + raise ValueError("Resource-class grid does not contain any valid classes") + n_classes = int(class_ids.max()) + 1 group_ids = (region_ids[valid] * n_classes + class_ids).astype(np.int32) region_names = extracted["region"].astype(str).to_numpy(dtype=str) diff --git a/workflow/scripts/build_crop_yields.py b/workflow/scripts/build_crop_yields.py index 820264e5..970fb8d0 100644 --- a/workflow/scripts/build_crop_yields.py +++ b/workflow/scripts/build_crop_yields.py @@ -31,34 +31,26 @@ gs_start_path: str = snakemake.input.growing_season_start_raster # type: ignore[name-defined] gs_length_path: str = snakemake.input.growing_season_length_raster # type: ignore[name-defined] crop_code: str = snakemake.wildcards.crop # type: ignore[name-defined] - conv_csv: str | None = getattr( # type: ignore[attr-defined] - snakemake.input, "yield_unit_conversions", None - ) - moisture_csv: str | None = getattr( # type: ignore[attr-defined] - snakemake.input, "moisture_content", None - ) + conv_csv: str = snakemake.input.yield_unit_conversions # type: ignore[name-defined] + moisture_csv: str = snakemake.input.moisture_content # type: ignore[name-defined] KG_TO_TONNE = 0.001 mapping = load_cell_mapping(mapping_path) - conversion_overrides: dict[str, float] = {} - if conv_csv: - conversion_overrides = ( - pd.read_csv(conv_csv, comment="#") - .set_index("code")["factor_to_t_per_ha"] - .to_dict() - ) + conversion_overrides: dict[str, float] = ( + pd.read_csv(conv_csv, comment="#") + .set_index("code")["factor_to_t_per_ha"] + .to_dict() + ) - use_actual_yields = bool(getattr(snakemake.params, "use_actual_yields", False)) # type: ignore[attr-defined] + use_actual_yields = bool(snakemake.params.use_actual_yields) # type: ignore[name-defined] - moisture_lookup: dict[str, float] = {} - if moisture_csv: - moisture_lookup = ( - pd.read_csv(moisture_csv, comment="#") - .set_index("crop")["moisture_fraction"] - .to_dict() - ) + moisture_lookup: dict[str, float] = ( + pd.read_csv(moisture_csv, comment="#") + .set_index("crop")["moisture_fraction"] + .to_dict() + ) def _yield_multiplier(crop: str) -> float: # GAEZ publishes RES05 potential yields in kg/ha but the historical diff --git a/workflow/scripts/crop_yield_aggregation.py b/workflow/scripts/crop_yield_aggregation.py index c358f8ba..429ed189 100644 --- a/workflow/scripts/crop_yield_aggregation.py +++ b/workflow/scripts/crop_yield_aggregation.py @@ -29,7 +29,7 @@ def load_cell_mapping(path: str) -> CellMapping: with np.load(path, allow_pickle=False) as data: return CellMapping( cell_ids=data["cell_ids"], - coverage=data["coverage"].astype(np.float64), + coverage=data["coverage"], group_ids=data["group_ids"], regions=data["regions"], n_classes=int(data["n_classes"]), From 1621aa2addc7515164ac50b79a759b4b4f4b835e Mon Sep 17 00:00:00 2001 From: Koen van Greevenbroek Date: Wed, 22 Jul 2026 12:27:17 -0700 Subject: [PATCH 4/4] perf: reuse cell coverage for harvested area --- CHANGELOG.md | 7 +- docs/crop_production.rst | 2 +- docs/workflow.rst | 15 ++- tests/test_crop_yield_aggregation.py | 64 --------- tests/test_region_class_aggregation.py | 123 ++++++++++++++++++ workflow/rules/crops.smk | 16 +-- workflow/scripts/build_crop_yields.py | 16 ++- workflow/scripts/build_harvested_area.py | 96 ++++---------- ....py => build_region_class_cell_mapping.py} | 2 + ...egation.py => region_class_aggregation.py} | 28 ++++ 10 files changed, 214 insertions(+), 155 deletions(-) delete mode 100644 tests/test_crop_yield_aggregation.py create mode 100644 tests/test_region_class_aggregation.py rename workflow/scripts/{build_crop_yield_cell_mapping.py => build_region_class_cell_mapping.py} (98%) rename workflow/scripts/{crop_yield_aggregation.py => region_class_aggregation.py} (66%) diff --git a/CHANGELOG.md b/CHANGELOG.md index 837b783f..d089f7b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -86,9 +86,10 @@ introduce breaking changes to configuration and outputs. ### Changed -- Crop-yield preparation now computes exact region and resource-class cell - coverage once per configuration and reuses it across crops, substantially - reducing build time and peak memory without changing crop-yield outputs. +- 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 diff --git a/docs/crop_production.rst b/docs/crop_production.rst index d2901295..02b97e9e 100644 --- a/docs/crop_production.rst +++ b/docs/crop_production.rst @@ -115,7 +115,7 @@ 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``. Exact region-to-cell coverage fractions are computed once per configuration by ``workflow/scripts/build_crop_yield_cell_mapping.py`` and reused for every crop and water supply. +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 ~~~~~~~~~~~~~~~~~~~ diff --git a/docs/workflow.rst b/docs/workflow.rst index 34c947ee..27a433cf 100644 --- a/docs/workflow.rst +++ b/docs/workflow.rst @@ -71,11 +71,11 @@ Data Preparation Rules * **Script**: ``workflow/scripts/aggregate_class_areas.py`` * **Purpose**: Compute available land area per (region, class, water, crop) -**build_crop_yield_cell_mapping** +**build_region_class_cell_mapping** * **Input**: Resource classes, regions - * **Output**: ``processing/{name}/crop_yield_cell_mapping.npz`` - * **Script**: ``workflow/scripts/build_crop_yield_cell_mapping.py`` - * **Purpose**: Cache exact region and resource-class coverage by GAEZ grid cell + * **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") @@ -84,6 +84,13 @@ Data Preparation Rules * **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) diff --git a/tests/test_crop_yield_aggregation.py b/tests/test_crop_yield_aggregation.py deleted file mode 100644 index 11e7520f..00000000 --- a/tests/test_crop_yield_aggregation.py +++ /dev/null @@ -1,64 +0,0 @@ -# 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 shapely.geometry import box -import xarray as xr - -from workflow.scripts.build_crop_yield_cell_mapping import build_cell_mapping -from workflow.scripts.crop_yield_aggregation import ( - load_cell_mapping, - 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.5, 2.0), box(1.5, 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, 15.0, 0.0, 15.0]) - - -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)) diff --git a/tests/test_region_class_aggregation.py b/tests/test_region_class_aggregation.py new file mode 100644 index 00000000..a7804c62 --- /dev/null +++ b/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 diff --git a/workflow/rules/crops.smk b/workflow/rules/crops.smk index 80a3803c..1354a7e7 100644 --- a/workflow/rules/crops.smk +++ b/workflow/rules/crops.smk @@ -149,29 +149,29 @@ def yield_and_suitability_for_crop(w): return inputs -rule build_crop_yield_cell_mapping: +rule build_region_class_cell_mapping: input: classes="/{name}/resource_classes.nc", regions="/{name}/regions.geojson", output: - mapping="/{name}/crop_yield_cell_mapping.npz", + mapping="/{name}/region_class_cell_mapping.npz", group: "prep" resources: runtime="1m", mem_mb=600, log: - "/{name}/build_crop_yield_cell_mapping.log", + "/{name}/build_region_class_cell_mapping.log", benchmark: - "/{name}/build_crop_yield_cell_mapping.tsv" + "/{name}/build_region_class_cell_mapping.tsv" script: - "../scripts/build_crop_yield_cell_mapping.py" + "../scripts/build_region_class_cell_mapping.py" rule build_crop_yields: input: unpack(yield_and_suitability_for_crop), - cell_mapping="/{name}/crop_yield_cell_mapping.npz", + cell_mapping="/{name}/region_class_cell_mapping.npz", yield_unit_conversions="data/curated/yield_unit_conversions.csv", moisture_content="data/curated/crop_moisture_content.csv", params: @@ -455,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"/{w.name}/resource_classes.nc", + "cell_mapping": f"/{w.name}/region_class_cell_mapping.npz", "regions": f"/{w.name}/regions.geojson", "crop_mapping": "data/curated/gaez_crop_code_mapping.csv", "faostat_production": f"/{w.name}/faostat_crop_production.csv", @@ -499,7 +499,7 @@ rule build_harvested_area_gaez: "prep" resources: runtime="1m", - mem_mb=700, + mem_mb=400, log: "/{name}/build_harvested_area_gaez_{crop}_{water_supply}.log", benchmark: diff --git a/workflow/scripts/build_crop_yields.py b/workflow/scripts/build_crop_yields.py index 970fb8d0..4319492f 100644 --- a/workflow/scripts/build_crop_yields.py +++ b/workflow/scripts/build_crop_yields.py @@ -9,16 +9,17 @@ import numpy as np import pandas as pd -from workflow.scripts.crop_yield_aggregation import ( - load_cell_mapping, - weighted_mean_by_group, - weighted_sum_by_group, -) from workflow.scripts.raster_utils import ( calculate_all_cell_areas, read_raster_float, scale_fraction, ) +from workflow.scripts.region_class_aggregation import ( + load_cell_mapping, + validate_raster_grid, + weighted_mean_by_group, + weighted_sum_by_group, +) if __name__ == "__main__": # Inputs @@ -69,6 +70,7 @@ def _yield_multiplier(crop: str) -> float: return base_scale * (override / KG_TO_TONNE) y_raw, y_src = read_raster_float(yield_path) + validate_raster_grid(y_raw, y_src, mapping) y_tpha = y_raw * _yield_multiplier(crop_code) if use_actual_yields: moisture_fraction = float(moisture_lookup[crop_code]) @@ -79,6 +81,7 @@ def _yield_multiplier(crop: str) -> float: del y_raw, y_tpha s_raw, s_src = read_raster_float(suit_path) + validate_raster_grid(s_raw, s_src, mapping) s_src.close() s_frac = scale_fraction(s_raw) area_ha = s_frac * cell_area_ha_1d[:, np.newaxis] @@ -87,6 +90,7 @@ def _yield_multiplier(crop: str) -> float: if water_path: water_raw_mm, water_src = read_raster_float(water_path) + validate_raster_grid(water_raw_mm, water_src, mapping) water_src.close() water_m3_per_ha = water_raw_mm * 10.0 # 1 mm over 1 ha equals 10 m3 water_by_group = weighted_mean_by_group(water_m3_per_ha, mapping) @@ -101,10 +105,12 @@ def _yield_multiplier(crop: str) -> float: water_by_group[weight != 0] = 0.0 gs_start_raw, gs_start_src = read_raster_float(gs_start_path) + validate_raster_grid(gs_start_raw, gs_start_src, mapping) gs_start_src.close() gs_start_by_group = weighted_mean_by_group(gs_start_raw, mapping) del gs_start_raw gs_length_raw, gs_length_src = read_raster_float(gs_length_path) + validate_raster_grid(gs_length_raw, gs_length_src, mapping) gs_length_src.close() gs_length_by_group = weighted_mean_by_group(gs_length_raw, mapping) del gs_length_raw diff --git a/workflow/scripts/build_harvested_area.py b/workflow/scripts/build_harvested_area.py index ecc29c93..ab4aecb7 100644 --- a/workflow/scripts/build_harvested_area.py +++ b/workflow/scripts/build_harvested_area.py @@ -7,78 +7,42 @@ import logging from pathlib import Path -from osgeo import gdal, osr +import geopandas as gpd +import numpy as np +import pandas as pd -gdal.UseExceptions() -osr.UseExceptions() - -from exactextract import exact_extract # noqa: E402 -from exactextract.raster import NumPyRasterSource # noqa: E402 -import geopandas as gpd # noqa: E402 -import numpy as np # noqa: E402 -import pandas as pd # noqa: E402 -import xarray as xr # noqa: E402 - -from workflow.scripts.harvested_area_shares import ( # noqa: E402 +from workflow.scripts.harvested_area_shares import ( RES06_HAR_SCALE_TO_HA, apply_country_shares, load_mapping, shares_for_crop, shares_from_fdd, ) -from workflow.scripts.raster_utils import raster_bounds, read_raster_float # noqa: E402 +from workflow.scripts.raster_utils import read_raster_float +from workflow.scripts.region_class_aggregation import ( + CellMapping, + load_cell_mapping, + validate_raster_grid, + weighted_sum_by_group, +) logger = logging.getLogger(__name__) def _extract_harvested_area( raster: np.ndarray, - transform, - crs_wkt: str | None, - class_labels: np.ndarray, - regions: gpd.GeoDataFrame, + mapping: CellMapping, ) -> pd.DataFrame: - xmin, ymin, xmax, ymax = raster_bounds(transform, raster.shape[1], raster.shape[0]) - - regions_for_extract = regions.reset_index(drop=True) - - records: list[pd.DataFrame] = [] - n_classes = ( - int(np.nanmax(class_labels)) + 1 if np.isfinite(class_labels).any() else 0 + values = weighted_sum_by_group(raster, mapping) + return pd.DataFrame( + { + "region": np.tile(mapping.regions, mapping.n_classes), + "resource_class": np.repeat( + np.arange(mapping.n_classes), len(mapping.regions) + ), + "value": values.reshape(len(mapping.regions), mapping.n_classes).T.ravel(), + } ) - for cls in range(n_classes): - mask = class_labels == cls - if not np.any(mask): - continue - masked = np.where(mask, raster, np.nan) - raster_src = NumPyRasterSource( - masked, - xmin=xmin, - ymin=ymin, - xmax=xmax, - ymax=ymax, - nodata=np.nan, - srs_wkt=crs_wkt, - ) - stats = exact_extract( - raster_src, - regions_for_extract, - ["sum"], - include_cols=["region"], - output="pandas", - ) - if stats.empty: - continue - stats = stats.rename(columns={"sum": "value"}) - stats["resource_class"] = cls - records.append(stats) - - if not records: - return pd.DataFrame(columns=["region", "resource_class", "value"]) - - combined = pd.concat(records, ignore_index=True) - combined["resource_class"] = combined["resource_class"].astype(int) - return combined def _optional_path(value) -> Path: @@ -144,7 +108,7 @@ def _yield_weighted_residual_addition( if __name__ == "__main__": - classes_nc = Path(snakemake.input.classes) # type: ignore[name-defined] + cell_mapping_path = Path(snakemake.input.cell_mapping) # type: ignore[name-defined] raster_path = Path(snakemake.input.harvested_area_raster) # type: ignore[name-defined] regions_path = Path(snakemake.input.regions) # type: ignore[name-defined] mapping_path = Path(snakemake.input.crop_mapping) # type: ignore[name-defined] @@ -158,14 +122,12 @@ def _yield_weighted_residual_addition( output_path = Path(snakemake.output[0]) # type: ignore[name-defined] crop = str(snakemake.wildcards.crop) # type: ignore[name-defined] - ds = xr.load_dataset(classes_nc) - class_labels = ds["resource_class"].values.astype(np.int16) + cell_mapping = load_cell_mapping(cell_mapping_path) harvested_raw, src = read_raster_float(raster_path) try: - harvested_raw = harvested_raw * RES06_HAR_SCALE_TO_HA - transform = src.transform - crs_wkt = src.crs.to_wkt() if src.crs else None + validate_raster_grid(harvested_raw, src, cell_mapping) + harvested_raw *= RES06_HAR_SCALE_TO_HA finally: src.close() @@ -199,13 +161,7 @@ def _yield_weighted_residual_addition( non_food_crops=non_food_crops, ) - extracted = _extract_harvested_area( - harvested_raw, - transform, - crs_wkt, - class_labels, - regions, - ) + extracted = _extract_harvested_area(harvested_raw, cell_mapping) extracted = extracted.merge(regions[["region", "country"]], on="region", how="left") extracted = apply_country_shares(extracted, shares_lookup, fallback_share) diff --git a/workflow/scripts/build_crop_yield_cell_mapping.py b/workflow/scripts/build_region_class_cell_mapping.py similarity index 98% rename from workflow/scripts/build_crop_yield_cell_mapping.py rename to workflow/scripts/build_region_class_cell_mapping.py index 6bf22a13..e37d839f 100644 --- a/workflow/scripts/build_crop_yield_cell_mapping.py +++ b/workflow/scripts/build_region_class_cell_mapping.py @@ -98,6 +98,8 @@ def build_cell_mapping(classes_path: str, regions_path: str, output_path: str) - n_classes=np.array(n_classes, dtype=np.int32), height=np.array(height, dtype=np.int32), width=np.array(width, dtype=np.int32), + transform=transform, + crs_wkt=np.array(crs_wkt), ) diff --git a/workflow/scripts/crop_yield_aggregation.py b/workflow/scripts/region_class_aggregation.py similarity index 66% rename from workflow/scripts/crop_yield_aggregation.py rename to workflow/scripts/region_class_aggregation.py index 429ed189..65f6feb9 100644 --- a/workflow/scripts/crop_yield_aggregation.py +++ b/workflow/scripts/region_class_aggregation.py @@ -18,6 +18,8 @@ class CellMapping(NamedTuple): regions: np.ndarray n_classes: int shape: tuple[int, int] + transform: tuple[float, ...] + crs_wkt: str @property def n_groups(self) -> int: @@ -34,9 +36,35 @@ def load_cell_mapping(path: str) -> CellMapping: regions=data["regions"], n_classes=int(data["n_classes"]), shape=(int(data["height"]), int(data["width"])), + transform=tuple(data["transform"]), + crs_wkt=str(data["crs_wkt"]), ) +def validate_raster_grid(values: np.ndarray, source, mapping: CellMapping) -> None: + """Fail if a raster does not use the grid represented by ``mapping``.""" + if values.shape != mapping.shape: + raise ValueError( + f"Raster shape {values.shape} does not match cell mapping " + f"shape {mapping.shape}" + ) + transform = source.transform.to_gdal() + if not np.allclose(transform, mapping.transform, rtol=0.0, atol=1e-12): + raise ValueError( + f"Raster transform {transform} does not match cell mapping " + f"transform {mapping.transform}" + ) + if source.crs is None: + raise ValueError("Raster CRS does not match cell mapping CRS") + actual_crs_wkt = source.crs.to_wkt() + if actual_crs_wkt != mapping.crs_wkt: + # Avoid importing pyproj when equivalent WKT strings already match. + from pyproj import CRS + + if CRS.from_wkt(actual_crs_wkt) != CRS.from_wkt(mapping.crs_wkt): + raise ValueError("Raster CRS does not match cell mapping CRS") + + def _mapped_values(values: np.ndarray, mapping: CellMapping) -> np.ndarray: if values.shape != mapping.shape: raise ValueError(