Python | Rossby波源

import warningswarnings.filterwarnings("ignore")import xarray as xrimport numpy as npnp.int = intimport matplotlib.pyplot as pltimport cartopy.crs as ccrsfrom windspharm.xarray import VectorWindimport cmapsfrom cartopy.util import add_cyclic_pointimport matplotlib.ticker as mtickerfrom cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatterfrom mpl_toolkits.axes_grid1.inset_locator import inset_axesfrom tqdm import tqdm# ------------------------------------------------------------------# 1. 读数据# ------------------------------------------------------------------u_file = 'uwnd.mon.mean.2.5.nc'v_file = 'vwnd.mon.mean.2.5.nc'ds_u = xr.open_dataset(u_file, chunks={'time': 12})ds_v = xr.open_dataset(v_file, chunks={'time': 12})level_sel = 200          # hPau_all = ds_u['u'].sel(level=level_sel, expver=1)v_all = ds_v['v'].sel(level=level_sel, expver=1)time_slice = slice('1979-01', '2022-12')u_all = u_all.sel(time=time_slice)v_all = v_all.sel(time=time_slice)# ------------------------------------------------------------------# 2. 计算 RWS# ------------------------------------------------------------------Omega = 7.2921150e-5f = 2 * Omega * np.sin(np.deg2rad(u_all.latitude))f2d = f.broadcast_like(u_all.isel(time=0))   # (lat,lon)rws_list = []for t in tqdm(u_all.time, desc='Computing RWS'):    u = u_all.sel(time=t).fillna(0)    v = v_all.sel(time=t).fillna(0)    vw = VectorWind(u, v)    zeta = vw.vorticity()    div  = vw.divergence()    u_chi, v_chi = vw.irrotationalcomponent()    vw_chi = VectorWind(u_chi, v_chi)    div_vchi = vw_chi.divergence()    try:        eta = vw.absolutevorticity()    except Exception:        eta = zeta + f2d    grad_x, grad_y = vw.gradient(eta)    rws = - (u_chi * grad_x + v_chi * grad_y) - eta * div_vchi    rws_list.append(rws.assign_coords(time=t))# 合并成一个 DataArrayRWS_all = xr.concat(rws_list, dim='time')# coslat = np.cos(np.deg2rad(RWS_all.latitude))   # 1D# RWS_weighted = RWS_all * coslat# ------------------------------------------------------------------# 3. 夏季 (6-7-8) RWS# ------------------------------------------------------------------jja = RWS_all.sel(time=RWS_all.time.dt.month.isin([6, 7, 8]))RWS_jja_clim = jja.mean('time')# ------------------------------------------------------------------# 4. 夏季平均的不可旋转风 (u_chi, v_chi) 矢量# ------------------------------------------------------------------u_chi_clim = []v_chi_clim = []for t in u_all.time:    if t.dt.month.item() in (6, 7, 8):        u = u_all.sel(time=t).fillna(0)        v = v_all.sel(time=t).fillna(0)        vw = VectorWind(u, v)        u_chi, v_chi = vw.irrotationalcomponent()        u_chi_clim.append(u_chi.assign_coords(time=t))        v_chi_clim.append(v_chi.assign_coords(time=t))u_chi_clim = xr.concat(u_chi_clim, dim='time').mean('time')v_chi_clim = xr.concat(v_chi_clim, dim='time').mean('time')# ------------------------------------------------------------------# 5. 绘图# ------------------------------------------------------------------# 处理0经度处断点问题def add_cyclic_da(da):    vals, lon = add_cyclic_point(da.values, coord=da.longitude)    return xr.DataArray(vals,                        coords={'longitude': lon, 'latitude': da.latitude},                        dims=da.dims)RWS_plot = add_cyclic_da(RWS_jja_clim)u_chi_plot = add_cyclic_da(u_chi_clim)v_chi_plot = add_cyclic_da(v_chi_clim)# ------------------------------------------------------------------fig = plt.figure(figsize=(12, 6))ax = plt.axes(projection=ccrs.Robinson())ax.set_global()im = ax.contourf(RWS_plot.longitude, RWS_plot.latitude, RWS_plot*1e11,                 transform=ccrs.PlateCarree(),                 cmap=cmaps.MPL_PuOr_r,                 levels=np.arange(-20, 22, 2), extend='both')# 辐散风step = max(1, int(u_chi_plot.sizes['longitude']//40))lon = RWS_plot.longitude.valueslat = RWS_plot.latitude.valueslon2d, lat2d = np.meshgrid(lon, lat)uv = ax.quiver(lon2d[::step, ::step], lat2d[::step, ::step],               u_chi_plot.values[::step, ::step],               v_chi_plot.values[::step, ::step],               transform=ccrs.PlateCarree(), scale=100, width=0.002)ax.quiverkey(uv, X=0.82, Y=-0.09, U=5, angle=0,             label='5 m/s', labelpos='E', color='k', labelcolor='k')ax.coastlines(linewidth=1.5, color='grey')gl = ax.gridlines(draw_labels=True, linewidth=0.5, color='gray',                  alpha=0.8, linestyle='--')# gl.top_labels   = False# gl.right_labels = Falsegl.xlocator = mticker.FixedLocator(np.arange(-180, 181, 60))gl.ylocator = mticker.FixedLocator(np.arange(-90, 91, 30))gl.xformatter = LongitudeFormatter()gl.yformatter = LatitudeFormatter()gl.xlabel_style = {'size': 14}gl.ylabel_style = {'size': 14}# colorbarcax = inset_axes(ax, width='50%', height='3%', loc='lower center',                 bbox_to_anchor=(0, -0.12, 1, 1), bbox_transform=ax.transAxes)cbar = plt.colorbar(im, cax=cax, orientation='horizontal')cbar.set_label('10$^{-11}$s$^{-2}$')ax.set_title('RWS Clim. JJA (1979-2022)',             fontdict={'family': 'sans-serif', 'size': 18}, fontweight='bold')plt.savefig('RWS_JJA_clim.png', dpi=600, bbox_inches='tight')plt.show()