diff --git a/CHANGELOG.md b/CHANGELOG.md index d089f7b2..9274aa31 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/tests/test_compute_resource_classes.py b/tests/test_compute_resource_classes.py new file mode 100644 index 00000000..12523858 --- /dev/null +++ b/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)) diff --git a/tests/test_luc_carbon.py b/tests/test_luc_carbon.py index a653a2cd..52f5cf38 100644 --- a/tests/test_luc_carbon.py +++ b/tests/test_luc_carbon.py @@ -7,12 +7,10 @@ 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 ( @@ -20,9 +18,11 @@ _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), diff --git a/tests/test_luicube_grassland_yields.py b/tests/test_luicube_grassland_yields.py new file mode 100644 index 00000000..eff2e9e8 --- /dev/null +++ b/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(), + ) diff --git a/tests/test_process_huang_irrigation_water.py b/tests/test_process_huang_irrigation_water.py new file mode 100644 index 00000000..4f764adb --- /dev/null +++ b/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] diff --git a/workflow/rules/crops.smk b/workflow/rules/crops.smk index 1354a7e7..b3109ecc 100644 --- a/workflow/rules/crops.smk +++ b/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: "/{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": "/{name}/resource_classes.nc", - "regions": "/{name}/regions.geojson", + "cell_mapping": "/{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: "/{name}/build_multi_cropping.log", benchmark: @@ -737,15 +736,14 @@ rule build_grassland_yields: rule build_luicube_grassland_yields: input: luicube="/shared/luc/luicube_grassland.nc", - classes="/{name}/resource_classes.nc", - regions="/{name}/regions.geojson", + cell_mapping="/{name}/region_class_cell_mapping.npz", output: "/{name}/luicube_grassland_yields.csv", group: "prep" resources: runtime="1m", - mem_mb=2600, + mem_mb=1000, log: "/{name}/build_luicube_grassland_yields.log", benchmark: diff --git a/workflow/rules/geography.smk b/workflow/rules/geography.smk index bafcf8b1..120bb38f 100644 --- a/workflow/rules/geography.smk +++ b/workflow/rules/geography.smk @@ -132,7 +132,7 @@ rule compute_resource_classes: "prep" resources: runtime="1m", - mem_mb=1900, + mem_mb=900, log: "/{name}/compute_resource_classes.log", benchmark: @@ -143,11 +143,10 @@ rule compute_resource_classes: rule aggregate_class_areas: input: - classes="/{name}/resource_classes.nc", + cell_mapping="/{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="/{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: "/{name}/aggregate_class_areas.log", benchmark: diff --git a/workflow/rules/luc.smk b/workflow/rules/luc.smk index 47a7a7f5..7218e4cf 100644 --- a/workflow/rules/luc.smk +++ b/workflow/rules/luc.smk @@ -304,7 +304,7 @@ rule build_grazing_only_land: rule build_luc_carbon_coefficients: input: classes="/{name}/resource_classes.nc", - regions="/{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, diff --git a/workflow/rules/model.smk b/workflow/rules/model.smk index 26bbdd1e..146939f0 100644 --- a/workflow/rules/model.smk +++ b/workflow/rules/model.smk @@ -230,7 +230,7 @@ rule build_model: "build_model" resources: runtime="1m", - mem_mb=900, + mem_mb=1500, log: "/{name}/build_model.log", benchmark: diff --git a/workflow/scripts/aggregate_class_areas.py b/workflow/scripts/aggregate_class_areas.py index abeef222..1c5c5f75 100644 --- a/workflow/scripts/aggregate_class_areas.py +++ b/workflow/scripts/aggregate_class_areas.py @@ -6,17 +6,18 @@ from pathlib import Path -from exactextract import exact_extract -from exactextract.raster import NumPyRasterSource -import geopandas as gpd import numpy as np import pandas as pd import rasterio from rasterio.enums import Resampling +from rasterio.env import set_gdal_config from rasterio.warp import reproject -import xarray as xr from workflow.scripts.raster_utils import calculate_all_cell_areas, scale_fraction +from workflow.scripts.region_class_aggregation import ( + load_cell_mapping, + weighted_sum_by_group, +) def read_raster_float(path: str): @@ -74,18 +75,11 @@ def load_scaled_fraction( return scale_fraction(arr) -def raster_bounds(transform, width: int, height: int): - xmin = transform.c - ymax = transform.f - xmax = xmin + width * transform.a - ymin = ymax + height * transform.e - return xmin, ymin, xmax, ymax - - if __name__ == "__main__": + set_gdal_config("GDAL_CACHEMAX", 128 * 1024**2) + # Inputs - regions_path: str = snakemake.input.regions # type: ignore[name-defined] - classes_nc: str = snakemake.input.classes # type: ignore[name-defined] + cell_mapping_path: str = snakemake.input.cell_mapping # type: ignore[name-defined] # Suitability/area inputs as lists of file paths sr_files: list[str] = list(snakemake.input.sr) # type: ignore[attr-defined] si_files: list[str] = list(snakemake.input.si) # type: ignore[attr-defined] @@ -93,9 +87,7 @@ def raster_bounds(transform, width: int, height: int): irrigated_area_source: str = snakemake.params.irrigated_area_source # type: ignore[name-defined] - # Load classes - ds = xr.load_dataset(classes_nc) - classes = ds["resource_class"].values.astype(np.int16) + cell_mapping = load_cell_mapping(cell_mapping_path) # Reference grid parameters from a suitability raster (rainfed) # Use first rainfed suitability file as reference @@ -106,18 +98,10 @@ def raster_bounds(transform, width: int, height: int): height, width = sr0.shape transform = src0.transform crs = src0.crs - xmin, ymin, xmax, ymax = raster_bounds(transform, width, height) - crs_wkt = crs.to_wkt() if crs else None cell_area_rows = calculate_all_cell_areas(src0, repeat=False) finally: src0.close() - # 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() - # Cell areas cell_area_rows = cell_area_rows.astype(np.float32, copy=False) @@ -151,47 +135,14 @@ def max_suitability( area_r_raw = sr_max def aggregate_area(area: np.ndarray, ws: str) -> pd.DataFrame: - out = [] - valid_mask = classes >= 0 - if not np.any(valid_mask): - return pd.DataFrame( - columns=["region", "resource_class", "water_supply", "area_ha"] - ) - class_ids = np.unique(classes[valid_mask]) - work_arr = np.empty_like(area, dtype=np.float32) - for cls in class_ids: - mask = classes == cls - if not np.any(mask): - continue - work_arr.fill(np.nan) - work_arr[mask] = area[mask] - a_src = NumPyRasterSource( - work_arr, - xmin=xmin, - ymin=ymin, - xmax=xmax, - ymax=ymax, - nodata=np.nan, - srs_wkt=crs_wkt, - ) - a_stats = exact_extract( - a_src, - regions_for_extract, - ["sum"], - include_cols=["region"], - output="pandas", - ) - if a_stats.empty: - continue - a_stats = a_stats.rename(columns={"sum": "area_ha"}) - a_stats["resource_class"] = cls - a_stats["water_supply"] = ws - out.append(a_stats) - if not out: - return pd.DataFrame( - columns=["region", "resource_class", "water_supply", "area_ha"] - ) - return pd.concat(out, ignore_index=True) + area_ha = weighted_sum_by_group(area, cell_mapping) + index = pd.MultiIndex.from_product( + [cell_mapping.regions, range(cell_mapping.n_classes)], + names=["region", "resource_class"], + ) + result = pd.DataFrame({"area_ha": area_ha}, index=index).reset_index() + result["water_supply"] = ws + return result if irrigated_area_source == "potential": area_i_raw = max_suitability(si_files) @@ -211,9 +162,11 @@ def aggregate_area(area: np.ndarray, ws: str) -> pd.DataFrame: # rainfed gets the remainder. This keeps the model's land budget faithful # to the underlying physical cell area regardless of how the two # suitability rasters overlap. - area_i = np.minimum(area_i_raw, area_r_raw) - area_r = np.maximum(area_r_raw - area_i, 0.0) - del area_i_raw, area_r_raw + np.minimum(area_i_raw, area_r_raw, out=area_i_raw) + np.subtract(area_r_raw, area_i_raw, out=area_r_raw) + np.maximum(area_r_raw, 0.0, out=area_r_raw) + area_i = area_i_raw + area_r = area_r_raw df_r = aggregate_area(area_r, "r") del area_r diff --git a/workflow/scripts/build_luc_carbon_coefficients.py b/workflow/scripts/build_luc_carbon_coefficients.py index 8a1bb574..cc1536b8 100644 --- a/workflow/scripts/build_luc_carbon_coefficients.py +++ b/workflow/scripts/build_luc_carbon_coefficients.py @@ -12,13 +12,15 @@ from pathlib import Path # noqa: E402 from affine import Affine # noqa: E402 -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.region_class_aggregation import ( # noqa: E402 + CellMapping, + load_cell_mapping, +) + CO2_PER_C = 44.0 / 12.0 ZONE_ORDER = ["tropical", "temperate", "boreal"] @@ -143,9 +145,38 @@ def _ensure_mode_zero(mode: str) -> None: ) +def _weighted_mean_by_group( + values: np.ndarray, + weights: np.ndarray, + mapping: CellMapping, +) -> np.ndarray: + """Return weighted means for every exact region/resource-class group.""" + mapped_values = values.ravel()[mapping.cell_ids] + mapped_weights = weights.ravel()[mapping.cell_ids] * mapping.coverage + valid = ~np.isnan(mapped_values) & ~np.isnan(mapped_weights) + group_ids = mapping.group_ids[valid] + mapped_weights = mapped_weights[valid] + numerator = np.bincount( + group_ids, + weights=mapped_values[valid] * mapped_weights, + minlength=mapping.n_groups, + ) + denominator = np.bincount( + group_ids, + weights=mapped_weights, + minlength=mapping.n_groups, + ) + return np.divide( + numerator, + denominator, + out=np.full(mapping.n_groups, np.nan), + where=denominator != 0, + ) + + def main() -> None: classes_path: str = snakemake.input.classes # type: ignore[name-defined] - regions_path: str = snakemake.input.regions # type: ignore[name-defined] + cell_mapping_path: str = snakemake.input.cell_mapping # type: ignore[name-defined] agb_path: str = snakemake.input.agb # type: ignore[name-defined] soc_path: str = snakemake.input.soc # type: ignore[name-defined] regrowth_path: str = snakemake.input.regrowth # type: ignore[name-defined] @@ -163,11 +194,10 @@ def main() -> None: Path(coeffs_out).parent.mkdir(parents=True, exist_ok=True) - classes_ds = xr.load_dataset(classes_path) - transform, height, width, lon, lat = _load_transform(classes_ds) - resource_class = classes_ds["resource_class"].astype(np.int16).values + with xr.open_dataset(classes_path) as classes_ds: + _transform, height, width, lon, lat = _load_transform(classes_ds) - zone_idx = _zone_index(lat, width) + zone_idx = _zone_index(lat, 1) params = _zone_parameters(zone_params_path) agb = xr.load_dataset(agb_path)["agb_tc_per_ha"].astype(np.float32).values @@ -179,11 +209,11 @@ def main() -> None: ) lc_masks_path: str = snakemake.input.lc_masks # type: ignore[name-defined] - lc_ds = xr.load_dataset(lc_masks_path) - cropland_frac = lc_ds["cropland_fraction"].astype(np.float32).values - pasture_frac = lc_ds["pasture_fraction"].astype(np.float32).values - forest_frac = lc_ds["forest_fraction"].astype(np.float32).values - grazing_intensity = lc_ds["grazing_intensity"].astype(np.float32).values + with xr.open_dataset(lc_masks_path) as lc_ds: + cropland_frac = lc_ds["cropland_fraction"].load().values + pasture_frac = lc_ds["pasture_fraction"].load().values + forest_frac = lc_ds["forest_fraction"].load().values + grazing_intensity = lc_ds["grazing_intensity"].load().values # Dominant-cover partition (see ``prepare_luc_inputs``): ``pasture_frac`` is # the open grazed land (LUIcube grass not overlapping forest/cropland), so # genuinely grazed open land is pasture -- not natural land -- for carbon, @@ -297,6 +327,39 @@ def main() -> None: natural_frac > 0, nonforest_frac / natural_frac, 0.0 ).astype(np.float32) + del ( + agb, + soc_0_30, + regrowth_tc, + grazing_intensity, + zone_idx, + bgb_ratio_nat, + bgb_ratio_nonforest, + soc_depth_factor, + agb_crop, + bgb_ratio_crop, + agb_past, + bgb_ratio_past, + soc_factor_crop, + soc_factor_past, + agb_nonforest_zone, + soc_factor_past_eff, + agb_past_eff, + agb_obs, + soc_0_30_nat, + agb_forest, + agb_nonforest, + soc_nat, + bgb_forest, + s_forest, + bgb_nonforest, + s_nonforest, + bgb_crop, + s_ag_crop, + bgb_past, + s_ag_past, + ) + pulses_ds = xr.Dataset( { "P_crop_forest_tCO2_per_ha": ( @@ -322,6 +385,14 @@ def main() -> None: pulses_out, encoding={v: {"zlib": True, "dtype": "float32"} for v in pulses_ds.data_vars}, ) + del pulses_ds + del ( + p_crop_forest, + p_crop_nonforest, + p_past_forest, + p_past_nonforest, + regrowth, + ) use_names = [ "cropland_forest", @@ -357,30 +428,17 @@ def main() -> None: annual_out, encoding={"LEF_tCO2_per_ha_yr": {"zlib": True, "dtype": "float32"}}, ) + del annual_ds, lef_stack # --- Aggregate per-pixel LEFs to per-region/class coefficients --- - # Uses exact_extract with region polygons and class masks so that tiny - # regions that don't cover a full grid cell still get correct - # area-weighted LEFs via fractional cell overlaps. - - regions_gdf = gpd.read_file(regions_path) - crs_wkt = classes_ds.attrs.get("crs_wkt") - if crs_wkt: - regions_gdf = regions_gdf.to_crs(crs_wkt) - regions_for_extract = regions_gdf.reset_index() - - xmin = float(transform.c) - ymax = float(transform.f) - xmax = xmin + width * transform.a - ymin = ymax + height * transform.e - raster_kwargs = { - "xmin": xmin, - "ymin": ymin, - "xmax": xmax, - "ymax": ymax, - "nodata": np.nan, - "srs_wkt": crs_wkt, - } + # Reuse the exact fractional cell coverage shared by crop raster rules. + mapping = load_cell_mapping(cell_mapping_path) + if mapping.shape != (height, width): + raise ValueError( + f"Cell mapping shape {mapping.shape} does not match resource-class " + f"shape {(height, width)}" + ) + unit_share = np.ones_like(share_forest) # For conversion uses, the LEF is weighted by the relevant land-cover # fraction (forest or nonforest), and the conversion_share tracks how @@ -388,34 +446,34 @@ def main() -> None: weighted_uses: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]] = { # (lef_array, area_weight, conversion_share) "cropland_forest": ( - lef_crop_forest.astype(np.float32), + lef_crop_forest, forest_frac, share_forest, ), "cropland_nonforest": ( - lef_crop_nonforest.astype(np.float32), + lef_crop_nonforest, nonforest_frac, share_nonforest, ), "pasture_forest": ( - lef_past_forest.astype(np.float32), + lef_past_forest, forest_frac, share_forest, ), "pasture_nonforest": ( - lef_past_nonforest.astype(np.float32), + lef_past_nonforest, nonforest_frac, share_nonforest, ), "spared_cropland": ( - lef_spared.astype(np.float32), + lef_spared, cropland_frac, - np.ones_like(share_forest), + unit_share, ), "spared_grassland": ( - lef_spared.astype(np.float32), + lef_spared, pasture_frac, - np.ones_like(share_forest), + unit_share, ), } water_options = { @@ -427,64 +485,22 @@ def main() -> None: "spared_grassland": ("r",), } - n_classes = ( - int(np.nanmax(resource_class)) + 1 - if np.isfinite(resource_class.astype(float)).any() - else 0 - ) - frames: list[pd.DataFrame] = [] - for cls in range(n_classes): - mask_float = (resource_class == cls).astype(np.float32) - if not np.any(mask_float > 0): - continue - - # Area-weighted mean LEF and conversion share per region for each use type - for use, (lef_arr, lc_weight, conv_share) in weighted_uses.items(): - composite_weight = mask_float * lc_weight - composite_src = NumPyRasterSource( - composite_weight, - xmin=xmin, - ymin=ymin, - xmax=xmax, - ymax=ymax, - srs_wkt=crs_wkt, - ) - lef_src = NumPyRasterSource(lef_arr, **raster_kwargs) - lef_stats = exact_extract( - lef_src, - regions_for_extract, - ["weighted_mean"], - weights=composite_src, - include_cols=["region"], - output="pandas", - ) - - # Aggregate conversion_share using the same weight (nonag area) - share_src = NumPyRasterSource( - conv_share.astype(np.float32), **raster_kwargs - ) - # Weight shares by natural_frac * mask to get area-weighted mean share - nonag_weight = mask_float * natural_frac - nonag_weight_src = NumPyRasterSource( - nonag_weight.astype(np.float32), - xmin=xmin, - ymin=ymin, - xmax=xmax, - ymax=ymax, - srs_wkt=crs_wkt, + region_ids = np.repeat(np.arange(len(mapping.regions)), mapping.n_classes) + class_ids = np.tile(np.arange(mapping.n_classes), len(mapping.regions)) + group_regions = mapping.regions[region_ids] + for use, (lef_arr, lc_weight, conv_share) in weighted_uses.items(): + lef_stats = _weighted_mean_by_group(lef_arr, lc_weight, mapping) + share_stats = _weighted_mean_by_group(conv_share, natural_frac, mapping) + for cls in range(mapping.n_classes): + class_mask = class_ids == cls + merged = pd.DataFrame( + { + "region": group_regions[class_mask], + "LEF_tCO2_per_ha_yr": lef_stats[class_mask], + "conversion_share": share_stats[class_mask], + } ) - share_stats = exact_extract( - share_src, - regions_for_extract, - ["weighted_mean"], - weights=nonag_weight_src, - include_cols=["region"], - output="pandas", - ) - - merged = lef_stats.rename(columns={"weighted_mean": "LEF_tCO2_per_ha_yr"}) - merged["conversion_share"] = share_stats["weighted_mean"] merged["resource_class"] = cls merged["use"] = use merged = merged.dropna(subset=["LEF_tCO2_per_ha_yr"]) diff --git a/workflow/scripts/build_luicube_grassland_yields.py b/workflow/scripts/build_luicube_grassland_yields.py index ab6fafc1..5e19dc75 100644 --- a/workflow/scripts/build_luicube_grassland_yields.py +++ b/workflow/scripts/build_luicube_grassland_yields.py @@ -5,7 +5,7 @@ """Compute grassland yields, suitable area and grazing intensity from LUIcube. Reads the resampled LUIcube grassland NetCDF and aggregates per -region/resource_class using exactextract zonal statistics. +region/resource_class using the shared exact cell-coverage mapping. Output CSV columns: region, resource_class, yield, suitable_area, grazing_intensity @@ -19,45 +19,33 @@ from pathlib import Path -from affine import Affine -from exactextract import exact_extract -from exactextract.raster import NumPyRasterSource -import geopandas as gpd import numpy as np import pandas as pd import xarray as xr -from workflow.scripts.raster_utils import raster_bounds +from workflow.scripts.region_class_aggregation import ( + CellMapping, + load_cell_mapping, + weighted_sum_by_group, +) # Carbon content of dry matter (tC per tDM) C_FRACTION = 0.45 -if __name__ == "__main__": - luicube_path: str = snakemake.input.luicube # type: ignore[name-defined] - classes_path: str = snakemake.input.classes # type: ignore[name-defined] - regions_path: str = snakemake.input.regions # type: ignore[name-defined] - output_path = Path(snakemake.output[0]) # type: ignore[name-defined] +def aggregate_grassland_yields( + area_km2: np.ndarray, + npp_act: np.ndarray, + hanpp_harv: np.ndarray, + mapping: CellMapping, +) -> pd.DataFrame: + """Aggregate LUIcube grassland variables by region and resource class.""" + if area_km2.shape != mapping.shape: + raise ValueError("LUIcube grid does not match region/class cell mapping") + if npp_act.shape != mapping.shape or hanpp_harv.shape != mapping.shape: + raise ValueError("LUIcube variables do not share one grid") - # Load resource classes grid - ds_classes = xr.load_dataset(classes_path) - class_labels = ds_classes["resource_class"].values.astype(np.int16) - region_id = ds_classes["region_id"].values.astype(np.int32) - transform = Affine.from_gdal(*ds_classes.attrs["transform"]) - height, width = class_labels.shape - crs_wkt = ds_classes.attrs.get("crs_wkt") - xmin, ymin, xmax, ymax = raster_bounds(transform, width, height) - - # Load LUIcube grassland data - ds = xr.load_dataset(luicube_path) - area_km2 = ds["area_km2"].values.astype(np.float64) - npp_act = ds["npp_act_tc_yr"].values.astype(np.float64) - hanpp_harv = ds["hanpp_harv_tc_yr"].values.astype(np.float64) - - if area_km2.shape != (height, width): - raise ValueError("LUIcube grid does not match resource_classes grid") - - # Convert area to hectares: 1 km² = 100 ha + # Convert area to hectares: 1 km2 = 100 ha area_ha = area_km2 * 100.0 # Compute per-cell grazing intensity = HANPP_harv / NPP_act, clipped [0, 1] @@ -68,141 +56,50 @@ # Managed pasture area: total grassland scaled by grazing intensity managed_area_ha = area_ha * gi_cell - # Load regions - regions_gdf = gpd.read_file(regions_path) - if regions_gdf.crs and regions_gdf.crs.to_epsg() != 4326: - regions_gdf = regions_gdf.to_crs("EPSG:4326") - regions_for_extract = regions_gdf.reset_index() - - valid_classes = sorted( - int(c) for c in np.unique(class_labels) if np.isfinite(c) and c >= 0 + sums = pd.DataFrame( + { + "hanpp_sum": weighted_sum_by_group(hanpp_harv, mapping), + "managed_area": weighted_sum_by_group(managed_area_ha, mapping), + "suitable_area": weighted_sum_by_group(area_ha, mapping), + "npp_sum": weighted_sum_by_group(npp_act, mapping), + "gi_weighted_sum": weighted_sum_by_group(gi_cell * npp_act, mapping), + }, + index=pd.MultiIndex.from_product( + [mapping.regions, range(mapping.n_classes)], + names=["region", "resource_class"], + ), ) - data_frames: list[pd.DataFrame] = [] - for cls in valid_classes: - mask = class_labels == cls - if not np.any(mask): - continue - - # Mask arrays to this resource class - hanpp_masked = np.where(mask, hanpp_harv, np.nan) - managed_area_masked = np.where(mask, managed_area_ha, np.nan) - physical_area_masked = np.where(mask, area_ha, np.nan) - # NPP_act-weighted grazing intensity (diagnostics): weight = npp_act - npp_masked = np.where(mask, npp_act, np.nan) - gi_weighted = np.where(mask, gi_cell * npp_act, np.nan) - - raster_kwargs = { - "xmin": xmin, - "ymin": ymin, - "xmax": xmax, - "ymax": ymax, - "nodata": np.nan, - "srs_wkt": crs_wkt, - } - hanpp_src = NumPyRasterSource(hanpp_masked, **raster_kwargs) - managed_area_src = NumPyRasterSource(managed_area_masked, **raster_kwargs) - physical_area_src = NumPyRasterSource(physical_area_masked, **raster_kwargs) - npp_src = NumPyRasterSource(npp_masked, **raster_kwargs) - gi_w_src = NumPyRasterSource(gi_weighted, **raster_kwargs) - - hanpp_stats = exact_extract( - hanpp_src, - regions_for_extract, - ["sum"], - include_cols=["region"], - output="pandas", - ) - managed_area_stats = exact_extract( - managed_area_src, - regions_for_extract, - ["sum"], - include_cols=["region"], - output="pandas", - ) - physical_area_stats = exact_extract( - physical_area_src, - regions_for_extract, - ["sum"], - include_cols=["region"], - output="pandas", - ) - npp_stats = exact_extract( - npp_src, - regions_for_extract, - ["sum"], - include_cols=["region"], - output="pandas", + # yield = sum(hanpp_harv) / sum(managed_area_ha) / C_FRACTION -> tDM/ha managed + with np.errstate(divide="ignore", invalid="ignore"): + sums["yield"] = np.where( + sums["managed_area"] > 0, + sums["hanpp_sum"] / sums["managed_area"] / C_FRACTION, + 0.0, ) - gi_w_stats = exact_extract( - gi_w_src, - regions_for_extract, - ["sum"], - include_cols=["region"], - output="pandas", + # grazing_intensity = sum(gi * npp) / sum(npp) (diagnostic) + with np.errstate(divide="ignore", invalid="ignore"): + sums["grazing_intensity"] = np.where( + sums["npp_sum"] > 0, + sums["gi_weighted_sum"] / sums["npp_sum"], + 0.0, ) + sums["grazing_intensity"] = sums["grazing_intensity"].clip(0.0, 1.0) + return sums[["yield", "suitable_area", "grazing_intensity"]].sort_index() - if hanpp_stats.empty or managed_area_stats.empty: - continue - - merged = ( - hanpp_stats.rename(columns={"sum": "hanpp_sum"}) - .merge( - managed_area_stats.rename(columns={"sum": "managed_area"}), - on="region", - ) - .merge( - physical_area_stats.rename(columns={"sum": "suitable_area"}), - on="region", - ) - .merge(npp_stats.rename(columns={"sum": "npp_sum"}), on="region") - .merge(gi_w_stats.rename(columns={"sum": "gi_weighted_sum"}), on="region") - ) - # yield = sum(hanpp_harv) / sum(managed_area_ha) / C_FRACTION → tDM/ha managed - with np.errstate(divide="ignore", invalid="ignore"): - merged["yield"] = np.where( - merged["managed_area"] > 0, - merged["hanpp_sum"] / merged["managed_area"] / C_FRACTION, - 0.0, - ) - # grazing_intensity = sum(gi * npp) / sum(npp) (diagnostic) - with np.errstate(divide="ignore", invalid="ignore"): - merged["grazing_intensity"] = np.where( - merged["npp_sum"] > 0, - merged["gi_weighted_sum"] / merged["npp_sum"], - 0.0, - ) - merged["grazing_intensity"] = merged["grazing_intensity"].clip(0.0, 1.0) - merged["resource_class"] = cls - data_frames.append( - merged[ - [ - "region", - "resource_class", - "yield", - "suitable_area", - "grazing_intensity", - ] - ] - ) +if __name__ == "__main__": + luicube_path: str = snakemake.input.luicube # type: ignore[name-defined] + mapping_path: str = snakemake.input.cell_mapping # type: ignore[name-defined] + output_path = Path(snakemake.output[0]) # type: ignore[name-defined] - if data_frames: - out_df = ( - pd.concat(data_frames, ignore_index=True) - .set_index(["region", "resource_class"]) - .sort_index() - ) - else: - out_df = pd.DataFrame( - columns=[ - "region", - "resource_class", - "yield", - "suitable_area", - "grazing_intensity", - ] - ).set_index(["region", "resource_class"]) + mapping = load_cell_mapping(mapping_path) + with xr.open_dataset(luicube_path) as ds: + area_km2 = ds["area_km2"].load().values.astype(np.float64) + npp_act = ds["npp_act_tc_yr"].load().values.astype(np.float64) + hanpp_harv = ds["hanpp_harv_tc_yr"].load().values.astype(np.float64) + + out_df = aggregate_grassland_yields(area_km2, npp_act, hanpp_harv, mapping) output_path.parent.mkdir(parents=True, exist_ok=True) out_df.to_csv(output_path) diff --git a/workflow/scripts/build_model/crops.py b/workflow/scripts/build_model/crops.py index 504642e2..d7215ed9 100644 --- a/workflow/scripts/build_model/crops.py +++ b/workflow/scripts/build_model/crops.py @@ -53,9 +53,10 @@ def multi_crop_cycle_multiplicities(links: pd.DataFrame) -> pd.DataFrame: "to account for harvested-cycle multiplicity. Rebuild the model." ) - records: list[dict[str, object]] = [] - for link, row in multi.iterrows(): - raw = row["crop_cycles"] + records: list[tuple[object, ...]] = [] + for link, raw, country, baseline_area_mha in multi[ + ["crop_cycles", "country", "baseline_area_mha"] + ].itertuples(index=True, name=None): try: cycles = json.loads(str(raw)) except (TypeError, ValueError) as exc: @@ -74,13 +75,13 @@ def multi_crop_cycle_multiplicities(links: pd.DataFrame) -> pd.DataFrame: ) for crop, multiplicity in Counter(cycles).items(): records.append( - { - "link": link, - "crop": crop, - "country": str(row["country"]), - "multiplicity": int(multiplicity), - "baseline_area_mha": float(row["baseline_area_mha"]), - } + ( + link, + crop, + str(country), + int(multiplicity), + float(baseline_area_mha), + ) ) return pd.DataFrame.from_records(records, columns=columns) @@ -295,6 +296,24 @@ def _apply_bounded_cost_calibration( ) +def _crop_cost_values( + crops: pd.Series, + countries: pd.Series, + crop_costs: pd.Series, + global_median_cost: pd.Series, +) -> np.ndarray: + """Look up country crop costs, falling back only when the key is absent.""" + keys = pd.MultiIndex.from_arrays( + [crops.astype(str).to_numpy(), countries.astype(str).to_numpy()] + ) + values = crop_costs.reindex(keys).to_numpy(dtype=float) + missing = ~keys.isin(crop_costs.index) + if missing.any(): + fallback = crops.astype(str).map(global_median_cost).fillna(0.0).to_numpy() + values[missing] = fallback[missing] + return values + + def add_regional_crop_production_links( n: pypsa.Network, crop_list: list, @@ -548,9 +567,10 @@ def add_regional_crop_production_links( # Look up per-(crop, country) cost, falling back to global median cost_keys = list(zip(all_df["crop"].astype(str), all_df["country"].astype(str))) per_link_cost = pd.Series( - [crop_costs.get(k, global_median_cost.get(k[0], 0.0)) for k in cost_keys], + _crop_cost_values( + all_df["crop"], all_df["country"], crop_costs, global_median_cost + ), index=all_df.index, - dtype=float, ) # Convert USD/ha to bnUSD/Mha all_df["marginal_cost"] = per_link_cost * 1e6 * constants.USD_TO_BNUSD @@ -904,20 +924,47 @@ def add_multi_cropping_links( cycle_df = cycle_df[~low_yield_mask] merged = cycle_df.merge(area_df, on=key_cols, how="inner") - valid_keys: list[tuple[object, ...]] = [] - for key, cycles in merged.groupby(key_cols, sort=False): - combo = str(key[0]) - expected = [str(crop) for crop in combinations[combo]["crops"]] - ordered = cycles.sort_values("cycle_index") - got_indices = ordered["cycle_index"].astype(int).tolist() - got_crops = ordered["crop"].astype(str).tolist() - if got_indices == list(range(1, len(expected) + 1)) and got_crops == expected: - valid_keys.append(key) - - if valid_keys: - valid_index = pd.MultiIndex.from_tuples(valid_keys, names=key_cols) + expected_cycles = pd.Series( + { + (str(combination), cycle_index): str(crop) + for combination, settings in combinations.items() + for cycle_index, crop in enumerate(settings["crops"], start=1) + } + ) + expected_counts = pd.Series( + { + str(combination): len(settings["crops"]) + for combination, settings in combinations.items() + } + ) + cycle_index_int = merged["cycle_index"].astype(int) + cycle_keys = pd.MultiIndex.from_arrays( + [merged["combination"].astype(str), cycle_index_int] + ) + merged["_cycle_index_int"] = cycle_index_int + merged["_valid_cycle"] = ( + merged["crop"].astype(str).to_numpy() + == expected_cycles.reindex(cycle_keys).to_numpy() + ) + sequence_stats = merged.groupby(key_cols, sort=False).agg( + cycle_count=("cycle_index", "size"), + distinct_cycles=("_cycle_index_int", "nunique"), + valid_cycles=("_valid_cycle", "all"), + ) + sequence_stats["expected_count"] = sequence_stats.index.get_level_values( + "combination" + ).map(expected_counts) + valid_sequences = sequence_stats[ + sequence_stats["valid_cycles"] + & sequence_stats["cycle_count"].eq(sequence_stats["expected_count"]) + & sequence_stats["distinct_cycles"].eq(sequence_stats["expected_count"]) + ] + valid_keys = valid_sequences.index + merged = merged.drop(columns=["_cycle_index_int", "_valid_cycle"]) + + if not valid_keys.empty: merged_index = pd.MultiIndex.from_frame(merged[key_cols]) - merged = merged[merged_index.isin(valid_index)].copy() + merged = merged[merged_index.isin(valid_keys)].copy() else: merged = merged.iloc[0:0].copy() @@ -975,10 +1022,9 @@ def add_multi_cropping_links( base = base.join(crop_counts) # Look up per-(crop, country) cost and sum across crops in combination - merged["cost_usd_per_ha"] = [ - crop_costs.get((c, cc), global_median_cost.get(c, 0.0)) - for c, cc in zip(merged["crop"], merged["country"]) - ] + merged["cost_usd_per_ha"] = _crop_cost_values( + merged["crop"], merged["country"], crop_costs, global_median_cost + ) # Marketing markup per cycle: marketing_cost_per_t * yield (post seed/loss) marketing_per_t = merged["crop"].astype(str).map(crop_marketing_cost_usd_per_t) if marketing_per_t.isna().any(): @@ -1024,49 +1070,59 @@ def add_multi_cropping_links( if not anchor_df.empty else pd.Series(dtype=float) ) - unavailable_groups: list[dict[str, str]] = [] - unavailable_mha = 0.0 anchor_group_cols = ["combination", "country", "water_supply"] - for group_key, source in anchor_df.groupby(anchor_group_cols, sort=False): - combination, country, water_supply = map(str, group_key) - destination_mask = ( - (base.index.get_level_values("combination") == combination) - & (base["country"].astype(str).to_numpy() == country) - & (base.index.get_level_values("water_supply") == water_supply) - ) - destinations = base.index[destination_mask] - if destinations.empty: - unavailable_mha += float(source["baseline_area_mha"].sum()) - unavailable_groups.append( - { - "combination": combination, - "country": country, - "water_supply": water_supply, - } - ) - continue - - local = anchor_by_key.reindex(destinations).fillna(0.0) - base.loc[destinations, "baseline_area_mha"] += local.to_numpy() - local_total = float(source["baseline_area_mha"].sum()) - matched_total = float(local.sum()) - relocate = local_total - matched_total - if relocate > 1e-12: - weights = base.loc[destinations, "eligible_area_ha"].astype(float) - base.loc[destinations, "baseline_area_mha"] += ( - relocate * weights / float(weights.sum()) - ).to_numpy() - - if unavailable_groups: + anchor_group_totals = anchor_df.groupby(anchor_group_cols, sort=False)[ + "baseline_area_mha" + ].sum() + destination_groups = pd.MultiIndex.from_arrays( + [ + base.index.get_level_values("combination").astype(str), + base["country"].astype(str), + base.index.get_level_values("water_supply").astype(str), + ], + names=anchor_group_cols, + ) + available_groups = destination_groups.unique() + unavailable = anchor_group_totals[~anchor_group_totals.index.isin(available_groups)] + if not unavailable.empty: + unavailable_groups = [ + dict(zip(anchor_group_cols, map(str, group), strict=True)) + for group in unavailable.index[:10] + ] logger.warning( "Retaining %.3f Mha of MIRCA anchors on single-crop baselines because " "%d combination/country/water groups have no valid full-sequence " "GAEZ target (examples: %s)", - unavailable_mha, - len(unavailable_groups), - unavailable_groups[:10], + float(unavailable.sum()), + len(unavailable), + unavailable_groups, ) + local_anchor = anchor_by_key.reindex(base.index).fillna(0.0) + matched_group_totals = ( + pd.Series(local_anchor.to_numpy(), index=destination_groups) + .groupby(level=anchor_group_cols, sort=False) + .sum() + ) + relocation = anchor_group_totals.sub(matched_group_totals, fill_value=0.0).clip( + lower=0.0 + ) + relocation_by_link = pd.Series( + destination_groups.map(relocation).fillna(0.0), index=base.index + ) + eligible_group_totals = ( + pd.Series( + base["eligible_area_ha"].to_numpy(dtype=float), index=destination_groups + ) + .groupby(level=anchor_group_cols, sort=False) + .transform("sum") + ) + base["baseline_area_mha"] = local_anchor + ( + relocation_by_link + * base["eligible_area_ha"].astype(float) + / eligible_group_totals.to_numpy() + ) + base["p_nom_max"] = np.maximum( base["eligible_area_ha"] / constants.HA_PER_MHA, base["baseline_area_mha"], @@ -1679,10 +1735,10 @@ def reconcile_single_crop_baselines( float(link_scales.min()), ) - # Re-expand the scaled anchors so all subsequent accounting uses exactly the - # metadata and baselines now stored on the network. + # Refresh the expanded anchors so all subsequent accounting uses exactly the + # baselines now stored on the network. multi = n.links.static[n.links.static["carrier"] == "crop_production_multi"] - cycle_long = multi_crop_cycle_multiplicities(n.links.static) + cycle_long["baseline_area_mha"] = cycle_long["link"].map(multi["baseline_area_mha"]) cycle_long["region"] = cycle_long["link"].map(multi["region"]) cycle_long["resource_class"] = ( cycle_long["link"].map(multi["resource_class"]).astype(int) @@ -1710,21 +1766,25 @@ def reconcile_single_crop_baselines( remaining_required = total_required.sub(locally_taken, fill_value=0.0).clip( lower=0.0 ) - for group, remainder in remaining_required.items(): - if remainder <= 1e-12: - continue - idx = single.index[group_index == group] - available = new_baseline.loc[idx] - total_available = float(available.sum()) - if remainder > total_available + 1e-9: - raise ValueError( - "Scaled multi-cropping anchors still exceed the FAOSTAT crop " - f"budget for {group}: need {remainder:.6g} Mha but only " - f"{total_available:.6g} Mha remains" - ) - new_baseline.loc[idx] -= ( - available / total_available * min(float(remainder), total_available) + remaining_by_link = pd.Series( + group_index.map(remaining_required).fillna(0.0), index=single.index + ) + available_by_group = new_baseline.groupby(group_index).transform("sum") + invalid = remaining_by_link > available_by_group + 1e-9 + if invalid.any(): + link = invalid[invalid].index[0] + group = tuple(group_index[single.index.get_loc(link)]) + raise ValueError( + "Scaled multi-cropping anchors still exceed the FAOSTAT crop " + f"budget for {group}: need {remaining_by_link[link]:.6g} Mha but " + f"only {available_by_group[link]:.6g} Mha remains" ) + active = remaining_by_link > 1e-12 + new_baseline.loc[active] -= ( + new_baseline.loc[active] + / available_by_group.loc[active] + * np.minimum(remaining_by_link.loc[active], available_by_group.loc[active]) + ) new_baseline = new_baseline.clip(lower=0.0) n.links.static.loc[new_baseline.index, "baseline_area_mha"] = new_baseline diff --git a/workflow/scripts/build_multi_cropping.py b/workflow/scripts/build_multi_cropping.py index e749a6e1..f4ac592b 100644 --- a/workflow/scripts/build_multi_cropping.py +++ b/workflow/scripts/build_multi_cropping.py @@ -6,21 +6,17 @@ from pathlib import Path -from exactextract import exact_extract -from exactextract.raster import NumPyRasterSource -import geopandas as gpd import numpy as np import pandas as pd -import xarray as xr +import rasterio from workflow.scripts.multi_cropping_combinations import effective_combinations from workflow.scripts.raster_utils import ( calculate_all_cell_areas, load_raster_array, - raster_bounds, - read_raster_float, scale_fraction, ) +from workflow.scripts.region_class_aggregation import load_cell_mapping ZONE_CAPABILITIES: dict[int, dict[str, int | bool]] = { 0: {"valid": False, "max_cycles": 0, "max_wetland_rice": 0}, @@ -53,109 +49,6 @@ WETLAND_RICE_CROPS = {"wetland-rice"} -# exactextract's NumPyRasterSource retains, at the C++ level, the numpy array it -# is handed, and that memory is never reclaimed even after the source is GC'd. -# Because this module wraps a fresh full-grid array on every call (once per combo -# x class x period), that leak accumulates to tens of GB. Reusing a single buffer -# per grid shape means only one array is ever retained -- copying the caller's -# transient array into the buffer lets the transient be freed normally. -_EXTRACT_BUFFERS: dict[tuple[int, ...], np.ndarray] = {} - - -def get_extract_buffer(shape: tuple[int, ...]) -> np.ndarray: - """Return the reused float64 exact_extract buffer for ``shape``. - - Callers may fill the buffer themselves (e.g. accumulate directly into it) - and pass it to the aggregation functions, which skip the defensive copy - when handed the buffer itself. - """ - buffer = _EXTRACT_BUFFERS.get(shape) - if buffer is None: - buffer = np.empty(shape, dtype=np.float64) - _EXTRACT_BUFFERS[shape] = buffer - return buffer - - -def aggregate_raster_by_region( - data_array: np.ndarray, - regions_gdf: gpd.GeoDataFrame, - xmin: float, - ymin: float, - xmax: float, - ymax: float, - crs_wkt: str | None, - stat: str = "sum", -) -> pd.DataFrame: - """Aggregate raster data by regions using exact_extract. - - The data is copied into a reused per-shape buffer before being handed to - ``NumPyRasterSource`` to avoid a per-call memory leak (see ``_EXTRACT_BUFFERS``). - """ - buffer = get_extract_buffer(data_array.shape) - if data_array is not buffer: - buffer[:] = data_array - src = NumPyRasterSource( - buffer, - xmin=xmin, - ymin=ymin, - xmax=xmax, - ymax=ymax, - nodata=np.nan, - srs_wkt=crs_wkt, - ) - return exact_extract( - src, - regions_gdf, - [stat], - include_cols=["region"], - output="pandas", - ) - - -def region_coverage_entries( - regions_gdf: gpd.GeoDataFrame, - xmin: float, - ymin: float, - xmax: float, - ymax: float, - crs_wkt: str | None, - shape: tuple[int, int], -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """Per-region cell coverage fractions for a grid, extracted once. - - Returns ``(region_names, rows, cells, fracs)`` where ``rows`` indexes into - ``region_names``, ``cells`` are flat row-major indices into the grid, and - ``fracs`` is the fraction of each cell covered by the region. For an - all-finite value grid ``v``, - ``np.bincount(rows, weights=v.ravel()[cells] * fracs, minlength=len(region_names))`` - reproduces exact_extract's ``sum`` op, so any number of rasters sharing the - grid can be aggregated without recomputing the polygon coverage. - """ - buffer = get_extract_buffer(shape) - buffer.fill(0.0) # cell values are irrelevant here but must be finite - src = NumPyRasterSource( - buffer, - xmin=xmin, - ymin=ymin, - xmax=xmax, - ymax=ymax, - nodata=np.nan, - srs_wkt=crs_wkt, - ) - res = exact_extract( - src, - regions_gdf, - ["cell_id", "coverage"], - include_cols=["region"], - output="pandas", - ) - lengths = res["cell_id"].map(len).to_numpy() - rows = np.repeat(np.arange(len(res)), lengths) - cells = np.concatenate([np.asarray(a, dtype=np.int64) for a in res["cell_id"]]) - fracs = np.concatenate([np.asarray(a, dtype=np.float64) for a in res["coverage"]]) - return res["region"].to_numpy(), rows, cells, fracs - - def compute_eligibility_mask( crop_sequence: list[str], ws: str, @@ -245,8 +138,7 @@ def compute_eligibility_mask( for ws in ("r", "i") if f"multiple_cropping_zone_{ws}" in inputs } - classes_nc = inputs.pop("classes") - regions_path = inputs.pop("regions") + mapping_path = inputs.pop("cell_mapping") inputs.pop("combinations") conv_csv = inputs.pop("yield_unit_conversions") moisture_csv = inputs.pop("moisture_content") @@ -307,29 +199,15 @@ def compute_eligibility_mask( .to_dict() ) - ds = xr.load_dataset(classes_nc) - if "resource_class" not in ds: - raise ValueError("resource_classes.nc is missing 'resource_class' data") - class_labels = ds["resource_class"].values.astype(np.int16) + mapping = load_cell_mapping(mapping_path) # Use any available crop/ws pair to get raster dimensions sample_crop, sample_ws = next(iter(crop_files.keys())) - yield_arr_ref, yield_src = read_raster_float( - crop_files[(sample_crop, sample_ws)]["yield"] - ) - try: - height, width = yield_arr_ref.shape - if class_labels.shape != (height, width): - raise ValueError( - "Resource class grid does not match GAEZ raster dimensions for multiple cropping" - ) - transform = yield_src.transform - crs = yield_src.crs - crs_wkt = crs.to_wkt() if crs else None - xmin, ymin, xmax, ymax = raster_bounds(transform, width, height) + with rasterio.open(crop_files[(sample_crop, sample_ws)]["yield"]) as yield_src: + height, width = yield_src.shape + if mapping.shape != (height, width): + raise ValueError("Cell mapping does not match GAEZ raster dimensions") cell_area_ha = calculate_all_cell_areas(yield_src) - finally: - yield_src.close() zone_arrays: dict[str, np.ndarray] = {} for ws, path in zone_paths.items(): @@ -340,11 +218,6 @@ def compute_eligibility_mask( ) zone_arrays[ws] = zone_arr.astype(np.int16, copy=False) - 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() - def conversion_factor(crop: str) -> float: base_scale = 1.0 if use_actual_yields else KG_TO_TONNE if crop in conv_df.index: @@ -393,21 +266,6 @@ def conversion_factor(crop: str) -> float: # with the single-crop path in build_crop_yields. water_requirement_data[(crop, ws)] = water_arr * 10.0 - valid_classes = [ - int(cls) - for cls in np.unique(class_labels[np.isfinite(class_labels)]) - if int(cls) >= 0 - ] - - # Region coverage fractions are identical for every raster on this grid, so - # extract them once and aggregate all combo/class/period/cycle rasters with - # gathers + bincount instead of one exact_extract call per raster. - region_names, cov_rows, cov_cells, cov_fracs = region_coverage_entries( - regions_for_extract, xmin, ymin, xmax, ymax, crs_wkt, (height, width) - ) - n_regions = len(region_names) - class_at_entry = class_labels.ravel()[cov_cells] - eligible_records: list[pd.DataFrame] = [] cycle_records: list[pd.DataFrame] = [] @@ -432,90 +290,81 @@ def conversion_factor(crop: str) -> float: eligible_fraction = np.where(combined_mask, min_fraction, np.nan) eligible_area = eligible_fraction * cell_area_ha - # Coverage entries restricted to the combo's eligible cells; per class, - # every aggregation is a gather over these entries + bincount against - # the same coverage fractions exact_extract would apply. - combo_at_entry = combined_mask.ravel()[cov_cells] - ea_flat = eligible_area.ravel() - - for cls in valid_classes: - sel = combo_at_entry & (class_at_entry == cls) - if not sel.any(): - continue - rows_s = cov_rows[sel] - cells_s = cov_cells[sel] - fracs_s = cov_fracs[sel] - ea_entries = ea_flat[cells_s] * fracs_s + selected = combined_mask.ravel()[mapping.cell_ids] + cells = mapping.cell_ids[selected] + groups = mapping.group_ids[selected] + area_entries = eligible_area.ravel()[cells] * mapping.coverage[selected] + area_by_group = np.bincount( + groups, weights=area_entries, minlength=mapping.n_groups + ) + positive = area_by_group > 0 + if not positive.any(): + continue + positive_groups = np.flatnonzero(positive) + region_ids, resource_classes = np.divmod(positive_groups, mapping.n_classes) + area_stats = pd.DataFrame( + { + "region": mapping.regions[region_ids], + "eligible_area_ha": area_by_group[positive], + } + ) - area_by_region = np.bincount( - rows_s, weights=ea_entries, minlength=n_regions + # Annual irrigation requirement (m3/ha): aggregate the summed-cycle + # demand*area numerator against the same eligible-area denominator. + if ws == "i" and total_water_arr is not None: + volume = np.bincount( + groups, + weights=total_water_arr.ravel()[cells] * area_entries, + minlength=mapping.n_groups, ) - pos = area_by_region > 0 - if not pos.any(): - continue - area_stats = pd.DataFrame( - { - "region": region_names[pos], - "eligible_area_ha": area_by_region[pos], - } + area_stats["water_requirement_m3_per_ha"] = ( + volume[positive] / area_by_group[positive] ) - - # Annual irrigation requirement (m3/ha): aggregate the summed-cycle - # demand*area numerator against the same eligible-area denominator. - if ws == "i" and total_water_arr is not None: - volume = np.bincount( - rows_s, - weights=total_water_arr.ravel()[cells_s] * ea_entries, - minlength=n_regions, - ) - area_stats["water_requirement_m3_per_ha"] = ( - volume[pos] / area_by_region[pos] - ) - else: - area_stats["water_requirement_m3_per_ha"] = 0.0 - - area_stats["resource_class"] = cls - area_stats["combination"] = combo_name - area_stats["water_supply"] = ws - eligible_records.append( - area_stats[ - [ - "combination", - "region", - "resource_class", - "water_supply", - "eligible_area_ha", - "water_requirement_m3_per_ha", - ] + else: + area_stats["water_requirement_m3_per_ha"] = 0.0 + + area_stats["resource_class"] = resource_classes + area_stats["combination"] = combo_name + area_stats["water_supply"] = ws + eligible_records.append( + area_stats[ + [ + "combination", + "region", + "resource_class", + "water_supply", + "eligible_area_ha", + "water_requirement_m3_per_ha", ] - ) + ] + ) - # Calculate yields for each crop cycle - for idx, (crop_name, yield_arr) in enumerate( - zip(crop_sequence, yield_stack), start=1 - ): - numerator = np.bincount( - rows_s, - weights=yield_arr.ravel()[cells_s] * ea_entries, - minlength=n_regions, - ) - yield_t_per_ha = numerator[pos] / area_by_region[pos] - keep = yield_t_per_ha > 0 - if not keep.any(): - continue - cycle_records.append( - pd.DataFrame( - { - "combination": combo_name, - "region": region_names[pos][keep], - "resource_class": cls, - "water_supply": ws, - "cycle_index": idx, - "crop": crop_name, - "yield_t_per_ha": yield_t_per_ha[keep], - } - ) + # Calculate yields for each crop cycle + for idx, (crop_name, yield_arr) in enumerate( + zip(crop_sequence, yield_stack), start=1 + ): + numerator = np.bincount( + groups, + weights=yield_arr.ravel()[cells] * area_entries, + minlength=mapping.n_groups, + ) + yield_t_per_ha = numerator[positive] / area_by_group[positive] + keep = yield_t_per_ha > 0 + if not keep.any(): + continue + cycle_records.append( + pd.DataFrame( + { + "combination": combo_name, + "region": mapping.regions[region_ids][keep], + "resource_class": resource_classes[keep], + "water_supply": ws, + "cycle_index": idx, + "crop": crop_name, + "yield_t_per_ha": yield_t_per_ha[keep], + } ) + ) if eligible_records: eligible_df = pd.concat(eligible_records, ignore_index=True) diff --git a/workflow/scripts/compute_resource_classes.py b/workflow/scripts/compute_resource_classes.py index 7fb2d0e4..b4ccbc4c 100644 --- a/workflow/scripts/compute_resource_classes.py +++ b/workflow/scripts/compute_resource_classes.py @@ -34,6 +34,22 @@ def read_raster_float(path: str): return arr, src +def read_raster_cells_float( + path: str, + cell_ids: np.ndarray, + expected_shape: tuple[int, int], +) -> np.ndarray: + with rasterio.open(path) as src: + if src.shape != expected_shape: + raise ValueError( + f"Raster shape mismatch for {path}: {src.shape} != {expected_shape}" + ) + arr = src.read(1).ravel()[cell_ids].astype(float) + if src.nodata is not None: + arr = np.where(arr == src.nodata, np.nan, arr) + return arr + + def weighted_median(values: np.ndarray, weights: np.ndarray) -> float: order = np.argsort(values) values = values[order] @@ -120,21 +136,6 @@ def shares_by_region( return countries.map(lambda country: lookup.get(country, fallback)).to_numpy(float) -def sum_by_region( - values: np.ndarray, - region_raster: np.ndarray, - n_regions: int, -) -> np.ndarray: - valid = (region_raster >= 0) & np.isfinite(values) - if not np.any(valid): - return np.zeros(n_regions, dtype=float) - return np.bincount( - region_raster[valid].ravel(), - weights=values[valid].ravel(), - minlength=n_regions, - ) - - def compute_max_yield_score( yield_paths: list[str], pairs: list[tuple[str, str]], @@ -180,19 +181,20 @@ def compute_regional_harvested_area_score( mapping_df = load_mapping(Path(mapping_path)) production_df = pd.read_csv(production_path) n_regions = len(regions_gdf) - numerator = np.zeros(expected_shape, dtype=float) - denominator = np.zeros(expected_shape, dtype=float) - region_valid = region_raster >= 0 - share_grid_cache: dict[str, np.ndarray] = {} + region_cell_ids = np.flatnonzero(region_raster.ravel() >= 0) + region_ids = region_raster.ravel()[region_cell_ids] + numerator = np.zeros(len(region_cell_ids), dtype=float) + denominator = np.zeros(len(region_cell_ids), dtype=float) + region_share_cache: dict[str, np.ndarray] = {} for yield_path, harvested_path, (crop, _water_supply) in zip( yield_paths, harvested_paths, pairs, strict=True ): - y_raw, y_src = read_raster_float(yield_path) - try: - validate_raster_shape(y_raw, expected_shape, yield_path) - finally: - y_src.close() + y_raw = read_raster_cells_float( + yield_path, + region_cell_ids, + expected_shape, + ) y = scale_yield( y_raw, crop, @@ -201,19 +203,18 @@ def compute_regional_harvested_area_score( moisture=moisture, ) - harvested_raw, harvested_src = read_raster_float(harvested_path) - try: - validate_raster_shape(harvested_raw, expected_shape, harvested_path) - finally: - harvested_src.close() - + harvested = read_raster_cells_float( + harvested_path, + region_cell_ids, + expected_shape, + ) harvested = np.where( - np.isfinite(harvested_raw) & (harvested_raw > 0.0), - harvested_raw * RES06_HAR_SCALE_TO_HA, + np.isfinite(harvested) & (harvested > 0.0), + harvested * RES06_HAR_SCALE_TO_HA, 0.0, ) - if crop not in share_grid_cache: - region_shares = shares_by_region( + if crop not in region_share_cache: + region_share_cache[crop] = shares_by_region( crop, regions_gdf, mapping_df, @@ -221,11 +222,8 @@ def compute_regional_harvested_area_score( fdd_shares_path, non_food_crops, ) - share_grid = np.zeros(expected_shape, dtype=float) - share_grid[region_valid] = region_shares[region_raster[region_valid]] - share_grid_cache[crop] = share_grid - share_grid = share_grid_cache[crop] - crop_area = harvested * share_grid + region_shares = region_share_cache[crop] + crop_area = harvested * region_shares[region_ids] scale_mask = np.isfinite(y) & (y > 0.0) & (crop_area > 0.0) if not np.any(scale_mask): @@ -234,22 +232,59 @@ def compute_regional_harvested_area_score( if not np.isfinite(scale) or scale <= 0.0: continue - regional_area = sum_by_region(crop_area, region_raster, n_regions) - region_weight = np.zeros(expected_shape, dtype=float) - region_weight[region_valid] = regional_area[region_raster[region_valid]] + regional_area = np.bincount( + region_ids, + weights=crop_area, + minlength=n_regions, + ) + region_weight = regional_area[region_ids] normalized = y / scale - valid = region_valid & np.isfinite(normalized) & (normalized > 0.0) + valid = np.isfinite(normalized) & (normalized > 0.0) valid &= region_weight > 0.0 numerator[valid] += region_weight[valid] * normalized[valid] denominator[valid] += region_weight[valid] - return np.divide( + score = np.full(expected_shape, np.nan, dtype=float) + score.ravel()[region_cell_ids] = np.divide( numerator, denominator, - out=np.full(expected_shape, np.nan, dtype=float), + out=np.full(len(region_cell_ids), np.nan, dtype=float), where=denominator > 0.0, ) + return score + + +def classify_by_region( + score: np.ndarray, + region_raster: np.ndarray, + quantiles: list[float], +) -> np.ndarray: + """Assign positive scores to unweighted within-region quantile bins.""" + flat_score = score.ravel() + flat_regions = region_raster.ravel() + valid_cell_ids = np.flatnonzero( + (flat_regions >= 0) & np.isfinite(flat_score) & (flat_score > 0.0) + ) + classes = np.full(flat_score.shape, -1, dtype=np.int8) + if not valid_cell_ids.size: + return classes.reshape(score.shape) + + order = np.argsort(flat_regions[valid_cell_ids], kind="stable") + sorted_cell_ids = valid_cell_ids[order] + sorted_regions = flat_regions[sorted_cell_ids] + starts = np.r_[0, np.flatnonzero(np.diff(sorted_regions)) + 1] + stops = np.r_[starts[1:], len(sorted_cell_ids)] + + for start, stop in zip(starts, stops, strict=True): + cell_ids = sorted_cell_ids[start:stop] + thresholds = np.quantile(flat_score[cell_ids], quantiles) + classes[cell_ids] = np.searchsorted( + thresholds[1:-1], + flat_score[cell_ids], + side="right", + ) + return classes.reshape(score.shape) if __name__ == "__main__": @@ -351,32 +386,19 @@ def compute_regional_harvested_area_score( raise ValueError(f"Unknown resource class score method: {score_method}") # Build xarray DataArrays - y_da = xr.DataArray(score, dims=("y", "x")) reg_da = xr.DataArray(region_raster, dims=("y", "x")) - # Vectorized per-region quantiles and class assignment - # Ignore cells with zero/negative scores so unsuitable or uncovered pixels - # do not collapse the quantile bins. - positive_y = xr.where((y_da > 0) & np.isfinite(y_da), y_da, np.nan) - reg_quantiles = positive_y.groupby(reg_da).quantile(quantiles) - thresholds = reg_quantiles.sel(group=reg_da).reset_coords(drop=True) - - class_da = xr.full_like(y_da, np.nan, dtype=float) - for ci in range(len(quantiles) - 1): - lo = thresholds.isel(quantile=ci) - hi = thresholds.isel(quantile=ci + 1) - if ci == len(quantiles) - 2: - sel = (reg_da >= 0) & np.isfinite(y_da) & (y_da >= lo) - else: - sel = (reg_da >= 0) & np.isfinite(y_da) & (y_da >= lo) & (y_da < hi) - class_da = xr.where(sel, float(ci), class_da) + class_da = xr.DataArray( + classify_by_region(score, region_raster, quantiles), + dims=("y", "x"), + ) ds = xr.Dataset( { "region_id": reg_da.astype(np.int32), - "resource_class": class_da.fillna(-1).astype(np.int8), + "resource_class": class_da, } - ) + ).assign_coords(quantile=quantiles[-2]) # Store transform/CRS/bounds as attrs for downstream use ds.attrs.update( { diff --git a/workflow/scripts/derive_mirca_multicropping.py b/workflow/scripts/derive_mirca_multicropping.py index 301218b9..7d29dd01 100644 --- a/workflow/scripts/derive_mirca_multicropping.py +++ b/workflow/scripts/derive_mirca_multicropping.py @@ -21,11 +21,13 @@ the active configuration's regions and resource classes. """ -from collections import Counter -from dataclasses import dataclass +from collections.abc import Callable +from dataclasses import dataclass, field from pathlib import Path from affine import Affine +from exactextract import exact_extract +from exactextract.raster import NumPyRasterSource import geopandas as gpd import numpy as np import pandas as pd @@ -33,11 +35,7 @@ from rasterio.crs import CRS import xarray as xr -from workflow.scripts.build_multi_cropping import ( - WETLAND_RICE_CROPS, - ZONE_CAPABILITIES, - region_coverage_entries, -) +from workflow.scripts.build_multi_cropping import WETLAND_RICE_CROPS, ZONE_CAPABILITIES from workflow.scripts.multi_cropping_combinations import load_catalog_combinations from workflow.scripts.raster_utils import raster_bounds @@ -79,13 +77,17 @@ def _assert_grid(actual: GridSpec, expected: GridSpec, label: str) -> None: ) -def load_tif(path: str, expected_grid: GridSpec | None = None) -> np.ndarray: - """Load a GeoTIFF as float64 and validate its complete grid identity.""" +def load_tif( + path: str, + expected_grid: GridSpec | None = None, + dtype: np.dtype = np.dtype(np.float64), +) -> np.ndarray: + """Load a GeoTIFF and validate its complete grid identity.""" with rasterio.open(path) as src: grid = _grid_from_raster(src, path) if expected_grid is not None: _assert_grid(grid, expected_grid, path) - arr = src.read(1).astype(np.float64) + arr = src.read(1).astype(dtype) nodata = src.nodata if nodata is not None: arr[arr == nodata] = 0.0 @@ -192,15 +194,15 @@ def candidate_capacity( rice2 = rice_support[ws] rice3 = rice_support[f"{ws}3"] if n_cycles == 2: - area = np.clip(rice2 - rice3, 0.0, None) + area = np.clip(np.subtract(rice2, rice3, dtype=np.float64), 0.0, None) elif n_cycles == 3: - area = rice3.copy() + area = rice3.astype(np.float64) else: raise ValueError(f"Unsupported repeated-rice cycle count: {n_cycles}") area[~zmask] = 0.0 return area - area = crop_area[(crops[0], ws)].copy() + area = crop_area[(crops[0], ws)].astype(np.float64) observed = area > 0 for crop in crops[1:]: support = crop_area[(crop, ws)] @@ -240,31 +242,31 @@ def allocate( if not capacities: return [], extra_cycle_area.copy() - total_physical = np.zeros_like(extra_cycle_area, dtype=np.float64) - total_extra = np.zeros_like(extra_cycle_area, dtype=np.float64) - required_by_crop = { - crop: np.zeros_like(extra_cycle_area, dtype=np.float64) - for crops in crop_sequences - for crop in crops - } + total = np.zeros_like(extra_cycle_area, dtype=np.float64) for capacity, crops in zip(capacities, crop_sequences, strict=True): - total_physical += capacity - total_extra += (len(crops) - 1) * capacity - for crop, multiplicity in Counter(crops).items(): - required_by_crop[crop] += multiplicity * capacity + total += (len(crops) - 1) * capacity scale = np.ones_like(extra_cycle_area, dtype=np.float64) - np.minimum(scale, _bounded_ratio(extra_cycle_area, total_extra), out=scale) - np.minimum(scale, _bounded_ratio(footprint_area, total_physical), out=scale) - for crop, required in required_by_crop.items(): - np.minimum(scale, _bounded_ratio(crop_support[crop], required), out=scale) + np.minimum(scale, _bounded_ratio(extra_cycle_area, total), out=scale) + total.fill(0.0) + for capacity in capacities: + total += capacity + np.minimum(scale, _bounded_ratio(footprint_area, total), out=scale) + + for crop in dict.fromkeys(crop for crops in crop_sequences for crop in crops): + total.fill(0.0) + for capacity, crops in zip(capacities, crop_sequences, strict=True): + multiplicity = crops.count(crop) + if multiplicity: + total += multiplicity * capacity + np.minimum(scale, _bounded_ratio(crop_support[crop], total), out=scale) areas: list[np.ndarray] = [] allocated_extra = np.zeros_like(extra_cycle_area, dtype=np.float64) for capacity, crops in zip(capacities, crop_sequences, strict=True): - area = capacity * scale - areas.append(area) - allocated_extra += (len(crops) - 1) * area + np.multiply(capacity, scale, out=capacity) + areas.append(capacity) + allocated_extra += (len(crops) - 1) * capacity residual = np.clip(extra_cycle_area - allocated_extra, 0.0, None) return areas, residual @@ -277,6 +279,7 @@ def run_derivation( rice_support: dict[str, np.ndarray], combos: list[dict], harvested_totals: dict[str, np.ndarray] | None = None, + area_sink: Callable[[tuple[str, str], np.ndarray], None] | None = None, ) -> tuple[dict[tuple[str, str], np.ndarray], np.ndarray, pd.DataFrame]: """Run independent rainfed and irrigated baseline attribution.""" if harvested_totals is None: @@ -314,9 +317,13 @@ def run_derivation( ) residual_total += residual + area = None for combo, area in zip(ws_combos, areas, strict=True): key = (combo["name"], ws) - area_rasters[key] = area + if area_sink is None: + area_rasters[key] = area + else: + area_sink(key, area) n_cycles = len(combo["crops"]) records.append( { @@ -329,6 +336,9 @@ def run_derivation( "system_residual_area_mha": float(residual.sum()) / 1e6, } ) + if area_sink is not None: + area = None + del areas return area_rasters, residual_total, pd.DataFrame.from_records(records) @@ -372,17 +382,28 @@ def _load_inputs( totals = {mws: np.zeros(grid.shape, dtype=np.float64) for mws in ("ir", "rf")} for row in mapping.itertuples(): for mws in ("ir", "rf"): - arr = load_tif(inp[_annual_key(row.mirca_crop, mws)], grid) + arr = load_tif( + inp[_annual_key(row.mirca_crop, mws)], grid, np.dtype(np.float32) + ) totals[mws] += arr if row.mirca_crop in mirca_to_glade: annual[(row.mirca_crop, mws)] = arr - footprint = {mws: load_tif(inp[f"footprint_{mws}"], grid) for mws in ("ir", "rf")} - zone = {ws: load_tif(inp[f"zone_{ws}"], grid) for ws in ("i", "r")} + footprint = { + mws: load_tif(inp[f"footprint_{mws}"], grid, np.dtype(np.float32)) + for mws in ("ir", "rf") + } + zone = { + ws: load_tif(inp[f"zone_{ws}"], grid, np.dtype(np.int16)) for ws in ("i", "r") + } rice_support: dict[str, np.ndarray] = {} for gws, mws in {"i": "ir", "r": "rf"}.items(): - rice_support[gws] = load_subcrop_maxmonth(inp[f"rice2_{mws}"], grid) - rice_support[f"{gws}3"] = load_subcrop_maxmonth(inp[f"rice3_{mws}"], grid) + rice_support[gws] = load_subcrop_maxmonth(inp[f"rice2_{mws}"], grid).astype( + np.float32 + ) + rice_support[f"{gws}3"] = load_subcrop_maxmonth( + inp[f"rice3_{mws}"], grid + ).astype(np.float32) return annual, totals, footprint, zone, rice_support @@ -411,66 +432,125 @@ def _load_resource_classes(path: str, grid: GridSpec) -> np.ndarray: return labels -def aggregate_baseline( - area_rasters: dict[tuple[str, str], np.ndarray], - class_labels: np.ndarray, +def region_coverage_entries( regions_gdf: gpd.GeoDataFrame, grid: GridSpec, -) -> pd.DataFrame: - """Aggregate physical bundle areas to region and resource class.""" - regions = regions_gdf - if regions.crs is None: - raise ValueError("regions.geojson is missing CRS information") - if regions.crs != grid.crs: - regions = regions.to_crs(grid.crs) - regions_for_extract = regions.reset_index() - if "region" not in regions_for_extract: - raise ValueError("regions.geojson must provide a 'region' index or column") - +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Return exact region/cell coverage entries for the MIRCA grid.""" height, width = grid.shape xmin, ymin, xmax, ymax = raster_bounds(grid.transform, width, height) - region_names, rows, cells, fractions = region_coverage_entries( - regions_for_extract, - xmin, - ymin, - xmax, - ymax, - grid.crs.to_wkt(), - grid.shape, + source = NumPyRasterSource( + np.zeros(grid.shape, dtype=np.float64), + xmin=xmin, + ymin=ymin, + xmax=xmax, + ymax=ymax, + nodata=np.nan, + srs_wkt=grid.crs.to_wkt(), ) - class_at_entry = class_labels.ravel()[cells] - valid_classes = sorted( - int(value) - for value in np.unique(class_labels[np.isfinite(class_labels)]) - if int(value) >= 0 + result = exact_extract( + source, + regions_gdf, + ["cell_id", "coverage"], + include_cols=["region"], + output="pandas", ) - records: list[pd.DataFrame] = [] - for (name, ws), area in area_rasters.items(): - values = area.ravel()[cells] * fractions - for resource_class in valid_classes: - selected = class_at_entry == resource_class + lengths = result["cell_id"].map(len).to_numpy() + rows = np.repeat(np.arange(len(result)), lengths) + cells = np.concatenate(result["cell_id"].to_numpy()).astype(np.int64) + fractions = np.concatenate(result["coverage"].to_numpy()).astype(np.float64) + return result["region"].to_numpy(), rows, cells, fractions + + +@dataclass +class BaselineAggregator: + """Accumulate exact region/class sums without retaining full rasters.""" + + region_names: np.ndarray + rows: np.ndarray + cells: np.ndarray + fractions: np.ndarray + class_at_entry: np.ndarray + valid_classes: list[int] + records: list[pd.DataFrame] = field(default_factory=list) + + @classmethod + def create( + cls, + class_labels: np.ndarray, + regions_gdf: gpd.GeoDataFrame, + grid: GridSpec, + ) -> "BaselineAggregator": + regions = regions_gdf + if regions.crs is None: + raise ValueError("regions.geojson is missing CRS information") + if regions.crs != grid.crs: + regions = regions.to_crs(grid.crs) + regions_for_extract = regions.reset_index() + if "region" not in regions_for_extract: + raise ValueError("regions.geojson must provide a 'region' index or column") + + region_names, rows, cells, fractions = region_coverage_entries( + regions_for_extract, grid + ) + valid_classes = sorted( + int(value) + for value in np.unique(class_labels[np.isfinite(class_labels)]) + if int(value) >= 0 + ) + return cls( + region_names, + rows, + cells, + fractions, + class_labels.ravel()[cells], + valid_classes, + ) + + def add(self, key: tuple[str, str], area: np.ndarray) -> None: + name, ws = key + values = area.ravel()[self.cells] * self.fractions + for resource_class in self.valid_classes: + selected = self.class_at_entry == resource_class sums = np.bincount( - rows[selected], weights=values[selected], minlength=len(region_names) + self.rows[selected], + weights=values[selected], + minlength=len(self.region_names), ) positive = sums > 0 if positive.any(): - records.append( + self.records.append( pd.DataFrame( { "combination": name, - "region": region_names[positive], + "region": self.region_names[positive], "resource_class": resource_class, "water_supply": ws, "baseline_area_ha": sums[positive], } ) ) - if not records: - return pd.DataFrame(columns=BASELINE_COLUMNS) - return pd.concat(records, ignore_index=True)[BASELINE_COLUMNS].sort_values( - ["combination", "water_supply", "region", "resource_class"], - ignore_index=True, - ) + + def result(self) -> pd.DataFrame: + if not self.records: + return pd.DataFrame(columns=BASELINE_COLUMNS) + return pd.concat(self.records, ignore_index=True)[BASELINE_COLUMNS].sort_values( + ["combination", "water_supply", "region", "resource_class"], + ignore_index=True, + ) + + +def aggregate_baseline( + area_rasters: dict[tuple[str, str], np.ndarray], + class_labels: np.ndarray, + regions_gdf: gpd.GeoDataFrame, + grid: GridSpec, +) -> pd.DataFrame: + """Aggregate physical bundle areas to region and resource class.""" + aggregator = BaselineAggregator.create(class_labels, regions_gdf, grid) + for key, area in area_rasters.items(): + aggregator.add(key, area) + return aggregator.result() def main() -> None: @@ -504,7 +584,10 @@ def main() -> None: for name, entry in catalog.items() for ws in entry["water_supplies"] ] - areas, residual, stats = run_derivation( + classes = _load_resource_classes(inp["classes"], grid) + regions = gpd.read_file(inp["regions"]) + aggregator = BaselineAggregator.create(classes, regions, grid) + _areas, residual, stats = run_derivation( annual, footprint, crop_area, @@ -512,11 +595,9 @@ def main() -> None: rice_support, combos, harvested_totals=totals, + area_sink=aggregator.add, ) - - classes = _load_resource_classes(inp["classes"], grid) - regions = gpd.read_file(inp["regions"]) - baseline = aggregate_baseline(areas, classes, regions, grid) + baseline = aggregator.result() baseline_path = Path(snakemake.output.baseline) # type: ignore[name-defined] residual_path = Path(snakemake.output.residual) # type: ignore[name-defined] diff --git a/workflow/scripts/process_huang_irrigation_water.py b/workflow/scripts/process_huang_irrigation_water.py index 95ae7e40..57d6c042 100644 --- a/workflow/scripts/process_huang_irrigation_water.py +++ b/workflow/scripts/process_huang_irrigation_water.py @@ -143,81 +143,20 @@ def aggregate_gridded_to_regions( def load_crop_growing_seasons(crop_files: Iterable[str]) -> pd.DataFrame: - """Load and aggregate crop growing seasons from yield files. - - This is a copy of the function from build_region_water_availability.py - to ensure consistency. - """ - records = [] - for path_str in crop_files: - path = Path(path_str) - stem = path.stem - if "_" not in stem: - continue - crop, water_supply = stem.split("_", 1) - - df = pd.read_csv(path) - - pivot = ( - df.pivot( - index=["region", "resource_class"], columns="variable", values="value" - ) - .rename_axis(columns=None) - .reset_index() - ) - - # CROPGRIDS-backed crops (config["cropgrids_crops"]) don't carry - # GAEZ growing-season rasters and contribute no irrigation demand - # (they're rainfed-only). Skip them rather than dropna-erroring. - required = { - "suitable_area", - "growing_season_start_day", - "growing_season_length_days", - } - if not required.issubset(pivot.columns): - continue - - pivot = pivot.dropna( - subset=[ - "region", - "suitable_area", - "growing_season_start_day", - "growing_season_length_days", - ] - ) - pivot = pivot[pivot["suitable_area"] > 0] - if pivot.empty: - continue - - pivot["resource_class"] = pivot["resource_class"].astype(int) - for column in [ - "suitable_area", - "growing_season_start_day", - "growing_season_length_days", - ]: - pivot[column] = pd.to_numeric(pivot[column], errors="coerce") - - grouped = pivot.groupby("region") - for region, group in grouped: - weight = group["suitable_area"].sum() - if weight <= 0: - continue - start = ( - group["growing_season_start_day"] * group["suitable_area"] - ).sum() / weight - length = ( - group["growing_season_length_days"] * group["suitable_area"] - ).sum() / weight - records.append( - { - "region": region, - "crop": crop, - "water_supply": water_supply, - "total_area": weight, - "growing_season_start_day": start, - "growing_season_length_days": length, - } - ) + """Load area-weighted crop growing seasons from yield files.""" + crop_files = list(crop_files) + irrigated_files = [path for path in crop_files if Path(path).stem.endswith("_i")] + records = [ + record + for path in irrigated_files + for record in _load_crop_growing_season_file(path) + ] + if not records: + records = [ + record + for path in crop_files + for record in _load_crop_growing_season_file(path) + ] if not records: return pd.DataFrame( columns=[ @@ -232,6 +171,71 @@ def load_crop_growing_seasons(crop_files: Iterable[str]) -> pd.DataFrame: return pd.DataFrame(records) +def _load_crop_growing_season_file(path_str: str) -> list[dict]: + """Load area-weighted growing seasons from one crop-yield file.""" + path = Path(path_str) + stem = path.stem + if "_" not in stem: + return [] + crop, water_supply = stem.split("_", 1) + + required = { + "suitable_area", + "growing_season_start_day", + "growing_season_length_days", + } + df = pd.read_csv( + path, + usecols=["region", "resource_class", "variable", "value"], + ) + df = df[df["variable"].isin(required)] + if not required.issubset(df["variable"].unique()): + return [] + + pivot = ( + df.pivot(index=["region", "resource_class"], columns="variable", values="value") + .rename_axis(columns=None) + .reset_index() + ) + pivot = pivot.dropna( + subset=[ + "region", + "suitable_area", + "growing_season_start_day", + "growing_season_length_days", + ] + ) + pivot = pivot[pivot["suitable_area"] > 0] + if pivot.empty: + return [] + + regions = pivot["region"].to_numpy() + area = pivot["suitable_area"].to_numpy() + start = pivot["growing_season_start_day"].to_numpy() + length = pivot["growing_season_length_days"].to_numpy() + group_starts = np.flatnonzero(np.r_[True, regions[1:] != regions[:-1]]) + group_ends = np.r_[group_starts[1:], len(regions)] + + records = [] + for first, last in zip(group_starts, group_ends, strict=True): + weight = area[first:last].sum() + records.append( + { + "region": regions[first], + "crop": crop, + "water_supply": water_supply, + "total_area": weight, + "growing_season_start_day": (start[first:last] * area[first:last]).sum() + / weight, + "growing_season_length_days": ( + length[first:last] * area[first:last] + ).sum() + / weight, + } + ) + return records + + def compute_region_growing_water( region_month_water: pd.DataFrame, crop_seasons: pd.DataFrame, @@ -277,14 +281,14 @@ def compute_region_growing_water( } region_total_area = dict.fromkeys(crop_seasons["region"].unique(), 0.0) - for _, row in irrigated.iterrows(): - region = row["region"] + for row in irrigated.itertuples(index=False): + region = row.region overlaps = compute_month_overlaps( - row["growing_season_start_day"], row["growing_season_length_days"] + row.growing_season_start_day, row.growing_season_length_days ) if overlaps.sum() <= 0: continue - area = row["total_area"] + area = row.total_area region_total_area[region] = region_total_area.get(region, 0.0) + area fraction = overlaps / MONTH_LENGTHS region_month_demand[region] = ( @@ -386,8 +390,6 @@ def process_huang_irrigation( regions_gdf = gpd.read_file(regions_path)[["region", "geometry"]] regions_list = regions_gdf["region"].tolist() - monthly_records = [] - lon_unique = np.sort(np.unique(lon.astype(float))) lat_unique = np.sort(np.unique(lat.astype(float))) lon_diffs = np.diff(lon_unique) @@ -421,6 +423,7 @@ def process_huang_irrigation( # The dataset spans 1971-2010, with monthly data = 480 time steps year_start_idx = (reference_year - 1971) * 12 + monthly_rasters = [] for month in range(1, 13): time_idx = year_start_idx + month - 1 monthly_values = np.asarray(data.isel({time_dim: time_idx}).values, dtype=float) @@ -431,25 +434,36 @@ def process_huang_irrigation( # before aggregating; a coverage-weighted sum of a depth is meaningless. monthly_volume_m3 = monthly_data * MM_TO_M * cell_area_m2 - # Aggregate to regions (coverage-weighted sum of per-cell volumes) - result = aggregate_gridded_to_regions( - monthly_volume_m3, lon_values, lat_values, regions_gdf + monthly_rasters.append( + NumPyRasterSource( + monthly_volume_m3, + xmin=lon_min - lon_res / 2, + ymin=lat_min - lat_res / 2, + xmax=lon_max + lon_res / 2, + ymax=lat_max + lat_res / 2, + nodata=np.nan, + name=f"month_{month}", + srs_wkt='GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["WGS 84",6378137,298.257223563]],PRIMEM["Greenwich",0],UNIT["degree",0.0174532925199433]]', + ) ) - for _, row in result.iterrows(): - water_m3 = float(row["value"]) - monthly_records.append( - { - "region": row["region"], - "month": month, - "water_available_m3": water_m3, - } - ) + result = exact_extract( + monthly_rasters, + regions_gdf.reset_index(), + ["sum"], + include_cols=["region"], + output="pandas", + ) ds.close() # Build monthly dataframe - monthly_df = pd.DataFrame(monthly_records) + monthly_df = result.melt( + id_vars="region", var_name="month", value_name="water_available_m3" + ) + monthly_df["month"] = ( + monthly_df["month"].str.removesuffix("_sum").str[6:].astype(int) + ) monthly_df = monthly_df.sort_values(["region", "month"]).reset_index(drop=True) # Load crop growing seasons and compute growing season water