Spaces:
Running
Running
| """HF-backed parameter-grid adapter for the AQUAH->CREST_demo integration. | |
| Clips the 1 km CONUS CREST/KW parameter COGs hosted in `vincewin/CREST_data` | |
| (param/crest/*, param/kw/*) to a basin bbox and returns the EF5 control-file | |
| `<param>_grid=` mapping. Used when running with spatially-distributed params | |
| (grid x scalar-multiplier) instead of AQUAH's uniform scalars. | |
| Grid<-COG mapping mirrors the calibrated conus-crest-hydro controls: | |
| wm_grid<-wm_usa im_grid<-im_usa fc_grid<-ksat_usa b_grid<-b_usa | |
| leaki_grid<-leaki_usa alpha_grid<-alpha_usa beta_grid<-beta_usa alpha0_grid<-alpha0_usa | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import time | |
| import rasterio | |
| from rasterio.windows import from_bounds | |
| HF_REPO = "vincewin/CREST_data" | |
| _RESOLVE = f"https://huggingface.co/datasets/{HF_REPO}/resolve/main" | |
| # GDAL-level retries on 429/5xx range reads — several gauges clip in parallel | |
| # and HF occasionally rate-limits the burst (process-wide; also covers the | |
| # terrain + snow-grid vsicurl reads) | |
| os.environ.setdefault("GDAL_HTTP_MAX_RETRY", "3") | |
| os.environ.setdefault("GDAL_HTTP_RETRY_DELAY", "2") | |
| # EF5 control key -> COG path in CREST_data -> local filename | |
| PARAM_GRIDS = { | |
| "wm_grid": ("param/crest/wm_usa.tif", "wm.tif"), | |
| "im_grid": ("param/crest/im_usa.tif", "im.tif"), | |
| "fc_grid": ("param/crest/ksat_usa.tif", "fc.tif"), | |
| "b_grid": ("param/crest/b_usa.tif", "b.tif"), | |
| "leaki_grid": ("param/kw/leaki_usa.tif", "leaki.tif"), | |
| "alpha_grid": ("param/kw/alpha_usa.tif", "alpha.tif"), | |
| "beta_grid": ("param/kw/beta_usa.tif", "beta.tif"), | |
| "alpha0_grid": ("param/kw/alpha0_usa.tif", "alpha0.tif"), | |
| } | |
| def _vsicurl(cog: str) -> str: | |
| return f"/vsicurl/{_RESOLVE}/{cog}" | |
| def _download_cog(cog: str) -> str: | |
| """Fetch the whole (small, public) param COG to the local HF cache and | |
| return its path. Fallback when windowed /vsicurl reads keep failing — | |
| reading the local file bypasses /vsicurl (and its chunk cache) entirely.""" | |
| from huggingface_hub import hf_hub_download | |
| return hf_hub_download(HF_REPO, cog, repo_type="dataset", | |
| token=os.environ.get("HF_TOKEN")) | |
| def clip_param_grids(bbox, out_dir: str, keys=None, unsafe_ssl: bool | None = None) -> dict: | |
| """Clip requested param COGs to bbox (W,S,E,N). Returns {control_key: local_path}.""" | |
| if unsafe_ssl is None: | |
| unsafe_ssl = os.name == "nt" | |
| if unsafe_ssl: | |
| os.environ.setdefault("GDAL_HTTP_UNSAFESSL", "YES") | |
| os.makedirs(out_dir, exist_ok=True) | |
| W, S, E, N = bbox | |
| keys = keys or list(PARAM_GRIDS) | |
| out, skipped = {}, {} | |
| for key in keys: | |
| cog, fname = PARAM_GRIDS[key] | |
| cached = os.path.join(out_dir, fname) | |
| if os.path.exists(cached): # shared-store hit: no remote read | |
| out[key] = cached | |
| continue | |
| # /vsicurl reads fail transiently (HF rate limits parallel gauge runs; | |
| # GDAL then reports "not recognized as being in a supported file | |
| # format", or a truncated tile read "got 0 bytes, expected N") — retry | |
| # with backoff like basic.clip_basic_data | |
| err = None | |
| for attempt in range(3): | |
| try: | |
| _clip_one_param(_vsicurl(cog), fname, bbox, out_dir, key, out, skipped) | |
| err = None | |
| break | |
| except Exception as e: | |
| err = e | |
| time.sleep(2 * (attempt + 1)) | |
| if err is not None: | |
| # An HF read hiccup can outlast the in-process retries and can | |
| # poison GDAL's process-wide /vsicurl chunk cache, so re-reads keep | |
| # returning the same bad bytes. Fall back to fetching the whole | |
| # (small, public) COG once and clipping from the local copy — | |
| # bypasses /vsicurl entirely and is cached for later gauges. | |
| try: | |
| _clip_one_param(_download_cog(cog), fname, bbox, out_dir, key, | |
| out, skipped) | |
| err = None | |
| except Exception as e: | |
| err = e | |
| if err is not None: | |
| raise err | |
| if skipped: | |
| for k, why in skipped.items(): | |
| print(f" WARN param grid {k} skipped: {why} (falls back to scalar)") | |
| return out | |
| def _clip_one_param(src_path: str, fname: str, bbox, out_dir: str, key: str, | |
| out: dict, skipped: dict) -> None: | |
| W, S, E, N = bbox | |
| with rasterio.open(src_path) as src: | |
| b = src.bounds | |
| # guard: bbox must intersect the grid (catches mis-georeferenced grids) | |
| if not (b.left < E and b.right > W and b.bottom < N and b.top > S): | |
| skipped[key] = f"bbox outside grid extent {tuple(round(x,1) for x in b)}" | |
| return | |
| win = from_bounds(W, S, E, N, src.transform).round_offsets().round_lengths() | |
| data = src.read(1, window=win, boundless=True, fill_value=src.nodata) | |
| if data.size == 0 or 0 in data.shape: | |
| skipped[key] = "empty clip window" | |
| return | |
| # standardized clip: Float32, strip-organized, nodata=-9999, WGS84 | |
| data = data.astype("float32") | |
| if src.nodata is not None: | |
| data[data == float(src.nodata)] = -9999.0 | |
| profile = src.profile.copy() | |
| profile.update(driver="GTiff", height=data.shape[0], width=data.shape[1], | |
| transform=src.window_transform(win), | |
| tiled=False, blockysize=1, dtype="float32", | |
| nodata=-9999.0, crs="EPSG:4326", interleave="pixel") | |
| profile.pop("compress", None) # EF5 TifGrid-safe: plain 1-row strips | |
| path = os.path.join(out_dir, fname) | |
| tmp = path + ".tmp" # atomic: parallel runs share this store | |
| with rasterio.open(tmp, "w", **profile) as dst: | |
| dst.write(data, 1) | |
| os.replace(tmp, path) | |
| out[key] = path | |
| if __name__ == "__main__": | |
| import sys, time | |
| out = sys.argv[1] if len(sys.argv) > 1 else "._param_test" | |
| bbox = (-69.83, 46.32, -68.33, 47.82) # Allagash | |
| t = time.time() | |
| grids = clip_param_grids(bbox, out) | |
| print(f"clipped {len(grids)} param grids in {time.time()-t:.1f}s") | |
| for k, p in grids.items(): | |
| with rasterio.open(p) as ds: | |
| a = ds.read(1) | |
| print(f" {k:12s} {os.path.basename(p):10s} {ds.width}x{ds.height} " | |
| f"range[{float(a.min()):.3f},{float(a.max()):.3f}]") | |