ingest, attribute: --modulus/--residues shard the tiles, as harmonize and emit do

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

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

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())

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
+                    )
+                ),
             )
         ),
     )