diff --git a/docs/architecture.md b/docs/architecture.md index 8c882bb..2c0ab5d 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -82,16 +82,18 @@ All rasters are harmonized onto the **GLAD tile grid**: global `EPSG:4326` at 0. A second, coarser grid (the ~10 km MapSPAM resolution) exists for the statistical attribution leg, which downsamples per-pixel emissions to match the resolution of the MapSPAM crop statistics rather than implying a precision the crop data does not have. -### One VRT per band per tile — chunked end to end +### One VRT per dataset per tile — chunked end to end -Harmonization is **one GDAL VRT per source band per tile**. Each VRT declares the destination tile's grid — `EPSG:4326`, with a `` anchored on the absolute 10° graticule — and names the ingested COG as a single `` whose `SrcRect`, `DstRect`, and `resampling` attribute rescale the source onto that grid. +Harmonization is **one GDAL VRT per dataset per tile**, with one band for each of the dataset's source bands. Each VRT declares the destination tile's grid — `EPSG:4326`, with a `` anchored on the absolute 10° graticule — and each band names the ingested COG as a single `` whose `SrcRect`, `DstRect`, and `resampling` attribute rescale the source onto that grid. + +Each chunk of a VRT is read with all of its bands at once. The multi-year sources, GPW grassland and GACED30, store 25 years pixel by pixel, so reading any one year decompresses all 25; reading them together decompresses each block once. There is no reprojection step and no mosaic step: - **No reprojection.** Every ingested COG is already `EPSG:4326` — `geo.validate_geotiff` asserts it at ingest — so the VRT only ever rescales and offsets within a single CRS. An earlier design added a second VRT pass (`rasterio.vrt.WarpedVRT`) to reproject; it was removed once the 4326-at-ingest invariant was enforced. (A `WarpedVRT` does survive in `geo.py`, where the statistical leg downsamples to the MapSPAM grid; it too is a same-CRS rescale.) - **No mosaic.** The unit of work is one tile, so there is never more than one source per band to combine. Two source layouts are handled: ten-degree-tiled datasets (GLAD, GNW peatlands, Harris AGB, Huang BGB, SoilGrids, CDL, TCL, GACED30, GPW grassland, GPW livestock, Descals oil palm) map a whole source tile onto the whole destination tile, while whole-world datasets (IPCC climate zones, MapSPAM) use `SrcRect` to window the tile's footprint out of the global raster. A publisher whose own tiling or projection does not match the 10° graticule is mosaicked or warped onto it during ingest instead — GACED30 ships three-degree tiles, and its dataset module resamples the sixteen covering a tile onto that tile; Descals oil palm ships hundred-kilometre cells on a lattice that divides neither the graticule nor its own pixel grid evenly, and merges up to 66 of them per tile; the FDaP masks ship as one-degree files on their own 10 m grid, so their dataset module pieces up to 100 of them into a tile at that resolution, 111,319 or 111,320 pixels a side, without resampling; GPW livestock ships one global raster per species and year on Interrupted Goode Homolosine, and its dataset module warps each tile's footprint onto a 30-arc-second grid — so harmonization still sees exactly one source per band. -A tile whose source was never published is tolerated when the caller opts in: the VRT is emitted with a band header and no source, which GDAL reads as all-nodata. A minority of tiles need this for the core inputs, and the emissions core coerces those nodatas to zero, so that "no data here" does not propagate as NaN. For a dataset whose publisher maps one crop rather than the whole land surface it is the majority case — Descals oil palm covers 48 of the 280 tiles — and there the band header carries no `NoDataValue` at all, because 0 is a mapped class, so the sourceless band reads back as zero directly. +A tile whose source was never published is tolerated when the caller opts in: each of the dataset's bands is emitted with a header and no source, which GDAL reads as all-nodata. A minority of tiles need this for the core inputs, and the emissions core coerces those nodatas to zero, so that "no data here" does not propagate as NaN. For a dataset whose publisher maps one crop rather than the whole land surface it is the majority case — Descals oil palm covers 48 of the 280 tiles — and there the band header carries no `NoDataValue` at all, because 0 is a mapped class, so the sourceless band reads back as zero directly. The emissions core harmonizes **two dataset lists and merges the results**: the primary stack that every conversion depends on, and the per-crop masks that widen its destination evidence (`emit.CROP_SUPPLEMENT_DATASET_NAMES`). They are kept apart because `harmonize.workflow` hashes its `dataset_names` into the cache key, so a single list would re-materialize every primary band across every tile each time a mask is admitted; separated, admitting one costs only its own cache. Both are harmonized onto the same tile grid, so the merge is an exact join. diff --git a/docs/cdl_glad_comparison_supplement.md b/docs/cdl_glad_comparison_supplement.md index fcc0568..56569b9 100644 --- a/docs/cdl_glad_comparison_supplement.md +++ b/docs/cdl_glad_comparison_supplement.md @@ -51,10 +51,10 @@ Because lost pixels are excluded from both the emissions numerator and the produ ## The statistical leg -The comparison above concerns the jurisdictional-direct leg, where the GLAD-cropland restriction is symmetric: the same per-pixel crop mask gates both the emissions numerator and the crop-area denominator, so a lost pixel drops out of both together. +The comparison above concerns the jurisdictional-direct leg, where the GLAD-cropland restriction it measures was symmetric: the same per-pixel crop mask gated both the emissions numerator and the crop-area denominator, so a lost pixel dropped out of both together. -The statistical leg applies the same GLAD-2020-cropland restriction, but only to the emissions numerator (`get_downscaled_luc_emissions` in `statistical.py`). This is deliberate, not an oversight. Its denominator is MapSPAM crop production — a crop quantity, and therefore zero wherever there is no cropland — so the denominator is already cropland-restricted by construction, to MapSPAM's delineation of cropland rather than GLAD's. Masking the numerator to GLAD-2020 cropland is what brings the two into alignment. Without it, conversion emissions whose destination is non-agricultural (e.g. forest or grassland lost to built-up) would enter the emissions pool and be spread across crops by expansion share — charging, say, a town's expansion to soybeans. +The statistical leg applies no cropland mask to its emissions numerator (`get_downscaled_luc_emissions` in `statistical.py`), and needs none to keep non-agricultural conversions out of it. `emit` charges a conversion's carbon only where a destination layer claims the pixel as cropland or pasture at the assessment year (see [`methodology.md`](methodology.md)). Carbon whose pixel no layer claims is reported on its own as `dropped-emissions`, so it never enters the pool that is spread across crops by expansion share: a town's expansion is not charged to soybeans. The denominator is MapSPAM crop production — a crop quantity, and therefore zero wherever there is no cropland — so it is cropland-restricted by construction, to MapSPAM's delineation of cropland. -The only residual inconsistency is that GLAD's and MapSPAM's cropland footprints do not perfectly coincide: a pixel MapSPAM credits with crop area but GLAD does not call cropland in 2020 sits in the denominator yet is excluded from the numerator. This is the same dataset-delineation disagreement quantified above for CDL, now between GLAD and MapSPAM. At the ~10 km MapSPAM resolution the two footprints cannot be co-registered any more tightly, and the mismatch is not removable by dropping the mask — that would only add the non-agricultural emissions described above. +The residual inconsistency is that the destination layers and MapSPAM do not delineate cropland identically: a cell MapSPAM credits with crop area, but where no destination layer claims the converted pixels, sits in the denominator with none of their emissions in the numerator. This is the same dataset-delineation disagreement quantified above for CDL, here between the destination layers and MapSPAM. At the ~10 km MapSPAM resolution the two footprints cannot be co-registered any more tightly. See [`analyses/cdl_glad_glc_comparison.ipynb`](../analyses/cdl_glad_glc_comparison.ipynb) for the full analysis. diff --git a/docs/further_research.md b/docs/further_research.md index 2997774..92ae5c8 100644 --- a/docs/further_research.md +++ b/docs/further_research.md @@ -35,7 +35,7 @@ Sampled over 1°×1° windows at 30 m on GLAD's land-cover map, not on GACED30, **Status**: open. No stage reads the masks yet, so this bounds what they can carry once one does. -The masks threshold FDaP's per-pixel probabilities per country and year (`jdluc/datasets/fdap.py`). Where a country has a fitted threshold, it is set so the masked area meets FAOSTAT's area harvested; that matches the national total and says nothing about which pixels carry it. Other producers use the model's precision/recall crossover, so their area is not matched to FAOSTAT's at all, and countries with no FAOSTAT production mask nothing. The 2024 masks are fitted to FAOSTAT's 2023 area. +The masks threshold FDaP's per-pixel probabilities per country and year (`jdluc/datasets/fdap.py`). Each commodity has a floor, the model's precision/recall crossover, and no threshold goes below it. Where a country's fitted threshold is above the floor, it is set so the masked area meets FAOSTAT's area harvested; that matches the national total and says nothing about which pixels carry it. Where the fit is below the floor, the floor applies and the mask maps less than FAOSTAT's area: 34 coffee entries across 11 countries, from Colombia to Vietnam, and Venezuela's cocoa in 2020. Other producers use the floor alone, so their area is not matched to FAOSTAT's at all, and countries with no FAOSTAT production mask nothing. The 2024 masks are fitted to FAOSTAT's 2023 area. Where even the highest threshold maps more than FAOSTAT's area harvested, the mask stays larger than FAOSTAT's figure: cocoa in Ghana (2020–2024), Liberia (2018), Fiji (2020 and 2022), Guyana (2022 and 2024) and Suriname, and coffee in Burundi (2018, 2022 and 2024), Kenya (2022), Malawi, Paraguay and Zimbabwe. Mozambique's coffee is excluded nationally instead, since its mask stays over 100 times FAOSTAT's area even at that threshold. diff --git a/jdluc/__tests__/geo_test.py b/jdluc/__tests__/geo_test.py index fb5fad2..a59ad0e 100644 --- a/jdluc/__tests__/geo_test.py +++ b/jdluc/__tests__/geo_test.py @@ -1,4 +1,5 @@ import contextlib +import math import pathlib import numpy @@ -13,6 +14,7 @@ from jdluc.geo import ( clip_dset, + convert_geotiff_to_cog, convert_vector_to_flatgeobuf, downscale_darray, get_chunk_size, @@ -101,6 +103,98 @@ def test_get_overview_level( ) +# NB: 450 is not a whole multiple of the coarsest overview factor, 4 +FOUR_BAND_SIZE = 450 + + +def get_four_band_data() -> numpy.ndarray: + data = numpy.random.default_rng(seed=0).integers( + 0, 100, dtype="uint8", size=(4, FOUR_BAND_SIZE, FOUR_BAND_SIZE) + ) + # Band 4 is all zero, which an RGBA reading would take for a transparent alpha band + data[3] = 0 + # A pixel missing from band 1 only + data[0, 0, 0] = 255 + return data + + +def convert_four_band_geotiff( + data: numpy.ndarray, interleave: str, tmp_path: pathlib.Path +) -> str: + path_to_geotiff = str(tmp_path / f"{interleave:s}-geotiff.tif") + with rasterio.open( + path_to_geotiff, + "w", + count=4, + crs=4326, + driver="GTiff", + dtype="uint8", + height=FOUR_BAND_SIZE, + interleave=interleave, + photometric="MINISBLACK", + tiled=True, + transform=rasterio.transform.from_origin(0, 10, 0.01, 0.01), + width=FOUR_BAND_SIZE, + ) as dataset: + dataset.write(data) + path_to_cog = str(tmp_path / f"{interleave:s}-cog.tif") + convert_geotiff_to_cog( + metadata={"cornerstone-source-name": "test"}, + no_data=255, + path_to_cog=path_to_cog, + path_to_geotiff=path_to_geotiff, + ) + return path_to_cog + + +@pytest.mark.parametrize( + "interleave", + ( + pytest.param("band", id="band-interleaved"), + pytest.param("pixel", id="pixel-interleaved"), + ), +) +def test_convert_geotiff_to_cog_keeps_bands_values_and_metadata( + interleave: str, tmp_path: pathlib.Path +) -> None: + data = get_four_band_data() + path_to_cog = convert_four_band_geotiff( + data=data, interleave=interleave, tmp_path=tmp_path + ) + + with rasterio.open(path_to_cog) as dataset: + assert rasterio.enums.ColorInterp.alpha not in dataset.colorinterp + assert dataset.nodata == 255 + assert dataset.tags()["cornerstone-source-name"] == "test" + assert dataset.tags(ns="IMAGE_STRUCTURE")["LAYOUT"] == "COG" + assert [dataset.overviews(idx) for idx in dataset.indexes] == [[2, 4]] * 4 + numpy.testing.assert_array_equal(dataset.read(), data) + with rasterio.open(path_to_cog, overview_level=0) as dataset: + numpy.testing.assert_array_equal(dataset.read(), data[:, ::2, ::2]) + # Each level rounds its size up from the full resolution, rather than down from the level + # before it + with rasterio.open(path_to_cog, overview_level=1) as dataset: + assert dataset.shape == (math.ceil(FOUR_BAND_SIZE / 4),) * 2 + + +def test_convert_geotiff_to_cog_samples_overviews_alike_for_either_interleave( + tmp_path: pathlib.Path, +) -> None: + data = get_four_band_data() + path_to_band_cog = convert_four_band_geotiff( + data=data, interleave="band", tmp_path=tmp_path + ) + path_to_pixel_cog = convert_four_band_geotiff( + data=data, interleave="pixel", tmp_path=tmp_path + ) + + with ( + rasterio.open(path_to_band_cog, overview_level=1) as band, + rasterio.open(path_to_pixel_cog, overview_level=1) as pixel, + ): + numpy.testing.assert_array_equal(band.read(), pixel.read()) + + @pytest.mark.parametrize( ("dtypes", "number_of_dimensions", "chunk_size"), ( diff --git a/jdluc/attribute.py b/jdluc/attribute.py index 8108789..448ecbb 100644 --- a/jdluc/attribute.py +++ b/jdluc/attribute.py @@ -142,7 +142,9 @@ def main() -> int: level=logging.INFO, ) - parser = argparse.ArgumentParser(description=__doc__) + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) parser.add_argument( "iso_3166s", nargs=argparse.ZERO_OR_MORE, diff --git a/jdluc/emit.py b/jdluc/emit.py index 59ffbeb..f4af306 100644 --- a/jdluc/emit.py +++ b/jdluc/emit.py @@ -805,7 +805,9 @@ def main() -> int: level=logging.INFO, ) - parser = argparse.ArgumentParser(description=__doc__) + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) parser.add_argument( "iso_3166s", help="cover exactly the tiles these countries' boundaries touch", diff --git a/jdluc/geo.py b/jdluc/geo.py index de6b36a..4696a30 100644 --- a/jdluc/geo.py +++ b/jdluc/geo.py @@ -7,6 +7,7 @@ import numpy import rasterio import rasterio.enums +import rasterio.shutil import rasterio.transform import shapely import xarray @@ -36,30 +37,39 @@ def convert_geotiff_to_cog( path_to_cog: str, path_to_geotiff: str, ) -> None: - import rio_cogeo.cogeo - import rio_cogeo.profiles - - logger.info(f"Getting overview level from {path_to_geotiff=:s}") - with rasterio.open(fp=path_to_geotiff) as dataset: - overview_level = get_overview_level( - height=dataset.height, - width=dataset.width, + logger.info(f"Setting {no_data=:} and {metadata=:} for {path_to_geotiff=:s}") + # NB: build the overviews in a separate file that stores the bands pixel by pixel. Inside a + # GeoTIFF that stores each band on its own, as FDaP's does, GDAL picks different pixels for + # an overview whose size does not divide the tile's evenly, so the same values would get + # different overviews depending on how their GeoTIFF was written + with ( + rasterio.Env( + COMPRESS_OVERVIEW="DEFLATE", + GDAL_NUM_THREADS="ALL_CPUS", + INTERLEAVE_OVERVIEW="PIXEL", + TIFF_USE_OVR=True, + ), + rasterio.open(fp=path_to_geotiff, mode="r+") as dataset, + ): + if no_data is not None: + dataset.nodata = no_data + dataset.update_tags(OVERVIEW_RESAMPLING="NEAREST", **metadata) + overview_level = get_overview_level(height=dataset.height, width=dataset.width) + logger.info(f"Building {overview_level=:d} overviews for {path_to_geotiff=:s}") + dataset.build_overviews( + factors=[1 << level for level in range(1, overview_level + 1)], + resampling=rasterio.enums.Resampling.nearest, ) logger.info(f"Converting from {path_to_geotiff=:s} to {path_to_cog=:s}") - rio_cogeo.cogeo.cog_translate( - additional_cog_metadata=metadata, - config={ - "GDAL_NUM_THREADS": "ALL_CPUS", - "GDAL_TIFF_INTERNAL_MASK": True, - }, - dst_kwargs=dict(rio_cogeo.profiles.DEFLATEProfile()) | {"BIGTIFF": "IF_SAFER"}, - dst_path=path_to_cog, - in_memory=False, - nodata=no_data, - overview_level=overview_level, - quiet=False, - source=path_to_geotiff, - use_cog_driver=True, + rasterio.shutil.copy( + path_to_geotiff, + path_to_cog, + BIGTIFF="IF_SAFER", + BLOCKSIZE=512, + COMPRESS="DEFLATE", + NUM_THREADS="ALL_CPUS", + OVERVIEWS="FORCE_USE_EXISTING", + driver="COG", ) diff --git a/jdluc/harmonize.py b/jdluc/harmonize.py index d8a40b4..923307f 100644 --- a/jdluc/harmonize.py +++ b/jdluc/harmonize.py @@ -1,8 +1,8 @@ """Place one ingested source tile per dataset onto a common grid and harmonize them. -For a single ten-degree tile, each raster dataset's tile becomes a per-band GDAL VRT on the -shared grid, returned as an xarray.Dataset with one variable per source band. Whole-world -datasets are windowed onto the same grid. The result is cached. +For a single ten-degree tile, each raster dataset's tile becomes one GDAL VRT on the shared +grid, with all of its bands, returned as an xarray.Dataset with one variable per source band. +Whole-world datasets are windowed onto the same grid. The result is cached. The tile set comes from the positional ISO 3166 alpha-3 codes -- or, with `--backfill`, from `tiling.GLOBAL_NATURE_WATCH_TILE_IDS`; `--stack` narrows which dataset stacks are built, @@ -74,9 +74,9 @@ def gdal_path(self) -> str: def iter_vrt_band_header( - band_name: str, dtype: str, no_data: int | float | None + band_idx: int, band_name: str, dtype: str, no_data: int | float | None ) -> collections.abc.Iterator[str]: - yield f' ' + yield f' ' yield f" {band_name:s}" if no_data is not None: yield f" {no_data}" @@ -189,9 +189,7 @@ def get_resampling_for_band_type( raise ValueError(band_type) -def get_vrt_for_dataset_band_tile_id( - band_idx: int, - band_name: str, +def get_vrt_for_dataset_tile_id( dataset: base.RasterDataset, grid: Grid, ignore_missing_tiles: bool, @@ -200,7 +198,7 @@ def get_vrt_for_dataset_band_tile_id( ) -> str: lines = list(grid.iter_preamble) - logger.info(f"Processing {dataset=} and {band_name=:s}") + logger.info(f"Processing {dataset=}") if dataset.partitioning == tiling.Partitioning.TEN_DEGREE_TILE: uri = storage.join_uri(prefix=dataset.get_prefix(tile_id=tile_id), root=root) if storage.path_exists(uri=uri): @@ -209,37 +207,50 @@ def get_vrt_for_dataset_band_tile_id( root=root, tile_id=tile_id, ) - lines.extend( - iter_vrt_band_header( - band_name=band_name, - dtype=tile.dtype, - no_data=tile.no_data, - ) + resampling = grid.get_resampling_for_band_type( + band_type=tile.band_type, + dest_resolution=grid.resolution, + src_resolution=tile.resolution, ) - lines.extend( - iter_vrt_band_content( - band_idx=band_idx, - dest_resolution=grid.resolution, - path_to_tile=tile.gdal_path, - resampling=grid.get_resampling_for_band_type( - band_type=tile.band_type, + for band_idx, band_name in enumerate( + dataset.fully_qualified_band_names, start=1 + ): + lines.extend( + iter_vrt_band_header( + band_idx=band_idx, + band_name=band_name, + dtype=tile.dtype, + no_data=tile.no_data, + ) + ) + lines.extend( + iter_vrt_band_content( + band_idx=band_idx, dest_resolution=grid.resolution, + path_to_tile=tile.gdal_path, + resampling=resampling, + src_offset=tiling.XY(x=0, y=0), src_resolution=tile.resolution, - ), - src_offset=tiling.XY(x=0, y=0), - src_resolution=tile.resolution, + ) ) - ) + lines.append(" ") elif ignore_missing_tiles: logger.warning( - f"{tile_id=:s} is missing for {dataset=} but due to {ignore_missing_tiles=} we are emitting an all-no-data band" + f"{tile_id=:s} is missing for {dataset=} but due to {ignore_missing_tiles=} we are emitting all-no-data bands" ) - # Yield an empty tile band - lines.extend( - iter_vrt_band_header( - band_name=band_name, dtype="Float32", no_data=dataset.no_data + # Bands with no source read as no-data everywhere + for band_idx, band_name in enumerate( + dataset.fully_qualified_band_names, start=1 + ): + lines.extend( + iter_vrt_band_header( + band_idx=band_idx, + band_name=band_name, + dtype="Float32", + no_data=dataset.no_data, + ) ) - ) + lines.append(" ") else: raise FileNotFoundError(uri) elif dataset.partitioning == tiling.Partitioning.WHOLE_WORLD: @@ -248,13 +259,6 @@ def get_vrt_for_dataset_band_tile_id( root=root, tile_id=tiling.WHOLE_WORLD_TILE_ID, ) - lines.extend( - iter_vrt_band_header( - band_name=band_name, - dtype=tile.dtype, - no_data=tile.no_data, - ) - ) world_span = tiling.XY(x=360, y=180) src_offset = grid.get_offset_for_world( resolution=tile.resolution, span=world_span @@ -262,23 +266,35 @@ def get_vrt_for_dataset_band_tile_id( src_resolution = grid.get_resolution_for_world( resolution=tile.resolution, span=world_span ) - lines.extend( - iter_vrt_band_content( - band_idx=band_idx, - dest_resolution=grid.resolution, - path_to_tile=tile.gdal_path, - resampling=grid.get_resampling_for_band_type( - band_type=tile.band_type, + resampling = grid.get_resampling_for_band_type( + band_type=tile.band_type, + dest_resolution=grid.resolution, + src_resolution=src_resolution, + ) + for band_idx, band_name in enumerate( + dataset.fully_qualified_band_names, start=1 + ): + lines.extend( + iter_vrt_band_header( + band_idx=band_idx, + band_name=band_name, + dtype=tile.dtype, + no_data=tile.no_data, + ) + ) + lines.extend( + iter_vrt_band_content( + band_idx=band_idx, dest_resolution=grid.resolution, + path_to_tile=tile.gdal_path, + resampling=resampling, + src_offset=src_offset, src_resolution=src_resolution, - ), - src_offset=src_offset, - src_resolution=src_resolution, + ) ) - ) + lines.append(" ") else: raise ValueError(dataset.partitioning) - lines.append(" ") lines.append("") with tempfile.NamedTemporaryFile(delete=False, mode="w", suffix=".vrt") as fp: @@ -288,23 +304,32 @@ def get_vrt_for_dataset_band_tile_id( return path_to_vrt -def get_dset_for_output(path_to_vrts: collections.abc.Sequence[str]) -> xarray.Dataset: +def get_dset_for_output( + path_to_vrt_to_band_names: dict[str, list[str]], +) -> xarray.Dataset: darrays: list[xarray.DataArray] = [] - chunk_size = geo.get_chunk_size(dtypes=[numpy.dtype("float32")] * len(path_to_vrts)) - for path_to_vrt in path_to_vrts: + chunk_size = geo.get_chunk_size( + dtypes=[numpy.dtype("float32")] + * sum(map(len, path_to_vrt_to_band_names.values())) + ) + for path_to_vrt, band_names in path_to_vrt_to_band_names.items(): logger.debug(f"Opening {path_to_vrt=:s} with {chunk_size=:d}") + # NB: every band in each chunk, because a pixel-interleaved source decompresses all of + # its bands to read any one of them darray = rioxarray.open_rasterio( - chunks=chunk_size, + chunks={"band": -1, "x": chunk_size, "y": chunk_size}, filename=path_to_vrt, lock=False, ) assert isinstance(darray, xarray.DataArray) - darray = darray.isel(band=0, drop=True) - darrays.append( - geo.unify_dtype_and_no_data( - darray=darray.rename(darray.attrs.pop("long_name")) + assert darray.sizes["band"] == len(band_names) + darray.attrs.pop("long_name") + for band_idx, band_name in enumerate(band_names): + darrays.append( + geo.unify_dtype_and_no_data( + darray=darray.isel(band=band_idx, drop=True).rename(band_name) + ) ) - ) return xarray.Dataset({darray.name: darray for darray in darrays}) @@ -335,23 +360,18 @@ def workflow( logger.info(f"Constructing common grid for {tile_resolution=:}") grid = Grid.from_tile_id_resolution(resolution=tile_resolution, tile_id=tile_id) - path_to_vrts = [ - get_vrt_for_dataset_band_tile_id( - band_idx=band_idx, - band_name=band_name, + path_to_vrt_to_band_names = { + get_vrt_for_dataset_tile_id( dataset=dataset, grid=grid, ignore_missing_tiles=ignore_missing_tiles, root=cfg.ingest_root, tile_id=tile_id, - ) + ): dataset.fully_qualified_band_names for dataset in datasets if isinstance(dataset, base.RasterDataset) - for band_idx, band_name in enumerate( - dataset.fully_qualified_band_names, start=1 - ) - ] - return get_dset_for_output(path_to_vrts=path_to_vrts) + } + return get_dset_for_output(path_to_vrt_to_band_names=path_to_vrt_to_band_names) @enum.unique @@ -388,7 +408,9 @@ def main() -> int: level=logging.INFO, ) - parser = argparse.ArgumentParser(description=__doc__) + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) parser.add_argument( "iso_3166s", help="cover exactly the tiles these countries' boundaries touch", diff --git a/jdluc/ingest.py b/jdluc/ingest.py index c1a0dd9..7bfd5d6 100644 --- a/jdluc/ingest.py +++ b/jdluc/ingest.py @@ -9,7 +9,8 @@ The tile set comes from the positional ISO 3166 alpha-3 codes -- or, with `--backfill`, from `tiling.GLOBAL_NATURE_WATCH_TILE_IDS`. Tiles are ingested one at a time unless -`--concurrency` raises the bound. +`--concurrency` raises the bound. `--modulus` and `--residues` split the tile set into disjoint +shards, so several processes can share a backfill. Example invocations: uv run python -m jdluc.ingest IPCC_CLIMATE_ZONES USA @@ -20,6 +21,7 @@ uv run python -m jdluc.ingest GPW_GRASSLAND BRA uv run python -m jdluc.ingest DESCALS_OIL_PALM IDN MYS uv run python -m jdluc.ingest FDAP_COCOA CIV GHA + uv run python -m jdluc.ingest FDAP_COCOA --backfill --modulus=2 --residues=0 --concurrency=2 """ import argparse @@ -27,7 +29,9 @@ import concurrent.futures import logging -from jdluc import config, datasets, tiling +import rasterio + +from jdluc import config, datasets, tiling, utils from jdluc.datasets import base, worldbank_jurisdictions logger = logging.getLogger(__name__) @@ -49,7 +53,12 @@ def workflow( if not all(map(tile_id_is_valid_func, tile_ids)): raise ValueError(f"{tile_ids=:} are not valid for {dataset=:}") results: dict[str, str | Exception] = {} - with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor: + with ( + # NB: Each process has its own cache and can OOM; reducing the cache size hasn't + # shown any slowdowns + rasterio.Env(GDAL_CACHEMAX=1 << 29), # 512 MiB + concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor, + ): futures: dict[concurrent.futures.Future[str], str] = {} for tile_id in sorted(tile_ids): futures[ @@ -77,9 +86,19 @@ def main() -> int: level=logging.INFO, ) - parser = argparse.ArgumentParser(description=__doc__) + dataset_names = sorted(e.name for e in datasets.DatasetName) + parser = argparse.ArgumentParser( + description=__doc__, + epilog="\n".join( + ("dataset names:", *(f" {name:s}" for name in dataset_names)) + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) parser.add_argument( - "dataset_name", choices=sorted(e.name for e in datasets.DatasetName) + "dataset_name", + choices=dataset_names, + help="one of the dataset names listed below", + metavar="dataset_name", ) parser.add_argument( "iso_3166s", @@ -88,24 +107,54 @@ def main() -> int: type=worldbank_jurisdictions.iso_3166_str, ) parser.add_argument("--backfill", action="store_true", help="cover all GNW tiles") - parser.add_argument("--concurrency", default=1, type=int) - parser.add_argument("--overwrite", action="store_true") + parser.add_argument( + "--concurrency", + default=1, + help="ingest up to this many tiles at once (default 1)", + type=int, + ) + parser.add_argument( + "--modulus", + default=1, + help="split the tiles into this many disjoint shards", + type=int, + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="re-ingest tiles that already exist under INGEST_ROOT", + ) + parser.add_argument( + "--residues", + action="append", + help="ingest only these shards, each in [0, --modulus); repeatable, all by default", + type=int, + ) args = parser.parse_args() assert bool(args.iso_3166s) ^ bool(args.backfill), ( "pass either one-or-more iso_3166s or --backfill" ) dataset = datasets.NAME_TO_CLS[datasets.DatasetName[str(args.dataset_name)]] + assert not ( + dataset.partitioning == tiling.Partitioning.WHOLE_WORLD and args.modulus > 1 + ), f"{args.dataset_name:s} is whole-world, so it cannot be split with --modulus" tile_id_to_path_or_exc: dict[str, str | Exception] = workflow( concurrency=int(args.concurrency), dataset=dataset, overwrite=args.overwrite, root=config.Config.from_dot_env().ingest_root, - tile_ids=( - tiling.GLOBAL_NATURE_WATCH_TILE_IDS - if args.backfill - else worldbank_jurisdictions.get_ten_degree_tile_ids_for_iso_3166s( - iso_3166s=args.iso_3166s + tile_ids=tuple( + utils.iter_sharded( + modulus=int(args.modulus), + residues=args.residues, + values=( + tiling.GLOBAL_NATURE_WATCH_TILE_IDS + if args.backfill + else worldbank_jurisdictions.get_ten_degree_tile_ids_for_iso_3166s( + iso_3166s=args.iso_3166s + ) + ), ) ), ) diff --git a/jdluc/trace.py b/jdluc/trace.py index bb414b4..335fcab 100644 --- a/jdluc/trace.py +++ b/jdluc/trace.py @@ -176,7 +176,9 @@ def main() -> int: level=logging.INFO, ) - parser = argparse.ArgumentParser(description=__doc__) + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) parser.add_argument( "iso_3166s", nargs=argparse.ZERO_OR_MORE, diff --git a/validation/__main__.py b/validation/__main__.py index 4b4782c..f760a6a 100644 --- a/validation/__main__.py +++ b/validation/__main__.py @@ -133,7 +133,9 @@ def get_document( def main() -> int: logging.basicConfig(format="%(levelname)s - %(message)s", level=logging.INFO) - parser = argparse.ArgumentParser(description=__doc__) + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) parser.add_argument( "--stage", choices=tuple(Stage),