diff --git a/nos_utils/forcing/adcirc_met.py b/nos_utils/forcing/adcirc_met.py index 825cd23..6dbce21 100644 --- a/nos_utils/forcing/adcirc_met.py +++ b/nos_utils/forcing/adcirc_met.py @@ -26,6 +26,7 @@ from __future__ import annotations +import functools import hashlib import json import calendar @@ -606,19 +607,54 @@ def sfcf_path(comin_gfs, cycle: datetime, fhr: int) -> Path: "gfs.t{:02d}z.sfcf{:03d}.nc".format(cycle.hour, fhr) -def getges_sfcf(comin_gfs, valid: datetime) -> Path: +def sflux_path(comin_gfs, cycle: datetime, fhr: int) -> Path: + return Path(comin_gfs) / "gfs.{:%Y%m%d}".format(cycle) / "{:02d}".format(cycle.hour) / "atmos" / \ + "gfs.t{:02d}z.sfluxgrbf{:03d}.grib2".format(cycle.hour, fhr) + + +_SOURCE_PATH = {"sfcf": sfcf_path, "sflux": sflux_path} + +# wgrib2 match string -> sfcf variable name, in output order. MJ (10/14/26) +SFLUX_FIELDS = { + "PRES:surface": "pressfc", + "UGRD:10 m above ground": "ugrd10m", + "VGRD:10 m above ground": "vgrd10m", + "ICEC:surface": "icec", +} +# sfcf grid_yt and the GRIB2 Gaussian grid (scanning mode 0) both run north to south, so no flip. MJ (10/14/26) +_SFCF_LAT_ASCENDING = False +_SFLUX_FILL = np.float32(9.99e20) # sfcf _FillValue/missing_value. MJ (10/14/26) +_WGRIB2_TIMEOUT_S = 300 +_SFLUX_UNDEF_MIN = 9.0e20 # wgrib2 writes 9.999e20 at bitmap-undefined points. MJ (10/14/26) +# Attributes of the real sfcf files (netCDF4 types included); im/jm are set from the grid. MJ (10/14/26) +_SFLUX_GLOBAL = {"fhzero": np.int32(6), "ncld": np.int32(5), "nsoil": np.int32(4), "imp_physics": np.int32(11), + "dtp": np.float32(150.0), "source": "FV3GFS", "grid": "gaussian"} +_SFLUX_ATTRS = { + "grid_xt": {"cartesian_axis": "X", "long_name": "T-cell longitude", "units": "degrees_E"}, + "grid_yt": {"cartesian_axis": "Y", "long_name": "T-cell latiitude", "units": "degrees_N"}, # sic, as in sfcf. MJ (10/14/26) + "lat": {"long_name": "T-cell latitude", "units": "degrees_N"}, + "lon": {"long_name": "T-cell longitude", "units": "degrees_E"}, + "time": {"long_name": "time", "cartesian_axis": "T", "calendar_type": "JULIAN", "calendar": "JULIAN"}, + "pressfc": {"long_name": "surface pressure", "units": "Pa"}, + "ugrd10m": {"long_name": "10 meter u wind", "units": "m/s"}, + "vgrd10m": {"long_name": "10 meter v wind", "units": "m/s"}, + "icec": {"long_name": "surface ice concentration (ice=1; no ice=0)", "units": "fraction"}, +} + + +def getges_sfcf(comin_gfs, valid: datetime, source: str = "sfcf") -> Path: """getges.sh -t sfgges -n gfs: the first readable file walking back one hour of lead at a time. The script starts at fhbeg=max(NHOUR valid, 0); a valid time in the past, as in ops, gives 0. A missing f000 therefore falls back to the previous cycle's f006. MJ (10/06/26) """ for fh in range(_GETGES_FHEND + 1): - cand = sfcf_path(comin_gfs, valid - timedelta(hours=fh), fh) - if os.access(str(cand), os.R_OK): + cand = _SOURCE_PATH[source](comin_gfs, valid - timedelta(hours=fh), fh) + if os.access(str(cand), os.R_OK) and (source == "sfcf" or _sfcf_ready(cand)): # a half-downloaded sflux must not win. MJ (10/14/26) return cand raise FileNotAvailableError( - "FATAL ERROR: no GFS sfcf file for valid {:%Y%m%d%H} under {} (searched f000-f{})".format( - valid, comin_gfs, _GETGES_FHEND)) + "FATAL ERROR: no GFS {} file for valid {:%Y%m%d%H} under {} (searched f000-f{})".format( + source, valid, comin_gfs, _GETGES_FHEND)) def sfcf_valid_times(cycle: datetime, segment: str, start: Optional[datetime] = None) -> List[datetime]: @@ -653,6 +689,8 @@ def _sfcf_ready(path: Path) -> bool: """Non-empty and, unless old enough to be complete, its logf marker exists (GFS writes sfcf straight into COM).""" if not (path.is_file() and path.stat().st_size > 0): return False + if ".sfluxgrbf" in path.name: # GFS writes the .idx after the grib, so idx or age means complete. MJ (10/14/26) + return path.with_name(path.name + ".idx").exists() or time.time() - path.stat().st_mtime > _SFCF_SETTLED_S logf = path.with_name(path.name.replace("sfcf", "logf").replace(".nc", ".txt")) return logf.exists() or time.time() - path.stat().st_mtime > _SFCF_SETTLED_S @@ -670,9 +708,12 @@ def _wait_for(path: Path, deadline: float, poll_s: float, sleep, clock) -> bool: def build_sfcf_forcing(comin_gfs, cycle: datetime, segment: str, out_dir, prefix: str = "stofs_2d_glo", start: Optional[datetime] = None, wait_s: float = 0.0, poll_s: float = 10.0, - sleep=time.sleep, clock=time.monotonic, now=time.time) -> Dict[str, Path]: + sleep=time.sleep, clock=time.monotonic, now=time.time, source: str = "sfcf") -> Dict[str, Path]: """Write ``_.{221,222,225}.nc`` from the GFS sfcf files (ncst, or one hourly 180 h fcst). + ``source="sflux"`` reads the sfluxgrbf GRIB2 files (possibly four-message subsets) instead, through wgrib2, + into the same output layout. MJ (10/14/26) + Valid times at or before ``cycle`` use getges; later ones the current cycle's own sfcf file, waited for up to ``wait_s`` seconds in total. Output files appear under their final names only when all three are complete, 221 last (the ops skip test is on 221). Records are copied one at a time. In real time (within @@ -681,6 +722,11 @@ def build_sfcf_forcing(comin_gfs, cycle: datetime, segment: str, out_dir, prefix """ import netCDF4 as nc + if source not in _SOURCE_PATH: + raise ValueError("source must be one of {}, got {!r}".format(tuple(_SOURCE_PATH), source)) + spath = _SOURCE_PATH[source] + # the two bracket files repeat from one hourly record to the next, so keep two decodes. MJ (10/14/26) + opener = nc.Dataset if source == "sfcf" else functools.partial(_SfluxDataset, cache={}) out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) valids = sfcf_valid_times(cycle, segment, start) @@ -694,28 +740,28 @@ def build_sfcf_forcing(comin_gfs, cycle: datetime, segment: str, out_dir, prefix for rec, valid in enumerate(valids): br = _interp_bracket(cycle, valid) if segment == "fcst" else None if br: - srcs = [sfcf_path(comin_gfs, cycle, h) for h in br[:2]] + srcs = [spath(comin_gfs, cycle, h) for h in br[:2]] for sp in srcs: if not _wait_for(sp, deadline, poll_s, sleep, clock): raise FileNotAvailableError("GFS file did not appear before the wait expired: {}".format(sp)) src = srcs[0] elif valid <= cycle: if valid == cycle and now() - calendar.timegm(cycle.timetuple()) < _REALTIME_S: - _wait_for(sfcf_path(comin_gfs, cycle, 0), deadline, poll_s, sleep, clock) - src = getges_sfcf(comin_gfs, valid) + _wait_for(spath(comin_gfs, cycle, 0), deadline, poll_s, sleep, clock) + src = getges_sfcf(comin_gfs, valid, source) else: - src = sfcf_path(comin_gfs, cycle, int((valid - cycle).total_seconds() // 3600)) + src = spath(comin_gfs, cycle, int((valid - cycle).total_seconds() // 3600)) if not _wait_for(src, deadline, poll_s, sleep, clock): raise FileNotAvailableError("GFS file did not appear before the wait expired: {}".format(src)) used.extend((valid, p) for p in (srcs if br else [src])) # interpolated records list both brackets. MJ (10/11/26) log.info("record %d valid %s <- %s", rec, valid, src) - with nc.Dataset(str(src)) as ds: + with opener(str(src)) as ds: ds.set_auto_maskandscale(False) if rec == 0: for n, names in SFC_OUTPUTS: outs[n] = _open_sfcf_output(nc, part[n], ds, names) lat, lon = ds.variables["lat"][:], ds.variables["lon"][:] - hi = nc.Dataset(str(srcs[1])) if br else None + hi = opener(str(srcs[1])) if br else None try: if hi is not None: hi.set_auto_maskandscale(False) @@ -753,7 +799,14 @@ def build_sfcf_forcing(comin_gfs, cycle: datetime, segment: str, out_dir, prefix for p in part.values(): if p.exists(): p.unlink() - _write_sfcf_manifest(out_dir / "{}_{}.sfcf_manifest.json".format(prefix, segment), cycle, segment, used) + try: + from ._log import log_input_files + except ImportError: # loaded by file path outside the package (tools/stage_stofs_2d_glo_ops_forcing.py). MJ (10/14/26) + log_input_files = None + if log_input_files is not None: + log_input_files("GFS_" + source.upper(), list(dict.fromkeys(p for _, p in used)), source="GFS_" + source.upper(), + category="atmospheric", note="cycle={:%Y%m%d%H} segment={}".format(cycle, segment)) + _write_sfcf_manifest(out_dir / "{}_{}.sfcf_manifest.json".format(prefix, segment), cycle, segment, used, source) return final @@ -761,8 +814,8 @@ def _open_sfcf_output(nc, path: Path, src, names: Tuple[str, ...]): """One output file with the ncecat+ncwa layout: record unlimited, lat/lon/data (record, y, x), scalar time.""" out = nc.Dataset(str(path), "w", format="NETCDF4_CLASSIC") ny, nx = len(src.dimensions["grid_yt"]), len(src.dimensions["grid_xt"]) - for k, v in src.__dict__.items(): - out.setncattr(k, v) + for k in src.ncattrs(): + out.setncattr(k, src.getncattr(k)) out.createDimension("grid_xt", nx) out.createDimension("grid_yt", ny) out.createDimension("record", None) @@ -793,10 +846,172 @@ def clone(name, dims, **kw): return out -def _write_sfcf_manifest(path: Path, cycle: datetime, segment: str, used: List[Tuple[datetime, Path]]) -> None: - data = {"source_model": "gfs", "product": "sfcf", "cycle_time": cycle.isoformat(), "segment": segment, +def _write_sfcf_manifest(path: Path, cycle: datetime, segment: str, used: List[Tuple[datetime, Path]], + source: str = "sfcf") -> None: + data = {"source_model": "gfs", "product": source, "cycle_time": cycle.isoformat(), "segment": segment, "nfiles": len(used), "files": [{"valid_time": v.isoformat(), "source_location": str(p), "file_size": p.stat().st_size} for v, p in used]} with open(str(path), "w") as fh: json.dump(data, fh, indent=2) + + +def gfs_input_files(comin_gfs, cycle: datetime, segment: str, start: Optional[datetime] = None, + source: str = "sfcf") -> List[Path]: + """Every file ``build_sfcf_forcing`` would read if all existed: getges' first choice (f000 of the valid + hour) at or before the cycle, the cycle's own files after it, both brackets where interpolated. MJ (10/14/26)""" + spath = _SOURCE_PATH[source] + out = [] # type: List[Path] + for valid in sfcf_valid_times(cycle, segment, start): + br = _interp_bracket(cycle, valid) if segment == "fcst" else None + if br: + cand = [spath(comin_gfs, cycle, h) for h in br[:2]] + elif valid <= cycle: + fh = valid.hour % 6 # getges' first hit on a complete tree: GFS cycles are 6-hourly. MJ (10/14/26) + cand = [spath(comin_gfs, valid - timedelta(hours=fh), fh)] + else: + cand = [spath(comin_gfs, cycle, int((valid - cycle).total_seconds() // 3600))] + out.extend(p for p in cand if p not in out) + return out + + +def _decode_sflux(path: Path): + """(lat, lon, {sfcf name: 2D float32}) from one sflux GRIB2 file via wgrib2, in file (scan) order. MJ (10/14/26)""" + import re + import tempfile + exe = shutil.which("wgrib2") + if exe is None: + raise MetDependencyError("wgrib2 not found on PATH (load the wgrib2 module) to read {}".format(path)) + with tempfile.TemporaryDirectory() as td: + bins = {m: os.path.join(td, "f{}.bin".format(i)) for i, m in enumerate(SFLUX_FIELDS)} + cmd = [exe, str(path), "-order", "raw", "-no_header"] + for m, b in bins.items(): + cmd += ["-if", ":{}:".format(m), "-bin", b, "-fi"] + first = next(iter(SFLUX_FIELDS)) + cmd += ["-if", ":{}:".format(first), "-grid", "-fi"] + try: + r = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, + timeout=_WGRIB2_TIMEOUT_S) + except subprocess.TimeoutExpired: + raise MetForcingError("wgrib2 timed out after {} s on {}".format(_WGRIB2_TIMEOUT_S, path)) + if r.returncode != 0: + raise MetForcingError("wgrib2 failed on {}: {}".format(path, r.stderr.strip())) + g = re.search(r"\((\d+) x (\d+)\).*?lat ([-\d.]+) to ([-\d.]+).*?lon ([-\d.]+) to ([-\d.]+)", + r.stdout, re.S) + if not g: + raise MetForcingError("{}: cannot read the grid from wgrib2 -grid output: {!r}".format(path, r.stdout[:200])) + scan = re.search(r"input (\w\w:\w\w)", r.stdout) + if not scan or scan.group(1) not in ("WE:NS", "WE:SN"): # a j-consecutive scan would garble the reshape. MJ (10/14/26) + raise MetForcingError("{}: unsupported GRIB scan order {}".format(path, scan.group(1) if scan else None)) + nx, ny = int(g.group(1)), int(g.group(2)) + lat0, lon0 = float(g.group(3)), float(g.group(5)) + lat = np.sort(np.rad2deg(np.arcsin(np.polynomial.legendre.leggauss(ny)[0]))) # Gaussian latitudes. MJ (10/14/26) + if lat0 > 0: + lat = lat[::-1] + if abs(lat[0] - lat0) > 1e-3: + raise MetForcingError("{}: not a {}-latitude Gaussian grid (first lat {})".format(path, ny, lat0)) + lon = lon0 + np.arange(nx) * (360.0 / nx) # global regular grid; the header's lon1 is rounded. MJ (10/14/26) + fields = {} + for m, b in bins.items(): + if not os.path.exists(b) or os.path.getsize(b) != 4 * nx * ny: + raise MetForcingError("{}: field {!r} missing or the wrong size".format(path, m)) + fields[SFLUX_FIELDS[m]] = np.fromfile(b, dtype=np.float32).reshape(ny, nx) + return lat, lon, fields + + +def _to_sfcf_orientation(lat, fields): + """Flip latitude to the sfcf order when the decoder returned the other one (one decision). MJ (10/14/26)""" + if (lat[0] < lat[-1]) == _SFCF_LAT_ASCENDING: + return lat, fields + return lat[::-1], {k: v[::-1] for k, v in fields.items()} + + +def _resolve_undefined(name, a): + """GFS sflux has no bitmap and ice-free/land icec is a real 0.0; if a bitmap ever appears, icec gets 0.0 + there and any other field raises. MJ (10/14/26)""" + bad = a >= _SFLUX_UNDEF_MIN + if not bad.any(): + return np.ascontiguousarray(a, dtype=np.float32) + if name != "icec": + raise MetForcingError("sflux {} has {} undefined points".format(name, int(bad.sum()))) + return np.where(bad, np.float32(0.0), a).astype(np.float32) + + +class _SfluxVar(object): + """The slice of a netCDF4 Variable that build_sfcf_forcing and _open_sfcf_output use. MJ (10/14/26)""" + + def __init__(self, data, attrs): + self._data, self._attrs = np.asarray(data), dict(attrs) + self.dtype = self._data.dtype + + def __getitem__(self, k): + return self._data[k] + + def ncattrs(self): + return list(self._attrs) + + def getncattr(self, k): + return self._attrs[k] + + def __getattr__(self, k): + try: + return self.__dict__["_attrs"][k] + except KeyError: + raise AttributeError(k) + + +class _SfluxDataset(object): + """An sflux GRIB2 file presented like the sfcf netCDF it replaces (same variables, dims, attributes). MJ (10/14/26)""" + + def __init__(self, path, cache=None): + import re + m = re.search(r"gfs\.(\d{8})/(\d{2})/atmos/gfs\.t\d{2}z\.sfluxgrbf(\d{3})\.grib2$", str(path)) + if not m: + raise MetForcingError("not an sflux path: {}".format(path)) + cyc, fhr = datetime.strptime(m.group(1) + m.group(2), "%Y%m%d%H"), int(m.group(3)) + if cache is not None and str(path) in cache: + lat, lon, fields = cache[str(path)] + else: + lat, lon, fields = _decode_sflux(Path(path)) + if cache is not None: + if len(cache) >= 2: + cache.pop(next(iter(cache))) + cache[str(path)] = (lat, lon, fields) + lat, fields = _to_sfcf_orientation(lat, fields) + lat, lon = np.asarray(lat, dtype="f8"), np.asarray(lon, dtype="f8") + ny, nx = len(lat), len(lon) + a = _SFLUX_ATTRS + self.dimensions = {"grid_xt": range(nx), "grid_yt": range(ny), "time": range(1)} + self.variables = { + "grid_xt": _SfluxVar(lon, a["grid_xt"]), + "grid_yt": _SfluxVar(lat, a["grid_yt"]), + "lat": _SfluxVar(np.repeat(lat[:, None], nx, axis=1), a["lat"]), + "lon": _SfluxVar(np.repeat(lon[None, :], ny, axis=0), a["lon"]), + "time": _SfluxVar(np.array([fhr], dtype="f8"), + dict(a["time"], units="hours since {:%Y-%m-%d %H:%M:%S}".format(cyc))), + } + for _, name in SFLUX_FIELDS.items(): + if name not in fields: + raise MetForcingError("{}: decoder returned no {}".format(path, name)) + self.variables[name] = _SfluxVar( + _resolve_undefined(name, fields[name])[None], dict(a[name], missing_value=_SFLUX_FILL, _FillValue=_SFLUX_FILL, + cell_methods="time: point", output_file="sfc")) + self._globals = dict(_SFLUX_GLOBAL, im=np.int32(nx), jm=np.int32(ny)) + + def ncattrs(self): + return list(self._globals) + + def getncattr(self, k): + return self._globals[k] + + def set_auto_maskandscale(self, flag): + pass + + def close(self): + pass + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False diff --git a/tests/test_adcirc_met.py b/tests/test_adcirc_met.py index 022595a..d5bc5f0 100644 --- a/tests/test_adcirc_met.py +++ b/tests/test_adcirc_met.py @@ -608,3 +608,237 @@ def test_retrospective_cycle_record_does_not_wait(tmp_path): epoch = __import__("calendar").timegm(CYC.timetuple()) am.build_sfcf_forcing(tmp_path, CYC, "ncst", tmp_path / "o", start=CYC - timedelta(hours=1), wait_s=30, sleep=lambda s: pytest.fail("waited"), now=lambda: epoch + 3 * 86400) + + +# --- the same forcing from the sfluxgrbf GRIB2 files (decoder mocked) ------------------------- + +_SF_NAMES = {"ugrd10m": 0, "vgrd10m": 1, "pressfc": 2, "icec": 3} + + +def _mk_sflux(comin, cycle, fhr, seed=None): + p = am.sflux_path(comin, cycle, fhr) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes(b"GRIB") + p.with_name(p.name + ".idx").write_text("1:0") # the downloader keeps the idx next to the grib. MJ (10/14/26) + return p + + +def _fake_decode(descending=False, undef=False): + import re + + def dec(path): + m = re.search(r"gfs\.(\d{8})/(\d{2})/atmos/.*sfluxgrbf(\d{3})", str(path)) + cyc, fhr = datetime.strptime(m.group(1) + m.group(2), "%Y%m%d%H"), int(m.group(3)) + seed = _valid_seed(cyc, fhr) + lat, lon = np.arange(NY, dtype="f8"), np.arange(NX, dtype="f8") + ramp = np.arange(NY * NX, dtype="f4").reshape(NY, NX) + f = {k: np.full((NY, NX), seed * 10 + i, dtype="f4") + ramp for k, i in _SF_NAMES.items()} + if undef: + f["icec"][0, 0] = np.float32(9.999e20) + if descending: + lat, f = lat[::-1], {k: v[::-1] for k, v in f.items()} + return lat, lon, f + return dec + + +def _fill_sflux(tank, pairs): + for cycle, fhr in pairs: + _mk_sflux(tank, cycle, fhr) + + +def _same_structure(a, b): + import netCDF4 as nc + with nc.Dataset(str(a)) as x, nc.Dataset(str(b)) as y: + assert list(x.dimensions) == list(y.dimensions) + assert list(x.variables) == list(y.variables) + for k, v in x.variables.items(): + w = y.variables[k] + assert (v.dtype, v.dimensions, v.shape) == (w.dtype, w.dimensions, w.shape), k + assert v.filters() == w.filters() and v.chunking() == w.chunking(), k + assert getattr(v, "_FillValue", None) == getattr(w, "_FillValue", None), k + assert getattr(v, "cell_methods", None) == getattr(w, "cell_methods", None), k + assert np.array_equal(v[:], w[:]), k + assert x.variables["time"].units == y.variables["time"].units + + +def test_sflux_matches_sfcf_path_ncst_and_fcst(tmp_path, monkeypatch): + monkeypatch.setattr(am, "_SFCF_LAT_ASCENDING", True) # the fixture sfcf grid and fake decoder are both south-north. MJ (10/14/26) + monkeypatch.setattr(am, "_decode_sflux", _fake_decode()) + c18 = CYC - timedelta(hours=6) + ncst = [(CYC - timedelta(hours=12), 6)] + [(c18, h) for h in range(1, 6)] + [(CYC, 0)] + fc = [(CYC, h) for h in range(0, 121)] + [(CYC, h) for h in range(123, 181, 3)] + _fill_tank(tmp_path / "a", ncst + fc) + _fill_sflux(tmp_path / "b", ncst + fc) + for seg in ("ncst", "fcst"): + fa = am.build_sfcf_forcing(tmp_path / "a", CYC, seg, tmp_path / "oa") + fb = am.build_sfcf_forcing(tmp_path / "b", CYC, seg, tmp_path / "ob", source="sflux") + for n in fa: + _same_structure(fa[n], fb[n]) + man = json.loads((tmp_path / "ob" / "stofs_2d_glo_fcst.sfcf_manifest.json").read_text()) + assert man["product"] == "sflux" and len(man["files"]) == 181 + 40 + assert man["files"][121]["source_location"] == str(am.sflux_path(tmp_path / "b", CYC, 120)) + assert man["files"][122]["source_location"] == str(am.sflux_path(tmp_path / "b", CYC, 123)) + assert json.loads((tmp_path / "oa" / "stofs_2d_glo_fcst.sfcf_manifest.json").read_text())["product"] == "sfcf" + + +def test_sflux_orientation_is_one_decision(tmp_path, monkeypatch): + _fill_sflux(tmp_path, [(CYC, 0)]) + monkeypatch.setattr(am, "_decode_sflux", _fake_decode(descending=True)) + f = am.build_sfcf_forcing(tmp_path, CYC, "ncst", tmp_path / "o1", start=CYC, source="sflux") # no flip: N to S kept. MJ (10/14/26) + assert np.array_equal(_read(f["221"], "grid_yt"), np.arange(NY)[::-1]) + assert _read(f["221"], "pressfc")[0][0, 0] == _valid_seed(CYC, 0) * 10 + 2 + (NY - 1) * NX + monkeypatch.setattr(am, "_SFCF_LAT_ASCENDING", True) + f = am.build_sfcf_forcing(tmp_path, CYC, "ncst", tmp_path / "o2", start=CYC, source="sflux") + assert np.array_equal(_read(f["221"], "grid_yt"), np.arange(NY)) + assert _read(f["221"], "pressfc")[0][0, 0] == _valid_seed(CYC, 0) * 10 + 2 + + +def test_sflux_bitmap_points_icec_zero_others_raise(tmp_path, monkeypatch): + monkeypatch.setattr(am, "_SFCF_LAT_ASCENDING", True) + _fill_sflux(tmp_path, [(CYC, 0)]) + monkeypatch.setattr(am, "_decode_sflux", _fake_decode(undef=True)) + f = am.build_sfcf_forcing(tmp_path, CYC, "ncst", tmp_path / "o", start=CYC, source="sflux") + i = _read(f["225"], "icec")[0] + assert i[0, 0] == 0.0 and i[0, 1] > 0 + with pytest.raises(am.MetForcingError, match="undefined"): + am._resolve_undefined("pressfc", np.full((2, 2), 9.999e20, dtype="f4")) + + +def test_sflux_attributes_match_sfcf(tmp_path, monkeypatch): + import netCDF4 as nc + _fill_sflux(tmp_path, [(CYC, 0)]) + monkeypatch.setattr(am, "_decode_sflux", _fake_decode()) + f = am.build_sfcf_forcing(tmp_path, CYC, "ncst", tmp_path / "o", start=CYC, source="sflux") + with nc.Dataset(str(f["222"])) as d: + g = {k: d.getncattr(k) for k in d.ncattrs()} + assert g == {"fhzero": 6, "ncld": 5, "nsoil": 4, "imp_physics": 11, "dtp": 150.0, "source": "FV3GFS", + "grid": "gaussian", "im": NX, "jm": NY} + assert type(g["fhzero"]) is np.int32 and type(g["dtp"]) is np.float32 + u = d.variables["ugrd10m"] + assert (u.long_name, u.units, u.output_file) == ("10 meter u wind", "m/s", "sfc") + assert u.getncattr("missing_value") == np.float32(9.99e20) == u.getncattr("_FillValue") + assert d.variables["grid_yt"].long_name == "T-cell latiitude" + assert d.variables["lon"].dtype == np.dtype("f8") and d.variables["lat"].dimensions == ("record", "grid_yt", "grid_xt") + assert d.variables["time"].units == "hours since 2026-10-04 00:00:00" and d.variables["time"].calendar == "JULIAN" + with nc.Dataset(str(f["225"])) as d: + assert (d.variables["icec"].units, d.variables["icec"].cell_methods) == ("fraction", "time: point time: mean") + + +def test_gfs_input_files(tmp_path): + c18, c12 = CYC - timedelta(hours=6), CYC - timedelta(hours=12) + n = am.gfs_input_files(tmp_path, CYC, "ncst") + assert n == [am.sfcf_path(tmp_path, c18, h) for h in range(0, 6)] + [am.sfcf_path(tmp_path, CYC, 0)] + assert len(am.gfs_input_files(tmp_path, CYC, "ncst", start=CYC - timedelta(hours=12))) == 13 + sf = am.gfs_input_files(tmp_path, CYC, "ncst", source="sflux") # includes the cycle's own f000 analysis. MJ (10/14/26) + assert sf == [am.sflux_path(tmp_path, c18, h) for h in range(0, 6)] + [am.sflux_path(tmp_path, CYC, 0)] + f = am.gfs_input_files(tmp_path, CYC, "fcst", source="sflux") + assert f == [am.sflux_path(tmp_path, CYC, h) for h in list(range(0, 121)) + list(range(123, 181, 3))] + assert len(f) == len(set(f)) + assert c18 != c12 + + +def test_sflux_missing_file_errors(tmp_path, monkeypatch): + monkeypatch.setattr(am, "_decode_sflux", _fake_decode()) + with pytest.raises(am.FileNotAvailableError, match="no GFS sflux file"): + am.build_sfcf_forcing(tmp_path, CYC, "ncst", tmp_path / "o", start=CYC, source="sflux") + _fill_sflux(tmp_path, [(CYC, h) for h in range(0, 121)]) + with pytest.raises(am.FileNotAvailableError, match="did not appear"): + am.build_sfcf_forcing(tmp_path, CYC, "fcst", tmp_path / "o", source="sflux") + + +def test_sflux_without_wgrib2_is_a_clear_error(tmp_path, monkeypatch): + monkeypatch.setattr(am.shutil, "which", lambda _: None) + with pytest.raises(am.MetDependencyError, match="wgrib2"): + am._decode_sflux(am.sflux_path(tmp_path, CYC, 0)) + + +@pytest.mark.parametrize("source", ["sfcf", "sflux"]) +@pytest.mark.parametrize("segment", ["ncst", "fcst"]) +def test_manifest_files_equal_gfs_input_files(tmp_path, monkeypatch, source, segment): + import re + monkeypatch.setattr(am, "_SFCF_LAT_ASCENDING", True) + monkeypatch.setattr(am, "_decode_sflux", _fake_decode()) + tank = tmp_path / "t" + want = am.gfs_input_files(tank, CYC, segment, source=source) + for p in want: # a tank holding exactly the listed files, nothing else. MJ (10/14/26) + h = int(re.search(r"f(\d{3})\.grib2$|sfcf(\d{3})\.nc$", p.name).group(1, 2)[0 if source == "sflux" else 1]) + cyc = datetime.strptime(p.parts[-4][4:] + p.parts[-3], "%Y%m%d%H") + if source == "sfcf": + _mk_sfcf(tank, cyc, h, _valid_seed(cyc, h)) + else: + _mk_sflux(tank, cyc, h) + am.build_sfcf_forcing(tank, CYC, segment, tmp_path / "o", source=source) + man = json.loads((tmp_path / "o" / "stofs_2d_glo_{}.sfcf_manifest.json".format(segment)).read_text()) + assert list(dict.fromkeys(e["source_location"] for e in man["files"])) == [str(p) for p in want] + + +def test_sflux_ready_needs_idx_or_age(tmp_path): + import os + p = _mk_sflux(tmp_path, CYC, 0) + assert am._sfcf_ready(p) + p.with_name(p.name + ".idx").unlink() + assert not am._sfcf_ready(p) + q = _mk_sflux(tmp_path, CYC, 1) + q.with_name(q.name + ".idx").unlink() + old = q.stat().st_mtime - 3600 + os.utime(str(q), (old, old)) + assert am._sfcf_ready(q) + + +def test_sflux_decoder_rejects_bad_scan_and_times_out(tmp_path, monkeypatch): + monkeypatch.setattr(am.shutil, "which", lambda _: "/bin/wgrib2") + out = "gaussian grid: (4 x 3) units 1e-06 input {} output WE:SN lat 1 to -1 lon 0 to 270" + monkeypatch.setattr(am.subprocess, "run", lambda *a, **k: type( + "R", (), {"returncode": 0, "stdout": out.format("WE:EW"), "stderr": ""})()) + with pytest.raises(am.MetForcingError, match="scan order"): + am._decode_sflux(tmp_path / "x.grib2") + + def boom(*a, **k): + raise am.subprocess.TimeoutExpired("wgrib2", 1) + monkeypatch.setattr(am.subprocess, "run", boom) + with pytest.raises(am.MetForcingError, match="timed out"): + am._decode_sflux(tmp_path / "x.grib2") + + +def test_sflux_bracket_files_decoded_once_per_build(tmp_path, monkeypatch): + _fill_sflux(tmp_path, [(CYC, h) for h in range(0, 121)] + [(CYC, h) for h in range(123, 181, 3)]) + inner, calls = _fake_decode(), [] + monkeypatch.setattr(am, "_decode_sflux", lambda p: (calls.append(str(p)), inner(p))[1]) + monkeypatch.setattr(am, "_SFCF_LAT_ASCENDING", True) + am.build_sfcf_forcing(tmp_path, CYC, "fcst", tmp_path / "o", source="sflux") + assert len(calls) == len(set(calls)) == 121 + 20 + + +def test_getges_sflux_skips_a_file_that_is_not_ready(tmp_path): + c18 = CYC - timedelta(hours=6) + _mk_sflux(tmp_path, c18, 6) + half = _mk_sflux(tmp_path, CYC, 0) + half.with_name(half.name + ".idx").unlink() # fresh and no idx: still downloading. MJ (10/14/26) + assert am.getges_sfcf(tmp_path, CYC, "sflux") == am.sflux_path(tmp_path, c18, 6) + half.with_name(half.name + ".idx").write_text("1:0") + assert am.getges_sfcf(tmp_path, CYC, "sflux") == half + + +def test_getges_sfcf_needs_only_a_readable_file(tmp_path): + p = _mk_sfcf(tmp_path, CYC, 0, 1, logf=False) # no logf marker and fresh: sfcf still takes it. MJ (10/14/26) + assert am.getges_sfcf(tmp_path, CYC) == p + + +@pytest.mark.parametrize("source", ["sfcf", "sflux"]) +def test_build_logs_the_used_input_files(tmp_path, monkeypatch, source): + from nos_utils.forcing import _log + monkeypatch.setattr(am, "_SFCF_LAT_ASCENDING", True) + monkeypatch.setattr(am, "_decode_sflux", _fake_decode()) + tank = tmp_path / "t" + want = am.gfs_input_files(tank, CYC, "ncst", source=source) + for p in want: + cyc = datetime.strptime(p.parts[-4][4:] + p.parts[-3], "%Y%m%d%H") + h = int(p.name.rsplit("sfcf" if source == "sfcf" else "sfluxgrbf", 1)[1][:3]) + (_mk_sfcf(tank, cyc, h, _valid_seed(cyc, h)) if source == "sfcf" else _mk_sflux(tank, cyc, h)) + calls = [] + monkeypatch.setattr(_log, "log_input_files", lambda *a, **k: calls.append((a, k))) + am.build_sfcf_forcing(tank, CYC, "ncst", tmp_path / "o", source=source) + assert len(calls) == 1 + (name, files), kw = calls[0] + assert name == "GFS_" + source.upper() and kw["category"] == "atmospheric" + assert list(files) == want