代码实战 | Skyborn | 假相当位温计算代码修正与开源共享

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple skyborn

Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple
Collecting skyborn
  Downloading skyborn-0.3.19-cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (2.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.0/2.0 MB 1.6 MB/s eta 0:00:00
...
Successfully installed skyborn-0.3.19

import skyborn
help(skyborn.calc.calculations.calculate_theta_se)

Help on function calculate_theta_se in module skyborn.calc.calculations:

calculate_theta_se(
    temperature: Union[numpy.ndarray, xarray.core.dataarray.DataArray], 
    pressure: Union[numpy.ndarray, xarray.core.dataarray.DataArray], 
    mixing_ratio: Union[numpy.ndarray, xarray.core.dataarray.DataArray], 
    dewpoint: Union[numpy.ndarray, xarray.core.dataarray.DataArray]
) -> Union[numpy.ndarray, xarray.core.dataarray.DataArray]

    Calculate Pseudo-equivalent potential temperature (theta-se)
    
    The formula is: 
    θ_se = T * (1000 / (p - e)) ** (Ra / cpd) * exp(L * r / (cpd * Tc))
    
    From: http://stream1.cmatc.cn/cmatcvod/12/tqx/second_content.html
    
    Parameters
    ----------
    temperature : array-like
        Temperature in Kelvin.
    pressure : array-like
        Pressure in hPa.
    mixing_ratio : array-like
        Water vapor mixing ratio in kg/kg.
    dewpoint : array-like
        Dewpoint temperature in Celsius.
    
    Returns
    -------
    array-like
        Pseudo-Equivalent potential temperature in Kelvin.
    
    Notes
    -----
    This function uses MetPy's `vapor_pressure` and `lcl` functions internally.

"""
基于 Skyborn 的假相当位温计算与 WRF 数据可视化
来源: 和鲸社区 - 酷炫用户名
修正版本: 使用 skyborn.calc.calculations.calculate_theta_se
"""

from wrf import (
    getvar, to_np, vertcross, interplevel,
    get_cartopy, latlon_coords, CoordPair
)
import numpy as np
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
import matplotlib.colors as mcolors
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import os

# ============================================
# 核心更新：使用 Skyborn 计算假相当位温
# ============================================
import skyborn
from skyborn.calc.calculations import calculate_theta_se

# 验证函数文档
# help(skyborn.calc.calculations.calculate_theta_se)

# --- 1. 数据读取和准备 ---
wrfout_path = "/home/mw/input/typhoon9537"
filename_prefix = "wrfout_d01_"
wrf_files = sorted([
    os.path.join(wrfout_path, f) 
    for f in os.listdir(wrfout_path) 
    if f.startswith(filename_prefix)
])
wrf_list = [Dataset(f) for f in wrf_files]

time_idx = -1

# 提取 WRF 变量
# 注意：Skyborn 函数需要特定单位，WRF-python 自动处理
p = getvar(wrf_list, 'pressure', timeidx=time_idx)      # hPa，Skyborn 需要
t = getvar(wrf_list, 'tk', timeidx=time_idx)            # K，Skyborn 需要
td = getvar(wrf_list, 'td', timeidx=time_idx)            # degC，Skyborn 需要
r = getvar(wrf_list, 'QVAPOR', timeidx=time_idx)        # kg/kg，Skyborn 需要

# 提取风场和通用坐标
ua = getvar(wrf_list, "ua", timeidx=time_idx)
va = getvar(wrf_list, "va", timeidx=time_idx)
rh = getvar(wrf_list, "rh", timeidx=time_idx)
lats, lons = latlon_coords(p)

# --- 2. 使用 Skyborn 计算假相当位温（核心更新）---
print("使用 Skyborn 计算假相当位温...")
theta_se_skyborn = calculate_theta_se(
    temperature=t,      # 温度，单位 K
    pressure=p,         # 气压，单位 hPa  
    mixing_ratio=r,     # 混合比，单位 kg/kg
    dewpoint=td         # 露点温度，单位 degC
)

print(f"θ_se 计算完成，形状: {theta_se_skyborn.shape}")
print(f"θ_se 范围: {float(theta_se_skyborn.min()):.2f} ~ {float(theta_se_skyborn.max()):.2f} K")

# 保持 xarray 结构，便于后续 WRF-python 工具调用
# Skyborn 返回的已经是 xarray.DataArray（如果输入是 xarray）

# --- 3. 定义剖面线并提取数据 ---
cross_start = CoordPair(lat=lats[0, 0], lon=lons[0, 0])
cross_end = CoordPair(lat=lats[-1, -1], lon=lons[-1, -1])
v_levels = np.arange(1000, 100, -25)  # 垂直层次 1000-100 hPa，间隔 25

# 提取剖面数据
theta_se_cross = vertcross(
    theta_se_skyborn, p, 
    wrfin=wrf_list[0], 
    start_point=cross_start, 
    end_point=cross_end, 
    latlon=True, 
    levels=v_levels
)

rh_cross = vertcross(
    rh, p, 
    wrfin=wrf_list[0], 
    start_point=cross_start, 
    end_point=cross_end, 
    latlon=True, 
    levels=v_levels
)
ua_cross = vertcross(
    ua, p, 
    wrfin=wrf_list[0], 
    start_point=cross_start, 
    end_point=cross_end, 
    latlon=True, 
    levels=v_levels
)
va_cross = vertcross(
    va, p, 
    wrfin=wrf_list[0], 
    start_point=cross_start, 
    end_point=cross_end, 
    latlon=True, 
    levels=v_levels
)

# --- 4. 可视化 ---
fig = plt.figure(figsize=(8, 10), dpi=200)

# 布局：地图在上，剖面在下
gs = fig.add_gridspec(2, 1, height_ratios=[1, 1.2])
ax_map = fig.add_subplot(gs[0], projection=get_cartopy(p))
ax_prof = fig.add_subplot(gs[1])

# --- 4.1 绘制地图（850 hPa θ_se）---
theta_se_850 = interplevel(theta_se_skyborn, p, 850)

# 使用 nipy_spectral 色图展示 θ_se 分布
cf_map = ax_map.contourf(
    to_np(lons), to_np(lats), to_np(theta_se_850), 
    levels=20,
    cmap='nipy_spectral', 
    transform=ccrs.PlateCarree()
)

# 添加地图要素
ax_map.add_feature(cfeature.COASTLINE.with_scale('50m'), linewidth=0.8)
ax_map.add_feature(cfeature.BORDERS, linewidth=0.5)

# 绘制剖面线
ax_map.plot(
    [cross_start.lon, cross_end.lon], 
    [cross_start.lat, cross_end.lat], 
    color='red', marker='o', 
    transform=ccrs.PlateCarree(), 
    linewidth=2, markersize=8,
    label='Cross-section'
)
ax_map.set_title('850 hPa Pseudo-Equivalent Potential Temperature (Skyborn)', fontsize=14)
ax_map.legend(loc='upper right')
ax_map.gridlines(draw_labels=True, dms=True, x_inline=False, y_inline=False)

# 添加色标
cbar_map = plt.colorbar(cf_map, ax=ax_map, orientation='horizontal', pad=0.05, fraction=0.05)
cbar_map.set_label('θ_se (K)', fontsize=10)

# --- 4.2 绘制剖面图 ---
xs = np.arange(0, theta_se_cross.shape[-1], 1)
ys = to_np(theta_se_cross.coords['vertical'])

# 绘制填充等值线
cf_prof = ax_prof.contourf(
    xs, ys, to_np(theta_se_cross), 
    levels=20,
    cmap='jet', 
    extend='both'
)

# 添加相对湿度等值线（虚线）
rh_levels = [80, 90]
rh_contours = ax_prof.contour(
    xs, ys, to_np(rh_cross), 
    levels=rh_levels, 
    colors='white', 
    linewidths=1.2, 
    linestyles='dashed'
)
ax_prof.clabel(rh_contours, inline=1, fontsize=10, fmt='%i%%')

ax_prof.set_title('Vertical Cross-section of θ_se (Skyborn calculate_theta_se)', fontsize=14)
ax_prof.set_ylabel('Pressure (hPa)', fontsize=12)
ax_prof.set_xlabel('Horizontal Distance (grid points)', fontsize=12)
ax_prof.invert_yaxis()
ax_prof.grid(True, linestyle='--', alpha=0.5)

# 添加色标
cbar_prof = plt.colorbar(cf_prof, ax=ax_prof, orientation='vertical', pad=0.02)
cbar_prof.set_label('Pseudo-Equivalent Potential Temperature (K)', fontsize=10)

plt.suptitle(
    'WRF Typhoon Simulation: θ_se Analysis using Skyborn\n' + 
    'Formula: θ_se = T*(1000/(p-e))^(Ra/cpd)*exp(L*r/(cpd*Tc))', 
    fontsize=16, y=0.98
)
plt.tight_layout(rect=[0, 0.03, 1, 0.96])
plt.show()

使用 Skyborn 计算假相当位温...
θ_se 计算完成，形状: (34, 437, 447)
θ_se 范围: 324.22 ~ 484.30 K