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/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 39e6751..ab757d2 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 @@ -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/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__/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__/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/__tests__/export_test.py b/jdluc/__tests__/export_test.py new file mode 100644 index 0000000..3464511 --- /dev/null +++ b/jdluc/__tests__/export_test.py @@ -0,0 +1,232 @@ +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_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"] + + +@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_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: + 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 + 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( + 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/__tests__/storage_test.py b/jdluc/__tests__/storage_test.py new file mode 100644 index 0000000..34a3953 --- /dev/null +++ b/jdluc/__tests__/storage_test.py @@ -0,0 +1,44 @@ +import pathlib + +import numpy +import pytest +import xarray + +from jdluc import storage + + +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/__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/__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/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/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/jdluc/export.py b/jdluc/export.py new file mode 100644 index 0000000..8f70ec5 --- /dev/null +++ b/jdluc/export.py @@ -0,0 +1,250 @@ +"""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, 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: + +- 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 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. +""" + +import collections.abc +import enum +import itertools +import logging +import os +import tempfile + +import numpy +import pandas +import xarray + +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, 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_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=get_export_root(output_root=output_root, version=version), + 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. + + 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)) + 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. 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) + 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, + 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 := [ + 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, + version=version, + ) + metadata = get_metadata( + product_name=name, source_name=source_name or name, version=version + ) + 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, + version: str = "", +) -> str: + """`{output_root}[/{version}]/{deliverable}/{name}.parquet`; see `get_export_root`.""" + return storage.join_uri( + root=get_export_root(output_root=output_root, version=version), + 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, + 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, + version=version, + ) + df = df.copy() + 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) + 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/storage.py b/jdluc/storage.py index 1159d1e..9abbc91 100644 --- a/jdluc/storage.py +++ b/jdluc/storage.py @@ -4,9 +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 @@ -16,21 +19,29 @@ 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) 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://") @@ -42,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 @@ -93,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) @@ -114,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, 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 + + from jdluc import geo + + stacked = stack_bands(dset=dset) + with tempfile.TemporaryDirectory() as tmpdir: + path_to_geotiff = os.path.join(tmpdir, "geotiff.tif") + 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, + 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") @@ -150,6 +233,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 +249,20 @@ def uri(self) -> str: @property def exists_uri(self) -> str: - return f"{self.uri:s}/zarr.json" + 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) 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) + touch_file(uri=self.exists_uri) class CacherDecoratorProtocol(typing.Protocol[R]): 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/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: 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..e4eeed4 --- /dev/null +++ b/orchestration/__main__.py @@ -0,0 +1,256 @@ +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..70141a8 --- /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, such as its 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. + + ``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( + 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", 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())