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 concurrent.futures
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.
    Uploads run on a thread pool: a tile is ~1.2k small objects and sequential
    PUTs are pure round-trip latency (measured 2.5 obj/s -> 8 min/tile; the
    pool brings that to well under a minute). boto3 clients are thread-safe."""
    puts = []
    rows = []
    for (i, j), pdf in tile_df.group_by(["lat_idx", "lon_idx"]):
        pdf = pdf.sort("date")
        puts.append((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"]):
        puts.append((era5lake.daily_part_key(ti, tj, year, month),
                     mdf.drop(["year", "month"]).sort(["lat_idx", "lon_idx", "date"])))
    with concurrent.futures.ThreadPoolExecutor(max_workers=16) as pool:
        for f in [pool.submit(sink.put, k, df) for k, df in puts]:
            f.result()   # surface the first failure instead of hiding it
    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()