大气视热源的python计算尝试

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()