Spaces:
Running
Running
| """Auto-update the CREST_data temperature forcing from NARR (task: NARR temp). | |
| Source: NCEP North American Regional Reanalysis 2-m air temperature | |
| (GDEX d608000, https://gdex.ucar.edu/datasets/d608000/ — 32 km Lambert | |
| Conformal, 3-hourly, GRIB1). We pull it from NOAA PSL's netCDF mirror of the | |
| SAME dataset (one yearly file per variable, anonymous HTTPS, updated ~weekly): | |
| https://psl.noaa.gov/thredds/fileServer/Datasets/NARR/monolevel/air.2m.YYYY.nc | |
| (downloads.psl.noaa.gov served the same files until it began answering 403 | |
| for the whole NARR tree on 2026-10-03; it stays as a fallback mirror) | |
| because the GDEX form is 0.5–1 GB multi-variable GRIB tars per ~8 days, i.e. | |
| ~2.6 GB of download per month to extract a single band. | |
| Each 3-hourly Lambert field is bilinearly resampled onto the canonical | |
| NLDAS-heritage 0.125° temp grid (forcing.TEMP_GRID — so every temp PQF in the | |
| store shares one geometry, which the [TEMPForcing] DEM= extrapolation grid is | |
| built against), converted K→°C, and linearly interpolated to HOURLY members | |
| (EF5 zero-fills missing hourly files, which would freeze SNOW17 at 0 °C). | |
| Members use the generic name temp_YYYYMMDDHH.pqf; existing NLDAS members are | |
| never touched — the updater only APPENDS missing hours to each month tar. | |
| Run on the local machine (needs netCDF4/pyproj/numpy + HF token): | |
| python scripts/update_temp_narr.py [--months 2026-06,2026-07] [--dry-run] | |
| Default: scan from BACKFILL_START (or the last 4 months, whichever is later) | |
| through the current month. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import io | |
| import json | |
| import os | |
| import re | |
| import sys | |
| import tarfile | |
| import tempfile | |
| import time | |
| from datetime import datetime, timedelta, timezone | |
| import numpy as np | |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| import truststore | |
| truststore.inject_into_ssl() | |
| import urllib.request | |
| from hf_data.forcing import TEMP_GRID, VARS, _write_pqf # noqa: E402 | |
| HF_REPO = "vincewin/CREST_data" | |
| PSL_URLS = ("https://psl.noaa.gov/thredds/fileServer/Datasets/NARR/monolevel/air.2m.{year}.nc", | |
| "https://downloads.psl.noaa.gov/Datasets/NARR/monolevel/air.2m.{year}.nc") | |
| TOKEN_PATH = r"C:\Users\chenm\Documents\EF5\CREST_token.txt" | |
| DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_narr_cache") | |
| BACKFILL_START = (2026, 6) # NLDAS members end 2026-06-23 12z | |
| _MEMBER_RES = [re.compile(r"NLDAS_FORA0125_H\.A(\d{8})\.(\d{2})00\."), | |
| re.compile(r"temp_(\d{10})\.pqf$")] | |
| def _member_hours(names) -> set[datetime]: | |
| """Hours already present in a tar (NLDAS or generic member names).""" | |
| hours = set() | |
| for n in names: | |
| m = _MEMBER_RES[0].search(n) | |
| if m: | |
| hours.add(datetime.strptime(m.group(1) + m.group(2), "%Y%m%d%H")) | |
| continue | |
| m = _MEMBER_RES[1].search(n) | |
| if m: | |
| hours.add(datetime.strptime(m.group(1), "%Y%m%d%H")) | |
| return hours | |
| def _download_year(year: int) -> str | None: | |
| """Fetch air.2m.<year>.nc if PSL has a newer copy than our cache.""" | |
| os.makedirs(DATA_DIR, exist_ok=True) | |
| dest = os.path.join(DATA_DIR, f"air.2m.{year}.nc") | |
| meta = dest + ".meta.json" | |
| url = remote = None | |
| for u in (m.format(year=year) for m in PSL_URLS): | |
| try: | |
| req = urllib.request.Request(u, method="HEAD") | |
| h = urllib.request.urlopen(req, timeout=60) | |
| url, remote = u, {"len": h.headers.get("Content-Length"), | |
| "mod": h.headers.get("Last-Modified")} | |
| break | |
| except Exception as e: | |
| print(f" [narr] HEAD {u} failed: {e}") | |
| if url is None: | |
| # no mirror answered: the cached copy is all we have — say so loudly, | |
| # a silent fallback hid the 2026-10-03 source move | |
| print(f" [narr] WARNING: no NARR mirror reachable for {year}; " | |
| f"{'using cached copy' if os.path.exists(dest) else 'no cached copy'}") | |
| return dest if os.path.exists(dest) else None | |
| if os.path.exists(dest) and os.path.exists(meta): | |
| if json.load(open(meta)) == remote: | |
| return dest | |
| print(f" [narr] downloading {url} ({int(remote['len'] or 0)/1e6:.0f} MB)…") | |
| # PSL rate-limits bursts (HTTP 429 on 2026-10-07, 21 min after it | |
| # republished the year file) — back off and retry transient errors | |
| # instead of failing the whole self-heal until the next cooldown | |
| for attempt, wait in enumerate((30, 120, 300, None)): | |
| try: | |
| urllib.request.urlretrieve(url, dest + ".part") | |
| break | |
| except Exception as e: | |
| code = getattr(e, "code", None) | |
| transient = code is None or code == 429 or code >= 500 | |
| if wait is None or not transient: | |
| raise | |
| ra = getattr(e, "headers", None) and e.headers.get("Retry-After") | |
| if ra and str(ra).isdigit(): | |
| wait = max(wait, min(int(ra), 900)) | |
| print(f" [narr] download failed ({e}); retry {attempt + 1}/3 in {wait} s") | |
| time.sleep(wait) | |
| os.replace(dest + ".part", dest) | |
| json.dump(remote, open(meta, "w")) | |
| return dest | |
| class NarrYear: | |
| """air.2m.<year>.nc with lazy bilinear resampling onto TEMP_GRID.""" | |
| def __init__(self, path: str): | |
| import netCDF4 | |
| self.ds = netCDF4.Dataset(path) | |
| t = self.ds.variables["time"] | |
| import cftime # noqa: F401 (num2date backend) | |
| import netCDF4 as nc4 | |
| # PSL's 2026-10-07 rewrite of air.2m.2026.nc (HDF5) holds Jan-Jul as | |
| # DAILY records, 1,732 masked (empty) time slots where August belongs, | |
| # and only September 3-hourly. A masked stamp crashed num2date; a daily | |
| # record read as a 00Z instant would be silently wrong. Index only | |
| # records that sit inside a 3-hourly run. | |
| tv = t[:] | |
| mask = np.ma.getmaskarray(tv) | |
| vals = np.ma.getdata(tv).astype("float64") | |
| good = np.nonzero(~mask)[0] | |
| gv = vals[good] | |
| step3 = np.zeros(len(good), bool) | |
| if len(good) > 1: | |
| d = np.diff(gv) | |
| nxt3 = np.isclose(d, 3.0) | |
| step3[:-1] |= nxt3 # opens/continues a 3-h step | |
| # reached by a 3-h step and not opening a daily run: a daily record | |
| # stamped D 00Z sits exactly 3 h after D-1 21Z (residual ambiguity: | |
| # a lone daily record at the very END of a file looks the same) | |
| step3[1:] |= nxt3 & np.append(~np.isclose(d[1:], 24.0), True) | |
| dates = nc4.num2date(gv[step3], t.units) if step3.any() else [] | |
| self.hours = {datetime(d.year, d.month, d.day, d.hour): int(i) | |
| for i, d in zip(good[step3], dates)} | |
| n_off = int((~step3).sum()) | |
| if mask.any() or n_off: | |
| print(f" [narr] {os.path.basename(path)}: {len(self.hours)} 3-hourly records " | |
| f"usable, {int(mask.sum())} empty time slots, {n_off} non-3-hourly " | |
| f"(e.g. daily) records ignored") | |
| self.air = self.ds.variables["air"] | |
| gm = self.ds.variables["Lambert_Conformal"] | |
| x = self.ds.variables["x"][:].astype("float64") | |
| y = self.ds.variables["y"][:].astype("float64") | |
| from pyproj import CRS, Transformer | |
| crs = CRS.from_cf({k: getattr(gm, k) for k in gm.ncattrs()}) | |
| tr = Transformer.from_crs("EPSG:4326", crs, always_xy=True) | |
| xll, yll, cell, nr, nc = TEMP_GRID[:5] | |
| lons = xll + (np.arange(nc) + 0.5) * cell | |
| lats = yll + (np.arange(nr) + 0.5) * cell | |
| LON, LAT = np.meshgrid(lons, lats[::-1]) # row 0 = north | |
| X, Y = tr.transform(LON, LAT) | |
| self.fi = (Y - y[0]) / (y[1] - y[0]) | |
| self.fj = (X - x[0]) / (x[1] - x[0]) | |
| self._cache: dict[datetime, np.ndarray] = {} | |
| def field(self, when: datetime) -> np.ndarray | None: | |
| """°C field on TEMP_GRID at a 3-hourly NARR step, or None.""" | |
| if when in self._cache: | |
| return self._cache[when] | |
| idx = self.hours.get(when) | |
| if idx is None: | |
| return None | |
| a = np.ma.filled(self.air[idx, :, :].astype("float64"), np.nan) - 273.15 | |
| if not np.isfinite(a).any(): | |
| return None # empty record: no data, not a nodata grid | |
| i0 = np.clip(np.floor(self.fi).astype(int), 0, a.shape[0] - 2) | |
| j0 = np.clip(np.floor(self.fj).astype(int), 0, a.shape[1] - 2) | |
| di = np.clip(self.fi - i0, 0, 1) | |
| dj = np.clip(self.fj - j0, 0, 1) | |
| v = (a[i0, j0] * (1 - di) * (1 - dj) + a[i0 + 1, j0] * di * (1 - dj) | |
| + a[i0, j0 + 1] * (1 - di) * dj + a[i0 + 1, j0 + 1] * di * dj) | |
| out = np.where(np.isfinite(v), v, TEMP_GRID[5]).astype("float32") | |
| if len(self._cache) > 4: | |
| self._cache.clear() | |
| self._cache[when] = out | |
| return out | |
| def close(self): | |
| self.ds.close() | |
| def hourly_field(years: dict[int, NarrYear], when: datetime) -> np.ndarray | None: | |
| """Linear interpolation between the two bracketing 3-hourly NARR steps.""" | |
| t0 = when.replace(hour=when.hour - when.hour % 3) | |
| ny = years.get(t0.year) | |
| f0 = ny.field(t0) if ny else None | |
| if f0 is None: | |
| return None | |
| if when == t0: | |
| return f0 | |
| t1 = t0 + timedelta(hours=3) | |
| ny1 = years.get(t1.year) | |
| f1 = ny1.field(t1) if ny1 else None | |
| if f1 is None: | |
| return None # can't bracket -> leave for the next run | |
| w = (when - t0).total_seconds() / 10800.0 | |
| bad = (f0 == TEMP_GRID[5]) | (f1 == TEMP_GRID[5]) | |
| v = ((1 - w) * f0 + w * f1).astype("float32") | |
| v[bad] = TEMP_GRID[5] | |
| return v | |
| def update_month(year: int, month: int, years: dict[int, NarrYear], api, | |
| dry_run: bool) -> str: | |
| from huggingface_hub import hf_hub_download | |
| cfg = VARS["temp"] | |
| ref = datetime(year, month, 1) | |
| tar_repo_path = ref.strftime(cfg.month_fmt) | |
| old_members: list[tuple[tarfile.TarInfo, bytes]] = [] | |
| have: set[datetime] = set() | |
| try: | |
| local = hf_hub_download(HF_REPO, tar_repo_path, repo_type="dataset") | |
| with tarfile.open(local) as tf: | |
| for ti in tf.getmembers(): | |
| old_members.append((ti, tf.extractfile(ti).read())) | |
| have = _member_hours([ti.name for ti, _ in old_members]) | |
| except Exception: | |
| pass # brand-new month tar | |
| now = datetime.now(timezone.utc).replace(tzinfo=None) | |
| end = min(now, datetime(year + (month == 12), month % 12 + 1, 1)) | |
| want, t = [], ref | |
| while t < end: | |
| if t not in have: | |
| want.append(t) | |
| t += timedelta(hours=1) | |
| if not want: | |
| return f"{year}-{month:02d}: complete ({len(have)} hours), nothing to do" | |
| new: list[tuple[str, bytes]] = [] | |
| xll, yll, cell, nr, nc = TEMP_GRID[:5] | |
| top = yll + nr * cell | |
| for t in want: | |
| v = hourly_field(years, t) | |
| if v is None: | |
| continue | |
| buf = os.path.join(tempfile.gettempdir(), "narr_member.pqf") | |
| _write_pqf(buf, v, xll, yll, cell, TEMP_GRID[5]) | |
| new.append((t.strftime(cfg.out_fmt), open(buf, "rb").read())) | |
| if not new: | |
| src_end = max((max(ny.hours) for ny in years.values() if ny.hours), | |
| default=None) | |
| if src_end is not None and src_end >= end: | |
| return (f"{year}-{month:02d}: SOURCE GAP — {len(want)} hour(s) missing and " | |
| f"NARR has later data (to {src_end:%Y-%m-%d %H}Z) but none for them") | |
| return (f"{year}-{month:02d}: {len(want)} hour(s) missing but NARR has " | |
| f"no data for them yet (lag)") | |
| if dry_run: | |
| return f"{year}-{month:02d}: would add {len(new)} NARR hour(s) to {len(have)} existing" | |
| out_tar = os.path.join(tempfile.gettempdir(), f"temp_{year}_{month:02d}.tar") | |
| with tarfile.open(out_tar, "w") as tf: | |
| for ti, data in old_members: | |
| tf.addfile(ti, io.BytesIO(data)) | |
| for name, data in sorted(new): | |
| ti = tarfile.TarInfo(name=name) | |
| ti.size = len(data) | |
| ti.mtime = int(datetime.now(timezone.utc).timestamp()) | |
| tf.addfile(ti, io.BytesIO(data)) | |
| from huggingface_hub import CommitOperationAdd | |
| api.create_commit( | |
| repo_id=HF_REPO, repo_type="dataset", | |
| operations=[CommitOperationAdd(path_in_repo=tar_repo_path, | |
| path_or_fileobj=out_tar)], | |
| commit_message=f"temp {year}-{month:02d}: +{len(new)} NARR-derived " | |
| f"hourly member(s) (32 km air.2m -> 0.125deg, K->C, " | |
| f"3h->1h linear)") | |
| os.remove(out_tar) | |
| return f"{year}-{month:02d}: uploaded +{len(new)} NARR hour(s) ({len(have)} kept)" | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--months", help="comma list like 2026-06,2026-07") | |
| ap.add_argument("--dry-run", action="store_true") | |
| args = ap.parse_args() | |
| now = datetime.now(timezone.utc) | |
| if args.months: | |
| months = [tuple(int(x) for x in m.split("-")) for m in args.months.split(",")] | |
| else: | |
| months = [] | |
| back4 = now - timedelta(days=123) # rescan the last ~4 months | |
| y, m = max(BACKFILL_START, (back4.year, back4.month)) | |
| while (y, m) <= (now.year, now.month): | |
| months.append((y, m)) | |
| y, m = (y + 1, 1) if m == 12 else (y, m + 1) | |
| years: dict[int, NarrYear] = {} | |
| need_years = {y for y, _ in months} | {y + 1 for y, mo in months if mo == 12} | |
| for yr in sorted(need_years): | |
| p = _download_year(yr) | |
| if p: | |
| try: | |
| years[yr] = NarrYear(p) | |
| except Exception as e: | |
| print(f" [narr] cannot open {p}: {e}") | |
| if not years: | |
| print("no NARR data available — aborting") | |
| return 1 | |
| from huggingface_hub import HfApi | |
| tok = os.environ.get("HF_TOKEN") or ( | |
| open(TOKEN_PATH).read().strip() if os.path.exists(TOKEN_PATH) else None) | |
| api = HfApi(token=tok) | |
| for y, m in months: | |
| try: | |
| print(update_month(y, m, years, api, args.dry_run)) | |
| except Exception as e: | |
| print(f"{y}-{m:02d}: FAILED — {e}") | |
| for ny in years.values(): | |
| ny.close() | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |