From 31a2a86e2bfb20562252d50c707a6c3ac9f952c2 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Thu, 15 Oct 2026 19:38:57 +0200 Subject: [PATCH 1/8] storage: a zarr cache entry counts only once its write has finished For zarr files, set up a .complete marker file to check when data is fully written. This replaces relying on zarr.json -- that doesn't work because zarr.json is the first to be written. This might happen from a failed or aborted calculation, or because multiple processes are trying to write the same file (e.g. running 2 countries which share a tile) A partially-written zarr might be not cause any warning signs later, since it might be techncially valid in spite of missing data. Now an empty .complete marker is written beside the zarr after its last chunk, and the entry counts as cached only once the marker exists. A zarr found without its marker is a write that was killed part-way or is still running: serialize raises FileExistsError and leaves it for someone to resolve by hand, usually by deleting the zarr, rather than writing over it and risking a race with a live writer. Entries written before this change have no marker, so each raises once until deleted. --- jdluc/__tests__/storage_test.py | 58 +++++++++++++++++++++++++++++++++ jdluc/storage.py | 22 ++++++++++--- 2 files changed, 75 insertions(+), 5 deletions(-) create mode 100644 jdluc/__tests__/storage_test.py diff --git a/jdluc/__tests__/storage_test.py b/jdluc/__tests__/storage_test.py new file mode 100644 index 0000000..2e5e48e --- /dev/null +++ b/jdluc/__tests__/storage_test.py @@ -0,0 +1,58 @@ +import pathlib + +import numpy +import pytest +import xarray + +from jdluc import storage +from jdluc.config import Config + + +@pytest.fixture(autouse=True) +def config_without_dot_env(monkeypatch: pytest.MonkeyPatch) -> None: + # Writing a zarr reads number_of_dask_workers from .env, which CI does not have + config = Config( + harvard_dataverse_guestbook_json="", + ingest_root="", + number_of_dask_workers=2, + scratch_root="", + usda_nass_api_key="", + ) + monkeypatch.setattr(Config, "from_dot_env", classmethod(lambda cls: config)) + + +def get_dset() -> xarray.Dataset: + return xarray.Dataset( + {"band": (("y", "x"), numpy.arange(16, dtype="float32").reshape(4, 4))} + ).chunk({"x": 2, "y": 2}) + + +def assert_reads_back(cacher: storage.ZarrCacher) -> None: + xarray.testing.assert_equal( + cacher.deserialize()["band"].drop_vars("spatial_ref"), get_dset()["band"] + ) + + +def test_zarr_cacher_marks_a_finished_write_complete(tmp_path: pathlib.Path) -> None: + cacher = storage.ZarrCacher(hash_key="abc", root=str(tmp_path)) + assert not storage.path_exists(uri=cacher.exists_uri) + + cacher.serialize(get_dset()) + + assert storage.path_exists(uri=cacher.exists_uri) + assert_reads_back(cacher) + + +def test_zarr_cacher_refuses_to_write_over_an_unfinished_zarr( + tmp_path: pathlib.Path, +) -> None: + cacher = storage.ZarrCacher(hash_key="abc", root=str(tmp_path)) + # A write killed part-way leaves zarr.json but no marker + storage.write_dask_dataset_to_zarr(dset=get_dset(), path_to_zarr=cacher.uri) + assert storage.path_exists(uri=f"{cacher.uri:s}/zarr.json") + assert not storage.path_exists(uri=cacher.exists_uri) + + with pytest.raises(FileExistsError): + cacher.serialize(get_dset()) + + assert not storage.path_exists(uri=cacher.exists_uri) diff --git a/jdluc/storage.py b/jdluc/storage.py index 1159d1e..684fda5 100644 --- a/jdluc/storage.py +++ b/jdluc/storage.py @@ -7,6 +7,7 @@ import threading import typing +import fsspec # type: ignore[import-untyped] import pandas import xarray @@ -16,15 +17,11 @@ def path_exists(uri: str) -> bool: - import fsspec # type: ignore[import-untyped] - fs, path = fsspec.url_to_fs(url=uri) return fs.exists(path) def put_file(local_path: str, uri: str) -> None: - import fsspec # type: ignore[import-untyped] - logger.info(f"Writing file {local_path=:s} to {uri=:s}") fs, path = fsspec.url_to_fs(url=uri) fs.makedirs(os.path.dirname(p=path), exist_ok=True) @@ -150,6 +147,13 @@ def serialize(self, df: pandas.DataFrame) -> None: @dataclasses.dataclass class ZarrCacher: + """A zarr at `uri`, complete once the `.complete` marker beside it exists. + + The marker is written after the last chunk, so a zarr without one is a write that was + killed part-way or is still running. Writing over it could race a live writer, so it is + an error left for someone to resolve by hand, usually by deleting the zarr. + """ + hash_key: str root: str @@ -159,13 +163,21 @@ def uri(self) -> str: @property def exists_uri(self) -> str: - return f"{self.uri:s}/zarr.json" + return f"{self.uri:s}.complete" def deserialize(self) -> xarray.Dataset: return open_zarr_to_dask_dataset(path_to_zarr=self.uri) def serialize(self, dset: xarray.Dataset) -> None: + if path_exists(uri=self.uri): + raise FileExistsError( + f"{self.uri=:s} exists without {self.exists_uri=:s}: a write was killed " + "part-way or is still running. Once no process is writing it, delete the " + "zarr to recompute it." + ) write_dask_dataset_to_zarr(dset=dset, path_to_zarr=self.uri) + fs, path = fsspec.url_to_fs(url=self.exists_uri) + fs.touch(path) class CacherDecoratorProtocol(typing.Protocol[R]): From 03a630e1b8f35d1a1eaa5d8c1c0fb21d7b16a97e Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Thu, 15 Oct 2026 20:50:29 +0200 Subject: [PATCH 2/8] tests: unit tests run without a .env An autouse fixture in jdluc/__tests__/conftest.py stubs Config.from_dot_env for every test not marked integration. Each field gets its type's empty value, with two dask workers, so a unit test that reaches for a storage root fails instead of touching a bucket. Integration tests still read the real .env. This replaces the fixture storage_test carried for itself, so the next unit test that writes a zarr needs no copy of it, and a new Config field needs no edit to the tests. --- jdluc/__tests__/conftest.py | 27 +++++++++++++++++++++++++++ jdluc/__tests__/storage_test.py | 14 -------------- 2 files changed, 27 insertions(+), 14 deletions(-) create mode 100644 jdluc/__tests__/conftest.py diff --git a/jdluc/__tests__/conftest.py b/jdluc/__tests__/conftest.py new file mode 100644 index 0000000..f5bcde1 --- /dev/null +++ b/jdluc/__tests__/conftest.py @@ -0,0 +1,27 @@ +import dataclasses +import typing + +import pytest + +from jdluc.config import Config + +# The fields a unit test needs a real value for +TEST_ENV = {"number_of_dask_workers": 2} + + +def get_unit_test_config() -> Config: + """TEST_ENV over every other field left empty, so a unit test that reaches for a root fails rather than touching real storage.""" + types = typing.get_type_hints(Config) + empty = {field.name: types[field.name]() for field in dataclasses.fields(Config)} + return Config(**{**empty, **TEST_ENV}) + + +@pytest.fixture(autouse=True) +def config_without_dot_env( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch +) -> None: + """Unit tests need no .env; integration tests still read the real one.""" + if request.node.get_closest_marker("integration"): + return + config = get_unit_test_config() + monkeypatch.setattr(Config, "from_dot_env", classmethod(lambda cls: config)) diff --git a/jdluc/__tests__/storage_test.py b/jdluc/__tests__/storage_test.py index 2e5e48e..34a3953 100644 --- a/jdluc/__tests__/storage_test.py +++ b/jdluc/__tests__/storage_test.py @@ -5,20 +5,6 @@ import xarray from jdluc import storage -from jdluc.config import Config - - -@pytest.fixture(autouse=True) -def config_without_dot_env(monkeypatch: pytest.MonkeyPatch) -> None: - # Writing a zarr reads number_of_dask_workers from .env, which CI does not have - config = Config( - harvard_dataverse_guestbook_json="", - ingest_root="", - number_of_dask_workers=2, - scratch_root="", - usda_nass_api_key="", - ) - monkeypatch.setattr(Config, "from_dot_env", classmethod(lambda cls: config)) def get_dset() -> xarray.Dataset: From a71f6c82e0db42be01196e8578e231c7c064ab91 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Mon, 19 Oct 2026 10:45:10 +0200 Subject: [PATCH 3/8] orchestration: run the pipeline on Argo, one tile per pod An optional `orchestration` package with one subcommand per workflow task (plan, ingest-tile, compute, trace), and an example Argo WorkflowTemplate that runs each tile's ingest -> compute as a DAG of its own, so a failed tile stops only itself. jdluc runs unmodified in the pods; every cache-key argument passes straight through, and an import-linter contract keeps jdluc from importing orchestration. trace writes the emissions-factor table to a new EXPORT_ROOT, through a small `export.write_table` utility. The package is built to be wrapped rather than edited: `build_parser` and `run` let a wrapper add subcommands, `plan.get_task_to_items` returns the lists for it to add to before `write_task_to_items`, and the template's tasks can be called by templateRef. See docs/orchestration.md. Also adds the Dockerfile the workflow's image is built from. --- .dockerignore | 23 ++ .env.example | 4 + .pre-commit-config.yaml | 2 +- Dockerfile | 25 ++ README.md | 1 + docs/orchestration.md | 95 ++++++ jdluc/__tests__/trace_test.py | 36 ++ jdluc/config.py | 1 + jdluc/export.py | 62 ++++ jdluc/trace.py | 36 +- orchestration/__init__.py | 22 ++ orchestration/__main__.py | 255 ++++++++++++++ orchestration/__tests__/__init__.py | 0 orchestration/__tests__/conftest.py | 2 + orchestration/__tests__/plan_test.py | 313 ++++++++++++++++++ orchestration/plan.py | 230 +++++++++++++ .../workflows/cornerstone.example.yaml | 238 +++++++++++++ pyproject.toml | 13 +- 18 files changed, 1355 insertions(+), 3 deletions(-) create mode 100644 .dockerignore create mode 100644 Dockerfile create mode 100644 docs/orchestration.md create mode 100644 jdluc/export.py create mode 100644 orchestration/__init__.py create mode 100644 orchestration/__main__.py create mode 100644 orchestration/__tests__/__init__.py create mode 100644 orchestration/__tests__/conftest.py create mode 100644 orchestration/__tests__/plan_test.py create mode 100644 orchestration/plan.py create mode 100644 orchestration/workflows/cornerstone.example.yaml diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..40c24c3 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,23 @@ +# Claude worktrees and other local tooling state. +.claude/ + +# A host venv, at any depth, must not shadow the one built inside the image. +**/.venv/ + +# Exact-patch pin (e.g. 3.14.5) a base image's uv may not fetch; requires-python (>=3.14,<3.15) +# is satisfied by any 3.14.x, so let uv use the base image's own. +.python-version + +# Secrets: supply .env at runtime, never bake it into the image. `*.env` matches any env file +# but not the `*.env.example` templates, which stay. +.env +**/*.env + +# Caches / build noise. +**/__pycache__/ +*.pyc +.mypy_cache/ +.pytest_cache/ +.ruff_cache/ +.cache/ +.DS_Store diff --git a/.env.example b/.env.example index dff0bf2..b53768a 100644 --- a/.env.example +++ b/.env.example @@ -1,3 +1,7 @@ +# EXPORT_ROOT is where terminal deliverables land (the emissions-factor table): not a cache +# the pipeline reads back. Any fsspec URI, like the roots below. +# EXPORT_ROOT = '/absolute/path/to/export' +EXPORT_ROOT = 'gs://cornerstone-emissions-us-central1' HARVARD_DATAVERSE_GUESTBOOK_JSON = '{"email":Replace with your email address,"institution":Replace with your institution}' # INGEST_ROOT and SCRATCH_ROOT take any fsspec URI. A local root must be an absolute path # that already exists; a gs:// root needs application-default credentials. diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ade003c..4c69651 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,7 +18,7 @@ repos: - id: mypy name: mypy language: system - entry: uv run mypy jdluc validation + entry: uv run mypy jdluc orchestration validation pass_filenames: false - repo: https://github.com/astral-sh/uv-pre-commit rev: 0.12.7 diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..0ccee1b --- /dev/null +++ b/Dockerfile @@ -0,0 +1,25 @@ +# syntax=docker/dockerfile:1.19 +# Build from a full clone, not a git worktree: a worktree's .git is a file pointing outside the +# build context, and the image needs the repository for git provenance (utils.get_git_version). + +FROM ghcr.io/astral-sh/uv:python3.14-bookworm-slim + +WORKDIR /app + +# The rasterio/GDAL wheels vendor GDAL but still link libexpat, which the slim image omits. +# safe.directory lets git read /app when the container runs as a user other than its owner. +RUN apt-get update \ + && apt-get install -y --no-install-recommends git libexpat1 \ + && rm -rf /var/lib/apt/lists/* \ + && git config --system --add safe.directory /app + +COPY pyproject.toml uv.lock ./ +RUN uv sync --frozen --no-install-project --no-dev + +COPY --exclude=.git . . +RUN uv sync --frozen --no-dev + +# .git in its own last layer, so a commit that changes no source rebuilds only this layer. +COPY .git .git + +ENV PATH="/app/.venv/bin:${PATH}" diff --git a/README.md b/README.md index f60f5fe..fbe4e2c 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,7 @@ Once you're ready to look under the hood: - `docs/data.md` — how to get the data products, grids, and full schemas for the harmonized inputs, per-pixel emissions, and emissions-factor table. - `docs/coverage.md` — the full list of countries and crops the pipeline produces emissions factors for, by ISO 3166-1 alpha-3 code. - `docs/validation.md` — how the pipeline is measured against external datasets: how the (country, crop) targets are chosen, which relationships are held as regression controls, and what such a comparison can and cannot conclude. The tooling lives in `validation/` and reports against the gaps in `docs/further_research.md`. +- `docs/orchestration.md` — an optional way to run the pipeline on Kubernetes with Argo Workflows, one tile per pod, for country runs and global backfills in hours rather than days. ## Running and contributing diff --git a/docs/orchestration.md b/docs/orchestration.md new file mode 100644 index 0000000..b4ddb43 --- /dev/null +++ b/docs/orchestration.md @@ -0,0 +1,95 @@ +# Running the pipeline on a cluster + +The pipeline runs on a single host by design (see [`architecture.md`](architecture.md#single-host-by-design)), and nothing here changes that. The `orchestration` package is an optional layer for running it on Kubernetes with [Argo Workflows](https://argoproj.github.io/workflows/), one 10° tile per pod, so that a country or a global backfill takes hours instead of days. `jdluc` runs unmodified inside the pods. Argo takes over only the driver's job: sequencing, fan-out, retries and visibility. + +## Design + +**One tile per pod, and no Dask cluster.** The pipeline processes each 10° tile independently, and one tile's working set fits one high-memory node. So the unit of distribution is a whole tile: a pod runs a tile's stages, each with its own in-process `LocalCluster`, exactly as a laptop run does. Nothing crosses between tiles except trace's final merge, over kilobyte-scale parquets. A shared Dask scheduler would add version skew between driver and workers and a scheduler that a heavy tile can destabilise, and a Dask failure would then look like a pipeline failure. Tasks are hour-long with tens of GB working sets, and they're already parallel inside the pod, so a distributed Dask cluster is the wrong granularity anyway. + +**The cache is the resume mechanism.** Every stage is wrapped in `@storage.cache_to_zarr` / `@storage.cache_to_parquet`, writing idempotently to `SCRATCH_ROOT`. Retrying anything skips finished work, so the orchestration can stay simple. A zarr counts as cached only once its `.complete` marker is written, so a pod killed mid-write leaves nothing a later pod would mistake for a result. + +**`orchestration` may not change a result.** It is one subcommand per workflow task, and every cache-key argument goes straight through to `jdluc`, with `jdluc`'s own defaults. An import-linter contract keeps `jdluc` from importing it. + +**Ingest is separate from compute.** They are different workloads: + +| | Ingest | Compute | +| --------------------- | --------------------------------- | ------------------------- | +| Bound by | network and disk | memory and CPU | +| Needs source-API keys | **yes** (USDA, Harvard Dataverse) | **no** (reads warm cache) | +| Cadence | once per country × source version | every code change | + +So only the ingest pods carry the keys, a flaky third-party source fails an ingest task rather than a compute run, and one warm `INGEST_ROOT` serves many compute runs. A laptop run ingests lazily, from inside harmonize; a workflow ingests up front instead, and `compute` runs harmonize with `skip_ingest=True`. That needs the list of datasets a run reads, which no `jdluc` constant declares, so `orchestration/plan.py` keeps one by hand. When `jdluc` starts reading a new dataset, add it there. + +## The DAG + +[`orchestration/workflows/cornerstone.example.yaml`](../orchestration/workflows/cornerstone.example.yaml) is an Argo `WorkflowTemplate`, with placeholders for the image, service account and storage class. Each run covers its `scope`: one country, or `world` for every country and GNW tile, as `--backfill` does locally. + +``` +plan -+-> ingest-world[dataset] -+-> tile[tile] --> trace + +--------------------------+ + +tile: ingest[dataset] -------------------+-> compute + ingest-allowed-failures[dataset] --+ +``` + +| Task | Pods | Runs | Keys | +| ------------------------- | -------------------- | --------------------------------------------------------------------------------------------------------- | ---- | +| `plan` | 1 | `python -m orchestration plan`: ingests the boundaries if absent, then writes the JSON lists to loop over | no | +| `ingest-world` | 1 per dataset | `jdluc/ingest.py NAME --backfill`, for each whole-world dataset | yes | +| `ingest` | 1 per dataset × tile | `python -m orchestration ingest-tile`: one dataset on one tile | yes | +| `ingest-allowed-failures` | 1 per dataset × tile | The same, for the datasets in `allow-failures` | yes | +| `compute` | 1 per tile | `python -m orchestration compute`: harmonize → emit → attribute, for each country the tile touches | no | +| `trace` | 1 | `python -m orchestration trace`: the cross-tile merge, written as the emissions-factor table | no | + +- **The tile list is derived, never typed in.** `plan` lists the tiles the scope's boundaries touch, the same derivation the `jdluc` CLIs use. +- **A failed tile fails only itself.** Each tile is a DAG of its own. A failed required ingest leaves that tile's compute unscheduled, and the other tiles carry on. +- **Ingest fails loudly, with exceptions.** A missing tile would otherwise reach harmonize's `ignore_missing_tiles` and become an all-nodata band without complaint. The datasets in `allow-failures` are the exception: by default only `DESCALS_OIL_PALM`, which is published only where oil palm was found, so most tiles have none. Compute runs whether or not those ingests fail. +- **`compute` checks its inputs.** With `--check-ingests` it raises before harmonize if a required ingest is missing from the ingest root, as a second line of defence behind the DAG. +- **`trace` needs every tile.** A table missing tiles would pass for a whole one, so with any tile failed `trace` doesn't run, and the run ends Failed. Once the failed tiles are fixed, `argo retry` reruns them, then `trace`. +- **Tiles wait for `ingest-world`.** That takes minutes once the whole-world datasets are cached. + +The flat alternative, one fan-out per stage across every tile, is simpler to read in the Argo UI, but one failed (dataset, tile) ingest then stops compute on every tile. `plan` writes the flat lists too (`ingest-tiles`, `tiles`, …), so a flat workflow needs no further code. + +## Setup + +**The image.** The repo-root `Dockerfile` builds it. Build from a full clone, not a git worktree, since the image carries `.git` for provenance (`utils.get_git_version`). Tag the image with the commit's SHA, and set the template's `image` to it. + +**Identity.** The pods need read/write access to `INGEST_ROOT`, `SCRATCH_ROOT` and `EXPORT_ROOT`. On a cloud provider, bind the template's Kubernetes service account to a cloud identity (on GKE, Workload Identity), so that no storage keys go in the image or the Secrets. + +**The two Secrets.** `Config.from_dot_env()` reads a `.env` *file*, so each pod mounts a Secret at `/app/.env`. Start from `.env.example`, keep the real file outside the repo, and fill in every field. Then make a copy with the two API-key fields present but blank, so that only the ingest pods get the keys: + +```bash +kubectl -n argo create secret generic cornerstone-env --from-file=.env=cloud.env +kubectl -n argo create secret generic cornerstone-env-nokeys --from-file=.env=cloud-nokeys.env +``` + +**Disk is what limits speed.** Ingest writes tens of GB of uncompressed GeoTIFF locally before converting it to a COG. On a node's shared boot disk, those writes can throttle the whole node to a crawl. So each pod gets its own SSD-backed generic ephemeral volume at `/localtmp`, with `TMPDIR` and `CPL_TMPDIR` pointed at it, created and deleted with the pod (on GKE, the built-in `premium-rwo` class). Don't shrink these volumes: SSD throughput scales with size. Each container also has a small `ephemeral-storage` request and limit, which covers only its own overlay, so a stray writer evicts its own pod and not the node. + +**Pod sizing.** `compute` requests 56 GiB, with peaks near 40 GiB observed: one tile per 64 GiB node, with room for system pods. The other tasks fit in 24 GiB. Spot or preemptible nodes work well, since a reclaimed node only fails its pod and `argo retry` resumes from the cache. Add their `nodeSelector` and `tolerations` under `templateDefaults`. + +## Running + +```bash +argo template create orchestration/workflows/cornerstone.example.yaml +argo submit --from workflowtemplate/cornerstone -p scope=HND --generate-name cornerstone-hnd- +``` + +Parameters: `scope` (one ISO 3166 alpha-3 code, or `world`), `methodology` (`STATISTICAL`, the default, or `JURISDICTIONAL_DIRECT`, which is US-only) and `allow-failures` (comma-separated tiled datasets, by default `DESCALS_OIL_PALM`). The template's `parallelism` and `activeDeadlineSeconds` are sized for one country; raise them for `world`. + +`trace` writes `{EXPORT_ROOT}/emissions-factors/{scope}-{methodology}.parquet`. Intermediate results stay under `INGEST_ROOT` and `SCRATCH_ROOT`. + +**Reruns are cheap.** Rerunning a country, or `argo retry` after a failure, skips everything already cached. A pod killed mid-write leaves a zarr without its `.complete` marker, and the next attempt to write it raises rather than reading back missing chunks. Once no pod is writing that zarr, delete it by hand to let it be recomputed. + +**Gotchas:** + +- **`compute` and `trace` must use the same methodology.** It is a cache-key argument, so a mismatch doesn't fail: it silently recomputes every tile inside the single `trace` pod. The template passes one `methodology` to both, so this bites only when running `python -m orchestration` by hand. +- **Storage, not compute, is the cost that recurs.** Per-tile intermediates in `SCRATCH_ROOT` are hundreds of GiB. Before a large run, decide their retention, for example with bucket lifecycle rules. +- **A warm cache can fake a fast run.** When measuring timings, make sure the run does real work. + +## Extending it + +To add tasks of your own, such as a per-tile export after `compute`, wrap the package rather than editing it: + +- **Subcommands.** `orchestration.__main__.build_parser()` returns the parser and its subparsers. Add a subcommand with `add_subparser` and the `add_scope` / `add_methodology` / `add_tile_id` helpers, then call `run(parser)`. A subcommand's flags are exactly its function's parameters. +- **Plan lists.** `orchestration.plan.get_task_to_items()` returns the lists without writing them. Add lists of your own, or keys to each `tile-units` item, then write them with `write_task_to_items()`. Call `check_allow_failures()` and `ingest_boundaries()` first, as `write_plan` does. +- **The workflow.** Call the template's tasks with `templateRef: {name: cornerstone, template: compute}` (and likewise `ingest-tile`, `trace`, …) from DAGs of your own, instead of copying the file. diff --git a/jdluc/__tests__/trace_test.py b/jdluc/__tests__/trace_test.py index aa02b59..56050e5 100644 --- a/jdluc/__tests__/trace_test.py +++ b/jdluc/__tests__/trace_test.py @@ -1,6 +1,7 @@ import numpy import pandas +from jdluc import trace from jdluc.attribute import Methodology from jdluc.datasets.worldbank_jurisdictions import AdminLevel from jdluc.trace import ( @@ -314,3 +315,38 @@ def test_methodologies_share_a_schema() -> None: assert set(jurisdictional.reset_index().columns) == set( statistical.reset_index().columns ) + + +def test_export_workflow_writes_the_table_by_aoi_and_methodology( + tmp_path, monkeypatch +) -> None: + table = attach_ratios( + df=derive_statistical_production_kg(emissions=STATISTICAL_EMISSIONS) + ) + monkeypatch.setattr(trace, "workflow", lambda **kwargs: table) + uri = trace.export_workflow( + commodity_names=("MAIZE",), + iso_3166s=("GTM", "HND"), + methodology=Methodology.STATISTICAL, + output_root=str(tmp_path), + ) + assert uri == str(tmp_path / "emissions-factors" / "GTM-HND-STATISTICAL.parquet") + written = pandas.read_parquet(uri) + pandas.testing.assert_frame_equal(written, table) + assert written.attrs["cornerstone-product-name"] == "emissions-factors" + assert written.attrs["cornerstone-source-name"] == "jdluc-trace" + + +def test_export_workflow_takes_a_name_for_its_countries(tmp_path, monkeypatch) -> None: + table = attach_ratios( + df=derive_statistical_production_kg(emissions=STATISTICAL_EMISSIONS) + ) + monkeypatch.setattr(trace, "workflow", lambda **kwargs: table) + uri = trace.export_workflow( + commodity_names=("MAIZE",), + iso_3166s=("GTM", "HND"), + methodology=Methodology.STATISTICAL, + name="world", + output_root=str(tmp_path), + ) + assert uri == str(tmp_path / "emissions-factors" / "world-STATISTICAL.parquet") diff --git a/jdluc/config.py b/jdluc/config.py index 5eab326..e455673 100644 --- a/jdluc/config.py +++ b/jdluc/config.py @@ -11,6 +11,7 @@ @dataclasses.dataclass class Config: + export_root: str harvard_dataverse_guestbook_json: str ingest_root: str number_of_dask_workers: int diff --git a/jdluc/export.py b/jdluc/export.py new file mode 100644 index 0000000..5860d5e --- /dev/null +++ b/jdluc/export.py @@ -0,0 +1,62 @@ +"""Write a pipeline step's output as a deliverable under EXPORT_ROOT. + +A utility, like `storage`: it knows nothing about the steps. A tabular deliverable (see +`trace.export_workflow`) is one parquet per set of countries, written by `write_table` to +`{EXPORT_ROOT}/{deliverable}/{name}.parquet`, with provenance in its `attrs`. + +Unlike SCRATCH_ROOT, EXPORT_ROOT is not a cache that the pipeline reads back: what lands +there is a terminal output. +""" + +import logging +import os +import tempfile + +import pandas + +from jdluc import config, storage, utils + +logger = logging.getLogger(__name__) + + +def get_metadata(product_name: str, source_name: str) -> dict[str, str]: + """Provenance, as `datasets.base` stamps on ingested COGs.""" + return { + "cornerstone-processing-time": utils.get_utc_timestamp(), + "cornerstone-processing-version": utils.get_git_version(), + "cornerstone-product-name": product_name, + "cornerstone-remote-url": utils.get_git_remote_url(), + "cornerstone-source-name": source_name, + } + + +def get_table_output_uri( + deliverable_name: str, name: str, output_root: str | None = None +) -> str: + """`{output_root}/{deliverable}/{name}.parquet`; `output_root` defaults to EXPORT_ROOT.""" + if output_root is None: + output_root = config.Config.from_dot_env().export_root + return storage.join_uri( + root=output_root, prefix=f"{deliverable_name:s}/{name:s}.parquet" + ) + + +def write_table( + df: pandas.DataFrame, + deliverable_name: str, + name: str, + source_name: str, + output_root: str | None = None, +) -> str: + """Write a tabular deliverable as one parquet, provenance in its `attrs`; return the URI.""" + uri = get_table_output_uri( + deliverable_name=deliverable_name, name=name, output_root=output_root + ) + df = df.copy() + df.attrs |= get_metadata(product_name=deliverable_name, source_name=source_name) + with tempfile.TemporaryDirectory() as tmpdir: + local_path = os.path.join(tmpdir, os.path.basename(uri)) + df.to_parquet(path=local_path) + storage.put_file(local_path=local_path, uri=uri) + logger.info(f"Wrote {deliverable_name=:s} ({len(df):d} rows) to {uri=:s}") + return uri diff --git a/jdluc/trace.py b/jdluc/trace.py index bb414b4..289a70b 100644 --- a/jdluc/trace.py +++ b/jdluc/trace.py @@ -9,6 +9,7 @@ It then derives, identically for both, each commodity's yield (kg per hectare) and emissions factor (kgCO2e per kg), and rolls the provincials up to national totals. Returns a cached pandas.DataFrame indexed by (admin_level, admin_id, commodity_name, methodology). +`export_workflow` writes that table as the `emissions-factors` deliverable under EXPORT_ROOT. The countries come from the positional ISO 3166 alpha-3 codes, or -- with `--backfill` -- from every country in the World Bank admin-0 layer. @@ -24,11 +25,14 @@ import pandas -from jdluc import attribute, emit, storage +from jdluc import attribute, emit, export, storage from jdluc.datasets import usda_nass_quickstats, worldbank_jurisdictions logger = logging.getLogger(__name__) +# Named after the published `emissions-factors.parquet` +EMISSIONS_FACTORS = "emissions-factors" + NASS_YIELD_YEARS = (2017, 2018, 2019, 2020) KG_PER_TONNE = 1000 @@ -170,6 +174,36 @@ def workflow( ) +def export_workflow( + commodity_names: tuple[str, ...], + iso_3166s: tuple[str, ...], + methodology: attribute.Methodology, + name: str | None = None, + output_root: str | None = None, +) -> str: + """Write the countries' table from the cached `workflow` to + `{EXPORT_ROOT}/emissions-factors/{name}-{METHODOLOGY}.parquet`; return the URI. + + `name` defaults to the countries' codes joined, `iso[-iso...]`; pass one for a set too large + to spell out. The methodology is in the name so the two legs over one set of countries don't + overwrite each other. + """ + df = workflow( + # NB: not a cache-key argument; with the attribute partials warm, this only reads them + concurrency=1, + commodity_names=commodity_names, + iso_3166s=iso_3166s, + methodology=methodology, + ) + return export.write_table( + df=df, + deliverable_name=EMISSIONS_FACTORS, + name=f"{name or '-'.join(iso_3166s):s}-{methodology.name:s}", + output_root=output_root, + source_name="jdluc-trace", + ) + + def main() -> int: logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", diff --git a/orchestration/__init__.py b/orchestration/__init__.py new file mode 100644 index 0000000..0d8371b --- /dev/null +++ b/orchestration/__init__.py @@ -0,0 +1,22 @@ +"""The glue between a workflow engine (Argo) and jdluc: one subcommand per workflow task. + +What may live here: listing the units of work a run fans out over, and running exactly one +unit per call. What must not: anything that changes a result. Every argument that reaches a +cache key is passed straight through to jdluc, with jdluc's own defaults. + + python -m orchestration plan HND --methodology=STATISTICAL --output-dir=/plan + python -m orchestration ingest-tile --dataset=GNW_TREE_COVER_LOSS --tile-id=10N_090W + python -m orchestration compute HND --methodology=STATISTICAL --tile-id=10N_090W + python -m orchestration trace HND --methodology=STATISTICAL + +Each subcommand that takes countries takes `world` in their place, for every country and GNW +tile, as jdluc's `--backfill` covers. + +Whole-world ingest needs no wrapper: the workflow runs jdluc's own CLI for it. A tile's ingest +has one, `ingest-tile`, so that its pod names the tile rather than the run's countries. + +To add tasks of your own, wrap rather than edit: `__main__.build_parser` returns the parser and +its subparsers to add subcommands to, and `plan.get_task_to_items` the lists to add to before +`plan.write_task_to_items`. The example in `workflows/` shows how a workflow calls each +subcommand. See docs/orchestration.md. +""" diff --git a/orchestration/__main__.py b/orchestration/__main__.py new file mode 100644 index 0000000..6daf089 --- /dev/null +++ b/orchestration/__main__.py @@ -0,0 +1,255 @@ +import argparse +import collections.abc +import logging + +import orchestration +from jdluc import ( + attribute, + config, + emit, + harmonize, + ingest, + jurisdictional_direct, + statistical, + storage, + tiling, + trace, +) +from jdluc.datasets import NAME_TO_CLS, DatasetName, worldbank_jurisdictions +from orchestration import plan + +logger = logging.getLogger(__name__) + + +def ingest_tile(dataset: DatasetName, tile_id: str) -> None: + """One ten-degree dataset on one tile, into the ingest root. Raises if the tile fails.""" + (path_or_exc,) = ingest.workflow( + concurrency=1, + dataset=NAME_TO_CLS[dataset], + overwrite=False, + root=config.Config.from_dot_env().ingest_root, + tile_ids=(tile_id,), + ).values() + if isinstance(path_or_exc, Exception): + raise path_or_exc + logger.info(f"ingested {dataset.name:s} {tile_id=:s} to {path_or_exc:s}") + + +def get_missing_ingests( + allow_failures: tuple[str, ...], methodology: attribute.Methodology, tile_id: str +) -> list[str]: + """The URIs of the tile's inputs absent from the ingest root: its ten-degree datasets + other than ``allow_failures``, and every whole-world dataset.""" + root = config.Config.from_dot_env().ingest_root + name_and_tile_ids = [ + *( + (name, tile_id) + for name in plan.get_dataset_names( + methodology=methodology, + partitioning=tiling.Partitioning.TEN_DEGREE_TILE, + ) + if name.name not in allow_failures + ), + *( + (name, tiling.WHOLE_WORLD_TILE_ID) + for name in plan.get_dataset_names( + methodology=methodology, partitioning=tiling.Partitioning.WHOLE_WORLD + ) + ), + ] + uris = [ + storage.join_uri( + prefix=NAME_TO_CLS[name].get_prefix(tile_id=tile_id), root=root + ) + for name, tile_id in name_and_tile_ids + ] + return [uri for uri in uris if not storage.path_exists(uri=uri)] + + +def compute( + allow_failures: tuple[str, ...], + check_ingests: bool, + methodology: attribute.Methodology, + scope: tuple[str, ...], + tile_id: str, +) -> None: + """harmonize -> emit -> attribute for one tile, in one process. + + harmonize goes first with ``skip_ingest=True``: the later stages call it again with + ``skip_ingest=False``, and the cache key ignores the flag, so warming it here is what + keeps source credentials off compute pods. + + With ``check_ingests``, raises first if any input other than ``allow_failures`` is + missing. harmonize would otherwise read it as an all-nodata band, and cache that. + """ + if check_ingests: + missing = get_missing_ingests( + allow_failures=allow_failures, methodology=methodology, tile_id=tile_id + ) + if missing: + raise FileNotFoundError(f"{tile_id=:s} is missing ingests: {missing}") + for stack in harmonize.Stack: + logger.info(f"harmonize {stack.name:s} {tile_id=:s}") + harmonize.workflow( + dataset_names=stack.value, + ignore_missing_tiles=True, + skip_ingest=True, + tile_id=tile_id, + tile_resolution=tiling.TileResolution.GLAD, + ) + logger.info(f"emit {tile_id=:s}") + emit.workflow(tile_id=tile_id) + + workflow_for_tile = { + attribute.Methodology.JURISDICTIONAL_DIRECT: jurisdictional_direct.workflow, + attribute.Methodology.STATISTICAL: statistical.workflow, + }[methodology] + # A tile can straddle several of the run's countries; trace sums exactly these + # per-(tile, country) partials, so compute has to produce all of them + for iso_3166 in get_iso_3166s_touching_tile( + iso_3166s=plan.get_iso_3166s(scope=scope), tile_id=tile_id + ): + logger.info(f"attribute {methodology.name:s} {tile_id=:s} {iso_3166=:s}") + workflow_for_tile( + commodity_names=attribute.get_commodity_names(methodology=methodology), + iso_3166=iso_3166, + tile_id=tile_id, + ) + + +def get_iso_3166s_touching_tile( + iso_3166s: collections.abc.Iterable[str], tile_id: str +) -> list[str]: + return sorted( + iso_3166 + for iso_3166 in iso_3166s + if tile_id + in worldbank_jurisdictions.get_ten_degree_tile_ids_for_iso_3166s( + iso_3166s=(iso_3166,) + ) + ) + + +def trace_countries(methodology: attribute.Methodology, scope: tuple[str, ...]) -> None: + """The cross-tile merge, written as the emissions-factor table. Cheap only if compute ran + with the same methodology.""" + uri = trace.export_workflow( + commodity_names=attribute.get_commodity_names(methodology=methodology), + iso_3166s=plan.get_iso_3166s(scope=scope), + methodology=methodology, + name=plan.get_name(scope=scope), + ) + logger.info(f"exported emissions factors to {uri:s}") + + +def add_subparser( + subparsers: argparse._SubParsersAction[argparse.ArgumentParser], + name: str, + func: collections.abc.Callable[..., None], +) -> argparse.ArgumentParser: + """A subcommand that calls ``func`` with its flags, which must be exactly ``func``'s + parameters.""" + subparser = subparsers.add_parser(name, help=(func.__doc__ or "").split("\n")[0]) + subparser.set_defaults(func=func) + return subparser + + +def add_scope(subparser: argparse.ArgumentParser) -> None: + subparser.add_argument( + "scope", + help=f"ISO 3166 alpha-3 codes, or {plan.WORLD:s} for every country and GNW tile", + nargs=argparse.ONE_OR_MORE, + type=plan.scope_str, + ) + + +def add_methodology(subparser: argparse.ArgumentParser) -> None: + subparser.add_argument( + "--methodology", + choices=list(attribute.Methodology), + metavar="{" + ",".join(m.name for m in attribute.Methodology) + "}", + required=True, + type=attribute.Methodology.__getitem__, + ) + + +def add_tile_id(subparser: argparse.ArgumentParser) -> None: + subparser.add_argument("--tile-id", required=True) + + +def add_allow_failures(subparser: argparse.ArgumentParser) -> None: + subparser.add_argument( + "--allow-failures", + default=(), + help="comma-separated ten-degree datasets whose failed tiles must not stop the run", + type=lambda s: tuple(name for name in s.split(",") if name), + ) + + +def build_parser() -> tuple[ + argparse.ArgumentParser, argparse._SubParsersAction[argparse.ArgumentParser] +]: + """The CLI, and its subparsers for a wrapper to add subcommands to before `run`.""" + parser = argparse.ArgumentParser( + description=orchestration.__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + subparsers = parser.add_subparsers(required=True) + + subparser = add_subparser(subparsers, "plan", plan.write_plan) + add_scope(subparser) + add_methodology(subparser) + add_allow_failures(subparser) + subparser.add_argument("--output-dir", required=True) + + subparser = add_subparser(subparsers, "ingest-tile", ingest_tile) + subparser.add_argument( + "--dataset", + choices=list(DatasetName), + metavar="DATASET", + required=True, + # NB: not `DatasetName.__getitem__`, which a StrEnum resolves to `str.__getitem__` + type=lambda s: DatasetName[s], + ) + add_tile_id(subparser) + + subparser = add_subparser(subparsers, "compute", compute) + add_scope(subparser) + add_methodology(subparser) + add_tile_id(subparser) + add_allow_failures(subparser) + subparser.add_argument( + "--check-ingests", + action="store_true", + help="fail if an input other than --allow-failures is missing from the ingest root", + ) + + subparser = add_subparser(subparsers, "trace", trace_countries) + add_scope(subparser) + add_methodology(subparser) + + return parser, subparsers + + +def run( + parser: argparse.ArgumentParser, args: collections.abc.Sequence[str] | None = None +) -> int: + """Parse ``args`` (by default the command line) and call the chosen subcommand.""" + logging.basicConfig( + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + level=logging.INFO, + ) + # Each subcommand's flags are exactly its function's parameters + kwargs = vars(parser.parse_args(args)) + func = kwargs.pop("func") + func(**{k: tuple(v) if isinstance(v, list) else v for k, v in kwargs.items()}) + return 0 + + +def main() -> int: + parser, _ = build_parser() + return run(parser) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/orchestration/__tests__/__init__.py b/orchestration/__tests__/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/orchestration/__tests__/conftest.py b/orchestration/__tests__/conftest.py new file mode 100644 index 0000000..139f566 --- /dev/null +++ b/orchestration/__tests__/conftest.py @@ -0,0 +1,2 @@ +# Unit tests here, as in jdluc, need no .env and must never touch real storage +from jdluc.__tests__.conftest import config_without_dot_env # noqa: F401 diff --git a/orchestration/__tests__/plan_test.py b/orchestration/__tests__/plan_test.py new file mode 100644 index 0000000..3773942 --- /dev/null +++ b/orchestration/__tests__/plan_test.py @@ -0,0 +1,313 @@ +import json +import pathlib +import re +import sys +import types + +import pytest + +from jdluc import attribute, datasets, tiling, trace +from jdluc.datasets import worldbank_jurisdictions +from orchestration import __main__ as cli +from orchestration import plan + +TILE_IDS = {"10N_010W", "20N_010W", "30N_010W"} +# A stand-in for every GNW tile: more than the countries' tiles, as the real list is +WORLD_TILE_IDS = ("00N_010W", *sorted(TILE_IDS), "40N_010W") + + +@pytest.fixture +def ingested(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, str]]: + """Stub the boundaries and every dataset; return the (dataset, tile) pairs ingested.""" + calls: list[tuple[str, str]] = [] + monkeypatch.setattr( + worldbank_jurisdictions, + "get_ten_degree_tile_ids_for_iso_3166s", + lambda iso_3166s: TILE_IDS, + ) + monkeypatch.setattr( + worldbank_jurisdictions, "get_all_iso_3166s", lambda: {"AAA", "BBB"} + ) + monkeypatch.setattr(tiling, "GLOBAL_NATURE_WATCH_TILE_IDS", WORLD_TILE_IDS) + for name, cls in list(datasets.NAME_TO_CLS.items()): + + def ingest_a_tile( + overwrite: bool, root: str, tile_id: str, name: str = name.name + ) -> str: + calls.append((name, tile_id)) + return f"{tile_id:s}.tif" + + monkeypatch.setitem( + datasets.NAME_TO_CLS, + name, + types.SimpleNamespace( + get_prefix=lambda tile_id, name=name.name: f"{name:s}/{tile_id:s}", + ingest_a_tile=ingest_a_tile, + partitioning=cls.partitioning, + ), + ) + return calls + + +def write_plan( + tmp_path: pathlib.Path, + allow_failures: tuple[str, ...], + scope: tuple[str, ...] = ("AAA",), +) -> dict: + plan.write_plan( + allow_failures=allow_failures, + methodology=attribute.Methodology.STATISTICAL, + output_dir=str(tmp_path), + scope=scope, + ) + return {p.stem: json.loads(p.read_text()) for p in tmp_path.glob("*.json")} + + +def test_plan_ingests_the_boundaries_it_reads( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + task_to_items = write_plan(tmp_path=tmp_path, allow_failures=()) + assert sorted(ingested) == [ + ("WORLD_BANK_ADMIN_0", tiling.WHOLE_WORLD_TILE_ID), + ("WORLD_BANK_ADMIN_1", tiling.WHOLE_WORLD_TILE_ID), + ] + assert "WORLD_BANK_ADMIN_0" not in task_to_items["ingest-world"] + + +def test_allowed_failures_get_their_own_list( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + task_to_items = write_plan(tmp_path=tmp_path, allow_failures=("DESCALS_OIL_PALM",)) + allowed = task_to_items["ingest-tiles-allowed-failures"] + assert {u["dataset"] for u in allowed} == {"DESCALS_OIL_PALM"} + assert {u["tile_id"] for u in allowed} == TILE_IDS + assert "DESCALS_OIL_PALM" not in { + u["dataset"] for u in task_to_items["ingest-tiles"] + } + + +def test_an_unknown_allowed_failure_fails( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + with pytest.raises(AssertionError, match="DESCALS_OIL_PLAM"): + write_plan(tmp_path=tmp_path, allow_failures=("DESCALS_OIL_PLAM",)) + + +@pytest.mark.parametrize("scope", [("AAA",), (plan.WORLD,)]) +def test_each_planned_unit_ingests_exactly_its_tile( + ingested: list[tuple[str, str]], + monkeypatch: pytest.MonkeyPatch, + scope: tuple[str, ...], + tmp_path: pathlib.Path, +) -> None: + task_to_items = write_plan(tmp_path=tmp_path, allow_failures=(), scope=scope) + for unit in task_to_items["ingest-tiles"]: + ingested.clear() + monkeypatch.setattr( + sys, + "argv", + [ + "orchestration", + "ingest-tile", + f"--dataset={unit['dataset']:s}", + f"--tile-id={unit['tile_id']:s}", + ], + ) + assert cli.main() == 0 + assert ingested == [(unit["dataset"], unit["tile_id"])] + + +def test_a_failed_tile_fails_its_ingest(monkeypatch: pytest.MonkeyPatch) -> None: + def ingest_a_tile(overwrite: bool, root: str, tile_id: str) -> str: + raise RuntimeError(f"no {tile_id:s}") + + monkeypatch.setitem( + datasets.NAME_TO_CLS, + datasets.DatasetName.DESCALS_OIL_PALM, + types.SimpleNamespace( + ingest_a_tile=ingest_a_tile, + partitioning=tiling.Partitioning.TEN_DEGREE_TILE, + ), + ) + with pytest.raises(RuntimeError, match="no 10N_010W"): + cli.ingest_tile( + dataset=datasets.DatasetName.DESCALS_OIL_PALM, tile_id="10N_010W" + ) + + +def test_world_plans_every_gnw_tile( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + task_to_items = write_plan( + tmp_path=tmp_path, allow_failures=(), scope=(plan.WORLD,) + ) + assert [u["tile_id"] for u in task_to_items["tiles"]] == list(WORLD_TILE_IDS) + assert {u["tile_id"] for u in task_to_items["ingest-tiles"]} == set(WORLD_TILE_IDS) + + +@pytest.mark.parametrize("scope", [(), (plan.WORLD, "AAA")]) +def test_a_scope_is_countries_or_world_alone(scope: tuple[str, ...]) -> None: + with pytest.raises(AssertionError): + plan.get_tile_ids(scope=scope) + + +def test_world_names_its_run_wide_outputs_world( + ingested: list[tuple[str, str]], monkeypatch: pytest.MonkeyPatch +) -> None: + calls: list[dict] = [] + + def record(**kwargs: object) -> str: + calls.append(kwargs) + return "uri" + + monkeypatch.setattr(trace, "export_workflow", record) + monkeypatch.setattr( + sys, "argv", ["orchestration", "trace", "world", "--methodology=STATISTICAL"] + ) + assert cli.main() == 0 + + (trace_kwargs,) = calls + assert trace_kwargs["name"] == "world" + assert trace_kwargs["iso_3166s"] == ("AAA", "BBB") + + +def test_a_country_run_names_its_outputs_by_its_countries() -> None: + assert plan.get_name(scope=("HND", "GTM")) == "GTM-HND" + + +def test_get_task_to_items_ingests_nothing_and_matches_the_plan( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + task_to_items = plan.get_task_to_items( + allow_failures=(), + methodology=attribute.Methodology.STATISTICAL, + scope=("AAA",), + ) + assert ingested == [] + assert task_to_items == write_plan(tmp_path=tmp_path, allow_failures=()) + + +def test_a_wrapper_adds_a_list_and_a_subcommand( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + """The extension points a wrapper builds on: a list added to the plan, and a subcommand + added to the CLI, with no edit to this package.""" + task_to_items = plan.get_task_to_items( + allow_failures=(), + methodology=attribute.Methodology.STATISTICAL, + scope=("AAA",), + ) + task_to_items["extra-tiles"] = [u["tile_id"] for u in task_to_items["tiles"]] + plan.write_task_to_items(output_dir=str(tmp_path), task_to_items=task_to_items) + assert json.loads((tmp_path / "extra-tiles.json").read_text()) == sorted(TILE_IDS) + + calls: list[dict] = [] + + def extra(scope: tuple[str, ...], tile_id: str) -> None: + calls.append({"scope": scope, "tile_id": tile_id}) + + parser, subparsers = cli.build_parser() + subparser = cli.add_subparser(subparsers, "extra", extra) + cli.add_scope(subparser) + cli.add_tile_id(subparser) + assert cli.run(parser, ["extra", "AAA", "--tile-id=10N_010W"]) == 0 + assert calls == [{"scope": ("AAA",), "tile_id": "10N_010W"}] + + +@pytest.mark.parametrize("scope", [("AAA",), (plan.WORLD,)]) +def test_tile_units_hold_the_flat_lists_by_tile( + ingested: list[tuple[str, str]], scope: tuple[str, ...], tmp_path: pathlib.Path +) -> None: + task_to_items = write_plan( + tmp_path=tmp_path, + allow_failures=("DESCALS_OIL_PALM",), + scope=scope, + ) + tile_units = task_to_items["tile-units"] + assert [u["tile_id"] for u in tile_units] == [ + u["tile_id"] for u in task_to_items["tiles"] + ] + for key, task in ( + ("ingests", "ingest-tiles"), + ("ingests_allowed_failures", "ingest-tiles-allowed-failures"), + ): + assert sorted( + (dataset, u["tile_id"]) for u in tile_units for dataset in u[key] + ) == sorted((u["dataset"], u["tile_id"]) for u in task_to_items[task]) + + +def test_each_tile_lists_each_tiled_dataset_once( + ingested: list[tuple[str, str]], tmp_path: pathlib.Path +) -> None: + task_to_items = write_plan(tmp_path=tmp_path, allow_failures=("DESCALS_OIL_PALM",)) + tiled = { + name.name + for name in plan.get_dataset_names( + methodology=attribute.Methodology.STATISTICAL, + partitioning=tiling.Partitioning.TEN_DEGREE_TILE, + ) + } + for u in task_to_items["tile-units"]: + assert u["ingests_allowed_failures"] == ["DESCALS_OIL_PALM"] + assert sorted([*u["ingests"], *u["ingests_allowed_failures"]]) == sorted(tiled) + + +@pytest.fixture +def present_uris( + ingested: list[tuple[str, str]], monkeypatch: pytest.MonkeyPatch +) -> set[str]: + """Every input of tile 10N_010W, as the ingest root would hold them; remove some to + make them missing.""" + root = cli.config.Config.from_dot_env().ingest_root + uris = { + cli.storage.join_uri(prefix=f"{name.name:s}/{tile_id:s}", root=root) + for partitioning, tile_id in ( + (tiling.Partitioning.TEN_DEGREE_TILE, "10N_010W"), + (tiling.Partitioning.WHOLE_WORLD, tiling.WHOLE_WORLD_TILE_ID), + ) + for name in plan.get_dataset_names( + methodology=attribute.Methodology.STATISTICAL, partitioning=partitioning + ) + } + monkeypatch.setattr(cli.storage, "path_exists", lambda uri: uri in uris) + return uris + + +def test_compute_checks_ingests_before_harmonize( + monkeypatch: pytest.MonkeyPatch, present_uris: set[str] +) -> None: + (missing,) = [uri for uri in present_uris if "GPW_GRASSLAND/" in uri] + present_uris.remove(missing) + + def harmonize(**kwargs: object) -> None: + raise AssertionError("harmonize ran on a tile missing an ingest") + + monkeypatch.setattr(cli.harmonize, "workflow", harmonize) + with pytest.raises(FileNotFoundError, match=re.escape(missing)): + cli.compute( + allow_failures=("DESCALS_OIL_PALM",), + check_ingests=True, + methodology=attribute.Methodology.STATISTICAL, + scope=("AAA",), + tile_id="10N_010W", + ) + + +def test_an_allowed_failure_is_not_missing(present_uris: set[str]) -> None: + present_uris -= {uri for uri in present_uris if "DESCALS_OIL_PALM/" in uri} + assert ( + cli.get_missing_ingests( + allow_failures=("DESCALS_OIL_PALM",), + methodology=attribute.Methodology.STATISTICAL, + tile_id="10N_010W", + ) + == [] + ) + assert ( + cli.get_missing_ingests( + allow_failures=(), + methodology=attribute.Methodology.STATISTICAL, + tile_id="10N_010W", + ) + != [] + ) diff --git a/orchestration/plan.py b/orchestration/plan.py new file mode 100644 index 0000000..964a17e --- /dev/null +++ b/orchestration/plan.py @@ -0,0 +1,230 @@ +"""What a run consists of, as the JSON lists the workflow loops over. + +The dataset lists are a hand-kept claim about what jdluc reads: no jdluc constant declares it +all. A laptop run ingests lazily, from inside harmonize, so it never needs them. A workflow +ingests up front instead, so that compute runs with ``skip_ingest`` and a missing dataset fails +at ingest rather than surfacing as an all-nodata band (``ignore_missing_tiles``). In particular +``faostat.load`` and ``usda_nass_quickstats.load`` read the ingest root directly and never +ingest, so those tables appear here by name. When jdluc starts reading a new dataset, add it +here. +""" + +import collections.abc +import json +import logging +import os + +from jdluc import ( + attribute, + config, + harmonize, + jurisdictional_direct, + statistical, + tiling, +) +from jdluc.datasets import NAME_TO_CLS, DatasetName, worldbank_jurisdictions + +logger = logging.getLogger(__name__) + +# Read by every tile lookup and every per-jurisdiction rollup, but named by no workflow +BOUNDARY_DATASET_NAMES = ( + DatasetName.WORLD_BANK_ADMIN_0, + DatasetName.WORLD_BANK_ADMIN_1, +) + +# Read with `load()` straight from the ingest root, never ingested on demand +METHODOLOGY_TO_LOADED_DATASET_NAMES = { + attribute.Methodology.JURISDICTIONAL_DIRECT: (DatasetName.USDA_NASS_QUICKSTATS,), + attribute.Methodology.STATISTICAL: ( + DatasetName.FAOSTAT_LIVESTOCK_PATTERNS, + DatasetName.FAOSTAT_PRODUCTION_LIVESTOCK, + ), +} + +METHODOLOGY_TO_HARMONIZED_DATASET_NAMES = { + attribute.Methodology.JURISDICTIONAL_DIRECT: jurisdictional_direct.DATASET_NAMES, + attribute.Methodology.STATISTICAL: statistical.DATASET_NAMES, +} + + +# The scope of a run over every country and GNW tile, as jdluc's `--backfill`; also what its +# run-wide outputs are named, in place of every country's code joined +WORLD = "world" + + +def get_dataset_names( + methodology: attribute.Methodology, partitioning: tiling.Partitioning +) -> list[DatasetName]: + """Every dataset of one partitioning that a run of ``methodology`` reads.""" + dataset_names = { + *BOUNDARY_DATASET_NAMES, + *(name for stack in harmonize.Stack for name in stack.value), + *METHODOLOGY_TO_HARMONIZED_DATASET_NAMES[methodology], + *METHODOLOGY_TO_LOADED_DATASET_NAMES[methodology], + } + return sorted( + ( + name + for name in dataset_names + if NAME_TO_CLS[name].partitioning == partitioning + ), + key=lambda name: name.name, + ) + + +def scope_str(s: str) -> str: + """One element of a run's scope: an ISO 3166 alpha-3 code, or ``world``.""" + return s if s == WORLD else worldbank_jurisdictions.iso_3166_str(s=s) + + +def is_world(scope: tuple[str, ...]) -> bool: + assert scope, "the scope names no country" + assert scope == (WORLD,) or WORLD not in scope, ( + f"{scope=} mixes {WORLD} with countries" + ) + return scope == (WORLD,) + + +def get_iso_3166s(scope: tuple[str, ...]) -> tuple[str, ...]: + """The scope's countries: those named, or for ``world`` every country, as jdluc's + ``--backfill`` reads them.""" + return tuple( + sorted( + worldbank_jurisdictions.get_all_iso_3166s() if is_world(scope) else scope + ) + ) + + +def get_name(scope: tuple[str, ...]) -> str: + """What the run's run-wide outputs (its mosaics and emissions-factor table) are named.""" + return WORLD if is_world(scope) else "-".join(sorted(scope)) + + +def get_tile_ids(scope: tuple[str, ...]) -> list[str]: + """The tiles the scope's countries touch, or for ``world`` every GNW tile, as jdluc's + ``--backfill`` reads them: tiles no country touches included.""" + return sorted( + tiling.GLOBAL_NATURE_WATCH_TILE_IDS + if is_world(scope) + else worldbank_jurisdictions.get_ten_degree_tile_ids_for_iso_3166s( + iso_3166s=get_iso_3166s(scope=scope) + ) + ) + + +def ingest_boundaries() -> None: + """The tile list is read off the boundaries, so the plan ingests them first. + + A no-op when they are already in the ingest root, as after any earlier run. + """ + root = config.Config.from_dot_env().ingest_root + for name in BOUNDARY_DATASET_NAMES: + logger.info(f"Ingesting {name.name:s} if absent") + NAME_TO_CLS[name].ingest_a_tile( + overwrite=False, root=root, tile_id=tiling.WHOLE_WORLD_TILE_ID + ) + + +def check_allow_failures( + allow_failures: tuple[str, ...], methodology: attribute.Methodology +) -> None: + tiled = get_dataset_names( + methodology=methodology, partitioning=tiling.Partitioning.TEN_DEGREE_TILE + ) + unknown = sorted(set(allow_failures) - {name.name for name in tiled}) + assert not unknown, f"{unknown} are not ten-degree datasets this run ingests" + + +def get_task_to_items( + allow_failures: tuple[str, ...], + methodology: attribute.Methodology, + scope: tuple[str, ...], +) -> dict[str, list]: + """Each list the workflow loops over, by task name. Writes nothing, so that a wrapper can + add lists of its own, or keys to each tile's unit, before `write_task_to_items`. + + Reads the boundaries for the tile list: `ingest_boundaries` first, on a fresh ingest root. + + The flat workflow loops over each stage's list across every tile. ``tile-units`` holds + the same units grouped by tile, for the per-tile workflow, which runs each tile's ingest + and compute as a DAG of its own. + """ + check_allow_failures(allow_failures=allow_failures, methodology=methodology) + tiled = get_dataset_names( + methodology=methodology, partitioning=tiling.Partitioning.TEN_DEGREE_TILE + ) + tile_ids = get_tile_ids(scope=scope) + # One flat list of (dataset, tile), so that the workflow needs no nested loop + ingest_units = [ + {"dataset": name.name, "tile_id": tile_id} + for name in tiled + for tile_id in tile_ids + ] + return { + "ingest-world": [ + name.name + for name in get_dataset_names( + methodology=methodology, partitioning=tiling.Partitioning.WHOLE_WORLD + ) + if name not in BOUNDARY_DATASET_NAMES + ], + "ingest-tiles": [u for u in ingest_units if u["dataset"] not in allow_failures], + "ingest-tiles-allowed-failures": [ + u for u in ingest_units if u["dataset"] in allow_failures + ], + "tiles": [{"tile_id": tile_id} for tile_id in tile_ids], + "tile-units": [ + { + "tile_id": tile_id, + "ingests": [ + u["dataset"] + for u in ingest_units + if u["tile_id"] == tile_id and u["dataset"] not in allow_failures + ], + "ingests_allowed_failures": [ + u["dataset"] + for u in ingest_units + if u["tile_id"] == tile_id and u["dataset"] in allow_failures + ], + } + for tile_id in tile_ids + ], + } + + +def write_task_to_items( + output_dir: str, task_to_items: collections.abc.Mapping[str, list] +) -> None: + """Write each list to ``output_dir``, as ``.json``.""" + os.makedirs(output_dir, exist_ok=True) + for task, items in task_to_items.items(): + path = os.path.join(output_dir, f"{task:s}.json") + logger.info(f"Writing {len(items):d} items to {path=:s}") + with open(path, "w") as fp: + json.dump(items, fp) + + +def write_plan( + allow_failures: tuple[str, ...], + methodology: attribute.Methodology, + output_dir: str, + scope: tuple[str, ...], +) -> None: + """Write each list the workflow loops over to ``output_dir``, as ``.json``. + + The run covers ``scope``: one or more countries, or ``world`` for every country and GNW + tile. + + ``allow_failures`` names ten-degree datasets whose failed tiles must not stop the run. + Their units go to their own list, so that the workflow can run them as a task that + compute does not wait to succeed. + """ + # Before the boundaries' ingest, so that a typo fails fast + check_allow_failures(allow_failures=allow_failures, methodology=methodology) + ingest_boundaries() + write_task_to_items( + output_dir=output_dir, + task_to_items=get_task_to_items( + allow_failures=allow_failures, methodology=methodology, scope=scope + ), + ) diff --git a/orchestration/workflows/cornerstone.example.yaml b/orchestration/workflows/cornerstone.example.yaml new file mode 100644 index 0000000..015143a --- /dev/null +++ b/orchestration/workflows/cornerstone.example.yaml @@ -0,0 +1,238 @@ +# The cornerstone pipeline as an Argo WorkflowTemplate, one DAG per tile: an example with +# placeholder names. Replace every UPPER-CASE placeholder. docs/orchestration.md describes the +# setup: the image, the two Secrets and the service account. +# +# argo template create cornerstone.example.yaml +# argo submit --from workflowtemplate/cornerstone -p scope=HND --generate-name cornerstone-hnd- +# argo submit --from workflowtemplate/cornerstone -p scope=world --generate-name cornerstone-world- +# +# Each tile runs its own ingest -> compute, so a failed (dataset, tile) ingest stops only its +# own tile, and every other tile goes on to compute. trace needs every tile, so with any tile +# failed it doesn't run, and the run ends Failed. `argo retry` resumes from the cache. +# +# `plan` writes, as `tile-units`, one item per tile listing its ingests, which the tile's own +# DAG loops over. A tile's tasks start once the whole-world datasets are in. +# +# plan ─┬─ ingest-world ─┐ +# └────────────────┴─ tile[tile] ── trace +# +# tile: ingest ──────────────────┐ +# ingest-allowed-failures ─┴─ compute +# +# Each task template is reusable on its own. A workflow that adds tasks of its own (say, a +# per-tile export after compute) calls them by `templateRef: {name: cornerstone, template: ...}` +# from DAGs of its own, rather than copying this file. +# +# Sizing: compute requests 56 GiB, with peaks near 40 GiB observed; the other tasks fit in +# 24 GiB. Every pod writes its temp files to a volume of its own at /localtmp: on a cloud +# node's boot disk, GDAL's temp files can throttle the whole node. + +apiVersion: argoproj.io/v1alpha1 +kind: WorkflowTemplate +metadata: + name: cornerstone +spec: + entrypoint: cornerstone + # Needs read/write access to INGEST_ROOT, SCRATCH_ROOT and EXPORT_ROOT, for instance through + # workload identity rather than keys + serviceAccountName: SERVICE-ACCOUNT + # Backstops against a runaway fan-out: at most 25 pods at once, and 6 hours for the whole run. + # Sized for one country: raise both for a scope=world run + parallelism: 25 + activeDeadlineSeconds: 21600 + arguments: + parameters: + # The run's country, as one ISO 3166 alpha-3 code, or world for every country and GNW tile + - name: scope + # STATISTICAL or JURISDICTIONAL_DIRECT (the latter is US-only) + - name: methodology + value: STATISTICAL + # Ten-degree datasets whose failed tiles must not stop the run, comma-separated. Descals + # is published only where there is oil palm, and its ingest fails on tiles without any. + - name: allow-failures + value: DESCALS_OIL_PALM + + # To run on preemptible (spot) nodes, add a nodeSelector and tolerations here. A reclaimed + # node fails its pod; `argo retry` resumes from the cache, as a zarr counts only once written. + # templateDefaults: + # nodeSelector: { NODE-POOL-LABEL: NODE-POOL } + # tolerations: [{ key: NODE-POOL-TAINT, operator: Equal, value: "true", effect: NoSchedule }] + + templates: + - name: cornerstone + dag: + # A failed task must not stop the scheduling of tasks that wait on its failure, nor a + # failed tile the other tiles + failFast: false + tasks: + - name: plan + template: plan + + - name: ingest-world + template: ingest-world + depends: plan + withParam: "{{tasks.plan.outputs.parameters.ingest-world}}" + arguments: { parameters: [{ name: dataset, value: "{{item}}" }] } + + # Each list in an item reaches the tile's DAG as a JSON string, for it to loop over + - name: tile + template: tile + depends: plan && ingest-world + withParam: "{{tasks.plan.outputs.parameters.tile-units}}" + arguments: + parameters: + - { name: tile-id, value: "{{item.tile_id}}" } + - { name: ingests, value: "{{item.ingests}}" } + - { name: ingests-allowed-failures, value: "{{item.ingests_allowed_failures}}" } + + # Runs only if every tile succeeded: a table missing tiles would pass for a whole one + - name: trace + template: trace + depends: tile + + # One tile's ingest -> compute. A failed required ingest leaves compute unscheduled, which + # fails this tile alone + - name: tile + inputs: + parameters: + - name: tile-id + - name: ingests + - name: ingests-allowed-failures + dag: + failFast: false + tasks: + # Ingest exits non-zero if its tile fails, which fails the task, and so the tile + - name: ingest + template: ingest-tile + withParam: "{{inputs.parameters.ingests}}" + arguments: &ingest-tile-arguments + parameters: + - { name: dataset, value: "{{item}}" } + - { name: tile-id, value: "{{inputs.parameters.tile-id}}" } + + # Failed tiles here stay visible as failed pods, but compute runs regardless + - name: ingest-allowed-failures + template: ingest-tile + withParam: "{{inputs.parameters.ingests-allowed-failures}}" + arguments: *ingest-tile-arguments + + - name: compute + template: compute + depends: ingest && (ingest-allowed-failures || ingest-allowed-failures.Failed) + arguments: { parameters: [{ name: tile-id, value: "{{inputs.parameters.tile-id}}" }] } + + # Ingests the boundaries if absent, then writes one JSON list per fan-out, and the per-tile list + - name: plan + activeDeadlineSeconds: 1800 + container: + image: &image IMAGE-REGISTRY/cornerstone:COMMIT-SHA + command: [python, -m, orchestration, plan] + args: + - "{{workflow.parameters.scope}}" + - "--methodology={{workflow.parameters.methodology}}" + - --output-dir=/plan + - "--allow-failures={{workflow.parameters.allow-failures}}" + env: &tmpdir-env + - { name: TMPDIR, value: /localtmp } + - { name: CPL_TMPDIR, value: /localtmp } + resources: &light-resources + requests: { cpu: "2", memory: 24Gi, ephemeral-storage: 2Gi } + limits: { cpu: "4", memory: 24Gi, ephemeral-storage: 8Gi } + volumeMounts: &mounts + - { name: cloud-env, mountPath: /app/.env, subPath: .env } + - { name: localtmp, mountPath: /localtmp } + volumes: + - &nokeys { name: cloud-env, secret: { secretName: cornerstone-env-nokeys } } + - &localtmp-20 + name: localtmp + ephemeral: + volumeClaimTemplate: + spec: + accessModes: [ReadWriteOnce] + # An SSD-backed class: throughput scales with size, so don't shrink these + storageClassName: SSD-STORAGE-CLASS + resources: { requests: { storage: 20Gi } } + outputs: + parameters: + - { name: ingest-world, valueFrom: { path: /plan/ingest-world.json } } + - { name: tile-units, valueFrom: { path: /plan/tile-units.json } } + + # One whole-world dataset, with jdluc's own CLI + - name: ingest-world + inputs: { parameters: [{ name: dataset }] } + activeDeadlineSeconds: 5400 + container: + image: *image + command: [python, jdluc/ingest.py] + args: ["{{inputs.parameters.dataset}}", --backfill] + env: *tmpdir-env + resources: *light-resources + volumeMounts: *mounts + volumes: + - &keys { name: cloud-env, secret: { secretName: cornerstone-env } } + - name: localtmp + ephemeral: + volumeClaimTemplate: + spec: + accessModes: [ReadWriteOnce] + storageClassName: SSD-STORAGE-CLASS + resources: { requests: { storage: 100Gi } } + + # One dataset on one tile. Room for a global source download (GPW livestock, ~5 GB) plus + # one tile staged uncompressed for its COG (Liao, ~55 GB) + - name: ingest-tile + inputs: { parameters: [{ name: dataset }, { name: tile-id }] } + activeDeadlineSeconds: 10800 + container: + image: *image + command: [python, -m, orchestration, ingest-tile] + args: + - "--dataset={{inputs.parameters.dataset}}" + - "--tile-id={{inputs.parameters.tile-id}}" + env: *tmpdir-env + resources: *light-resources + volumeMounts: *mounts + volumes: + - *keys + - name: localtmp + ephemeral: + volumeClaimTemplate: + spec: + accessModes: [ReadWriteOnce] + storageClassName: SSD-STORAGE-CLASS + resources: { requests: { storage: 150Gi } } + + # harmonize -> emit -> attribute for one tile, from the warm ingest cache, without keys + - name: compute + inputs: { parameters: [{ name: tile-id }] } + activeDeadlineSeconds: 14400 + container: + image: *image + command: [python, -m, orchestration, compute] + args: + - "{{workflow.parameters.scope}}" + - "--methodology={{workflow.parameters.methodology}}" + - "--tile-id={{inputs.parameters.tile-id}}" + # A second line of defence: harmonize would read a missing ingest as nodata + - --check-ingests + - "--allow-failures={{workflow.parameters.allow-failures}}" + env: *tmpdir-env + resources: + requests: { cpu: "6", memory: 56Gi, ephemeral-storage: 2Gi } + limits: { cpu: "8", memory: 56Gi, ephemeral-storage: 8Gi } + volumeMounts: *mounts + volumes: [*nokeys, *localtmp-20] + + # The cross-tile merge over the scope; writes emissions-factors/{scope}-{methodology}.parquet + - name: trace + activeDeadlineSeconds: 3600 + container: + image: *image + command: [python, -m, orchestration, trace] + args: + - "{{workflow.parameters.scope}}" + - "--methodology={{workflow.parameters.methodology}}" + env: *tmpdir-env + resources: *light-resources + volumeMounts: *mounts + volumes: [*nokeys, *localtmp-20] diff --git a/pyproject.toml b/pyproject.toml index 6e37fee..a2ee865 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,7 +59,7 @@ follow_imports = "skip" implicit_reexport = true [tool.importlinter] -root_packages = ["jdluc", "validation"] +root_packages = ["jdluc", "orchestration", "validation"] [[tool.importlinter.contracts]] name = "The pipeline never imports the validator" @@ -71,6 +71,16 @@ forbidden_modules = [ "validation", ] +[[tool.importlinter.contracts]] +name = "The pipeline never imports its orchestration" +type = "forbidden" +source_modules = [ + "jdluc", +] +forbidden_modules = [ + "orchestration", +] + [[tool.importlinter.contracts]] name = "ETL" type = "layers" @@ -88,6 +98,7 @@ name = "Utilities are foundational" type = "forbidden" source_modules = [ "jdluc.config", + "jdluc.export", "jdluc.geo", "jdluc.storage", "jdluc.tiling", From c6e3b0a3c5e0e6e23026d18908c90c6a4492edfb Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Mon, 19 Oct 2026 12:16:11 +0200 Subject: [PATCH 4/8] orchestration: docstrings name nothing that lives only in the fork get_name no longer mentions mosaics, get_task_to_items describes the flat lists without a flat workflow to point at, and trace's --help is a whole sentence. --- orchestration/__main__.py | 5 +++-- orchestration/plan.py | 8 ++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/orchestration/__main__.py b/orchestration/__main__.py index 6daf089..e4eeed4 100644 --- a/orchestration/__main__.py +++ b/orchestration/__main__.py @@ -131,8 +131,9 @@ def get_iso_3166s_touching_tile( def trace_countries(methodology: attribute.Methodology, scope: tuple[str, ...]) -> None: - """The cross-tile merge, written as the emissions-factor table. Cheap only if compute ran - with the same methodology.""" + """The cross-tile merge, written as the emissions-factor table. + + Cheap only if compute ran with the same methodology.""" uri = trace.export_workflow( commodity_names=attribute.get_commodity_names(methodology=methodology), iso_3166s=plan.get_iso_3166s(scope=scope), diff --git a/orchestration/plan.py b/orchestration/plan.py index 964a17e..70141a8 100644 --- a/orchestration/plan.py +++ b/orchestration/plan.py @@ -96,7 +96,7 @@ def get_iso_3166s(scope: tuple[str, ...]) -> tuple[str, ...]: def get_name(scope: tuple[str, ...]) -> str: - """What the run's run-wide outputs (its mosaics and emissions-factor table) are named.""" + """What the run's run-wide outputs, such as its emissions-factor table, are named.""" return WORLD if is_world(scope) else "-".join(sorted(scope)) @@ -145,9 +145,9 @@ def get_task_to_items( Reads the boundaries for the tile list: `ingest_boundaries` first, on a fresh ingest root. - The flat workflow loops over each stage's list across every tile. ``tile-units`` holds - the same units grouped by tile, for the per-tile workflow, which runs each tile's ingest - and compute as a DAG of its own. + ``tile-units`` groups the units by tile, for the example workflow, which runs each tile's + ingest and compute as a DAG of its own. The other lists hold the same units flat, one per + stage across every tile, for a workflow that fans out stage by stage instead. """ check_allow_failures(allow_failures=allow_failures, methodology=methodology) tiled = get_dataset_names( From 01d7ff3fd469f8202dcce47dda8c2a6d29d13f0a Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Mon, 19 Oct 2026 13:54:39 +0200 Subject: [PATCH 5/8] export: write_tile writes a step's per-tile output as a COG or a zarr A Python API only: nothing calls it, a workflow opts in. Writes the chosen variables of a lazy dataset to {EXPORT_ROOT}/{name}/{tile_id}.{tif|zarr}, provenance from get_metadata. - COG: every variable as float32 2-D bands, NaN no-data; extra dims are flattened to one band per combination of values, named variable:value[:value...]. Staged locally and uploaded whole, through storage.write_dask_dataset_to_cog (moved into storage with stack_bands). - zarr: variables as they are, dtypes and extra dims kept. An export is re-runnable, so it writes over an existing zarr (write_dask_dataset_to_zarr gains mode=, default "w-" so the cache is unchanged), dropping the .complete marker first and touching it once the write finishes. --- docs/data.md | 2 +- jdluc/__tests__/export_test.py | 185 +++++++++++++++++++++++++++++++++ jdluc/export.py | 159 ++++++++++++++++++++++++++-- jdluc/storage.py | 95 ++++++++++++++++- 4 files changed, 428 insertions(+), 13 deletions(-) create mode 100644 jdluc/__tests__/export_test.py diff --git a/docs/data.md b/docs/data.md index 39e6751..cd0149a 100644 --- a/docs/data.md +++ b/docs/data.md @@ -10,7 +10,7 @@ Of the three artifacts below, one is published and two are not. **The emissions-factor table is published as an archived dataset.** Each data version is deposited as `emissions-factors.parquet` in a public data archive that mints a DOI per version, and the matching git tag names the commit behind it. The archive is still being set up; cite a version by its DOI, or follow the fallback in `ATTRIBUTION.md`. -**The two raster artifacts are pipeline outputs.** They are zarr stores far too large to deposit, so they exist only where a run writes them: under the `SCRATCH_ROOT` of whoever ran the pipeline, keyed by the cache hash `storage.py` derives (see [`architecture.md`](architecture.md#caching-versioning-and-provenance)). Sharing the underlying maps in a form that can be read remotely is open work — the schemas below are documented so a run's own output is legible in the meantime. +**The two raster artifacts are pipeline outputs.** They are zarr stores far too large to deposit, so they exist only where a run writes them: under the `SCRATCH_ROOT` of whoever ran the pipeline, keyed by the cache hash `storage.py` derives (see [`architecture.md`](architecture.md#caching-versioning-and-provenance)). Sharing the underlying maps in a form that can be read remotely is open work — the schemas below are documented so a run's own output is legible in the meantime. A run can also write a step's per-tile output as a COG or a zarr under `EXPORT_ROOT` with `export.write_tile`; that copies the maps out of the cache, but does not by itself publish them. ## Conventions diff --git a/jdluc/__tests__/export_test.py b/jdluc/__tests__/export_test.py new file mode 100644 index 0000000..a6d2400 --- /dev/null +++ b/jdluc/__tests__/export_test.py @@ -0,0 +1,185 @@ +import dataclasses +import pathlib + +import numpy +import pytest +import rasterio +import xarray + +from jdluc import config, export, storage + +TILE_ID = "00N_000E" + + +def get_darray(value: float) -> xarray.DataArray: + return xarray.DataArray( + coords={"y": [1.5, 0.5], "x": [0.5, 1.5, 2.5]}, # north up, as a raster + data=numpy.full((2, 3), value, dtype=numpy.float32), + dims=("y", "x"), + ) + + +def get_dset() -> xarray.Dataset: + """A step's output: an integer (y, x) variable, and one that varies by `scheme`.""" + by_scheme = xarray.concat( + [get_darray(value=1), get_darray(value=2), get_darray(value=3)], dim="scheme" + ).assign_coords(scheme=["c", "a", "b"]) + return xarray.Dataset( + { + "total:t": by_scheme, + "area:ha": get_darray(value=7).astype(numpy.int16), + "count": get_darray(value=4), + } + ).chunk({"x": 2}) + + +def write_tile(format: export.Format, tmp_path: pathlib.Path, **kwargs: object) -> str: + return export.write_tile( + dset=get_dset(), + format=format, + name="example", + output_root=str(tmp_path), + tile_id=TILE_ID, + **kwargs, # type: ignore[arg-type] + ) + + +def test_get_tile_output_uri_is_under_the_name_by_tile(tmp_path: pathlib.Path) -> None: + uri = export.get_tile_output_uri( + format=export.Format.COG, + name="example", + output_root=str(tmp_path), + tile_id=TILE_ID, + ) + assert uri == str(tmp_path / "example" / f"{TILE_ID:s}.tif") + + +def test_get_tile_output_uri_defaults_to_the_export_root( + tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch +) -> None: + with_root = dataclasses.replace( + config.Config.from_dot_env(), export_root=str(tmp_path) + ) + monkeypatch.setattr( + config.Config, "from_dot_env", classmethod(lambda cls: with_root) + ) + uri = export.get_tile_output_uri( + format=export.Format.ZARR, name="example", tile_id=TILE_ID + ) + assert uri == str(tmp_path / "example" / f"{TILE_ID:s}.zarr") + + +def test_iter_bands_flattens_extra_dims_in_coord_order() -> None: + names = [name for name, _ in export.iter_bands(darray=get_dset()["total:t"])] + assert names == ["total:t:c", "total:t:a", "total:t:b"] + + +@pytest.mark.parametrize("format", list(export.Format)) +def test_write_tile_rejects_a_missing_variable_before_writing( + format: export.Format, tmp_path: pathlib.Path +) -> None: + with pytest.raises(ValueError, match="missing"): + write_tile(format=format, tmp_path=tmp_path, variables=["count", "missing"]) + assert not list(tmp_path.iterdir()) + + +def test_write_tile_as_cog_writes_every_variable_by_default( + tmp_path: pathlib.Path, +) -> None: + uri = write_tile(format=export.Format.COG, tmp_path=tmp_path) + assert uri == str(tmp_path / "example" / f"{TILE_ID:s}.tif") + with rasterio.open(uri) as dataset: + assert list(dataset.descriptions) == [ + "total:t:c", + "total:t:a", + "total:t:b", + "area:ha", + "count", + ] + assert set(dataset.dtypes) == {"float32"} + assert numpy.isnan(dataset.nodata) + numpy.testing.assert_array_equal(dataset.read()[:, 0, 0], [1, 2, 3, 7, 4]) + tags = dataset.tags() + assert tags["cornerstone-product-name"] == "example" + assert tags["cornerstone-source-name"] == "example" + assert "cornerstone-processing-version" in tags + + +def test_write_tile_as_cog_writes_a_subset_in_the_given_order( + tmp_path: pathlib.Path, +) -> None: + uri = write_tile( + format=export.Format.COG, + source_name="emit", + tmp_path=tmp_path, + variables=["count", "area:ha"], + ) + with rasterio.open(uri) as dataset: + assert list(dataset.descriptions) == ["count", "area:ha"] + assert dataset.tags()["cornerstone-source-name"] == "emit" + + +def test_write_tile_as_cog_marks_no_data_as_nan(tmp_path: pathlib.Path) -> None: + dset = get_dset() + dset["count"] = dset["count"].where(dset["x"] > 1) + uri = export.write_tile( + dset=dset, + format=export.Format.COG, + name="example", + output_root=str(tmp_path), + tile_id=TILE_ID, + variables=["count"], + ) + with rasterio.open(uri) as dataset: + band = dataset.read(1, masked=True) + assert band.mask.tolist() == [[True, False, False], [True, False, False]] + + +def test_write_tile_as_zarr_keeps_dtypes_and_extra_dims( + tmp_path: pathlib.Path, +) -> None: + uri = write_tile(format=export.Format.ZARR, tmp_path=tmp_path) + assert uri == str(tmp_path / "example" / f"{TILE_ID:s}.zarr") + written = xarray.open_zarr(uri, consolidated=False) + # NB: a zarr group's members have no order; a reader lists them by name + assert set(written.data_vars) == {"total:t", "area:ha", "count"} + assert written["area:ha"].dtype == numpy.int16 + assert written["total:t"].dims == ("scheme", "y", "x") + assert list(written["scheme"].values) == ["c", "a", "b"] + assert written.attrs["cornerstone-product-name"] == "example" + assert storage.path_exists(uri=storage.get_zarr_complete_uri(path_to_zarr=uri)) + + +def test_write_tile_as_zarr_writes_only_a_subset( + tmp_path: pathlib.Path, +) -> None: + uri = write_tile( + format=export.Format.ZARR, tmp_path=tmp_path, variables=["count", "area:ha"] + ) + written = xarray.open_zarr(uri, consolidated=False) + assert set(written.data_vars) == {"count", "area:ha"} + assert "scheme" not in written.dims + + +def test_write_tile_as_zarr_writes_over_an_earlier_export( + tmp_path: pathlib.Path, +) -> None: + write_tile(format=export.Format.ZARR, tmp_path=tmp_path) + uri = write_tile(format=export.Format.ZARR, tmp_path=tmp_path, variables=["count"]) + assert list(xarray.open_zarr(uri, consolidated=False).data_vars) == ["count"] + assert storage.path_exists(uri=storage.get_zarr_complete_uri(path_to_zarr=uri)) + + +def test_write_tile_as_zarr_drops_the_marker_before_writing( + tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch +) -> None: + uri = write_tile(format=export.Format.ZARR, tmp_path=tmp_path) + + def fail(**kwargs: object) -> None: + raise RuntimeError("killed part-way") + + monkeypatch.setattr(storage, "write_dask_dataset_to_zarr", fail) + with pytest.raises(RuntimeError, match="killed part-way"): + write_tile(format=export.Format.ZARR, tmp_path=tmp_path) + + assert not storage.path_exists(uri=storage.get_zarr_complete_uri(path_to_zarr=uri)) diff --git a/jdluc/export.py b/jdluc/export.py index 5860d5e..0fbb364 100644 --- a/jdluc/export.py +++ b/jdluc/export.py @@ -1,23 +1,63 @@ -"""Write a pipeline step's output as a deliverable under EXPORT_ROOT. +"""Write a pipeline step's output under EXPORT_ROOT, a raster per tile or a table. -A utility, like `storage`: it knows nothing about the steps. A tabular deliverable (see -`trace.export_workflow`) is one parquet per set of countries, written by `write_table` to -`{EXPORT_ROOT}/{deliverable}/{name}.parquet`, with provenance in its `attrs`. +A utility, like `storage`: it knows nothing about the steps, and nothing calls it +automatically. A workflow opts in by passing its own lazy dataset to `write_tile`, which +writes the chosen variables of one tile to `{EXPORT_ROOT}/{name}/{tile_id}.{tif|zarr}`, +computing only those, as it writes: -Unlike SCRATCH_ROOT, EXPORT_ROOT is not a cache that the pipeline reads back: what lands -there is a terminal output. +- COG: every variable as float32 2-D bands, NaN no-data. A variable with dimensions beyond + the pixel grid is flattened to one band per combination of their values, named + `:[:...]`; +- zarr: the variables as they are, with their dtypes and extra dimensions. A rewrite + replaces the zarr, and the `.complete` marker beside it says the write finished. + +For example: + + export.write_tile( + dset=emit.workflow(tile_id=tile_id), + format=export.Format.COG, + name="emit", + tile_id=tile_id, + ) + export.write_tile( + dset=statistical.get_downscaled_luc_emissions(tile_id=tile_id), + format=export.Format.ZARR, + name="statistical-emissions", + tile_id=tile_id, + ) + +A tabular deliverable (see `trace.export_workflow`) is one parquet per set of countries, +written by `write_table` to `{EXPORT_ROOT}/{deliverable}/{name}.parquet`. + +Either way, provenance is stamped in the output's metadata. Unlike SCRATCH_ROOT, EXPORT_ROOT +is not a cache that the pipeline reads back: what lands there is a terminal output. """ +import collections.abc +import enum +import itertools import logging import os import tempfile +import numpy import pandas +import xarray -from jdluc import config, storage, utils +from jdluc import config, geo, storage, utils logger = logging.getLogger(__name__) +PIXEL_DIMS = ("y", "x") + + +class Format(enum.StrEnum): + COG = "cog" + ZARR = "zarr" + + +FORMAT_TO_EXTENSION = {Format.COG: ".tif", Format.ZARR: ".zarr"} + def get_metadata(product_name: str, source_name: str) -> dict[str, str]: """Provenance, as `datasets.base` stamps on ingested COGs.""" @@ -30,6 +70,111 @@ def get_metadata(product_name: str, source_name: str) -> dict[str, str]: } +def get_tile_output_uri( + format: Format, name: str, tile_id: str, output_root: str | None = None +) -> str: + """`{output_root}/{name}/{tile_id}.{tif|zarr}`; `output_root` defaults to EXPORT_ROOT.""" + if output_root is None: + output_root = config.Config.from_dot_env().export_root + return storage.join_uri( + root=output_root, + prefix=f"{name:s}/{tile_id:s}{FORMAT_TO_EXTENSION[format]:s}", + ) + + +def iter_bands( + darray: xarray.DataArray, +) -> collections.abc.Iterator[tuple[str, xarray.DataArray]]: + """Each 2-D band of a variable and its name: one per combination of values on its dims + beyond the pixel grid, taken in the variable's dim order and each dim's coord order.""" + dims = [dim for dim in map(str, darray.dims) if dim not in PIXEL_DIMS] + for indexes in itertools.product(*(range(darray.sizes[dim]) for dim in dims)): + point = dict(zip(dims, indexes, strict=True)) + values = (str(darray[dim].values[index]) for dim, index in point.items()) + yield ":".join((str(darray.name), *values)), darray.isel(point, drop=True) + + +def to_cog_dataset(dset: xarray.Dataset) -> xarray.Dataset: + """`dset`'s variables as 2-D (y, x) bands, float32, NaN no-data, in order. + + Chunked so one chunk of every band fits in memory together, as `emit` writes its output. + """ + name_and_darrays = [ + band for darray in dset.data_vars.values() for band in iter_bands(darray=darray) + ] + name_to_darray = dict(name_and_darrays) + if len(name_to_darray) < len(name_and_darrays): + raise ValueError(f"Band names repeat: {[name for name, _ in name_and_darrays]}") + chunk_size = geo.get_chunk_size( + dtypes=[numpy.dtype("float32")] * len(name_to_darray) + ) + return xarray.Dataset( + { + name: geo.unify_dtype_and_no_data(darray=darray.rename(None)).chunk( + chunks=chunk_size + ) + for name, darray in name_to_darray.items() + } + ) + + +def to_zarr_dataset(dset: xarray.Dataset) -> xarray.Dataset: + # NB: zarr v3 specifies variable-length strings, not numpy's fixed-length ones + return dset.assign_coords( + { + name: coord.astype(object) + for name, coord in dset.coords.items() + if coord.dtype.kind == "U" + } + ) + + +def write_zarr(dset: xarray.Dataset, uri: str) -> None: + """Write over any zarr at `uri`, marking it complete only once the write is.""" + complete_uri = storage.get_zarr_complete_uri(path_to_zarr=uri) + storage.remove_file(uri=complete_uri) + storage.write_dask_dataset_to_zarr(dset=dset, mode="w", path_to_zarr=uri) + storage.touch_file(uri=complete_uri) + + +def write_tile( + dset: xarray.Dataset, + format: Format, + name: str, + tile_id: str, + output_root: str | None = None, + source_name: str | None = None, + variables: collections.abc.Sequence[str] | None = None, +) -> str: + """Write `variables` of the (lazy) `dset` for one tile; return the URI. + + `variables` defaults to every data variable, in order. Computes only those, as it writes. + """ + if variables is None: + variables = list(map(str, dset.data_vars)) + if missing := [ + variable for variable in variables if variable not in dset.data_vars + ]: + raise ValueError(f"{name=:s} has no variables {missing}") + uri = get_tile_output_uri( + format=format, name=name, output_root=output_root, tile_id=tile_id + ) + metadata = get_metadata(product_name=name, source_name=source_name or name) + selected = dset[list(variables)] + logger.info(f"Writing {name=:s} for {tile_id=:s} as {format!s}") + match format: + case Format.COG: + storage.write_dask_dataset_to_cog( + dset=to_cog_dataset(dset=selected), metadata=metadata, uri=uri + ) + case Format.ZARR: + write_zarr( + dset=to_zarr_dataset(dset=selected).assign_attrs(metadata), uri=uri + ) + logger.info(f"Wrote {name=:s} for {tile_id=:s} to {uri=:s}") + return uri + + def get_table_output_uri( deliverable_name: str, name: str, output_root: str | None = None ) -> str: diff --git a/jdluc/storage.py b/jdluc/storage.py index 684fda5..5f1b317 100644 --- a/jdluc/storage.py +++ b/jdluc/storage.py @@ -4,10 +4,12 @@ import inspect import logging import os +import tempfile import threading import typing import fsspec # type: ignore[import-untyped] +import numpy import pandas import xarray @@ -28,6 +30,18 @@ def put_file(local_path: str, uri: str) -> None: fs.put_file(local_path, path) +def touch_file(uri: str) -> None: + fs, path = fsspec.url_to_fs(url=uri) + fs.touch(path) + + +def remove_file(uri: str) -> None: + """Remove the file at `uri`, if there is one.""" + fs, path = fsspec.url_to_fs(url=uri) + if fs.exists(path): + fs.rm_file(path) + + def to_gdal_path(uri: str) -> str: if uri.startswith("gs://"): return "/vsigs/" + uri.removeprefix("gs://") @@ -39,7 +53,14 @@ def join_uri(root: str, prefix: str) -> str: return f"{root.rstrip('/'):s}/{prefix:s}" -def write_dask_dataset_to_zarr(dset: xarray.Dataset, path_to_zarr: str) -> None: +def get_zarr_complete_uri(path_to_zarr: str) -> str: + """The marker beside a zarr, written once its last chunk is.""" + return f"{path_to_zarr:s}.complete" + + +def write_dask_dataset_to_zarr( + dset: xarray.Dataset, path_to_zarr: str, mode: typing.Literal["w", "w-"] = "w-" +) -> None: assert dset.chunks is not None, f"{dset=:} is not a chunked dask array" import dask.config @@ -90,7 +111,11 @@ def write_dask_dataset_to_zarr(dset: xarray.Dataset, path_to_zarr: str) -> None: for variable in to_write.variables.values(): variable.encoding.pop("chunks", None) delayed = to_write.to_zarr( - compute=False, consolidated=False, group=None, store=path_to_zarr + compute=False, + consolidated=False, + group=None, + mode=mode, + store=path_to_zarr, ) logger.info("Submitting task graph to dask scheduler") future = client.compute(delayed, retries=5) @@ -111,6 +136,67 @@ def open_zarr_to_dask_dataset(path_to_zarr: str) -> xarray.Dataset: return dset.rio.write_crs(4326) +COG_NO_DATA = float("nan") + + +def stack_bands(dset: xarray.Dataset) -> xarray.DataArray: + """`dset`'s variables as one float32 array along `band`, in the dataset's order.""" + import rioxarray # noqa: F401 + + band_names = list(map(str, dset.data_vars)) + stacked = xarray.concat([dset[name] for name in band_names], dim="band").astype( + numpy.float32 + ) + stacked = stacked.assign_coords(band=("band", band_names)) + crs = dset.rio.crs or "EPSG:4326" + return stacked.rio.write_crs(crs).rio.write_nodata(COG_NO_DATA) + + +def write_dask_dataset_to_cog( + dset: xarray.Dataset, metadata: dict[str, str], uri: str +) -> None: + """Write `dset` as a Cloud-Optimised GeoTIFF: one band per variable, named after it. + + EPSG:4326, float32, NaN no-data. Computes the (lazy) dataset as it writes. Stages as + `datasets.base` ingest does: a local GeoTIFF, band names, validation, COG, then one + upload of the whole file. + """ + import dask.diagnostics + + from jdluc import geo + + stacked = stack_bands(dset=dset) + with tempfile.TemporaryDirectory() as tmpdir: + path_to_geotiff = os.path.join(tmpdir, "geotiff.tif") + with dask.diagnostics.ProgressBar(dt=5, minimum=1): + logger.info(f"Writing staged GeoTIFF to {path_to_geotiff=:s}") + stacked.rio.to_raster( + path_to_geotiff, + blockxsize=512, + blockysize=512, + compress="ZSTD", + driver="GTiff", + dtype="float32", + lock=True, + num_threads="all_cpus", + tiled=True, + BIGTIFF="IF_SAFER", + ) + geo.set_band_names_for_geotiff( + band_names=list(map(str, stacked.band.values)), + path_to_geotiff=path_to_geotiff, + ) + geo.validate_geotiff(dtype="float32", path_to_geotiff=path_to_geotiff) + path_to_cog = os.path.join(tmpdir, "cog.tif") + geo.convert_geotiff_to_cog( + metadata=metadata, + no_data=COG_NO_DATA, + path_to_cog=path_to_cog, + path_to_geotiff=path_to_geotiff, + ) + put_file(local_path=path_to_cog, uri=uri) + + P = typing.ParamSpec("P") R = typing.TypeVar("R") @@ -163,7 +249,7 @@ def uri(self) -> str: @property def exists_uri(self) -> str: - return f"{self.uri:s}.complete" + return get_zarr_complete_uri(path_to_zarr=self.uri) def deserialize(self) -> xarray.Dataset: return open_zarr_to_dask_dataset(path_to_zarr=self.uri) @@ -176,8 +262,7 @@ def serialize(self, dset: xarray.Dataset) -> None: "zarr to recompute it." ) write_dask_dataset_to_zarr(dset=dset, path_to_zarr=self.uri) - fs, path = fsspec.url_to_fs(url=self.exists_uri) - fs.touch(path) + touch_file(uri=self.exists_uri) class CacherDecoratorProtocol(typing.Protocol[R]): From 0bd70119471bfffb20630c536d08dd360360c054 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Mon, 19 Oct 2026 15:33:52 +0200 Subject: [PATCH 6/8] export: an optional version puts an export under {EXPORT_ROOT}/{version}/ write_tile and write_table take version=, such as a run id or a release label. Blank, the default, writes directly under EXPORT_ROOT as before; set, the whole export goes under {EXPORT_ROOT}/{version}/ and the version is stamped as cornerstone-export-version. A version must be one path segment. Also from review: - write_tile rejects a bare str as variables, rather than reading its letters - the COG compute runs under COMPUTE_LOCK, as the zarr write does - iter_bands notes how datetime and float coords name their bands - to_zarr_dataset says why only coords become variable-length strings --- jdluc/__tests__/export_test.py | 47 ++++++++++++++++++++ jdluc/export.py | 81 ++++++++++++++++++++++++++-------- jdluc/storage.py | 8 ++-- 3 files changed, 113 insertions(+), 23 deletions(-) diff --git a/jdluc/__tests__/export_test.py b/jdluc/__tests__/export_test.py index a6d2400..3464511 100644 --- a/jdluc/__tests__/export_test.py +++ b/jdluc/__tests__/export_test.py @@ -69,6 +69,39 @@ def test_get_tile_output_uri_defaults_to_the_export_root( assert uri == str(tmp_path / "example" / f"{TILE_ID:s}.zarr") +def test_get_tile_output_uri_puts_a_version_above_the_name( + tmp_path: pathlib.Path, +) -> None: + uri = export.get_tile_output_uri( + format=export.Format.COG, + name="example", + output_root=str(tmp_path), + tile_id=TILE_ID, + version="run-1", + ) + assert uri == str(tmp_path / "run-1" / "example" / f"{TILE_ID:s}.tif") + + +def test_get_table_output_uri_puts_a_version_above_the_deliverable( + tmp_path: pathlib.Path, +) -> None: + uri = export.get_table_output_uri( + deliverable_name="factors", + name="BRA", + output_root=str(tmp_path), + version="run-1", + ) + assert uri == str(tmp_path / "run-1" / "factors" / "BRA.parquet") + + +@pytest.mark.parametrize("version", ["a/b", ".", ".."]) +def test_get_export_root_rejects_a_version_that_is_not_one_segment( + version: str, tmp_path: pathlib.Path +) -> None: + with pytest.raises(ValueError, match="one path segment"): + export.get_export_root(output_root=str(tmp_path), version=version) + + def test_iter_bands_flattens_extra_dims_in_coord_order() -> None: names = [name for name, _ in export.iter_bands(darray=get_dset()["total:t"])] assert names == ["total:t:c", "total:t:a", "total:t:b"] @@ -83,6 +116,12 @@ def test_write_tile_rejects_a_missing_variable_before_writing( assert not list(tmp_path.iterdir()) +def test_write_tile_rejects_one_name_as_variables(tmp_path: pathlib.Path) -> None: + with pytest.raises(TypeError, match="sequence of names"): + write_tile(format=export.Format.ZARR, tmp_path=tmp_path, variables="count") + assert not list(tmp_path.iterdir()) + + def test_write_tile_as_cog_writes_every_variable_by_default( tmp_path: pathlib.Path, ) -> None: @@ -103,6 +142,14 @@ def test_write_tile_as_cog_writes_every_variable_by_default( assert tags["cornerstone-product-name"] == "example" assert tags["cornerstone-source-name"] == "example" assert "cornerstone-processing-version" in tags + assert "cornerstone-export-version" not in tags + + +def test_write_tile_with_a_version_stamps_it(tmp_path: pathlib.Path) -> None: + uri = write_tile(format=export.Format.COG, tmp_path=tmp_path, version="run-1") + assert uri == str(tmp_path / "run-1" / "example" / f"{TILE_ID:s}.tif") + with rasterio.open(uri) as dataset: + assert dataset.tags()["cornerstone-export-version"] == "run-1" def test_write_tile_as_cog_writes_a_subset_in_the_given_order( diff --git a/jdluc/export.py b/jdluc/export.py index 0fbb364..8f70ec5 100644 --- a/jdluc/export.py +++ b/jdluc/export.py @@ -29,6 +29,10 @@ A tabular deliverable (see `trace.export_workflow`) is one parquet per set of countries, written by `write_table` to `{EXPORT_ROOT}/{deliverable}/{name}.parquet`. +Either may be given a `version`, such as a run id or a release label, which puts the whole +export under `{EXPORT_ROOT}/{version}/` and stamps it in the metadata. The default, blank, +writes directly under EXPORT_ROOT, and a rerun writes over the last one. + Either way, provenance is stamped in the output's metadata. Unlike SCRATCH_ROOT, EXPORT_ROOT is not a cache that the pipeline reads back: what lands there is a terminal output. """ @@ -59,25 +63,41 @@ class Format(enum.StrEnum): FORMAT_TO_EXTENSION = {Format.COG: ".tif", Format.ZARR: ".zarr"} -def get_metadata(product_name: str, source_name: str) -> dict[str, str]: - """Provenance, as `datasets.base` stamps on ingested COGs.""" +def get_metadata( + product_name: str, source_name: str, version: str = "" +) -> dict[str, str]: + """Provenance, as `datasets.base` stamps on ingested COGs, and the export's version.""" return { "cornerstone-processing-time": utils.get_utc_timestamp(), "cornerstone-processing-version": utils.get_git_version(), "cornerstone-product-name": product_name, "cornerstone-remote-url": utils.get_git_remote_url(), "cornerstone-source-name": source_name, - } + } | ({"cornerstone-export-version": version} if version else {}) -def get_tile_output_uri( - format: Format, name: str, tile_id: str, output_root: str | None = None -) -> str: - """`{output_root}/{name}/{tile_id}.{tif|zarr}`; `output_root` defaults to EXPORT_ROOT.""" +def get_export_root(output_root: str | None = None, version: str = "") -> str: + """`{output_root}/{version}`, or `output_root` itself for a blank version; + `output_root` defaults to EXPORT_ROOT.""" if output_root is None: output_root = config.Config.from_dot_env().export_root + if not version: + return output_root + if "/" in version or version in (".", ".."): + raise ValueError(f"{version=:s} is not one path segment") + return storage.join_uri(root=output_root, prefix=version) + + +def get_tile_output_uri( + format: Format, + name: str, + tile_id: str, + output_root: str | None = None, + version: str = "", +) -> str: + """`{output_root}[/{version}]/{name}/{tile_id}.{tif|zarr}`; see `get_export_root`.""" return storage.join_uri( - root=output_root, + root=get_export_root(output_root=output_root, version=version), prefix=f"{name:s}/{tile_id:s}{FORMAT_TO_EXTENSION[format]:s}", ) @@ -86,7 +106,10 @@ def iter_bands( darray: xarray.DataArray, ) -> collections.abc.Iterator[tuple[str, xarray.DataArray]]: """Each 2-D band of a variable and its name: one per combination of values on its dims - beyond the pixel grid, taken in the variable's dim order and each dim's coord order.""" + beyond the pixel grid, taken in the variable's dim order and each dim's coord order. + + A value is named by `str`, so it reads well for string or integer coords; a datetime or + float coord gives a name like `variable:2020-01-01T00:00:00.000000000`.""" dims = [dim for dim in map(str, darray.dims) if dim not in PIXEL_DIMS] for indexes in itertools.product(*(range(darray.sizes[dim]) for dim in dims)): point = dict(zip(dims, indexes, strict=True)) @@ -119,7 +142,9 @@ def to_cog_dataset(dset: xarray.Dataset) -> xarray.Dataset: def to_zarr_dataset(dset: xarray.Dataset) -> xarray.Dataset: - # NB: zarr v3 specifies variable-length strings, not numpy's fixed-length ones + # NB: zarr v3 specifies variable-length strings, not numpy's fixed-length ones. Only + # coords are converted: they are in memory, whereas an `object` dask variable is loaded + # whole to infer its dtype, so a string data variable is left fixed-length return dset.assign_coords( { name: coord.astype(object) @@ -145,11 +170,15 @@ def write_tile( output_root: str | None = None, source_name: str | None = None, variables: collections.abc.Sequence[str] | None = None, + version: str = "", ) -> str: """Write `variables` of the (lazy) `dset` for one tile; return the URI. `variables` defaults to every data variable, in order. Computes only those, as it writes. + A `version` puts the tile under `{output_root}/{version}/` (see `get_export_root`). """ + if isinstance(variables, str): + raise TypeError(f"{variables=:s} is one name; pass a sequence of names") if variables is None: variables = list(map(str, dset.data_vars)) if missing := [ @@ -157,9 +186,15 @@ def write_tile( ]: raise ValueError(f"{name=:s} has no variables {missing}") uri = get_tile_output_uri( - format=format, name=name, output_root=output_root, tile_id=tile_id + format=format, + name=name, + output_root=output_root, + tile_id=tile_id, + version=version, + ) + metadata = get_metadata( + product_name=name, source_name=source_name or name, version=version ) - metadata = get_metadata(product_name=name, source_name=source_name or name) selected = dset[list(variables)] logger.info(f"Writing {name=:s} for {tile_id=:s} as {format!s}") match format: @@ -176,13 +211,15 @@ def write_tile( def get_table_output_uri( - deliverable_name: str, name: str, output_root: str | None = None + deliverable_name: str, + name: str, + output_root: str | None = None, + version: str = "", ) -> str: - """`{output_root}/{deliverable}/{name}.parquet`; `output_root` defaults to EXPORT_ROOT.""" - if output_root is None: - output_root = config.Config.from_dot_env().export_root + """`{output_root}[/{version}]/{deliverable}/{name}.parquet`; see `get_export_root`.""" return storage.join_uri( - root=output_root, prefix=f"{deliverable_name:s}/{name:s}.parquet" + root=get_export_root(output_root=output_root, version=version), + prefix=f"{deliverable_name:s}/{name:s}.parquet", ) @@ -192,13 +229,19 @@ def write_table( name: str, source_name: str, output_root: str | None = None, + version: str = "", ) -> str: """Write a tabular deliverable as one parquet, provenance in its `attrs`; return the URI.""" uri = get_table_output_uri( - deliverable_name=deliverable_name, name=name, output_root=output_root + deliverable_name=deliverable_name, + name=name, + output_root=output_root, + version=version, ) df = df.copy() - df.attrs |= get_metadata(product_name=deliverable_name, source_name=source_name) + df.attrs |= get_metadata( + product_name=deliverable_name, source_name=source_name, version=version + ) with tempfile.TemporaryDirectory() as tmpdir: local_path = os.path.join(tmpdir, os.path.basename(uri)) df.to_parquet(path=local_path) diff --git a/jdluc/storage.py b/jdluc/storage.py index 5f1b317..9abbc91 100644 --- a/jdluc/storage.py +++ b/jdluc/storage.py @@ -157,9 +157,9 @@ def write_dask_dataset_to_cog( ) -> None: """Write `dset` as a Cloud-Optimised GeoTIFF: one band per variable, named after it. - EPSG:4326, float32, NaN no-data. Computes the (lazy) dataset as it writes. Stages as - `datasets.base` ingest does: a local GeoTIFF, band names, validation, COG, then one - upload of the whole file. + EPSG:4326, float32, NaN no-data. Computes the (lazy) dataset as it writes, under + COMPUTE_LOCK as `write_dask_dataset_to_zarr` does. Stages as `datasets.base` ingest + does: a local GeoTIFF, band names, validation, COG, then one upload of the whole file. """ import dask.diagnostics @@ -168,7 +168,7 @@ def write_dask_dataset_to_cog( stacked = stack_bands(dset=dset) with tempfile.TemporaryDirectory() as tmpdir: path_to_geotiff = os.path.join(tmpdir, "geotiff.tif") - with dask.diagnostics.ProgressBar(dt=5, minimum=1): + with COMPUTE_LOCK, dask.diagnostics.ProgressBar(dt=5, minimum=1): logger.info(f"Writing staged GeoTIFF to {path_to_geotiff=:s}") stacked.rio.to_raster( path_to_geotiff, From 7076a715c76e434fc16fad3801b53394ac9296b5 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Mon, 19 Oct 2026 15:47:35 +0200 Subject: [PATCH 7/8] utils: get_git_version names the commit on a detached HEAD A detached HEAD, as CI checks out a pull request, has no active branch, so get_git_version raised and with it every export that stamps provenance. It now returns the short commit SHA. No .git still fails, as before. --- jdluc/__tests__/utils_test.py | 19 +++++++++++++++++++ jdluc/utils.py | 3 +++ 2 files changed, 22 insertions(+) diff --git a/jdluc/__tests__/utils_test.py b/jdluc/__tests__/utils_test.py index 2bff881..ad538b2 100644 --- a/jdluc/__tests__/utils_test.py +++ b/jdluc/__tests__/utils_test.py @@ -3,15 +3,34 @@ import itertools import re import time +import types +import git import pytest +from jdluc import utils from jdluc.utils import ( iter_sharded, threadsafe_cache, ) +def test_get_git_version_on_a_detached_head_is_the_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + head = types.SimpleNamespace( + commit=types.SimpleNamespace(hexsha="0123456789abcdef"), is_detached=True + ) + monkeypatch.setattr( + git, "Repo", lambda *args, **kwargs: types.SimpleNamespace(head=head) + ) + utils.get_git_version.cache_clear() + try: + assert utils.get_git_version() == "01234567" + finally: + utils.get_git_version.cache_clear() + + def test_threadsafe_cache_serial() -> None: @threadsafe_cache def now() -> float: diff --git a/jdluc/utils.py b/jdluc/utils.py index 61e7ca0..20ff959 100644 --- a/jdluc/utils.py +++ b/jdluc/utils.py @@ -75,6 +75,9 @@ def save_remote_url_to_local_path( @functools.cache def get_git_version(default_branch_name: str = "main") -> str: repo = git.Repo(__file__, search_parent_directories=True) + # NB: a detached HEAD, as CI checks out a pull request, has no branch to name + if repo.head.is_detached: + return repo.head.commit.hexsha[:8] if (branch_name := repo.active_branch.name) == default_branch_name: return f"{branch_name:s}-{repo.head.commit.hexsha[:8]:s}" else: From bd087a1bf5b9a87c303e094bcfc994d1af6882c2 Mon Sep 17 00:00:00 2001 From: Dan O'Huiginn Date: Mon, 19 Oct 2026 20:47:07 +0200 Subject: [PATCH 8/8] emit: Crop-Independent Emissions bands, from which the existing bands derive emit writes nine `cie-` bands beside its output: the conversion source and year, and each carbon pool the conversion puts at risk, before any destination gate, destination-specific soil factor or discount. `derive_from_cie` rebuilds every existing band from them alone, and a test checks it does on a grid covering each event, destination and soil case. Additive: the existing bands, and every consumer of them, are unchanged. `CIE_BANDS` names the bands for `export.write_tile`. docs/crop_independent_emissions.md sets the two side by side; tools/compare-emit-layers.py checks the identities on a real tile. --- docs/crop_independent_emissions.md | 70 ++++++ docs/data.md | 14 ++ jdluc/__tests__/emit_test.py | 326 ++++++++++++++++++++++++++++ jdluc/emit.py | 328 +++++++++++++++++++++++++---- tools/compare-emit-layers.py | 110 ++++++++++ 5 files changed, 810 insertions(+), 38 deletions(-) create mode 100644 docs/crop_independent_emissions.md create mode 100644 tools/compare-emit-layers.py diff --git a/docs/crop_independent_emissions.md b/docs/crop_independent_emissions.md new file mode 100644 index 0000000..a01f8ba --- /dev/null +++ b/docs/crop_independent_emissions.md @@ -0,0 +1,70 @@ +# Crop-Independent Emissions (CIE) + +`emit` writes a second set of bands, prefixed `cie-`, beside its destination-gated output (see [`data.md`](data.md#2-per-pixel-emissions-zarr)). They describe the same land conversions and carbon pools, but stop **before** anything that depends on what the land became: + +- **No destination gate.** A pixel is charged as soon as it lost a source class (forest, natural grassland or pasture) inside the 20-year window, whether or not a destination dataset claims it in the assessment year. +- **No destination-specific factor.** Mineral soil is the full stock at risk. The cropland loss fraction is not applied. +- **No discount.** Each pool is a single undiscounted value. The span split and the 20-year linear discount are left to the consumer. + +This makes the `cie-` bands a common input from which any destination- or crop-specific accounting can be derived: a consumer that wants a different soil factor per crop, a different discount, or the emissions a conversion would carry under a destination other than the one observed, can compute it from these bands without rerunning `emit`. The existing bands are one such derivation (see [Deriving the existing bands](#deriving-the-existing-bands-from-cie) below). The code is `emit.get_crop_independent_emissions`. + +## Side by side + +| quantity | existing band(s) | CIE band | difference | +| ----------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------- | +| what happened | `conversion`: source × destination, 0–5 (`emit.Conversion`) | `cie-conversion-source`: 0 none · 1 forest · 2 natural grassland · 3 pasture (`emit.ConversionSource`) | source only; set whether or not a destination resolved | +| when | `conversion-year` | `cie-conversion-year` | copy: the year is already 0 wherever there is no source | +| vegetation | `vegetation-emissions:tco2e-per-ha:{span}` ×4 | `cie-vegetation-emissions-undiscounted:tco2e-per-ha` | one band rather than one per span; not gated on destination | +| soil | `soil-emissions:tco2e-per-ha:{span}` ×4, one band holding either term: on peat, the 621 tCO2e/ha pulse; off peat, SOC × 44/12 × the cropland loss fraction for a cropland destination, or 0 for a pasture destination | two bands: `cie-mineral-soil-carbon-at-risk:tco2e-per-ha` (SOC × 44/12 off peat, 0 on peat) and `cie-peat-transformation-emissions-undiscounted:tco2e-per-ha` (621 on peat, 0 off peat) | mineral and peat kept apart; no loss fraction on mineral; neither gated on destination | +| peat occupation | `cropland-peatland-occupation:tco2e-per-ha`, `pastureland-peatland-occupation:tco2e-per-ha`: 37.3 on peat under that destination | `cie-peat-occupation-emissions:tco2e-per-ha-per-year`: 37.3 on all peat | the potential, whatever the destination and whether or not the pixel was converted | +| discounted total | `emissions-per-hectare:tco2e-per-ha` | — | the discount is applied downstream | +| unclaimed source carbon | `dropped-emissions:tco2e-per-ha` | — | these pixels are carried in the CIE pools themselves | +| destination evidence | `destination-dataset`: bitmask, 1 Descals oil palm · 2 GACED30 cropland · 4 GPW cultivated grassland | `cie-destination-dataset` | copy | +| climate zone | — (input only) | `cie-climate-zone`: `ipcc_climate_zones.Zone`, 1–10 | new: needed to look up destination-specific soil factors | +| pixel area | `hectares-per-pixel:ha` | `cie-hectares-per-pixel:ha` | copy | + +`{span}` is one of `2000-2005`, `2005-2010`, `2010-2015` and `2015-2020`. A conversion falls in span `(a, b)` when `a < year ≤ b`. The copied bands are duplicated so that the `cie-` set is complete on its own. + +## Deriving the existing bands from CIE + +Every existing band can be computed from the `cie-` bands. `emit.derive_from_cie` does so, though `jurisdictional_direct` and `statistical` still read the existing bands directly. The only extra inputs are two lookup tables already in `emit`: the soil loss fraction and the discount weights. Below, `ConversionSource` and `Conversion` refer to the band values in the table above. + +**Destination.** Cropland outranks pasture: + +- `to_cropland = cie-destination-dataset & (1 | 2)` +- `to_pasture = ¬to_cropland ∧ cie-destination-dataset & 4` + +**Conversion.** It fired where `cie-conversion-source ≠ NONE ∧ (to_cropland ∨ to_pasture)`, except pasture to pasture, which is not a conversion. The `conversion` value is the pair of the two: forest to cropland, forest to pasture, natural grassland to cropland, natural grassland to pasture, and pasture to cropland. A pasture-to-pasture pixel appears in neither the per-span pools nor `dropped-emissions`. + +**Per-span pools.** Each pool is set only where a conversion fired and `cie-conversion-year` is in the span: + +- `vegetation-emissions:{span}` = `cie-vegetation-emissions-undiscounted` +- `soil-emissions:{span}` = `cie-peat-transformation-emissions-undiscounted` + `f(cie-climate-zone)` × `cie-mineral-soil-carbon-at-risk` for a cropland destination +- `soil-emissions:{span}` = `cie-peat-transformation-emissions-undiscounted` for a pasture destination +- the per-span total, `e(span)`, is vegetation + soil (not itself a band) + +Here `f` is the cropland soil loss fraction, 1 − retention (IPCC 2019, Vol 4, Table 5.5; `emit.CLIMATE_ZONE_TO_SOC_RETENTION_FRACTION`). + +**Occupation.** Each band is gated on its destination only: + +- `cropland-peatland-occupation` = `cie-peat-occupation-emissions` where `to_cropland` +- `pastureland-peatland-occupation` = `cie-peat-occupation-emissions` where `to_pasture` + +**Discounted values.** Both use the span weights `w` = 0.0125, 0.0375, 0.0625 and 0.0875, which sum to 0.2 (`emit.SPAN_TO_LINEAR_DISCOUNT_WEIGHT`): + +- `emissions-per-hectare` = Σspan `w` × `e(span)` + both occupation bands +- `dropped-emissions` = Σspan `w` × `cie-vegetation-emissions-undiscounted`, where there is a source, the year is in the span, and neither `to_cropland` nor `to_pasture` holds + +The vegetation, soil and year identities are unit-tested in `jdluc/__tests__/emit_test.py` (`test_cie_*`), and `test_derive_from_cie_reproduces_every_band_it_replaces` checks every existing band against `derive_from_cie` on a grid covering each event, destination and soil case. `tools/compare-emit-layers.py` checks the vegetation and year identities on a real tile, and the largest difference between each existing band and its derived counterpart. It also totals the old and new pools by source, and by whether a conversion fired, which shows how much carbon the CIE bands carry that the existing bands leave uncharged. + +## Encoding + +- Every band is float32 with NaN as no-data, including the enum and bitmask bands (`geo.unify_dtype_and_no_data`). +- Bands share the tile's 30 m grid. The units follow the band name after the first `:`. +- Adding the `cie-` bands takes each tile from 16 bands to 25. The `emit` cache version is 5. + +## Not yet covered + +- **Assessment year.** It is fixed at 2020 (`emit.ASSESSMENT_YEAR`), which sets the lookback window and the destination predicates. +- **Continent.** Some destination-specific factors vary by continent, which no ingested dataset provides yet. +- **Land cover at the start of the window.** There is no band for it. A forest bit would need a year-2000 tree-cover extent, which is not currently ingested. diff --git a/docs/data.md b/docs/data.md index cd0149a..ab757d2 100644 --- a/docs/data.md +++ b/docs/data.md @@ -72,6 +72,20 @@ The crop-agnostic, per-pixel, per-span LUC emissions produced by the `emit` stag | `emissions-per-hectare:tco2e-per-ha` | float32 | t CO₂e/ha | 20-year linearly-discounted LUC emissions + peatland occupation — the headline per-hectare layer | | `hectares-per-pixel:ha` | float32 | ha | Pixel area (varies with latitude) | +Beside these, the `cie-` bands carry Crop-Independent Emissions: the same conversions and pools, stopped before anything that depends on what the land became (no destination gate, no destination-specific soil factor, no discount). Every band above can be derived from them; see [`crop_independent_emissions.md`](crop_independent_emissions.md). + +| Variable | Type | Units | Description | +| ------------------------------------------------------------- | ------- | --------------- | ---------------------------------------------------------------------------------------------------------------- | +| `cie-conversion-source` | float32 | Source code | The class the pixel left inside the window, whatever it became: 0 none, 1 forest, 2 natural grassland, 3 pasture | +| `cie-conversion-year` | float32 | year | As `conversion-year` | +| `cie-vegetation-emissions-undiscounted:tco2e-per-ha` | float32 | t CO₂e/ha | Vegetation carbon lost with the source class, undiscounted and not gated on a destination | +| `cie-mineral-soil-carbon-at-risk:tco2e-per-ha` | float32 | t CO₂e/ha | The full mineral soil-carbon stock under a converted pixel off peat, before any land-use factor | +| `cie-peat-transformation-emissions-undiscounted:tco2e-per-ha` | float32 | t CO₂e/ha | The peatland drainage pulse under a converted pixel on peat, undiscounted | +| `cie-peat-occupation-emissions:tco2e-per-ha-per-year` | float32 | t CO₂e/ha/yr | Annual peatland-occupation emissions on all peat, whatever occupies it | +| `cie-climate-zone` | float32 | Zone code | IPCC climate zone, for looking up destination-specific soil factors | +| `cie-destination-dataset` | float32 | Dataset bitmask | As `destination-dataset` | +| `cie-hectares-per-pixel:ha` | float32 | ha | As `hectares-per-pixel:ha` | + ```python import xarray diff --git a/jdluc/__tests__/emit_test.py b/jdluc/__tests__/emit_test.py index 87dc466..c01032d 100644 --- a/jdluc/__tests__/emit_test.py +++ b/jdluc/__tests__/emit_test.py @@ -1,21 +1,27 @@ import collections import collections.abc +import itertools import numpy import pytest import xarray +from jdluc import emit from jdluc.datasets.descals_oil_palm import DATASET as OIL_PALM from jdluc.datasets.gnw_global_peatlands import DATASET as PEATLANDS +from jdluc.datasets.gnw_harris_agb import DATASET as ABOVEGROUND_BIOMASS from jdluc.datasets.gnw_tcl import DATASET as TREE_COVER_LOSS from jdluc.datasets.gnw_tcl import LOSS_YEAR_OFFSET from jdluc.datasets.gpw_grassland import DATASET as GRASSLAND from jdluc.datasets.gpw_grassland import YEARS as GRASSLAND_YEARS from jdluc.datasets.gpw_grassland import Grassland +from jdluc.datasets.huang_bgb import DATASET as BELOWGROUND_BIOMASS +from jdluc.datasets.ipcc_climate_zones import DATASET as CLIMATE_ZONES from jdluc.datasets.ipcc_climate_zones import Zone from jdluc.datasets.liao_gaced30 import DATASET as CROPLAND from jdluc.datasets.liao_gaced30 import YEARS as CROPLAND_YEARS from jdluc.datasets.liao_gaced30 import Cropland +from jdluc.datasets.soilgrids_ocs import DATASET as SOIL_ORGANIC_CARBON_STOCK from jdluc.emit import ( ASSESSMENT_YEAR, CARBON_PER_BIOMASS_LIVE_WOOD, @@ -29,12 +35,15 @@ Conversion, ConversionEmissions, ConversionRecord, + ConversionSource, DestinationDataset, EmissionComponent, SpanType, + derive_from_cie, get_belowground_carbon, get_conversion_emissions, get_conversion_record, + get_crop_independent_emissions, get_dead_organic_matter_carbon, get_dropped_emissions, get_dset_for_output, @@ -43,6 +52,7 @@ get_last_departure_year, get_linear_discounted_total, get_mineral_soil_emissions, + get_output_dset, get_peatland_occupation_emissions, get_span_to_charge, get_span_to_component_to_emissions, @@ -148,6 +158,17 @@ def test_get_mineral_soil_emissions() -> None: ) +CONVERSION_TO_SOURCE: dict[Conversion, ConversionSource] = { + Conversion.NONE: ConversionSource.NONE, + Conversion.FOREST_TO_CROPLAND: ConversionSource.FOREST, + Conversion.FOREST_TO_PASTURE: ConversionSource.FOREST, + Conversion.RANGELAND_TO_CROPLAND: ConversionSource.NATURAL_GRASSLAND, + Conversion.RANGELAND_TO_PASTURE: ConversionSource.NATURAL_GRASSLAND, + Conversion.PASTURE_TO_CROPLAND: ConversionSource.PASTURE, +} +assert set(CONVERSION_TO_SOURCE) == set(Conversion) + + def get_record_for(conversion: Conversion, width: int = 1) -> ConversionRecord: """A record holding one conversion, with its halves read off the axis tuples. @@ -184,6 +205,7 @@ def mask(members: tuple[Conversion, ...]) -> xarray.DataArray: from_rangeland=mask( (Conversion.RANGELAND_TO_CROPLAND, Conversion.RANGELAND_TO_PASTURE) ), + source=get_darray_for_data(data=[[CONVERSION_TO_SOURCE[conversion]] * width]), to_cropland=mask(to_cropland), to_pasture=mask(to_pasture), year=get_darray_for_data(data=[[ASSESSMENT_YEAR] * width]), @@ -223,6 +245,11 @@ def mask(value: bool) -> xarray.DataArray: from_forest=mask(from_forest), from_pasture=mask(False), from_rangeland=mask(False), + source=get_darray_for_data( + data=[ + [ConversionSource.FOREST if from_forest else ConversionSource.NONE] + ] + ), to_cropland=mask(to_cropland), to_pasture=mask(False), year=get_darray_for_data(data=[[ASSESSMENT_YEAR]]), @@ -1115,6 +1142,16 @@ def test_the_halves_a_conversion_record_resolves( ) assert bool(result.from_forest) == from_forest assert bool(result.from_rangeland | result.from_pasture) == from_grassland + # NB: the source is the same halves, read before any destination gates them + assert bool(result.source == ConversionSource.FOREST) == from_forest + assert ( + bool( + result.source.isin( + (ConversionSource.NATURAL_GRASSLAND, ConversionSource.PASTURE) + ) + ) + == from_grassland + ) assert bool(result.to_cropland | result.to_pasture) == has_destination @@ -1159,3 +1196,292 @@ def test_a_conversion_record_tells_pasture_from_rangeland_sources( ) assert bool(result.from_pasture) == from_pasture assert bool(result.from_rangeland) == from_rangeland + + +@pytest.mark.parametrize( + ("grassland_by_year", "loss_year", "source"), + ( + pytest.param(NO_GRASSLAND, 2012, ConversionSource.FOREST, id="forest loss"), + pytest.param( + [ + Grassland.NATURAL if year <= 2011 else Grassland.OTHER + for year in LOOKBACK_YEARS_RANGE + ], + None, + ConversionSource.NATURAL_GRASSLAND, + id="natural grassland departed", + ), + pytest.param( + [ + Grassland.CULTIVATED if year <= 2011 else Grassland.OTHER + for year in LOOKBACK_YEARS_RANGE + ], + None, + ConversionSource.PASTURE, + id="pasture departed, which CIE keeps apart from natural grassland", + ), + pytest.param(NO_GRASSLAND, None, ConversionSource.NONE, id="nothing departed"), + ), +) +def test_get_conversion_record_source( + grassland_by_year: list[float], loss_year: float | None, source: ConversionSource +) -> None: + # No destination claims the pixel, so the conversion is NONE; the source is reported anyway + result = get_conversion_record( + dset=get_dset_for_one_pixel( + grassland_by_year=grassland_by_year, + is_cropland=False, + loss_year=loss_year, + planting_year=0, + ) + ) + assert result.conversion.data == [[Conversion.NONE]] + assert result.source.data == [[source]] + # NB: which is what lets `cie-conversion-year` copy the year rather than gate it + assert bool(result.year == 0) == (source == ConversionSource.NONE) + + +def get_crop_independent_emissions_for( + conversion_record: ConversionRecord, is_peatland: bool = False +) -> dict[str, xarray.DataArray]: + return get_crop_independent_emissions( + climate_zones=get_darray_for_data(data=[[Zone.TROPICAL_DRY.value]]), + conversion_record=conversion_record, + forest_carbon=get_darray_for_data(data=[[FOREST_TCARBON_PER_HA]]), + grassland_carbon=get_darray_for_data(data=[[GRASSLAND_TCARBON_PER_HA]]), + hectares_per_pixel=get_darray_for_data(data=[[1.0]]), + is_peatland=get_darray_for_data(data=[[is_peatland]]), + soil_organic_carbon=get_darray_for_data(data=[[SOIL_ORGANIC_CARBON]]), + ) + + +@pytest.mark.parametrize("is_peatland", (False, True)) +@pytest.mark.parametrize("conversion", tuple(Conversion)) +def test_cie_reproduces_the_pools_a_fired_conversion_releases( + conversion: Conversion, is_peatland: bool +) -> None: + """Where a conversion fired, the old pools are CIE's with the destination applied. + + Vegetation is CIE's as it stands; soil is the peat pulse on peat, and otherwise the + mineral stock at risk scaled by the cropland loss fraction, or by zero for pasture. + """ + record = get_record_for(conversion=conversion) + old = get_pools_for(conversion=conversion, is_peatland=is_peatland) + new = get_crop_independent_emissions_for( + conversion_record=record, is_peatland=is_peatland + ) + + assert new["cie-conversion-source"].data == [[CONVERSION_TO_SOURCE[conversion]]] + if conversion == Conversion.NONE: + return + numpy.testing.assert_allclose( + old.vegetation.data, new["cie-vegetation-emissions-undiscounted"].data + ) + soil_loss_fraction = ( + 1 - CLIMATE_ZONE_TO_SOC_RETENTION_FRACTION[Zone.TROPICAL_DRY] + if bool(record.to_cropland) + else 0.0 + ) + numpy.testing.assert_allclose( + old.soil.data, + new["cie-peat-transformation-emissions-undiscounted"].data + + soil_loss_fraction * new["cie-mineral-soil-carbon-at-risk"].data, + rtol=1e-6, + ) + + +def test_cie_charges_a_source_no_destination_claimed() -> None: + # The old bands charge this pixel nothing and report its vegetation as dropped; CIE + # carries the whole of what is at risk, leaving the destination to the consumer + record = ConversionRecord( + conversion=get_darray_for_data(data=[[Conversion.NONE]]), + destination_dataset=get_darray_for_data(data=[[0]]), + from_forest=get_darray_for_data(data=[[True]]), + from_pasture=get_darray_for_data(data=[[False]]), + from_rangeland=get_darray_for_data(data=[[False]]), + source=get_darray_for_data(data=[[ConversionSource.FOREST]]), + to_cropland=get_darray_for_data(data=[[False]]), + to_pasture=get_darray_for_data(data=[[False]]), + year=get_darray_for_data(data=[[2012]]), + ) + result = get_crop_independent_emissions_for(conversion_record=record) + dropped = get_dropped_emissions( + forest_carbon=get_darray_for_data(data=[[FOREST_TCARBON_PER_HA]]), + from_forest=record.from_forest, + from_pasture=record.from_pasture, + from_rangeland=record.from_rangeland, + grassland_carbon=get_darray_for_data(data=[[GRASSLAND_TCARBON_PER_HA]]), + has_destination=get_darray_for_data(data=[[False]]), + ) + numpy.testing.assert_allclose( + result["cie-vegetation-emissions-undiscounted"].data, dropped.data + ) + numpy.testing.assert_allclose( + result["cie-mineral-soil-carbon-at-risk"].data, + [[SOIL_ORGANIC_CARBON * CO2E_PER_CARBON]], + ) + assert result["cie-conversion-year"].data == [[2012]] + + +def test_cie_without_a_source() -> None: + # With no source, CIE prices only the occupation potential of the peat underneath + result = get_crop_independent_emissions_for( + conversion_record=get_record_for(conversion=Conversion.NONE), is_peatland=True + ) + for name in ( + "cie-mineral-soil-carbon-at-risk", + "cie-peat-transformation-emissions-undiscounted", + "cie-vegetation-emissions-undiscounted", + ): + assert float(result[name].sum()) == 0.0, name + numpy.testing.assert_allclose( + result["cie-peat-occupation-emissions"].data, + [[PEATLAND_EMISSIONS_ANNUAL_TCO2E_PER_HA]], + ) + + +def test_cie_names_its_units() -> None: + dset = get_dset_for_output( + name_to_darray=get_crop_independent_emissions_for( + conversion_record=get_record_for(conversion=Conversion.FOREST_TO_CROPLAND) + ) + ) + assert set(dset) == { + "cie-climate-zone", + "cie-conversion-source", + "cie-conversion-year", + "cie-destination-dataset", + "cie-hectares-per-pixel", + "cie-mineral-soil-carbon-at-risk:tco2e-per-ha", + "cie-peat-occupation-emissions:tco2e-per-ha-per-year", + "cie-peat-transformation-emissions-undiscounted:tco2e-per-ha", + "cie-vegetation-emissions-undiscounted:tco2e-per-ha", + } + + +# Every (loss, grassland sequence, destination bits) class +LOSS_YEARS: tuple[int | None, ...] = ( + None, + LOOKBACK_YEARS_RANGE[0], + 2003, + 2012, + ASSESSMENT_YEAR, + ASSESSMENT_YEAR + 4, +) +GRASSLAND_SEQUENCES: tuple[collections.abc.Callable[[int], Grassland], ...] = ( + lambda year: Grassland.OTHER, + lambda year: Grassland.NATURAL if year <= 2011 else Grassland.OTHER, + lambda year: Grassland.CULTIVATED if year <= 2011 else Grassland.OTHER, + lambda year: ( + Grassland.CULTIVATED + if year <= 2005 + else Grassland.NATURAL + if year <= 2014 + else Grassland.OTHER + ), + lambda year: ( + Grassland.NATURAL + if year <= 2005 + else Grassland.CULTIVATED + if year <= 2014 + else Grassland.OTHER + ), +) + + +def get_harmonized_dset_for_every_case() -> xarray.Dataset: + """One column per event class; one row per soil case: mineral, peat, no peat data, no data. + + The last row has no soil carbon and no climate zone, which is where NaN handling differs. + """ + columns = list( + itertools.product( + LOSS_YEARS, + GRASSLAND_SEQUENCES, + itertools.product((False, True), repeat=3), + ) + ) + zones = list(Zone) + height, width = 4, len(columns) + grassland = numpy.empty((len(LOOKBACK_YEARS_RANGE), height, width)) + loss = numpy.empty((height, width)) + planting_year = numpy.empty((height, width)) + cropland = numpy.empty((height, width)) + for x, ( + loss_year, + get_class, + (is_oil_palm, is_cropland, is_cultivated), + ) in enumerate(columns): + for offset, year in enumerate(LOOKBACK_YEARS_RANGE): + grassland[offset, :, x] = get_class(year) + if is_cultivated: + grassland[-1, :, x] = Grassland.CULTIVATED + loss[:, x] = 0 if loss_year is None else loss_year - LOSS_YEAR_OFFSET + planting_year[:, x] = 2012 if is_oil_palm else 0 + cropland[:, x] = Cropland.CROPLAND if is_cropland else Cropland.NOT_CROPLAND + + def darray(data: numpy.ndarray) -> xarray.DataArray: + return xarray.DataArray( + coords={"y": numpy.arange(height) * -0.5, "x": numpy.arange(width) * 0.5}, + data=data, + dims=("y", "x"), + ) + + xs = numpy.arange(width) + climate_zone = numpy.tile([float(zones[x % len(zones)]) for x in xs], (height, 1)) + climate_zone[3] = numpy.nan + soil_organic_carbon = numpy.tile(50.0 + xs, (height, 1)) + soil_organic_carbon[3] = numpy.nan + peat = numpy.array([[0.0], [1.0], [numpy.nan], [0.0]]) * numpy.ones(width) + aboveground_biomass = numpy.tile(100.0 + 3 * xs, (height, 1)) + # NB: zero where the belowground fallback to the root-to-shoot ratio applies + belowground_biomass = numpy.tile(numpy.where(xs % 2, 20.0 + xs, 0.0), (height, 1)) + return xarray.Dataset( + { + GRASSLAND.fully_qualified_band_names[GRASSLAND_YEARS.index(year)]: darray( + grassland[offset] + ) + for offset, year in enumerate(LOOKBACK_YEARS_RANGE) + } + | { + ABOVEGROUND_BIOMASS.fully_qualified_band_name: darray(aboveground_biomass), + BELOWGROUND_BIOMASS.fully_qualified_band_name: darray(belowground_biomass), + CLIMATE_ZONES.fully_qualified_band_name: darray(climate_zone), + CROPLAND.fully_qualified_band_names[ + CROPLAND_YEARS.index(ASSESSMENT_YEAR) + ]: darray(cropland), + OIL_PALM.fully_qualified_band_name: darray(planting_year), + PEATLANDS.fully_qualified_band_name: darray(peat), + SOIL_ORGANIC_CARBON_STOCK.fully_qualified_band_name: darray( + soil_organic_carbon + ), + TREE_COVER_LOSS.fully_qualified_band_name: darray(loss), + } + ) + + +def test_derive_from_cie_reproduces_every_band_it_replaces() -> None: + output = get_output_dset(dset=get_harmonized_dset_for_every_case()).compute() + derived = derive_from_cie(dset=output).compute() + + assert set(derived) == { + name for name in map(str, output) if not name.startswith("cie-") + } + for name in derived: + numpy.testing.assert_allclose( + derived[name].data, + output[name].data, + atol=1e-4, + err_msg=name, + rtol=1e-6, + ) + # Every conversion, and both a charged and a dropped pixel, occur in the grid + assert set(numpy.unique(output["conversion"].data)) == set(map(float, Conversion)) + assert float(output["dropped-emissions:tco2e-per-ha"].sum()) > 0 + + +def test_cie_bands_names_every_cie_band() -> None: + dset = emit.get_output_dset(dset=get_harmonized_dset_for_every_case()) + assert list(emit.CIE_BANDS) == sorted( + name for name in map(str, dset.data_vars) if name.startswith("cie-") + ) diff --git a/jdluc/emit.py b/jdluc/emit.py index 59ffbeb..562a6e4 100644 --- a/jdluc/emit.py +++ b/jdluc/emit.py @@ -7,7 +7,9 @@ Returns a cached xarray.Dataset: the conversion, the year its source class ended and the datasets that claimed its destination; vegetation and soil emissions per span; the two occupation bands; the discounted per-hectare total; the source carbon no destination claimed; and a -hectares-per-pixel band for downstream area-scaling. +hectares-per-pixel band for downstream area-scaling. Beside them, the `cie-` bands carry +Crop-Independent Emissions (CIE), from which every band above can be derived; see +`get_crop_independent_emissions`, `derive_from_cie` and `docs/crop_independent_emissions.md`. Source carbon that no destination claimed is charged to nobody and reported on its own as `dropped-emissions`. @@ -285,6 +287,16 @@ def _generate_next_value_( assert set(TO_CROPLAND) ^ set(TO_PASTURE) == set(Conversion) - {Conversion.NONE} +@enum.unique +class ConversionSource(enum.IntEnum): + """The class a pixel left, whatever -- if anything -- it became: CIE's conversion source.""" + + NONE = 0 + FOREST = 1 + NATURAL_GRASSLAND = 2 + PASTURE = 3 + + class DestinationDataset(enum.IntFlag): """Which datasets (the can superpose) claim a pixel's land class at the assessment year.""" @@ -358,6 +370,8 @@ class ConversionRecord: from_forest: xarray.DataArray from_pasture: xarray.DataArray from_rangeland: xarray.DataArray + # NB: the source alone, before any destination gates it -- what CIE reports + source: xarray.DataArray to_cropland: xarray.DataArray to_pasture: xarray.DataArray year: xarray.DataArray @@ -436,6 +450,13 @@ def get_conversion_record(dset: xarray.Dataset) -> ConversionRecord: from_forest=from_forest, from_pasture=from_pasture, from_rangeland=from_rangeland, + source=( + from_forest * numpy.uint8(ConversionSource.FOREST) + + from_rangeland * numpy.uint8(ConversionSource.NATURAL_GRASSLAND) + + from_pasture * numpy.uint8(ConversionSource.PASTURE) + ) + .astype(numpy.uint8) + .rename(None), to_cropland=to_cropland, to_pasture=to_pasture, # NB: a forest source is dated by its loss, a grassland one by its last departure @@ -538,6 +559,57 @@ def get_peatland_occupation_emissions( ).rename("tco2e-per-ha") +def get_crop_independent_emissions( + climate_zones: xarray.DataArray, + conversion_record: ConversionRecord, + forest_carbon: xarray.DataArray, + grassland_carbon: xarray.DataArray, + hectares_per_pixel: xarray.DataArray, + is_peatland: xarray.DataArray, + soil_organic_carbon: xarray.DataArray, +) -> dict[str, xarray.DataArray]: + """Crop-Independent Emissions (CIE); see `docs/crop_independent_emissions.md`. + + Gated on the source alone: no destination is asked, nothing is discounted, and no land-use + factor is applied, so each pool is the stock the conversion puts at risk. A consumer applies + its own destination- or crop-specific factors and discount, as `derive_from_cie` does for + the bands above. + """ + has_source = conversion_record.source != ConversionSource.NONE + return { + "cie-conversion-source": conversion_record.source, + # NB: already 0 wherever there is no source, so a copy + "cie-conversion-year": conversion_record.year, + # NB: peat replaces the mineral term, as it does in `get_conversion_emissions` + "cie-mineral-soil-carbon-at-risk": ( + CO2E_PER_CARBON * soil_organic_carbon.fillna(0) + ) + .where(has_source & ~is_peatland, other=0) + .rename("tco2e-per-ha"), + "cie-peat-transformation-emissions-undiscounted": ( + PEATLAND_EMISSIONS_PULSE_TCO2E_PER_HA + * (has_source & is_peatland).astype(numpy.float32) + ).rename("tco2e-per-ha"), + "cie-vegetation-emissions-undiscounted": ( + CO2E_PER_CARBON + * ( + forest_carbon.where(conversion_record.from_forest, other=0) + + grassland_carbon.where( + conversion_record.from_rangeland | conversion_record.from_pasture, + other=0, + ) + ) + ).rename("tco2e-per-ha"), + # NB: the potential, on all peat whether or not it was converted or is occupied + "cie-peat-occupation-emissions": ( + PEATLAND_EMISSIONS_ANNUAL_TCO2E_PER_HA * is_peatland.astype(numpy.float32) + ).rename("tco2e-per-ha-per-year"), + "cie-climate-zone": climate_zones.rename(None), + "cie-hectares-per-pixel": hectares_per_pixel, + "cie-destination-dataset": conversion_record.destination_dataset, + } + + def get_span_to_charge( conversion_year: xarray.DataArray, darray: xarray.DataArray ) -> dict[SpanType, xarray.DataArray]: @@ -675,26 +747,187 @@ def merge_name_units(name: str, units: typing.Hashable | None) -> str: ) -@storage.cache_to_zarr(version=1) -def workflow(tile_id: str) -> xarray.Dataset: - logger.info(f"Running the land conversion and emissions worflow for {tile_id=:s}") - dset = geo.exact_merge( - harmonize.workflow( - dataset_names=harmonize.Stack.LUC_AND_EMISSIONS.value, - ignore_missing_tiles=True, - skip_ingest=False, - tile_id=tile_id, - tile_resolution=tiling.TileResolution.GLAD, +def get_cropland_soil_loss_fraction( + climate_zones: xarray.DataArray, +) -> xarray.DataArray: + """1 - retention, per zone (IPCC 2019 Vol 4 Table 5.5); zero where the zone is missing.""" + lookup = numpy.zeros(256, dtype=numpy.float32) + for climate_zone, retention in CLIMATE_ZONE_TO_SOC_RETENTION_FRACTION.items(): + lookup[climate_zone.value] = 1 - retention + return xarray.apply_ufunc( + lookup.__getitem__, + climate_zones.fillna(0).astype(numpy.uint8), + dask="parallelized", + output_dtypes=[numpy.float32], + ) + + +def derive_from_cie(dset: xarray.Dataset) -> xarray.Dataset: + """The destination-gated, discounted bands, rebuilt from `emit`'s `cie-` bands alone. + + Reproduces the non-`cie-` bands of `workflow` (see `docs/crop_independent_emissions.md`): + gate on the destination, scale mineral soil by the cropland loss fraction (zero for + pasture), and charge each pool to its span. + """ + + def cie(name: str) -> xarray.DataArray: + return dset[f"cie-{name:s}"] + + source = cie("conversion-source").fillna(0).astype(numpy.uint8) + destination_dataset = cie("destination-dataset").fillna(0).astype(numpy.uint8) + year = cie("conversion-year") + to_cropland = (destination_dataset & CROPLAND_DATASETS).astype(bool) + to_pasture = ~to_cropland & (destination_dataset & PASTURE_DATASETS).astype(bool) + has_destination = to_cropland | to_pasture + + source_to_destination_conversion = { + ConversionSource.FOREST: ( + Conversion.FOREST_TO_CROPLAND, + Conversion.FOREST_TO_PASTURE, ), - harmonize.workflow( - dataset_names=harmonize.Stack.CROP_SUPPLEMENT.value, - ignore_missing_tiles=True, - skip_ingest=False, - tile_id=tile_id, - tile_resolution=tiling.TileResolution.GLAD, + ConversionSource.NATURAL_GRASSLAND: ( + Conversion.RANGELAND_TO_CROPLAND, + Conversion.RANGELAND_TO_PASTURE, ), + # NB: pasture to pasture is not a conversion + ConversionSource.PASTURE: (Conversion.PASTURE_TO_CROPLAND, Conversion.NONE), + } + conversion = sum( + (source == member) + * ( + to_cropland * numpy.uint8(to_cropland_conversion) + + to_pasture * numpy.uint8(to_pasture_conversion) + ) + for member, ( + to_cropland_conversion, + to_pasture_conversion, + ) in source_to_destination_conversion.items() + ) + assert isinstance(conversion, xarray.DataArray) + conversion = conversion.astype(numpy.uint8) + fired = conversion != Conversion.NONE + + vegetation = cie("vegetation-emissions-undiscounted:tco2e-per-ha") + peat_transformation = cie("peat-transformation-emissions-undiscounted:tco2e-per-ha") + soil = ( + peat_transformation + + ( + get_cropland_soil_loss_fraction(climate_zones=cie("climate-zone")) + * cie("mineral-soil-carbon-at-risk:tco2e-per-ha") + ).where(to_cropland, other=0) + ).where(fired, other=0) + span_to_vegetation_emissions = get_span_to_charge( + conversion_year=year, + darray=vegetation.where(fired, other=0).rename("tco2e-per-ha"), + ) + span_to_soil_emissions = get_span_to_charge( + conversion_year=year, darray=soil.rename("tco2e-per-ha") + ) + span_to_emissions = { + span: ( + span_to_vegetation_emissions[span] + span_to_soil_emissions[span] + ).rename("tco2e-per-ha") + for span in span_to_vegetation_emissions + } + + peat_occupation = cie("peat-occupation-emissions:tco2e-per-ha-per-year") + cropland_occupation_emissions = peat_occupation.where(to_cropland, other=0).rename( + "tco2e-per-ha" + ) + pastureland_occupation_emissions = peat_occupation.where( + to_pasture, other=0 + ).rename("tco2e-per-ha") + emissions_per_hectare = ( + get_linear_discounted_total(span_to_value=span_to_emissions) + + cropland_occupation_emissions + + pastureland_occupation_emissions + ).rename("tco2e-per-ha") + dropped_emissions = get_linear_discounted_total( + span_to_value=get_span_to_charge( + conversion_year=year, darray=vegetation.where(~has_destination, other=0) + ) + ).rename("tco2e-per-ha") + + return get_dset_for_output( + name_to_darray=get_name_to_legacy_darray( + conversion=conversion.rename(None), + conversion_year=year.rename(None), + cropland_occupation_emissions=cropland_occupation_emissions, + destination_dataset=cie("destination-dataset").rename(None), + dropped_emissions=dropped_emissions, + emissions_per_hectare=emissions_per_hectare, + hectares_per_pixel=cie("hectares-per-pixel:ha").rename("ha"), + pastureland_occupation_emissions=pastureland_occupation_emissions, + span_to_soil_emissions=span_to_soil_emissions, + span_to_vegetation_emissions=span_to_vegetation_emissions, + ) + ) + + +def get_name_to_legacy_darray( + conversion: xarray.DataArray, + conversion_year: xarray.DataArray, + cropland_occupation_emissions: xarray.DataArray, + destination_dataset: xarray.DataArray, + dropped_emissions: xarray.DataArray, + emissions_per_hectare: xarray.DataArray, + hectares_per_pixel: xarray.DataArray, + pastureland_occupation_emissions: xarray.DataArray, + span_to_soil_emissions: dict[SpanType, xarray.DataArray], + span_to_vegetation_emissions: dict[SpanType, xarray.DataArray], +) -> dict[str, xarray.DataArray]: + """The non-`cie-` bands, named as `workflow` writes them.""" + return ( + { + "conversion": conversion, + "conversion-year": conversion_year, + "destination-dataset": destination_dataset, + } + | { + f"soil-emissions:{before:d}-{after:d}": darray + for (before, after), darray in span_to_soil_emissions.items() + } + | { + f"vegetation-emissions:{before:d}-{after:d}": darray + for (before, after), darray in span_to_vegetation_emissions.items() + } + | { + f"{CROPLAND:s}-peatland-occupation": cropland_occupation_emissions, + # NB: charged to nobody, so it is reported beside the total rather than inside it + "dropped-emissions": dropped_emissions, + "emissions-per-hectare": emissions_per_hectare, + "hectares-per-pixel": hectares_per_pixel, + f"{PASTURELAND:s}-peatland-occupation": pastureland_occupation_emissions, + } + ) + + +# NB: not 2-4, which unmerged branches have already written +@storage.cache_to_zarr(version=5) +def workflow(tile_id: str) -> xarray.Dataset: + logger.info(f"Running the land conversion and emissions worflow for {tile_id=:s}") + return get_output_dset( + dset=geo.exact_merge( + harmonize.workflow( + dataset_names=harmonize.Stack.LUC_AND_EMISSIONS.value, + ignore_missing_tiles=True, + skip_ingest=False, + tile_id=tile_id, + tile_resolution=tiling.TileResolution.GLAD, + ), + harmonize.workflow( + dataset_names=harmonize.Stack.CROP_SUPPLEMENT.value, + ignore_missing_tiles=True, + skip_ingest=False, + tile_id=tile_id, + tile_resolution=tiling.TileResolution.GLAD, + ), + ) ) + +def get_output_dset(dset: xarray.Dataset) -> xarray.Dataset: + """`workflow` on an already-harmonized dataset.""" logger.info("Resolving which conversion each pixel underwent, and when") conversion_record: ConversionRecord = get_conversion_record(dset=dset) @@ -774,31 +1007,50 @@ def workflow(tile_id: str) -> xarray.Dataset: ) ).rename("tco2e-per-ha") + logger.info("Assembling the Crop-Independent Emissions (CIE)") + hectares_per_pixel = get_hectares_per_pixel(darray=emissions_per_hectare) + cie = get_crop_independent_emissions( + climate_zones=climate_zones, + conversion_record=conversion_record, + forest_carbon=forest_carbon, + grassland_carbon=grassland_carbon, + hectares_per_pixel=hectares_per_pixel, + is_peatland=is_peatland, + soil_organic_carbon=dset[soilgrids_ocs.DATASET.fully_qualified_band_name], + ) + return get_dset_for_output( - name_to_darray={ - "conversion": conversion_record.conversion, - "conversion-year": conversion_record.year, - "destination-dataset": conversion_record.destination_dataset, - } - | { - f"soil-emissions:{before:d}-{after:d}": darray - for (before, after), darray in span_to_soil_emissions.items() - } - | { - f"vegetation-emissions:{before:d}-{after:d}": darray - for (before, after), darray in span_to_vegetation_emissions.items() - } - | { - f"{CROPLAND:s}-peatland-occupation": cropland_occupation_emissions, - # NB: charged to nobody, so it is reported beside the total rather than inside it - "dropped-emissions": dropped_emissions, - "emissions-per-hectare": emissions_per_hectare, - "hectares-per-pixel": get_hectares_per_pixel(darray=emissions_per_hectare), - f"{PASTURELAND:s}-peatland-occupation": pastureland_occupation_emissions, - } + name_to_darray=get_name_to_legacy_darray( + conversion=conversion_record.conversion, + conversion_year=conversion_record.year, + cropland_occupation_emissions=cropland_occupation_emissions, + destination_dataset=conversion_record.destination_dataset, + dropped_emissions=dropped_emissions, + emissions_per_hectare=emissions_per_hectare, + hectares_per_pixel=hectares_per_pixel, + pastureland_occupation_emissions=pastureland_occupation_emissions, + span_to_soil_emissions=span_to_soil_emissions, + span_to_vegetation_emissions=span_to_vegetation_emissions, + ) + | cie ) +# The Crop-Independent Emissions, as `get_crop_independent_emissions` names them in the output: +# the variables to pass to `export.write_tile` +CIE_BANDS = ( + "cie-climate-zone", + "cie-conversion-source", + "cie-conversion-year", + "cie-destination-dataset", + "cie-hectares-per-pixel:ha", + "cie-mineral-soil-carbon-at-risk:tco2e-per-ha", + "cie-peat-occupation-emissions:tco2e-per-ha-per-year", + "cie-peat-transformation-emissions-undiscounted:tco2e-per-ha", + "cie-vegetation-emissions-undiscounted:tco2e-per-ha", +) + + def main() -> int: logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", diff --git a/tools/compare-emit-layers.py b/tools/compare-emit-layers.py new file mode 100644 index 0000000..29c253a --- /dev/null +++ b/tools/compare-emit-layers.py @@ -0,0 +1,110 @@ +"""Contrast emit's existing bands with its `cie-` bands (Crop-Independent Emissions) on one tile. + +Pixels are grouped by CIE conversion source and by whether the old record fired a conversion +(i.e. a destination claimed the pixel). For each group: pixel count, hectares, and the hectare- +weighted totals of the old undiscounted pools (summed over spans) beside CIE's. Where a +conversion fired, the old vegetation must equal CIE's and the year must match; the residuals +are printed so any disagreement is visible. Where none fired, CIE carries emissions the old +bands charge to nobody (and report, discounted, as `dropped-emissions`). + +Last, every existing band is rebuilt from the `cie-` bands by `emit.derive_from_cie`, and the +largest absolute difference from the stored band is printed; each should be ~0. + +Must run from the repo root, so Config finds .env. + + uv run python tools/compare-emit-layers.py 10N_010W + uv run python tools/compare-emit-layers.py --zarr /path/to/emit.zarr +""" + +import argparse + +import dask +import xarray + +from jdluc import emit, storage + +SPANS = [ + f"{before:d}-{after:d}" for before, after in emit.SPAN_TO_LINEAR_DISCOUNT_WEIGHT +] +TCO2E = "tco2e-per-ha" + + +def get_rows_checks( + dset: xarray.Dataset, +) -> tuple[list[dict[str, object]], dict[str, float]]: + source = dset["cie-conversion-source"] + fired = dset["conversion"] != emit.Conversion.NONE + ha = dset["hectares-per-pixel:ha"] + old_vegetation = sum(dset[f"vegetation-emissions:{TCO2E}:{span}"] for span in SPANS) + old_soil = sum(dset[f"soil-emissions:{TCO2E}:{span}"] for span in SPANS) + new_vegetation = dset[f"cie-vegetation-emissions-undiscounted:{TCO2E}"] + new_mineral = dset[f"cie-mineral-soil-carbon-at-risk:{TCO2E}"] + new_peat = dset[f"cie-peat-transformation-emissions-undiscounted:{TCO2E}"] + + lazy: list[dict[str, object]] = [] + for member in emit.ConversionSource: + for is_fired in (True, False): + mask = (source == member) & (fired if is_fired else ~fired) + lazy.append( + { + "source": member.name.lower(), + "fired": is_fired, + "pixels": mask.sum(), + "ha": ha.where(mask).sum(), + "old veg t": (ha * old_vegetation).where(mask).sum(), + "cie veg t": (ha * new_vegetation).where(mask).sum(), + "old soil t": (ha * old_soil).where(mask).sum(), + "cie mineral t": (ha * new_mineral).where(mask).sum(), + "cie peat t": (ha * new_peat).where(mask).sum(), + } + ) + checks = { + "max |old veg - cie veg| where fired": abs(old_vegetation - new_vegetation) + .where(fired) + .max(), + "year mismatches where fired": ( + dset["conversion-year"] != dset["cie-conversion-year"] + ) + .where(fired, other=False) + .sum(), + "fired pixels with no CIE source": ( + fired & (source == emit.ConversionSource.NONE) + ).sum(), + } + derived = emit.derive_from_cie(dset=dset) + checks |= { + f"max |derived - existing| {name:s}": abs(derived[name] - dset[name]).max() + for name in map(str, derived.data_vars) + } + rows, checks = dask.compute(lazy, checks) + return rows, {name: float(value) for name, value in checks.items()} + + +def print_rows_checks(rows: list[dict[str, object]], checks: dict[str, float]) -> None: + print("| " + " | ".join(rows[0]) + " |") + print("|" + "---|" * len(rows[0])) + for row in rows: + cells = (v if isinstance(v, str | bool) else f"{v:,.0f}" for v in row.values()) + print("| " + " | ".join(map(str, cells)) + " |") + print() + for name, value in checks.items(): + print(f"- {name}: {value:g}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + group = parser.add_mutually_exclusive_group(required=True) + group.add_argument("tile_id", nargs="?", help="run (or read the cache of) emit") + group.add_argument("--zarr", help="an emit zarr already on disk") + args = parser.parse_args() + dset = ( + storage.open_zarr_to_dask_dataset(path_to_zarr=args.zarr) + if args.zarr + else emit.workflow(tile_id=args.tile_id) + ) + print_rows_checks(*get_rows_checks(dset=dset)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())