代码实战 | 如何基于fnl计算可降水量并绘图

pip install xarray cfgrib wrf-python matplotlib cartopy meteva

import argparse
from pathlib import Path

import numpy as np
import xarray as xr
import wrf
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from meteva.base.tool.plot_tools import add_china_map_2basemap

# --- 变量名映射表 ---
PRESSURE_LEVEL_KEYS = ["isobaricInhPa", "level", "pressure", "lev"]
TEMP_KEYS = ["t", "temperature", "T"]
Q_KEYS = ["q", "qv", "specific_humidity", "QV", "qvapor"]
GHT_KEYS = ["gh", "ght", "geopotential_height", "HGT", "z"]
SURFACE_PS_KEYS = ["sp", "ps", "psfc", "surface_pressure", "PRES"]

def pick_var(ds, candidates, label):
    """从数据集中自动识别变量"""
    for name in candidates:
        if name in ds.data_vars:
            return ds[name]
    raise KeyError(f"未找到{label}变量，候选: {candidates}，当前变量: {list(ds.data_vars)}")


def pick_coord(ds, candidates, label):
    """从数据集中自动识别坐标"""
    for name in candidates:
        if name in ds.coords:
            return ds.coords[name]
    raise KeyError(f"未找到{label}坐标，候选: {candidates}，当前坐标: {list(ds.coords)}")


def open_fnl_datasets(file_path):
    """读取 FNL GRIB2 文件（等压面+地面）"""
    ds_pl = xr.open_dataset(
        file_path,
        engine="cfgrib",
        backend_kwargs={"filter_by_keys": {"typeOfLevel": "isobaricInhPa"}},
    )
    ds_sfc = xr.open_dataset(
        file_path,
        engine="cfgrib",
        backend_kwargs={"filter_by_keys": {"typeOfLevel": "surface"}},
    )
    return ds_pl, ds_sfc

def preprocess_data(ds_pl, ds_sfc):
    """提取变量、层对齐、单位转换、垂直方向调整"""
    t = pick_var(ds_pl, TEMP_KEYS, "温度")
    q = pick_var(ds_pl, Q_KEYS, "比湿")
    ght = pick_var(ds_pl, GHT_KEYS, "位势高度")

    # 关键：共同层对齐
    t, q, ght = xr.align(t, q, ght, join="inner")

    # 构造 3D 气压场 [Pa]
    z_dim = t.dims[0]
    common_pres_hpa = t.coords[z_dim]
    pres_3d, _ = xr.broadcast(common_pres_hpa, t)
    pres_pa = pres_3d * 100.0

    # 垂直顺序：地面 -> 高空（高压到低压）
    if pres_pa.isel({z_dim: 0}).values.mean() < pres_pa.isel({z_dim: -1}).values.mean():
        print("检测到气压层为高空到地面，正在进行垂直翻转...")
        pres_pa = pres_pa.isel({z_dim: slice(None, None, -1)})
        t = t.isel({z_dim: slice(None, None, -1)})
        q = q.isel({z_dim: slice(None, None, -1)})
        ght = ght.isel({z_dim: slice(None, None, -1)})

    print(f"垂直层对齐完成: pres/t/q/ght = {pres_pa.shape}/{t.shape}/{q.shape}/{ght.shape}")
    return pres_pa, t, q, ght

def calculate_pw(pres_pa, tkel, qv, height, meta=True):
    z_dim = pres_pa.dims[0]
    level_coord = pres_pa.coords[z_dim]

    # 兼容 wrf.pw 对 height staggered 维度的要求
    if height.sizes[z_dim] == pres_pa.sizes[z_dim]:
        if pres_pa.sizes[z_dim] < 2:
            raise ValueError("垂直层数不足，无法构造 staggered height")

        top_diff = (height.isel({z_dim: -1}) - height.isel({z_dim: -2})).expand_dims({z_dim: [0]})
        top_diff = top_diff.assign_coords({z_dim: [0]})

        top_height = (height.isel({z_dim: -1}) + top_diff.isel({z_dim: 0})).expand_dims({z_dim: [0]})
        top_height = top_height.assign_coords({z_dim: [float(level_coord.values[-1]) - 1e-3]})

        height_for_wrf = xr.concat([height, top_height], dim=z_dim)
    else:
        height_for_wrf = height

    return wrf.pw(pres_pa, tkel, qv, height_for_wrf, meta=meta)

def plot_pw(lon, lat, pw, extent, output_file, title="PW"):
    fig = plt.figure(figsize=(16, 10))
    ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())

    is_global = extent == [-180, 180, -90, 90]
    if is_global:
        ax.set_global()
    else:
        ax.set_extent(extent, crs=ccrs.PlateCarree())

    levels = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
    cf = ax.contourf(lon, lat, pw, levels=levels, cmap="Blues", extend="max", transform=ccrs.PlateCarree())

    try:
        if is_global:
            add_china_map_2basemap(ax, name="world", edgecolor='black', lw=0.8, encoding='gbk', grid0=None)
        else:
            add_china_map_2basemap(ax, name="river", edgecolor='dodgerblue', lw=0.5, encoding='gbk', grid0=None)
            add_china_map_2basemap(ax, name="nation", edgecolor='black', lw=1.0, encoding='gbk', grid0=None)
            add_china_map_2basemap(ax, name="province", edgecolor='gray', lw=0.5, encoding='gbk', grid0=None)
    except Exception as e:
        ax.coastlines(resolution='110m'if is_global else'50m', lw=0.8)
        ax.add_feature(cfeature.BORDERS, linestyle=':', alpha=0.5)
        print(f"提示: meteva 地图加载失败 ({e})，已切换至 cartopy 海岸线。")

    gl = ax.gridlines(draw_labels=True, linestyle="--", alpha=0.5)
    gl.top_labels = False
    gl.right_labels = False

    cbar = plt.colorbar(cf, ax=ax, label="PW (kg m$^{-2}$)", orientation='vertical', pad=0.02, shrink=0.7)
    cbar.ax.set_title("mm", fontsize=9, pad=5)
    plt.title(f"Precipitable Water (PW)\n{title}", loc="left", fontsize=10)
#    plt.savefig(output_file, dpi=200, bbox_inches="tight")
    print(f"绘图成功，输出文件: {output_file}")
    plt.show()

def main(file_path, output_file, extent):
    file_path = Path(file_path)
    ifnot file_path.exists():
        raise FileNotFoundError(f"文件不存在: {file_path}")

    print("加载数据...")
    ds_pl, ds_sfc = open_fnl_datasets(str(file_path))

    print("预处理...")
    pres_pa, t, q, ght = preprocess_data(ds_pl, ds_sfc)

    print("计算 PW...")
    pw = calculate_pw(pres_pa, t, q, ght, meta=True)
    print(f"PW 范围: {pw.min().values:.2f} - {pw.max().values:.2f} kg/m²")

    lon = pick_coord(ds_pl, ["longitude", "lon"], "经度")
    lat = pick_coord(ds_pl, ["latitude", "lat"], "纬度")

    if lon.max() > 181:
        lon_data = lon.values.copy()
        lon_data[lon_data > 180] -= 360
        lon = xr.DataArray(lon_data, dims=lon.dims, coords=lon.coords, attrs=lon.attrs)
        sort_idx = np.argsort(lon.values)
        lon = lon.isel({lon.dims[0]: sort_idx})
        pw = pw.isel({pw.dims[-1]: sort_idx})

    print("绘图...")
    plot_pw(lon, lat, pw, extent, output_file, title=f"Source: {file_path.name}")
    print(f"输出: {output_file}")

    return pw

# 示例：全球范围
pw = main(
    file_path="fnl_20220101_00_00.grib2",
    output_file="pw_global_wrfpw_from_notebook.png",
    extent=[-180, 180, -90, 90],
)

加载数据...
预处理...
垂直层对齐完成: pres/t/q/ght = (33, 181, 360)/(33, 181, 360)/(33, 181, 360)/(33, 181, 360)
计算 PW...
PW 范围: 0.72 - 87.94 kg/m²
绘图...
绘图成功，输出文件: pw_global_wrfpw_from_notebook.png
输出: pw_global_wrfpw_from_notebook.png