ebirdst_geotiff_to_parquet.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "cyclopts>=4.23",
#     "flox>=0.11.2",
#     "h3ronpy>=0.22",
#     "numpy>=2",
#     "pandas>=2.3",
#     "pyarrow>=24",
#     "pyproj>=3.7.2",
#     "rasterio>=1.5.1",
#     "rich>=14",
#     "rioxarray>=0.23",
#     "xarray>=2026.7",
# ]
# [tool.uv]
# exclude-newer = "2026-10-07T00:00:00Z"
# ///
"""Average a weekly eBird Status & Trends GeoTIFF into H3 cells and write parquet.

Each band must be described by its ISO week date, as in the weekly products
(`<species>_<product>_{median,mean}_<res>_<year>.tif`). Seasonal rasters are
rejected. Nothing else is read: species, product and review status are left to
downstream code (e.g. from `tiff_basename` and the species `config.json`).

Examples:

    uv run https://gist.github.com/gcr/4a18224cd05f26c88b08a62e2e607ef0/raw/8a45589f8fcfd82a81c9a55efa453d873ade481d/ebirdst_geotiff_to_parquet.py --help
    uv run ebirdst_geotiff_to_parquet.py abetow_occurrence_median_3km_2023.tif out.parquet

Output columns, one row per (H3 cell, band) with at least one non-NaN pixel,
sorted by `h3_level`, `h3_cell_id`, `band`:

- `band`: 1-based band number in the GeoTIFF.
- `week_of_year`: the band's ISO date (string).
- `value`: mean of the non-NaN pixels whose centre falls in the cell.
- `h3_cell_id` (int64), `h3_cell_id_str` (hex), `h3_level`.
- `tiff_basename`: input file name.
"""

from datetime import date
from pathlib import Path
from typing import Annotated

import flox.xarray
import numpy as np
import pandas as pd
import rioxarray
import xarray as xr
from cyclopts import App, Parameter
from h3ronpy.vector import coordinates_to_cells
from pyproj import Transformer
from rich import status

app = App(help_format="markdown")


def load_weekly_geotiff(src: Path) -> xr.DataArray:
    """Open a weekly GeoTIFF with nodata as NaN and a `week_of_year` band coordinate.

    Raises:
        ValueError: If any band description is missing or not an ISO date.
    """
    data_arr = rioxarray.open_rasterio(src, masked=True, num_threads="all_cpus")
    names = data_arr.attrs.get("long_name")
    names = [names] if isinstance(names, str) else list(names or [])
    try:
        if len(names) != data_arr.sizes["band"]:
            raise ValueError
        weeks = [date.fromisoformat(name).isoformat() for name in names]
    except (TypeError, ValueError):
        data_arr.close()
        raise ValueError(f"Band descriptions must be ISO week dates: {src}") from None
    return data_arr.assign_coords(week_of_year=("band", weeks))


def h3_cell_indexes(data_arr: xr.DataArray, h3_levels: list[int]) -> xr.DataArray:
    """Return the uint64 H3 cell of each pixel centre, with dims `(h3_level, y, x)`."""
    native_y, native_x = xr.broadcast(data_arr.y, data_arr.x)
    transformer = Transformer.from_crs(data_arr.rio.crs, "EPSG:4326", always_xy=True)
    longitude, latitude = transformer.transform(native_x.values, native_y.values)
    cells = [
        np.asarray(
            coordinates_to_cells(latitude.ravel(), longitude.ravel(), level, radians=False)
        ).reshape(longitude.shape)
        for level in h3_levels
    ]
    return xr.DataArray(
        np.stack(cells),
        dims=("h3_level", "y", "x"),
        coords={"h3_level": h3_levels, "y": data_arr.y, "x": data_arr.x},
    )


def h3_means(data_arr: xr.DataArray, h3_levels: list[int]) -> pd.DataFrame:
    """Average non-NaN pixels per (H3 cell, band) at each level; see module docstring."""
    cells = h3_cell_indexes(data_arr, h3_levels)
    data_arr = data_arr.expand_dims(h3_level=cells.h3_level).assign_coords(cell_ids=cells)
    # Most pixels are NaN in every band; drop them before grouping to save memory.
    data_arr = data_arr.stack(pixel=("y", "x"), create_index=False).dropna(
        "pixel", how="all"
    )
    columns = ["band", "week_of_year", "value", "h3_cell_id", "h3_cell_id_str", "h3_level"]
    if not data_arr.sizes["pixel"]:
        return pd.DataFrame(columns=columns)
    # All levels share one flattened cell_ids axis: the level is encoded in the id.
    grouped = flox.xarray.xarray_reduce(
        data_arr.drop_vars("spatial_ref", errors="ignore"),
        data_arr.cell_ids,
        dim=("h3_level", "pixel"),
        func="mean",
        skipna=True,
        fill_value=np.nan,
    )
    frame = grouped.rename("value").to_dataframe().reset_index().dropna(subset=["value"])
    ids = frame["cell_ids"].to_numpy(dtype=np.uint64)
    frame["h3_cell_id"] = ids.astype("int64")
    frame["h3_cell_id_str"] = pd.Series(ids, index=frame.index).map("{:x}".format)
    # H3 index bits 52-55 hold the resolution.
    frame["h3_level"] = ((ids >> np.uint64(52)) & np.uint64(0xF)).astype("int64")
    frame = frame[columns]
    return frame.sort_values(["h3_level", "h3_cell_id", "band"]).reset_index(drop=True)


@app.default
def convert(
    tiff: Path,
    output: Path | None = None,
    *,
    h3_levels: Annotated[tuple[int, ...], Parameter(consume_multiple=True)] = (2, 3, 4, 5),
    overwrite: Annotated[bool, Parameter(name=["--overwrite", "-f"])] = False,
) -> None:
    """Average one weekly eBird S&T GeoTIFF into H3 cells and write it as parquet.

    Parameters
    ----------
    tiff
        Weekly eBird S&T GeoTIFF; band descriptions must be ISO dates.
    output
        Parquet path. Default: `tiff` with a `.parquet` suffix, next to the tiff.
    h3_levels
        H3 resolutions (0-15) to aggregate into.
    overwrite
        Replace an existing output.
    """
    tiff = tiff.expanduser()
    if not tiff.exists():
        raise SystemExit(f"No such file: {tiff}")
    if not h3_levels or len(set(h3_levels)) != len(h3_levels) or any(
        not 0 <= level <= 15 for level in h3_levels
    ):
        raise SystemExit("h3_levels must be distinct resolutions from 0 to 15")
    levels = sorted(h3_levels)
    output = (output or tiff.with_suffix(".parquet")).expanduser()
    if output.exists() and not overwrite:
        raise SystemExit(f"{output} exists; pass --overwrite/-f to replace it")

    try:
        raster = load_weekly_geotiff(tiff)
    except ValueError as e:
        raise SystemExit(str(e)) from None
    with raster, status.Status(f"Averaging {tiff.name} into H3 levels {levels}..."):
        frame = h3_means(raster, levels)
    frame["tiff_basename"] = tiff.name

    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_suffix(".parquet.partial")
    try:
        frame.to_parquet(temporary, index=False)
        temporary.replace(output)
    finally:
        temporary.unlink(missing_ok=True)
    print(f"Wrote {len(frame):,} rows (H3 levels {levels}) to {output}")


if __name__ == "__main__":
    app()