"""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()