WRF | 如何使用WRF计算K指数

pip install easyclimate xarray cartopy pooch

import os
import numpy as np
import matplotlib.pyplot as plt
from netCDF4 import Dataset
from matplotlib.cm import get_cmap
import cartopy.crs as ccrs
from cartopy.feature import NaturalEarthFeature
from wrf import getvar, interplevel, latlon_coords, get_cartopy, cartopy_xlim, cartopy_ylim
from meteva.base.tool.plot_tools import add_china_map_2basemap

# ==================== 配置参数 ====================
WRF_FILE = "/home/mw/input/typhoon9537/wrfout_d01_2019-08-08_20_00_00"# 替换为你的wrfout文件路径
OUTPUT_DIR = "./output"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ==================== 核心函数 ====================
def k_index_calculation(wrf_file):
    """
    从WRF输出文件计算K指数
    参数:
        wrf_file: wrfout文件路径
    返回:
        k_index: K指数数组
        lats, lons: 经纬度坐标
        time: 时间对象
    """
    ncfile = Dataset(wrf_file)

    print("正在提取变量...")
    temp = getvar(ncfile, "tc", timeidx=0)      # 温度(°C)
    td = getvar(ncfile, "td", timeidx=0)        # 露点温度(°C)
    pressure = getvar(ncfile, "pressure", timeidx=0)  # 气压(hPa)

    print("正在插值到标准等压面...")
    t_850 = interplevel(temp, pressure, 850.0)
    t_700 = interplevel(temp, pressure, 700.0)
    t_500 = interplevel(temp, pressure, 500.0)
    td_850 = interplevel(td, pressure, 850.0)
    td_700 = interplevel(td, pressure, 700.0)

    print("正在计算K指数...")
    k_index = (t_850.values - t_500.values) + td_850.values - (t_700.values - td_700.values)

    lats, lons = latlon_coords(t_850)
    time = t_850.Time
    ncfile.close()
    return k_index, lats, lons, time

def plot_k_index(k_index, lats, lons, time, output_path):
    cart_proj = ccrs.PlateCarree()
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(1, 1, 1, projection=cart_proj)
    ax.set_extent([115, 130, 20, 35], crs=cart_proj)

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

    contour_levels = np.arange(-10, 51, 5)
    ax.contour(lons, lats, k_index, levels=contour_levels,
               colors="black", linewidths=0.5, transform=ccrs.PlateCarree())
    cf = ax.contourf(lons, lats, k_index, levels=contour_levels,
                     cmap=get_cmap("RdYlBu_r"), extend="both",
                     transform=ccrs.PlateCarree())

    cbar = fig.colorbar(cf, ax=ax, shrink=0.7, pad=0.02)
    cbar.set_label("K指数 (°C)", fontsize=12)

    plt.title(f"WRF模式K指数分析\n{str(time.values)[:19]} UTC", fontsize=14, fontweight="bold", pad=15)

    gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
                      linewidth=1, color='gray', alpha=0.5, linestyle='--')
    gl.top_labels = False
    gl.right_labels = False

    ax.text(0.02, 0.02,
            "K指数等级:\n<20 无对流\n20-25 弱\n26-30 中等\n31-35 强\n>35 极强",
            transform=ax.transAxes, fontsize=9,
            bbox=dict(boxstyle="round", facecolor="white", alpha=0.8))

    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.show()

# ==================== 主程序 ====================
print("=" * 50)
print("WRF K指数计算程序启动")
print("=" * 50)

k_index, lats, lons, time = k_index_calculation(WRF_FILE)

print("\nK指数统计信息:")
print(f"  最小值: {np.nanmin(k_index):.2f} °C")
print(f"  最大值: {np.nanmax(k_index):.2f} °C")
print(f"  平均值: {np.nanmean(k_index):.2f} °C")

output_path = os.path.join(OUTPUT_DIR, "k_index_analysis.png")
plot_k_index(k_index, lats, lons, time, output_path)
print("\n程序执行完毕！")