"""Concrete GRIB -> canonical preprocessing (Phase 1, against the FROZEN schema). Implemented ONLY after inspecting a real ECDS s2s-reforecasts GRIB; see docs/canonical_schema.md and metadata/grib_inspection.md. Path (D6): GRIB2 --cdo -f nc copy--> CF NetCDF --xarray--> standardize --> CanonicalField Standardization (frozen rules): flip latitude to ascending; K -> degC for temperature; keep native coarse grid; valid_time = init + endStep (UTC). """ import os import re import subprocess import tempfile import numpy as np from .interfaces import CanonicalField, Preprocessor from utils.timealign import valid_time, local_calendar_day def grib_to_netcdf(grib_path, nc_path=None): """Convert GRIB -> NetCDF with cdo (no cfgrib). Returns the NetCDF path.""" if nc_path is None: _fd, nc_path = tempfile.mkstemp(suffix=".nc") os.close(_fd) # mkstemp leaves the fd open -> close to avoid an fd leak r = subprocess.run(["cdo", "-s", "-f", "nc", "copy", grib_path, nc_path], capture_output=True, text=True, timeout=300) if r.returncode != 0 or not os.path.exists(nc_path): raise RuntimeError(f"cdo conversion failed: {r.stderr[:500]}") return nc_path def _grib_key(grib_path, key): """Read one eccodes key from the first GRIB message.""" r = subprocess.run(["grib_get", "-w", "count=1", "-p", key, grib_path], capture_output=True, text=True, timeout=60) return r.stdout.strip() class S2SReforecastPreprocessor(Preprocessor): """ECDS s2s-reforecasts control-forecast GRIB -> CanonicalField(s).""" def __init__(self, tz_hours=7): self.tz_hours = tz_hours # ---- interface methods -------------------------------------------- def extract_valid_times(self, init_dt, leads_hours): """valid = init + lead (hours), UTC. Aligned by real date.""" return valid_time(init_dt, np.asarray(leads_hours), "h") def standardize_coordinates(self, ds): """Return (lat_asc, lon, lat_ascending_flag). Flip lat to ascending.""" lat = ds["lat"].values if lat[0] > lat[-1]: ds = ds.isel(lat=slice(None, None, -1)) return ds def standardize_units(self, values_K): """K -> degC for temperature fields.""" return np.asarray(values_K, dtype="float64") - 273.15 # ---- driver ------------------------------------------------------- def load_canonical(self, grib_path, to_degC=True): """Read one reforecast GRIB into a CanonicalField (coarse grid).""" import xarray as xr # dates from eccodes (authoritative), independent of cdo time decoding init_s = _grib_key(grib_path, "dataDate") # e.g. 20220105 mv_s = _grib_key(grib_path, "mars.date") or init_s # model version date init_dt = np.datetime64(f"{init_s[:4]}-{init_s[4:6]}-{init_s[6:8]}T00") mv_dt = np.datetime64(f"{mv_s[:4]}-{mv_s[4:6]}-{mv_s[6:8]}") short = _grib_key(grib_path, "shortName") units = _grib_key(grib_path, "units") level_type = "sfc" if "heightAboveGround" in _grib_key(grib_path, "typeOfLevel") \ or "surface" in _grib_key(grib_path, "typeOfLevel") else "pl" nc = grib_to_netcdf(grib_path) try: ds = xr.open_dataset(nc) ds = self.standardize_coordinates(ds) var = [v for v in ds.data_vars][0] da = ds[var] if "height" in da.dims: da = da.isel(height=0) # end-step (lead) hours = valid_time - init, from cdo's time coord vt = da["time"].values.astype("datetime64[h]") leads = (vt - init_dt.astype("datetime64[h]")).astype("timedelta64[h]").astype(int) values = da.values.astype("float64") if to_degC and units == "K": values = self.standardize_units(values) field = CanonicalField( name=short, values=values.astype("float32"), lat=ds["lat"].values, lon=ds["lon"].values, init_time=init_dt, model_version_date=mv_dt, valid_time=self.extract_valid_times(init_dt, leads), lead=leads, lead_unit="h", level_type=level_type, member_type="cf", units=("degC" if (to_degC and units == "K") else units), provenance={"grib": os.path.basename(grib_path), "shortName": short, "stepType": _grib_key(grib_path, "stepType")}, ) ds.close() return field finally: if nc.startswith(tempfile.gettempdir()): os.remove(nc) # ---- multi-hindcast files (20 init dates per model-version date) ---- def split_by_init(self, grib_path, shortname=None): """Split a multi-hindcast GRIB into one temp GRIB per dataDate (init), optionally keeping only messages of a given shortName. Returns {dataDate: path}.""" out = subprocess.run(["grib_get", "-p", "dataDate", grib_path], capture_output=True, text=True, timeout=180).stdout.split() dds = sorted({d for d in out if d.isdigit() and len(d) == 8}) where = "dataDate={}" + (f",shortName={shortname}" if shortname else "") parts = {} for dd in dds: _fd, tmp = tempfile.mkstemp(suffix=f"_{dd}.grib") os.close(_fd) subprocess.run(["grib_copy", "-w", where.format(dd), grib_path, tmp], capture_output=True, text=True, timeout=180) if os.path.exists(tmp) and os.path.getsize(tmp) > 0: parts[dd] = tmp return parts def load_canonical_by_init(self, grib_path, to_degC=True, shortname=None): """Return a list of CanonicalField, one per hindcast init date. If ``shortname`` is given, keep only that variable (for multi-var files).""" parts = self.split_by_init(grib_path, shortname=shortname) fields = [] try: for dd in sorted(parts): fields.append(self.load_canonical(parts[dd], to_degC=to_degC)) finally: for p in parts.values(): if os.path.exists(p): os.remove(p) return fields def daily_max_from_6h(field, tz_hours=7): """Collapse 6-hourly mx2t6 windows into daily-max Tx on the LOCAL (UTC+tz) calendar day. Returns (local_days, tx_daily[day,lat,lon], coverage[day]). coverage = number of 6 h windows contributing to each local day (a day with <4 is partial, e.g. LT0's day). """ local_day = local_calendar_day(field.valid_time, tz_hours) days = np.unique(local_day) Y, X = field.values.shape[1:] out = np.full((len(days), Y, X), np.nan, dtype="float64") cov = np.zeros(len(days), dtype=int) for k, d in enumerate(days): sel = local_day == d cov[k] = int(sel.sum()) out[k] = np.nanmax(field.values[sel], axis=0) return days, out, cov