Python计算单层、整层水汽物理量(水汽通量和水汽通量散度)

import xarray as xrimport numpy as npimport cfgribimport matplotlib.pyplot as pltimport matplotlib.ticker as mtickerimport cartopy.crs as ccrsimport cartopy.feature as cfeaturefrom cartopy.io.shapereader import Reader  # 读取 .shpimport metpy.calc as mpcalcimport geocat.compimport cmapsfrom metpy.units import units# 读取省界shpName = 'Province_9.shp'reader = Reader(shpName)provinces = cfeature.ShapelyFeature(reader.geometries(), ccrs.PlateCarree(), edgecolor='k', facecolor='none')# 需要的层次levels = [1000.,  975.,  950.,  925.,  900.,  850.,  800.,  750.,  700.,  650.,          600.,  550.,  500.,  450.,  400.,  350.,  300.,  250.,  200.]# 常数ptop = 200*100 # 垂直积分层顶，单位Pag = 9.8# 读取文件files = 'ERA5/' # 这里改成自己的数据路径file = 'pl_2014070506.grib'file_sl = "sl_2014070506.grib"ds = cfgrib.open_datasets(files+file)ds_sl = cfgrib.open_datasets(files+file_sl)# 读取变量u_levels = ds[0].u.sel(isobaricInhPa=levels)v_levels = ds[0].v.sel(isobaricInhPa=levels)q_levels = ds[0].q.sel(isobaricInhPa=levels) # kg/kgpsfc = ds_sl[4].sp # 地表气压，单位Pa# 接下来是数据的处理：plev = u_levels.coords["isobaricInhPa"] * 100 # 单位变为Paplev.attrs['units'] = 'Pa'q_levels = q_levels * 1000 # kg/kg -> g/kgq_levels.attrs['units'] = 'g/kg'# 垂直积分的关键，计算各层次厚度，效果类似ncl中的dpres_pleveldp = geocat.comp.dpres_plevel(plev, psfc, ptop)# Layer Mass Weightingdpg = dp / gdpg.attrs["units"] = "kg/m2"# 水汽通量计算：# 水汽通量uq = u_levels * q_levelsvq = v_levels * q_levelsuq.attrs["units"] = "("+u_levels.units+")("+q_levels.units+")" vq.attrs["units"] = "("+v_levels.units+")("+q_levels.units+")" uq_dpg = uq * dpg.datavq_dpg = vq * dpg.data# 整层水汽通量iuq = uq_dpg.sum(axis=0)ivq = vq_dpg.sum(axis=0)iuq.attrs["units"] = "[m/s][g/kg]"


    
ivq.attrs["units"] = "[m/s][g/kg]"# 水汽通量散度：# 计算散度要用到dx，dylons = u_levels.coords['longitude'][:]lats = v_levels.coords['latitude'][:]dx, dy = mpcalc.lat_lon_grid_deltas(lons, lats)# 水汽通量散度duvq = mpcalc.divergence(uq/g, vq/g, dx=dx[np.newaxis,:,:], dy=dy[np.newaxis,:,:]) duvq.attrs["units"] = "g/(kg-s)"duvq_dpg = duvq * dpg.data# 整层水汽通量散度iduvq = duvq_dpg.sum(axis=0)iduvq = iduvq * units('kg/m^2')iduvq.attrs["units"] = "g/(m2-s)"# 画图：leftlon, rightlon, lowerlat, upperlat = 80, 100, 20, 40img_extent = [leftlon, rightlon, lowerlat, upperlat]levs = np.arange(-0.3, 0.3+0.05, 0.05)fig = plt.figure(figsize=(12,9))ax = plt.subplot(111, projection=ccrs.PlateCarree())ax.set_extent(img_extent)ax.coastlines('50m', linewidth=0.8)# 整层水汽通量散度填色mfc_contourf = ax.contourf(lons, lats,                            iduvq,                           cmap=cmaps.MPL_seismic,                            levels=levs,                           extend='both', transform=ccrs.PlateCarree())# 整层水汽通量流线mf_stream = ax.streamplot(lons, lats,                          iuq.data, ivq.data,                          color='black', transform=ccrs.PlateCarree())gl = ax.gridlines(draw_labels=True, dms=False,                   linestyle='--', color='gray',                  x_inline=False, y_inline=False)gl.top_labels = Falsegl.right_labels = Falseax.add_feature(provinces, linewidth=1)position = fig.add_axes([ax.get_position().x0,                         ax.get_position().y0-0.08,                         ax.get_position().width,                         0.02])cb = fig.colorbar(mfc_contourf, orientation='horizontal', cax=position)cb.set_label('%s'%iduvq.units, fontsize=12)# 保存图片plt.savefig('./plo99.png',dpi=800,bbox_inches='tight',pad_inches=0)plt.show()