File size: 4,862 Bytes
1a437b8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
#!/usr/bin/env python3
"""
Fetch NOAA GFSWave Arctic 9km GRIB2 and extract wave height, direction, period
without pygrib. Uses xarray+cfgrib (ecCodes backend).
"""
import argparse
import json
import os
import sys
import tempfile
from datetime import datetime
from typing import Dict, Any

import numpy as np
import requests
import xarray as xr


def build_url(date_str: str, run: str, fh: int) -> str:
    base = "https://nomads.ncep.noaa.gov/pub/data/nccf/com/gfs/prod"
    fname = f"gfswave.t{run}z.arctic.9km.f{fh:03d}.grib2"
    return f"{base}/gfs.{date_str}/{run}/wave/gridded/{fname}"


def download(url: str) -> str:
    r = requests.get(url, timeout=300)
    r.raise_for_status()
    tf = tempfile.NamedTemporaryFile(delete=False, suffix=".grib2")
    tf.write(r.content)
    tf.close()
    return tf.name


def open_grib(path: str) -> xr.Dataset:
    # cfgrib uses ecCodes; ensure eccodes is installed on the system
    return xr.open_dataset(path, engine="cfgrib", decode_timedelta=True)


def pick_var(ds: xr.Dataset, cands) -> str | None:
    for c in cands:
        if c in ds.variables:
            return c
    # heuristic fallback
    for v in ds.variables:
        lv = v.lower()
        if ("wave" in lv and "height" in lv) or v in ("HTSGW", "htsgw"):
            return v
    return None


def sample_points(ds: xr.Dataset, n: int = 1000) -> list[Dict[str, Any]]:
    lat_name = "latitude" if "latitude" in ds.variables else "lat"
    lon_name = "longitude" if "longitude" in ds.variables else "lon"
    lats = ds[lat_name].values
    lons = ds[lon_name].values

    wave_var = pick_var(ds, ["swh", "HTSGW", "htsgw"])  # significant wave height
    if wave_var is None:
        raise RuntimeError("No wave height variable found in GRIB")
    wh = ds[wave_var].values

    dir_var = pick_var(ds, [
        "dirpw", "DIRPW", "dp", "wvdir", "WVDIR", "dir", "mwd", "MWD", "MWDIR",
    ])
    per_var = pick_var(ds, [
        "perpw", "PERPW", "tp", "wvper", "WVPER", "per", "pp1d", "PP1D", "mwp", "MWP",
    ])
    wd = ds[dir_var].values if dir_var else None
    wp = ds[per_var].values if per_var else None

    lon_grid, lat_grid = np.meshgrid(lons, lats)
    flat_lats = lat_grid.flatten()
    flat_lons = lon_grid.flatten()
    flat_wh = wh.flatten()

    mask = np.isfinite(flat_wh)
    if wd is not None:
        mask &= np.isfinite(wd.flatten())
    idx_all = np.where(mask)[0]
    if idx_all.size == 0:
        return []
    choose = np.random.choice(idx_all, size=min(n, idx_all.size), replace=False)

    out = []
    for i in choose:
        item: Dict[str, Any] = {
            "lat": float(flat_lats[i]),
            "lon": float(flat_lons[i]),
            "wave_height": float(flat_wh[i]),
        }
        if wd is not None:
            item["wave_direction"] = float(wd.flatten()[i])
        if wp is not None:
            item["wave_period"] = float(wp.flatten()[i])
        out.append(item)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--date", help="YYYYMMDD (default: today UTC)")
    ap.add_argument("--run", default=None, help="Cycle hour: 00,06,12,18 (default: best guess)")
    ap.add_argument("--fh", type=int, default=0, help="Forecast hour (0-240)")
    ap.add_argument("--out-json", default=None, help="Write sampled points JSON")
    # ap.add_argument("--out-nc", default=None, help="Write NetCDF copy of the GRIB")
    args = ap.parse_args()

    now = datetime.utcnow()
    date_str = args.date or now.strftime("%Y%m%d")
    # pick a likely available run by current hour if not provided
    if args.run is None:
        hr = now.hour
        if hr >= 18:
            run = "18"
        elif hr >= 12:
            run = "12"
        elif hr >= 6:
            run = "06"
        else:
            run = "00"
    else:
        run = args.run.zfill(2)

    url = build_url(date_str, run, args.fh)
    print(f"Downloading: {url}")
    path = download(url)
    print(f"Saved: {path}")

    ds = open_grib(path)
    pts = sample_points(ds, n=1000)

    # Summaries
    wh_name = pick_var(ds, ["swh", "HTSGW", "htsgw"]) or "swh"
    wh = ds[wh_name].values
    print(json.dumps({
        "date": date_str,
        "run": run,
        "forecast_hour": args.fh,
        "points": len(pts),
        "height_min": float(np.nanmin(wh)),
        "height_max": float(np.nanmax(wh)),
        "height_mean": float(np.nanmean(wh)),
    }, indent=2))

    if args.out_json:
        with open(args.out_json, "w") as f:
            json.dump({"type": "points", "points": pts}, f)
        print(f"Wrote JSON: {args.out_json}")

    # if args.out_nc:
    #     ds.to_netcdf(args.out_nc)
    #     print(f"Wrote NetCDF: {args.out_nc}")

    ds.close()

    # Clean up temp file
    try:
        os.unlink(path)
    except Exception:
        pass


if __name__ == "__main__":
    sys.exit(main())