代码实战 | FNL 数据处理与 MCAPE 计算教程

pip install xarray cfgrib wrf-python matplotlib cartopy meteva  

#!/usr/bin/env python3
import argparse
import sys
from pathlib import Path

import numpy as np
import xarray as xr
from wrf import cape_2d
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

[1000, 975, 950, 925, 900, 850, 800, 750, 700, 650, 600, 550, 500,  
 450, 400, 350, 300, 250, 200, 150, 100, 70, 50, 40, 30, 20, 15,  
 10, 7, 5, 3, 2, 1] hPa  

# 变量名映射表（用于自动识别不同命名规范的变量）
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"]
TERRAIN_KEYS = ["orog", "terrain", "HGT_surface", "z_surface", "geopotential_height"]

def open_fnl_datasets(file_path):
    """
    读取 FNL GRIB2 文件
    
    Parameters:
        file_path: FNL GRIB2 文件路径
    
    Returns:
        ds_pl: 等压面数据 (3D)
        ds_sfc: 地面数据 (2D)
    """
    # 读取等压面数据 (3D)
    ds_pl = xr.open_dataset(
        file_path,
        engine="cfgrib",
        backend_kwargs={"filter_by_keys": {"typeOfLevel": "isobaricInhPa"}},
    )
    
    # 读取地面数据 (2D)
    ds_sfc = xr.open_dataset(
        file_path,
        engine="cfgrib",
        backend_kwargs={"filter_by_keys": {"typeOfLevel": "surface"}},
    )
    
    return ds_pl, ds_sfc

# 使用示例
ds_pl, ds_sfc = open_fnl_datasets("fnl_20220101_00_00.grib2")

# 等压面变量（3D）
print("等压面变量:", list(ds_pl.data_vars))
# 输出: ['gh', 't', 'r', 'q', 'w', 'wz', 'u', 'v', 'absv', 'o3mr']

# 地面变量（2D）
print("地面变量:", list(ds_sfc.data_vars))
# 输出包含: ['sp', 'orog', 't', 'cape', 'cin', ...]

# 其中 orog 是地形数据，单位为米

wrf.cape_2d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow, meta=True)  

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 preprocess_data(ds_pl, ds_sfc):
    """
    数据预处理：提取变量、单位转换、维度调整
    """
    # 1. 提取变量
    t = pick_var(ds_pl, TEMP_KEYS, "温度")          # 温度 [K]
    q = pick_var(ds_pl, Q_KEYS, "比湿")            # 比湿 [kg/kg]
    ght = pick_var(ds_pl, GHT_KEYS, "位势高度")    # 位势高度 [m]
    ps = pick_var(ds_sfc, SURFACE_PS_KEYS, "地面气压")  # 地面气压 [Pa]
    cape = pick_var(ds_sfc, SURFACE_PS_KEYS, "地面气压")
    # 提取地形数据
    try:
        terrain = pick_var(ds_sfc, TERRAIN_KEYS, "地形")
    except KeyError:
        terrain = xr.zeros_like(ps)
    
    # 2. 单位转换
    # 温度：确保单位为 K
    if t.max() < 100:
        t = t + 273.15
    
    # 构造 3D 气压场并转换为 hPa
    pres_coord = pick_coord(ds_pl, PRESSURE_LEVEL_KEYS, "气压层")
    pres_3d, _ = xr.broadcast(pres_coord, t)
    pres_hpa = pres_3d  # 已经是 hPa
    
    # 地面气压：Pa -> hPa
    ps_hpa = ps / 100.0
    
    # 3. 垂直层顺序调整（从地面向上：1000 -> 1 hPa）
    # 如果第一层气压小于最后一层（如 1hPa < 1000hPa），需要翻转
    if pres_hpa.isel(isobaricInhPa=0).values.mean() < pres_hpa.isel(isobaricInhPa=-1).values.mean():
        print("检测到气压层为高空到地面，正在进行垂直翻转...")
        pres_hpa = pres_hpa.isel(isobaricInhPa=slice(None, None, -1))
        t = t.isel(isobaricInhPa=slice(None, None, -1))
        q = q.isel(isobaricInhPa=slice(None, None, -1))
        ght = ght.isel(isobaricInhPa=slice(None, None, -1))
    
    return pres_hpa, t, q, ght, terrain, ps_hpa

def calculate_mcape(pres_hpa, t, q, ght, terrain, ps_hpa):
    """
    使用 wrf-python 计算 MCAPE
    
    Parameters:
        pres_hpa: 3D 气压场 [hPa]
        t: 3D 温度 [K]
        q: 3D 比湿 [kg/kg]
        ght: 3D 位势高度 [m]
        terrain: 2D 地形高度 [m]
        ps_hpa: 2D 地面气压 [hPa]
    
    Returns:
        mcape: 2D MCAPE 场 [J/kg]
        mcin: 2D MCIN 场 [J/kg]
        lcl: 2D LCL 场 [m]
        lfc: 2D LFC 场 [m]
    """
    print("正在计算 MCAPE (这可能需要一点时间)...")
    
    cape_result = cape_2d(
        pres_hpa, 
        t, 
        q, 
        ght, 
        terrain, 
        ps_hpa, 
        ter_follow=False,  # FNL 是等压面坐标，非地形跟随
        meta=True
    )
    
    # 提取结果
    mcape = cape_result[0]  # MCAPE
    mcin = cape_result[1]   # MCIN
    lcl = cape_result[2]    # LCL
    lfc = cape_result[3]    # LFC
    
    return mcape, mcin, lcl, lfc

add_china_map_2basemap(ax, name="province", edgecolor='k', lw=0.5, encoding='gbk', grid0=None)  

def plot_mcape(lon, lat, mcape, extent, output_file, title="MCAPE"):
    """
    绘制 MCAPE 分布图
    
    Parameters:
        lon: 经度坐标
        lat: 纬度坐标
        mcape: MCAPE 数据
        extent: 绘图范围 [lon_min, lon_max, lat_min, lat_max]
        output_file: 输出文件名
        title: 图表标题
    """
    # 创建图形
    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, 100, 250, 500, 750, 1000, 1500, 2000, 3000, 4000]
    cf = ax.contourf(lon, lat, mcape, levels=levels, cmap="Spectral_r", 
                     extend="max", transform=ccrs.PlateCarree())
    
    # 添加地图要素
    try:
        if is_global:
            # 全球范围使用 world 参数
            add_china_map_2basemap(ax, name="world", edgecolor='black', lw=0.8, encoding='gbk', grid0=None)
            print("成功加载 meteva 世界地图")
        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)
            print("成功加载 meteva 中国地图")
    except Exception as e:
        # 如果 meteva 加载失败，使用 cartopy 作为备用
        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
    
    # 添加色标和标题
    plt.colorbar(cf, ax=ax, label="MCAPE (J/kg)", orientation='vertical', pad=0.02, shrink=0.7)
    plt.title(f"Maximum Convective Available Potential Energy (MCAPE)\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):
    """
    主函数：完整的 MCAPE 计算和绘图流程
    
    Parameters:
        file_path: FNL GRIB2 文件路径
        output_file: 输出图片路径
        extent: 绘图范围 [lon_min, lon_max, lat_min, lat_max]
    """
    file_path = Path(file_path)
    ifnot file_path.exists():
        raise FileNotFoundError(f"未找到文件: {file_path}")
    
    # 1. 加载数据
    print("步骤 1: 加载 FNL 数据...")
    ds_pl, ds_sfc = open_fnl_datasets(str(file_path))
    
    # 2. 数据预处理
    print("步骤 2: 数据预处理...")
    pres_hpa, t, q, ght, terrain, ps_hpa = preprocess_data(ds_pl, ds_sfc)
    
    # 3. 计算 MCAPE
    print("步骤 3: 计算 MCAPE...")
    mcape, mcin, lcl, lfc = calculate_mcape(pres_hpa, t, q, ght, terrain, ps_hpa)
    
    # 4. 获取经纬度坐标
    lon = pick_coord(ds_pl, ["longitude", "lon"], "经度")
    lat = pick_coord(ds_pl, ["latitude", "lat"], "纬度")
    
    # 5. 经度处理（从 0-360 转换到 -180-180）
    if lon.max() > 181:
        lon_data = lon.values.copy()
        lon_idx = lon_data > 180
        lon_data[lon_idx] -= 360
        lon = xr.DataArray(lon_data, dims=lon.dims, coords=lon.coords, attrs=lon.attrs)
        # 重新排序以防绘图断裂
        sort_idx = np.argsort(lon.values)
        lon = lon.isel(longitude=sort_idx)
        mcape = mcape.isel(longitude=sort_idx)

    # 6. 绘图
    print("步骤 4: 绘制地图...")
    plot_mcape(lon, lat, mcape, extent, output_file, title=f"Source: {file_path.name}")
    
    print("\n完成！")
    return mcape, mcin, lcl, lfc

# 示例 1: 绘制全球 MCAPE 分布
mcape_global, mcin_global, lcl_global, lfc_global = main(
    file_path="/home/mw/project/fnl_20220101_00_00.grib2",
    output_file="mcape_global.png",
    extent=[-180, 180, -90, 90]
)

# 示例 2: 绘制中国区域 MCAPE 分布
mcape_china, mcin_china, lcl_china, lfc_china = main(
    file_path="/home/mw/project/fnl_20220101_00_00.grib2",
    output_file="mcape_china.png",
    extent=[70, 140, 15, 55]
)

def read_fnl_cape(ds_sfc):
    """
    从 FNL 地面层数据中读取自带的 CAPE 和 CIN
    
    Parameters:
        ds_sfc: FNL 地面层数据集

    Returns:
        cape: 对流有效位能 [J/kg]
        cin: 对流抑制 [J/kg]
    """
    # FNL 中 CAPE 的变量名可能是 'cape' 或类似名称
    cape_keys = ["cape", "CAPE", "convective_available_potential_energy"]
    cin_keys = ["cin", "CIN", "convective_inhibition"]

    cape = pick_var(ds_sfc, cape_keys, "CAPE")
    cin = pick_var(ds_sfc, cin_keys, "CIN")

    return cape, cin

   # ============================================
   # 绘制 FNL 自带的 CAPE
   # ============================================

def plot_fnl_cape(lon, lat, cape, extent, output_file, title="FNL CAPE"):
    """
    绘制 FNL 自带的 CAPE 分布图

    Parameters:
        lon: 经度坐标
        lat: 纬度坐标
        cape: CAPE 数据 (来自FNL文件)
        extent: 绘图范围 [lon_min, lon_max, lat_min, lat_max]
        output_file: 输出文件名
        title: 图表标题
    """
    # 创建图形
    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())

    # 绘制填色图 - CAPE 使用与 MCAPE 相同的色标
    levels = [0, 100, 250, 500, 750, 1000, 1500, 2000, 3000, 4000]
    cf = ax.contourf(lon, lat, cape, levels=levels, cmap="Spectral_r",
                    extend="max", transform=ccrs.PlateCarree())

    # 添加地图要素
    try:
        if is_global:
            add_china_map_2basemap(ax, name="world", edgecolor='black', lw=0.8, encoding='gbk')
        else:
            add_china_map_2basemap(ax, name="river", edgecolor='dodgerblue', lw=0.5, encoding='gbk')
            add_china_map_2basemap(ax, name="nation", edgecolor='black', lw=1.0, encoding='gbk')
            add_china_map_2basemap(ax, name="province", edgecolor='gray', lw=0.5, encoding='gbk')
    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)

    # 添加网格线
    gl = ax.gridlines(draw_labels=True, linestyle="--", alpha=0.5)
    gl.top_labels = False
    gl.right_labels = False

    # 添加色标和标题
    plt.colorbar(cf, ax=ax, label="CAPE (J/kg)", orientation='vertical', pad=0.02, shrink=0.7)
    plt.title(f"Convective Available Potential Energy (CAPE) from FNL\n{title}",
                loc="left", fontsize=10)

    # 保存图片
    plt.savefig(output_file, dpi=200, bbox_inches="tight")
    print(f"绘图成功，输出文件: {output_file}")
    plt.show()

   # ============================================
   # 主函数：读取并绘制 FNL 自带的 CAPE
   # ============================================

def plot_fnl_cape_main(file_path, output_file, extent):
    """
    主函数：读取 FNL 自带的 CAPE 并绘图

    Parameters:
        file_path: FNL GRIB2 文件路径
        output_file: 输出图片路径
        extent: 绘图范围 [lon_min, lon_max, lat_min, lat_max]
    """
    file_path = Path(file_path)
    ifnot file_path.exists():
        raise FileNotFoundError(f"未找到文件: {file_path}")

    # 1. 加载数据
    print("步骤 1: 加载 FNL 数据...")
    ds_pl, ds_sfc = open_fnl_datasets(str(file_path))

    # 2. 读取 FNL 自带的 CAPE
    print("步骤 2: 读取 FNL 自带的 CAPE...")
    cape, cin = read_fnl_cape(ds_sfc)

    # 3. 获取经纬度坐标
    lon = pick_coord(ds_pl, ["longitude", "lon"], "经度")
    lat = pick_coord(ds_pl, ["latitude", "lat"], "纬度")

    # 4. 经度处理（从 0-360 转换到 -180-180）
    if lon.max() > 181:
        lon_data = lon.values.copy()
        lon_idx = lon_data > 180
        lon_data[lon_idx] -= 360
        lon = xr.DataArray(lon_data, dims=lon.dims, coords=lon.coords, attrs=lon.attrs)
        sort_idx = np.argsort(lon.values)
        lon = lon.isel(longitude=sort_idx)
        cape = cape.isel(longitude=sort_idx)

    # 5. 绘图
    print("步骤 3: 绘制 CAPE 地图...")
    plot_fnl_cape(lon, lat, cape, extent, output_file, title=f"Source: {file_path.name}")

    print("\n完成！")
    return cape, cin

   # ============================================
   # 使用示例
   # ============================================

#    示例 1: 绘制全球 CAPE 分布
cape_global, cin_global = plot_fnl_cape_main(
    file_path="fnl_20220101_00_00.grib2",
    output_file="fnl_cape_global.png",
    extent=[-180, 180, -90, 90]
)

#    示例 2: 绘制中国区域 CAPE 分布
#    cape_china, cin_china = plot_fnl_cape_main(
#        file_path="fnl_20220101_00_00.grib2",
#        output_file="fnl_cape_china.png",
#        extent=[70, 140, 15, 55]
#    )

lon_data = lon.values.copy()  
lon_data[lon_data > 180] -= 360  
lon = xr.DataArray(lon_data, dims=lon.dims, coords=lon.coords)  

if pres_hpa.isel(isobaricInhPa=0).values.mean() < pres_hpa.isel(isobaricInhPa=-1).values.mean():  
    # 需要翻转（从高空向地面 -> 从地面向高空）  
    pres_hpa = pres_hpa.isel(isobaricInhPa=slice(None, None, -1))  
    t = t.isel(isobaricInhPa=slice(None, None, -1))  
    # ... 其他变量同样处理  

cape_result = cape_2d(pres_hpa, t, q, ght, terrain, ps_hpa, ter_follow=False, meta=True)  

mcape = cape_result[0]  # 最大对流有效位能  
mcin = cape_result[1]   # 最大对流抑制  
lcl = cape_result[2]    # 抬升凝结高度  
lfc = cape_result[3]    # 自由对流高度  

import glob  

files = sorted(glob.glob("fnl_*.grib2"))  
for file in files:  
    output = file.replace(".grib2", "_mcape.png")  
    main(file, output, extent=[70, 140, 15, 55])  

import imageio  

images = []  
for png_file in sorted(glob.glob("*_mcape.png")):  
    images.append(imageio.imread(png_file))  

imageio.mimsave("mcape_animation.gif", images, duration=0.5)