1
#!/usr/bin/env python
2
# -*- coding: utf-8 -*-
3
"""
4
从 NOMADS 自动下载 GFS 0.25° GRIB2 文件的简单脚本。
5

6
- 支持指定日期、起报时间、预报小时
7
- 默认下载“最新起报时间 + f000 分析场”
8
"""
9

10
import os
11
import sys
12
import datetime as dt
13
import requests
14

15
NOMADS_BASE = "https://nomads.ncep.noaa.gov/pub/data/nccf/com/gfs/prod"
16

17

18
def build_gfs_0p25_url(date, cycle, fhour):
19
    """
20
    构造 GFS 0.25° GRIB2 文件的下载 URL
21

22
    :param date: dt.date 对象
23
    :param cycle: 起报小时字符串, e.g. "00", "06", "12", "18"
24
    :param fhour: 预报小时整数, e.g. 0, 6, 12 ...
25
    """
26
    ymd = date.strftime("%Y%m%d")
27
    # 文件名形如: gfs.t00z.pgrb2.0p25.f000
28
    fname = f"gfs.t{cycle}z.pgrb2.0p25.f{fhour:03d}"
29

30
    url = f"{NOMADS_BASE}/gfs.{ymd}/{cycle}/atmos/{fname}"
31
    return url, fname
32

33

34
def download_file(url, out_path, chunk_size=1024 * 1024):
35
    """
36
    通用下载函数，带简单的进度打印
37
    """
38
    print(f"[INFO] Downloading: {url}")
39
    resp = requests.get(url, stream=True, timeout=60)
40
    resp.raise_for_status()
41

42
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
43

44
    total = int(resp.headers.get("Content-Length", 0))
45
    downloaded = 0
46

47
    with open(out_path, "wb") as f:
48
        for chunk in resp.iter_content(chunk_size=chunk_size):
49
            if not chunk:
50
                continue
51
            f.write(chunk)
52
            downloaded += len(chunk)
53
            if total > 0:
54
                done = int(50 * downloaded / total)
55
                sys.stdout.write(
56
                    "\r[INFO] Progress: [{}{}] {:.1f}%".format(
57
                        "#" * done,
58
                        "." * (50 - done),
59
                        100 * downloaded / total,
60
                    )
61
                )
62
                sys.stdout.flush()
63

64
    print(f"\n[INFO] Saved to {out_path}")
65

66

67
def download_latest_gfs(out_dir, fhour=0, cycle=None):
68
    """
69
    下载最新一套 GFS 0.25° 数据中的一个预报时次
70

71
    :param out_dir: 输出目录
72
    :param fhour: 预报小时 (0 = 分析场)
73
    :param cycle: 可指定 "00"/"06"/"12"/"18"，不指定则自动选最近的一个
74
    """
75
    now_utc = dt.datetime.utcnow()
76

77
    if cycle is None:
78
        # 简单根据当前 UTC 小时选一个最近的 cycle
79
        hour = now_utc.hour
80
        if hour >= 18:
81
            cycle = "18"
82
        elif hour >= 12:
83
            cycle = "12"
84
        elif hour >= 6:
85
            cycle = "06"
86
        else:
87
            cycle = "00"
88

89
    # 如果当前时间还没到该 cycle 完整产出，就往前推一天
90
    date = now_utc.date()
91
    if now_utc.hour < int(cycle):
92
        date = date - dt.timedelta(days=1)
93

94
    url, fname = build_gfs_0p25_url(date, cycle, fhour)
95
    out_path = os.path.join(out_dir, fname)
96

97
    download_file(url, out_path)
98
    return out_path
99

100

101
if __name__ == "__main__":
102
    # 示例用法：python download_gfs_grib2.py [out_dir] [fhour]
103
    if len(sys.argv) >= 2:
104
        out_dir = sys.argv[1]
105
    else:
106
        out_dir = "./data/raw"
107

108
    if len(sys.argv) >= 3:
109
        fhour = int(sys.argv[2])
110
    else:
111
        fhour = 0  # 默认下载分析场
112

113
    download_latest_gfs(out_dir, fhour=fhour)

1
"""
2
grib2_to_wind3d.py
3
==================
4
Convert GFS GRIB2 file to Wind3D JSON format (full global, no region crop).
5
Extracts U (east-west wind), V (north-south wind), W (vertical velocity)
6
at multiple pressure levels.
7

8
Usage:
9
  .venv\\Scripts\\python.exe scripts\\grib2_to_wind3d.py
10

11
Dependencies:
12
  pip install xarray cfgrib eccodes numpy
13
"""
14

15
import os
16
import sys
17
import io
18
import json
19
import numpy as np
20
import xarray as xr
21

22
# Fix Windows console encoding
23
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
24
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")
25

26
# ============================================================
27
# Configuration
28
# ============================================================
29

30
INPUT_FILE = os.path.join("public", "gfs.t00z.pgrb2.1p00 (1).anl")
31
OUTPUT_FILE = os.path.join("public", "wind3d.json")
32

33
# Target pressure levels (hPa) - will match any available in the file
34
TARGET_LEVELS = [1000, 975, 950, 925, 900, 850, 800, 700, 600, 500, 400, 300, 250, 200, 150, 100]
35

36
# Approximate height (meters) for each pressure level
37
LEVEL_HEIGHTS = {
38
    1000: 100, 975: 300, 950: 500, 925: 750,
39
    900: 1000, 850: 1500, 800: 2000, 700: 3000,
40
    600: 4200, 500: 5500, 400: 7200, 300: 9000,
41
    250: 10400, 200: 12000, 150: 13500, 100: 16000,
42
}
43

44

45
# ============================================================
46
# Helper Functions
47
# ============================================================
48

49
def open_grib_variable(filepath, short_name):
50
    """Read a single variable from GRIB2 at isobaric levels."""
51
    try:
52
        ds = xr.open_dataset(
53
            filepath,
54
            engine="cfgrib",
55
            backend_kwargs={
56
                "filter_by_keys": {
57
                    "typeOfLevel": "isobaricInhPa",
58
                    "shortName": short_name,
59
                }
60
            },
61
        )
62
        return ds
63
    except Exception as e:
64
        print(f"  [WARN] Cannot read variable '{short_name}': {e}")
65
        return None
66

67

68
def get_levels_from_dataset(ds):
69
    """Extract available pressure levels from dataset."""
70
    level_dim = "isobaricInhPa"
71
    if level_dim in ds.dims:
72
        return sorted(ds[level_dim].values.tolist()), True
73
    elif level_dim in ds.coords:
74
        val = float(ds[level_dim].values)
75
        return [val], False
76
    else:
77
        return [], False
78

79

80
def select_levels(da, level_dim, levels, is_multi_level):
81
    """Select pressure levels from a DataArray (no geographic cropping)."""
82
    if is_multi_level and level_dim in da.dims:
83
        return da.sel(**{level_dim: levels})
84
    return da
85

86

87
# ============================================================
88
# Main
89
# ============================================================
90

91
def main():
92
    print("=" * 60)
93
    print("  GFS GRIB2 -> Wind3D JSON Converter (Full Global)")
94
    print("=" * 60)
95

96
    if not os.path.exists(INPUT_FILE):
97
        print(f"\n[ERROR] File not found: {INPUT_FILE}")
98
        sys.exit(1)
99

100
    file_size_mb = os.path.getsize(INPUT_FILE) / (1024 * 1024)
101
    print(f"\nInput:  {INPUT_FILE} ({file_size_mb:.1f} MB)")
102
    print(f"Output: {OUTPUT_FILE}")
103
    print(f"Mode:   Full global (no region crop)")
104

105
    # ----------------------------------------------------------
106
    # 1. Read U, V, W separately
107
    # ----------------------------------------------------------
108
    print("\n[1/6] Reading GRIB2 data...")
109

110
    print("  -> U component (UGRD)...")
111
    u_ds = open_grib_variable(INPUT_FILE, "u")
112

113
    print("  -> V component (VGRD)...")
114
    v_ds = open_grib_variable(INPUT_FILE, "v")
115

116
    print("  -> W component (VVEL/omega)...")
117
    w_ds = open_grib_variable(INPUT_FILE, "w")
118

119
    if u_ds is None or v_ds is None:
120
        print("\n[ERROR] Cannot read U or V wind components!")
121
        sys.exit(1)
122

123
    u_var = list(u_ds.data_vars)[0]
124
    v_var = list(v_ds.data_vars)[0]
125
    print(f"\n  U var: '{u_var}', dims: {list(u_ds[u_var].dims)}")
126
    print(f"  V var: '{v_var}', dims: {list(v_ds[v_var].dims)}")
127
    if w_ds is not None:
128
        w_var = list(w_ds.data_vars)[0]
129
        print(f"  W var: '{w_var}', dims: {list(w_ds[w_var].dims)}")
130

131
    # ----------------------------------------------------------
132
    # 2. Determine available pressure levels
133
    # ----------------------------------------------------------
134
    print("\n[2/6] Analyzing pressure levels...")
135

136
    available_levels, is_multi_level = get_levels_from_dataset(u_ds)
137
    print(f"  Available: {available_levels}")
138

139
    selected_levels = sorted([l for l in TARGET_LEVELS if l in available_levels])
140
    if not selected_levels:
141
        if available_levels:
142
            selected_levels = sorted([int(l) for l in available_levels])
143
            print(f"  [INFO] No target levels matched, using all: {selected_levels}")
144
        else:
145
            print(f"  [ERROR] No pressure levels found!")
146
            sys.exit(1)
147
    else:
148
        print(f"  Selected: {selected_levels}")
149

150
    single_level = len(selected_levels) == 1
151
    if single_level:
152
        print(f"  [INFO] Only 1 level, will replicate across layers.")
153

154
    # ----------------------------------------------------------
155
    # 3. Select levels (full spatial extent, no cropping)
156
    # ----------------------------------------------------------
157
    print("\n[3/6] Selecting levels (full spatial extent)...")
158

159
    level_dim = "isobaricInhPa"
160

161
    u_data = select_levels(u_ds[u_var], level_dim, selected_levels, is_multi_level)
162
    v_data = select_levels(v_ds[v_var], level_dim, selected_levels, is_multi_level)
163

164
    has_w = False
165
    if w_ds is not None:
166
        w_var = list(w_ds.data_vars)[0]
167
        try:
168
            w_data = select_levels(w_ds[w_var], level_dim, selected_levels, is_multi_level)
169
            has_w = True
170
        except Exception as e:
171
            print(f"  [WARN] W selection failed: {e}")
172

173
    # ----------------------------------------------------------
174
    # 4. Convert to numpy arrays
175
    # ----------------------------------------------------------
176
    print("\n[4/6] Processing data...")
177

178
    u_arr = u_data.values
179
    v_arr = v_data.values
180

181
    if u_arr.ndim == 2:
182
        u_arr = u_arr[np.newaxis, :, :]
183
        v_arr = v_arr[np.newaxis, :, :]
184

185
    if has_w:
186
        w_arr = w_data.values
187
        if w_arr.ndim == 2:
188
            w_arr = w_arr[np.newaxis, :, :]
189
        # Convert omega (Pa/s) -> vertical velocity (m/s)
190
        w_arr = -w_arr / (1.225 * 9.8)
191
        print("  [OK] W: converted from Pa/s to m/s")
192
    else:
193
        w_arr = np.zeros_like(u_arr)
194
        print("  [INFO] W: not found, using zeros")
195

196
    # Replace NaN with 0
197
    u_arr = np.nan_to_num(u_arr, nan=0.0)
198
    v_arr = np.nan_to_num(v_arr, nan=0.0)
199
    w_arr = np.nan_to_num(w_arr, nan=0.0)
200

201
    # Get coordinates
202
    lons = u_data.longitude.values
203
    lats = u_data.latitude.values
204

205
    # Ensure latitude ascending (south -> north)
206
    if len(lats) > 1 and lats[0] > lats[-1]:
207
        lats = lats[::-1]
208
        u_arr = u_arr[:, ::-1, :]
209
        v_arr = v_arr[:, ::-1, :]
210
        w_arr = w_arr[:, ::-1, :]
211

212
    nz_raw, ny, nx = u_arr.shape
213

214
    # If single level, replicate
215
    if single_level and len(TARGET_LEVELS) > 1:
216
        nz_target = len(TARGET_LEVELS)
217
        u_arr = np.repeat(u_arr, nz_target, axis=0)
218
        v_arr = np.repeat(v_arr, nz_target, axis=0)
219
        w_arr = np.repeat(w_arr, nz_target, axis=0)
220
        output_levels = TARGET_LEVELS
221
        print(f"  [INFO] Replicated single level to {nz_target} layers")
222
    else:
223
        output_levels = selected_levels
224

225
    nz, ny, nx = u_arr.shape
226

227
    print(f"\n  Grid: {nx} x {ny} x {nz} (lon x lat x levels)")
228
    print(f"  Lon:  {lons[0]:.2f} to {lons[-1]:.2f}")
229
    print(f"  Lat:  {lats[0]:.2f} to {lats[-1]:.2f}")
230
    print(f"  Levels: {output_levels}")
231
    print(f"  U: [{u_arr.min():.2f}, {u_arr.max():.2f}] m/s")
232
    print(f"  V: [{v_arr.min():.2f}, {v_arr.max():.2f}] m/s")
233
    print(f"  W: [{w_arr.min():.6f}, {w_arr.max():.6f}] m/s")
234

235
    # ----------------------------------------------------------
236
    # 5. Build output JSON
237
    # ----------------------------------------------------------
238
    print(f"\n[5/6] Building JSON...")
239

240
    output = {
241
        "header": {
242
            "nx": int(nx),
243
            "ny": int(ny),
244
            "nz": int(nz),
245
            "lo1": round(float(lons[0]), 4),
246
            "lo2": round(float(lons[-1]), 4),
247
            "la1": round(float(lats[0]), 4),    # south (min lat)
248
            "la2": round(float(lats[-1]), 4),    # north (max lat)
249
            "dx": round(float(lons[1] - lons[0]), 4) if nx > 1 else 1.0,
250
            "dy": round(float(lats[1] - lats[0]), 4) if ny > 1 else 1.0,
251
            "levels": [int(l) for l in output_levels],
252
            "levelHeights": [LEVEL_HEIGHTS.get(int(l), 0) for l in output_levels],
253
            "uMin": round(float(u_arr.min()), 4),
254
            "uMax": round(float(u_arr.max()), 4),
255
            "vMin": round(float(v_arr.min()), 4),
256
            "vMax": round(float(v_arr.max()), 4),
257
            "wMin": round(float(w_arr.min()), 6),
258
            "wMax": round(float(w_arr.max()), 6),
259
        },
260
        "data": {
261
            "u": [round(float(x), 2) for x in u_arr.flatten()],
262
            "v": [round(float(x), 2) for x in v_arr.flatten()],
263
            "w": [round(float(x), 4) for x in w_arr.flatten()],
264
        },
265
    }
266

267
    # ----------------------------------------------------------
268
    # 6. Write file
269
    # ----------------------------------------------------------
270
    print(f"\n[6/6] Writing {OUTPUT_FILE}...")
271
    with open(OUTPUT_FILE, "w") as f:
272
        json.dump(output, f, separators=(",", ":"))
273

274
    out_size = os.path.getsize(OUTPUT_FILE)
275
    unit = "KB" if out_size < 1024 * 1024 else "MB"
276
    size_val = out_size / 1024 if unit == "KB" else out_size / (1024 * 1024)
277
    print(f"  File size: {size_val:.1f} {unit}")
278

279
    total_points = nx * ny * nz
280
    print(f"  Total points: {total_points} ({nx}x{ny}x{nz})")
281

282
    print(f"\n{'=' * 60}")
283
    print(f"  DONE! Full global wind3d.json is ready.")
284
    print(f"{'=' * 60}")
285

286

287
if __name__ == "__main__":
288
    main()