from wrf import uvmet, to_np, getvar, interplevel, smooth2d, get_cartopy, cartopy_xlim, cartopy_ylim, latlon_coords,ALL_TIMES
import numpy as np
from netCDF4 import Dataset
import xarray as xr
from metpy.units import units
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
from matplotlib.colors import from_levels_and_colors
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import cartopy.io.shapereader as shpreader
from cartopy.feature import NaturalEarthFeature
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
from cartopy.mpl.gridliner import LATITUDE_FORMATTER, LONGITUDE_FORMATTER
import cmaps
from glob import glob
import metpy.calc as mpcalc
import metpy.constants as constants
import os

Warning: ecCodes 2.21.0 or higher is recommended. You are running version 2.14.1

# 定义 WRF 文件夹路径和文件名前缀
wrfout_path = "/home/mw/input/wrfout3385/"
filename_prefix = "wrfout_d02_"
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]
提取要素
lon = getvar(wrf_list, 'lon', timeidx=ALL_TIMES, method='cat')

lat = getvar(wrf_list, 'lat', timeidx=ALL_TIMES, method='cat')

u = getvar(wrf_list, 'ua', timeidx=ALL_TIMES, method='cat')

v = getvar(wrf_list, 'va', timeidx=ALL_TIMES, method='cat')

w = getvar(wrf_list, 'omg', timeidx=ALL_TIMES, method='cat')

t = getvar(wrf_list, 'tk', timeidx=ALL_TIMES, method='cat')

theta = getvar(wrf_list, 'theta', timeidx=ALL_TIMES, method='cat')

#q = getvar(wrf_file, 'QVAPOR', timeidx=0)*1000
提取WRF模拟的经纬度数组
lats, lons = latlon_coords(u)
提取WRF模拟的地图投影
wrf_proj = get_cartopy(u)

levels = [1000,925,850,700,600,500,400,300,200,100]

p = getvar(wrf_list, 'pressure', timeidx=ALL_TIMES, method='cat')

u_in = interplevel(u, p, levels)

v_in = interplevel(v, p, levels)

w_in = interplevel(w, p, levels)

theta_in = interplevel(theta, p, levels)

t_in = interplevel(t, p, levels)

#z = getvar(wrf_list, 'z', timeidx=ALL_TIMES, method='cat')

#z_in = interplevel(z, p, levels)

dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat)

p_in =interplevel(p, p, levels)

In [3]:

p_in.units,theta_in.units,t_in.units,u_in.units,v_in.units,w_in.units

('hPa', 'K', 'K', 'm s-1', 'm s-1', 'Pa s-1')

from metpy.calc import advection

cp = 1004 #j/(kg*k)

R = 287 #j/(kg*k)

p0 = 1000 # hpa

kappa = R/cp

w_in = w_in/100

lev = u_in.level

time = u_in.Time

计算 d theta/ dp
dthetadp = mpcalc.first_derivative(theta_in, axis=1, x=lev)

##时间处理注意单位

time_sec = (time - np.datetime64('2022-07-14T06:00:00.000000000')) / np.timedelta64(1, 's')
dT/dt
delta_t = time_sec[-1] - time_sec[0]
dT = mpcalc.first_derivative(t_in, axis=0,x=time_sec)

delta_t_bcast = delta_t.broadcast_like(t_in)

dTdt = dT/delta_t_bcast /units.sec
温度平流
adv_T = np.zeros((time.shape[0],lev.shape[0],lon.shape[1],lat.shape[2]))
for j in range(time.shape[0]):

for i in range(lev.shape[0]):

adv_T[j,i] = advection(to_np(t_in[j,i]),u=to_np(u_in[j,i]), v=to_np(v_in[j,i]), dx=dx[j], dy=dy[j])

adv_T=adv_T * units.K/ units.s
Q1
vert_adv_theta = w_in * dthetadp * (p0/p_in)**kappa   / units.s

Q1 = cp * (dTdt + adv_T+ vert_adv_theta)

Q1 = np.nan_to_num(Q1, nan=0)
积分，nan值需要去除才能积分
g=9.8 # m*s^-2

q1_int = np.trapz(Q1,lev,axis=1)/g

# 创建画布

fig = plt.figure(figsize=(10, 8))

设置地图投影
ax = plt.axes(projection=wrf_proj)
设置地图范围
ax.set_xlim(cartopy_xlim(u))

ax.set_ylim(cartopy_ylim(u))
读取国界线和九段线
province = shpreader.Reader(

'/home/mw/input/data5246/中国地图/China_provinces/China_provinces.shp')

nineline = shpreader.Reader(

'/home/mw/input/data5246/中国地图/China_10-dash_line/China_10-dash_line.shp')
绘制国界线和九段线
ax.add_geometries(province.geometries(),

crs=ccrs.PlateCarree(),

linewidth=0.5,

edgecolor='k',

facecolor='none',

zorder=2)

ax.add_geometries(nineline.geometries(),

crs=ccrs.PlateCarree(),

linewidth=0.5,

color='k',

zorder=2)
用contourf方法实现等值线填充
plt.contourf(to_np(lons),

to_np(lats),

to_np(q1_int[2]),
         extend='both',
         transform=ccrs.PlateCarree(),
         cmap=cmaps.MPL_BuPu_r)

添加colorbar
cbar = plt.colorbar(ax=ax, extend='both', shrink=1)

cbar.set_label('atmospheric apparent heat source', fontdict={'size': 20})

cbar.ax.tick_params(labelsize=20)
设置刻度标签的字体大小
plt.tick_params(labelsize=15)
添加标题
plt.title((str(u.Time.values)[0:16]+'(UTC)'), loc='left', fontsize=20)

plt.show()