diff --git a/jdluc/__tests__/attribute_test.py b/jdluc/__tests__/attribute_test.py new file mode 100644 index 0000000..1e6fd31 --- /dev/null +++ b/jdluc/__tests__/attribute_test.py @@ -0,0 +1,74 @@ +import types + +import pandas +import pytest + +from jdluc import attribute, statistical +from jdluc.datasets import worldbank_jurisdictions + +ISO_3166_TO_TILE_IDS = {"AAA": {"10N_010W", "20N_010W"}, "BBB": {"20N_010W"}} + + +@pytest.fixture +def attributed_tile_ids(monkeypatch: pytest.MonkeyPatch) -> list[str]: + """Stub the boundaries and the per-tile leg; return the tiles the leg was called for.""" + calls: list[str] = [] + + def workflow_for_tile( + commodity_names: tuple[str, ...], iso_3166: str, tile_id: str + ) -> pandas.DataFrame: + calls.append(tile_id) + return pandas.DataFrame.from_records( + [("PROVINCIAL", f"{iso_3166:s}001", "MAIZE", iso_3166, 1.0)], + columns=[ + "admin_level", + "admin_id", + "commodity_name", + "jurisdiction_name", + "emissions_mt", + ], + ).set_index(["admin_level", "admin_id", "commodity_name"]) + + monkeypatch.setattr( + worldbank_jurisdictions, + "get_ten_degree_tile_ids_for_admin_id", + lambda admin_id, admin_level: ISO_3166_TO_TILE_IDS[admin_id], + ) + monkeypatch.setattr( + worldbank_jurisdictions, + "iter_jurisdiction_for_iso_3166_tile_id", + lambda admin_level, iso_3166, tile_id: iter( + [ + types.SimpleNamespace( + id=f"{iso_3166:s}001", level="PROVINCIAL", name=iso_3166 + ) + ] + ), + ) + monkeypatch.setattr(statistical, "workflow", workflow_for_tile) + return calls + + +@pytest.mark.parametrize( + ("residues", "expected_tile_ids"), + [ + pytest.param(None, ["10N_010W", "20N_010W", "20N_010W"], id="every-tile"), + pytest.param([0], ["10N_010W"], id="first-shard"), + # A tile shared by two countries is one shard, so both its partials land together + pytest.param([1], ["20N_010W", "20N_010W"], id="second-shard"), + ], +) +def test_shards_select_tiles_by_their_sorted_index( + attributed_tile_ids: list[str], + expected_tile_ids: list[str], + residues: list[int] | None, +) -> None: + attribute.workflow( + commodity_names=("MAIZE",), + concurrency=1, + iso_3166s=("AAA", "BBB"), + methodology=attribute.Methodology.STATISTICAL, + modulus=2, + residues=residues, + ) + assert sorted(attributed_tile_ids) == expected_tile_ids diff --git a/jdluc/__tests__/ingest_test.py b/jdluc/__tests__/ingest_test.py new file mode 100644 index 0000000..73d753d --- /dev/null +++ b/jdluc/__tests__/ingest_test.py @@ -0,0 +1,67 @@ +import sys +import types + +import pytest + +from jdluc import config, datasets, ingest, tiling +from jdluc.datasets import worldbank_jurisdictions + +TILE_IDS = {"10N_010W", "20N_010W", "30N_010W"} + + +@pytest.fixture +def ingested_tile_ids(monkeypatch: pytest.MonkeyPatch) -> list[str]: + """Stub the config, the boundaries and a tiled dataset; return the tiles it was asked to ingest.""" + calls: list[str] = [] + + def ingest_a_tile(overwrite: bool, root: str, tile_id: str) -> str: + calls.append(tile_id) + return f"{tile_id:s}.tif" + + monkeypatch.setattr( + config.Config, + "from_dot_env", + lambda: types.SimpleNamespace(ingest_root="unused"), + ) + monkeypatch.setattr( + worldbank_jurisdictions, + "get_ten_degree_tile_ids_for_iso_3166s", + lambda iso_3166s: TILE_IDS, + ) + monkeypatch.setitem( + datasets.NAME_TO_CLS, + datasets.DatasetName.GNW_HARRIS_AGB, + types.SimpleNamespace( + ingest_a_tile=ingest_a_tile, + partitioning=tiling.Partitioning.TEN_DEGREE_TILE, + ), + ) + return calls + + +@pytest.mark.parametrize( + ("shard_args", "expected_tile_ids"), + [ + pytest.param([], ["10N_010W", "20N_010W", "30N_010W"], id="every-tile"), + pytest.param( + ["--modulus=2", "--residues=0"], ["10N_010W", "30N_010W"], id="first-shard" + ), + pytest.param(["--modulus=3", "--residues=1"], ["20N_010W"], id="one-tile"), + pytest.param( + ["--modulus=3", "--residues=0", "--residues=2"], + ["10N_010W", "30N_010W"], + id="two-shards", + ), + ], +) +def test_shards_select_tiles_by_their_sorted_index( + ingested_tile_ids: list[str], + monkeypatch: pytest.MonkeyPatch, + expected_tile_ids: list[str], + shard_args: list[str], +) -> None: + monkeypatch.setattr( + sys, "argv", ["ingest.py", "GNW_HARRIS_AGB", "AAA", *shard_args] + ) + assert ingest.main() == 0 + assert sorted(ingested_tile_ids) == expected_tile_ids diff --git a/jdluc/attribute.py b/jdluc/attribute.py index 8108789..f5c5e71 100644 --- a/jdluc/attribute.py +++ b/jdluc/attribute.py @@ -8,7 +8,9 @@ Both clip to provincial (World Bank admin-1) polygons and sum commodity area, peatland commodity area, peatland-occupation emissions, and total emissions. Returns a pandas.DataFrame indexed by (admin_level, admin_id, commodity_name, methodology); the per-(country, tile) sub-workflows are -cached and `--concurrency` bounds how many tiles are in flight at once. +cached and `--concurrency` bounds how many tiles are in flight at once. `--modulus` and +`--residues` restrict a run to a shard of the countries' sorted tiles, as in harmonize, emit and +ingest, so one tile's partials can be computed apart from the others. The countries come from the positional ISO 3166 alpha-3 codes, or -- with `--backfill` -- from every country in the World Bank admin-0 layer. @@ -29,6 +31,7 @@ from jdluc import ( jurisdictional_direct, statistical, + utils, ) from jdluc.datasets import worldbank_jurisdictions @@ -58,6 +61,8 @@ def workflow( commodity_names: tuple[str, ...], iso_3166s: collections.abc.Iterable[str], methodology: Methodology, + modulus: int = 1, + residues: collections.abc.Collection[int] | None = None, ) -> pandas.DataFrame: workflow_for_tile = ( jurisdictional_direct.workflow @@ -75,6 +80,12 @@ def workflow( ) ): tile_id_to_iso_3166s[tile_id].append(iso_3166) + tile_id_to_iso_3166s = { + tile_id: tile_id_to_iso_3166s[tile_id] + for tile_id in utils.iter_sharded( + modulus=modulus, residues=residues, values=tile_id_to_iso_3166s + ) + } def workflow_for_tile_id(tile_id: str) -> list[pandas.DataFrame]: return [ @@ -160,6 +171,18 @@ def main() -> int: choices=sorted(e.name for e in Methodology), default=Methodology.STATISTICAL.name, ) + parser.add_argument( + "--modulus", + default=1, + type=int, + help="split the tiles into this many disjoint shards", + ) + parser.add_argument( + "--residues", + action="append", + type=int, + help="attribute only these shards, each in [0, --modulus); repeatable, all by default", + ) args = parser.parse_args() assert bool(args.iso_3166s) ^ bool(args.backfill), ( "pass either one-or-more iso_3166s or --backfill" @@ -175,6 +198,8 @@ def main() -> int: else args.iso_3166s ), methodology=methodology, + modulus=int(args.modulus), + residues=args.residues, ) if args.display_results: print(df.to_string()) diff --git a/jdluc/ingest.py b/jdluc/ingest.py index 8a35db1..106169a 100644 --- a/jdluc/ingest.py +++ b/jdluc/ingest.py @@ -9,7 +9,8 @@ The tile set comes from the positional ISO 3166 alpha-3 codes -- or, with `--backfill`, from `tiling.GLOBAL_NATURE_WATCH_TILE_IDS`. Tiles are ingested one at a time unless -`--concurrency` raises the bound. +`--concurrency` raises the bound. `--modulus` and `--residues` restrict a run to a shard of +the sorted tiles; like the tile set, the shard only bounds tiled datasets. Example invocations: uv run python -m jdluc.ingest IPCC_CLIMATE_ZONES USA @@ -26,7 +27,7 @@ import concurrent.futures import logging -from jdluc import config, datasets, tiling +from jdluc import config, datasets, tiling, utils from jdluc.datasets import base, worldbank_jurisdictions logger = logging.getLogger(__name__) @@ -88,7 +89,19 @@ def main() -> int: ) parser.add_argument("--backfill", action="store_true", help="cover all GNW tiles") parser.add_argument("--concurrency", default=1, type=int) + parser.add_argument( + "--modulus", + default=1, + type=int, + help="split the tiles into this many disjoint shards", + ) parser.add_argument("--overwrite", action="store_true") + parser.add_argument( + "--residues", + action="append", + type=int, + help="ingest only these shards, each in [0, --modulus); repeatable, all by default", + ) args = parser.parse_args() assert bool(args.iso_3166s) ^ bool(args.backfill), ( "pass either one-or-more iso_3166s or --backfill" @@ -100,11 +113,17 @@ def main() -> int: dataset=dataset, overwrite=args.overwrite, root=config.Config.from_dot_env().ingest_root, - tile_ids=( - tiling.GLOBAL_NATURE_WATCH_TILE_IDS - if args.backfill - else worldbank_jurisdictions.get_ten_degree_tile_ids_for_iso_3166s( - iso_3166s=args.iso_3166s + tile_ids=tuple( + utils.iter_sharded( + modulus=int(args.modulus), + residues=args.residues, + values=( + tiling.GLOBAL_NATURE_WATCH_TILE_IDS + if args.backfill + else worldbank_jurisdictions.get_ten_degree_tile_ids_for_iso_3166s( + iso_3166s=args.iso_3166s + ) + ), ) ), )