thermograph/backend/gen_era5_lake.py

"""Extract the ERA5 lake: Earthmover Icechunk archive -> the Thermograph bucket.

Builds the lake that ``data/era5lake.py`` reads (layout documented there):
per-point serving files, the tile/year/month hive table, and the manifest.
Pulls are tile-aligned to the source's chunk grid (8736h x 12 x 12) — one
46-year pull (~19s measured) decodes all 144 points of a tile, so per-point
pulls would waste 144x the wire. Roughly: ~600 tiles cover the curated
cities (hours), ~3000 cover land (a day or two); both resume from the
manifest, so ctrl-C and re-run is always safe.

    pip install -r requirements-seed.txt 'numcodecs[pcodec]' boto3
    python gen_era5_lake.py --cities --dry-run      # one tile, no writes
    python gen_era5_lake.py --cities                # tiles covering cities.json
    python gen_era5_lake.py --tiles 20:30,40:60     # explicit tile-index window
    python gen_era5_lake.py --land                  # every land-containing tile

Destination: the bucket from THERMOGRAPH_LAKE_S3_* (see era5lake.s3_config),
or a local mirror when THERMOGRAPH_LAKE_LOCAL_DIR is set (tests/rehearsal).
Seed-only deps (icechunk/xarray/zarr) stay out of the app image, same policy
as seed_era5.py.
"""
import argparse
import datetime
import io
import os
import sys
import time

import polars as pl

from data import cities
from data import climate
from data import era5lake
from seed_era5 import ERA5_BUCKET, ERA5_GROUP, ERA5_PREFIX, ERA5_REGION, ERA5_VARS, \
    hourly_to_daily

# The lake keeps each location's WHOLE available record (the store reaches
# back to 1940); the app slices to its own grading window (climate.START_DATE)
# at read time, so the extra decades are analytical surface, not a grading
# change.
START_DATE = "1940-01-01"


def _open_store():
    import icechunk           # noqa: PLC0415 - seed-only dep
    import xarray as xr       # noqa: PLC0415

    storage = icechunk.s3_storage(
        bucket=ERA5_BUCKET, prefix=ERA5_PREFIX, region=ERA5_REGION, anonymous=True)
    repo = icechunk.Repository.open(storage)
    return xr.open_zarr(repo.readonly_session("main").store,
                        group=ERA5_GROUP, consolidated=False)


class _Sink:
    """Writes lake objects to the bucket (boto3) or a local mirror dir."""

    def __init__(self):
        self.local = era5lake.local_dir()
        self.cfg = era5lake.s3_config()
        if not self.local and self.cfg is None:
            sys.exit("no destination: set THERMOGRAPH_LAKE_S3_* or "
                     "THERMOGRAPH_LAKE_LOCAL_DIR")
        self._s3 = None
        if not self.local:
            import boto3  # noqa: PLC0415 - extractor-only dep
            self._s3 = boto3.client(
                "s3", endpoint_url=self.cfg["endpoint"],
                region_name=self.cfg["region"],
                aws_access_key_id=self.cfg["access_key"],
                aws_secret_access_key=self.cfg["secret_key"])

    def put(self, key: str, df: pl.DataFrame) -> None:
        if self.local:
            path = os.path.join(self.local, key)
            os.makedirs(os.path.dirname(path), exist_ok=True)
            df.write_parquet(path)
            return
        buf = io.BytesIO()
        df.write_parquet(buf)
        self._s3.put_object(Bucket=self.cfg["bucket"], Key=key,
                            Body=buf.getvalue())

    def read_manifest(self) -> "pl.DataFrame | None":
        try:
            if self.local:
                return pl.read_parquet(
                    os.path.join(self.local, era5lake.MANIFEST_KEY))
            obj = self._s3.get_object(Bucket=self.cfg["bucket"],
                                      Key=era5lake.MANIFEST_KEY)
            return pl.read_parquet(io.BytesIO(obj["Body"].read()))
        except Exception:  # noqa: BLE001 - a fresh lake has no manifest yet
            return None


def _tile_frame(ds, ti: int, tj: int, end: str) -> "pl.DataFrame | None":
    """Pull one 12x12 tile's hourly series and return the finalized daily frame
    for all its points (lat_idx, lon_idx, date, ...). None for all-sea tiles."""
    t = era5lake.TILE
    i0, j0 = ti * t, tj * t
    lat_sl = slice(i0, min(i0 + t, era5lake.LAT_N))
    lon_sl = slice(j0, j0 + t)
    if "lsm" in ds:            # land-sea mask: skip all-sea tiles outright
        lsm = ds["lsm"].isel(latitude=lat_sl, longitude=lon_sl)
        if "valid_time" in lsm.dims:   # static field, stored with a time axis
            lsm = lsm.isel(valid_time=-1)
        if float(lsm.max()) < 0.05:
            return None
    sub = (ds[list(ERA5_VARS)]
           .isel(latitude=lat_sl, longitude=lon_sl)
           .sel(valid_time=slice(START_DATE, end))
           .load())

    frames = []
    for a, lat in enumerate(sub.latitude.values):
        for b, lon in enumerate(sub.longitude.values):
            cols = {"time": sub.valid_time.values}
            for src, dst in ERA5_VARS.items():
                cols[dst] = sub[src].isel(latitude=a, longitude=b).values
            daily = climate._finalize_approximated(hourly_to_daily(pl.DataFrame(cols)))
            i, j = era5lake.to_idx(float(lat), float(lon))
            frames.append(daily.with_columns(
                pl.lit(i, dtype=pl.Int32).alias("lat_idx"),
                pl.lit(j, dtype=pl.Int32).alias("lon_idx")))
    return pl.concat(frames)


def _write_tile(sink: _Sink, ti: int, tj: int, tile_df: pl.DataFrame) -> pl.DataFrame:
    """Write one tile's point files + hive partitions; return its manifest rows."""
    rows = []
    for (i, j), pdf in tile_df.group_by(["lat_idx", "lon_idx"]):
        pdf = pdf.sort("date")
        sink.put(era5lake.point_key(i, j), pdf.drop(["lat_idx", "lon_idx"]))
        lat, lon = era5lake.to_coords(i, j)
        rows.append({"lat_idx": i, "lon_idx": j, "lat": lat, "lon": lon,
                     "tile": f"{ti}_{tj}", "rows": pdf.height,
                     "date_min": str(pdf["date"].min()),
                     "date_max": str(pdf["date"].max())})
    parts = tile_df.with_columns(
        pl.col("date").dt.year().alias("year"),
        pl.col("date").dt.month().alias("month"))
    for (year, month), mdf in parts.group_by(["year", "month"]):
        sink.put(era5lake.daily_part_key(ti, tj, year, month),
                 mdf.drop(["year", "month"]).sort(["lat_idx", "lon_idx", "date"]))
    return pl.DataFrame(rows)


def _city_tiles() -> list:
    from data import grid  # noqa: PLC0415

    tiles = set()
    for c in cities.all_cities():
        cell = grid.snap(c["lat"], c["lon"])
        tiles.add(era5lake.tile_of(
            *era5lake.to_idx(cell["center_lat"], cell["center_lon"])))
    return sorted(tiles)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--cities", action="store_true")
    ap.add_argument("--land", action="store_true")
    ap.add_argument("--tiles", help="ti0:ti1,tj0:tj1 tile-index window")
    ap.add_argument("--limit", type=int)
    ap.add_argument("--overwrite", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    if args.cities:
        todo = _city_tiles()
    elif args.tiles:
        (a, b), (c, d) = (p.split(":") for p in args.tiles.split(","))
        todo = [(ti, tj) for ti in range(int(a), int(b))
                for tj in range(int(c), int(d))]
    elif args.land:
        todo = [(ti, tj) for ti in range(era5lake.LAT_N // era5lake.TILE + 1)
                for tj in range(era5lake.LON_N // era5lake.TILE)]
    else:
        ap.error("pick one of --cities / --tiles / --land")
    if args.limit:
        todo = todo[:args.limit]

    sink = _Sink()
    ds = _open_store()
    end = (datetime.date.today()
           - datetime.timedelta(days=climate.ARCHIVE_LATENCY_DAYS)).isoformat()

    manifest = sink.read_manifest()
    done = set() if manifest is None or args.overwrite else \
        set(manifest["tile"].unique().to_list())

    written = skipped = empty = 0
    for n, (ti, tj) in enumerate(todo, 1):
        if f"{ti}_{tj}" in done:
            skipped += 1
            continue
        t0 = time.time()
        tile_df = _tile_frame(ds, ti, tj, end)
        if tile_df is None:
            empty += 1
            continue
        if args.dry_run:
            print(f"[dry-run] tile {ti}_{tj}: {tile_df.height} rows "
                  f"({time.time() - t0:.1f}s); nothing written")
            print(tile_df.head())
            return
        rows = _write_tile(sink, ti, tj, tile_df)
        manifest = rows if manifest is None else pl.concat(
            [manifest.filter(pl.col("tile") != f"{ti}_{tj}"), rows])
        sink.put(era5lake.MANIFEST_KEY, manifest)   # after every tile: resumable
        written += 1
        print(f"[{n}/{len(todo)}] tile {ti}_{tj}: {rows.height} points "
              f"in {time.time() - t0:.1f}s")
    print(f"done: tiles written={written} skipped(done)={skipped} sea={empty}")


if __name__ == "__main__":
    main()