读者答疑 | python怎么计算流函数

!pip install windspharm -i https://pypi.mirrors.ustc.edu.cn/simple/

!pip install xinvert -i https://pypi.mirrors.ustc.edu.cn/simple/

import cartopy.crs as ccrsimport cartopy.feature as cfeatureimport matplotlib.pyplot as pltimport numpy as npimport xarray as xr
import metpy.calc as mpcalcfrom metpy.cbook import get_test_datafrom metpy.plots import add_metpy_logo, add_timestampfrom metpy.units import units

import xarray as xrds = xr.open_dataset('/home/mw/input/xinvert2128/atmos3D.nc')print(ds)

print(list(ds.variables))

msf = mpcalc.montgomery_streamfunction(    ds['hgt'].sel(LEV=850),    ds['T'].sel(LEV=850)).data.to_base_units() * 1e-2lon =ds.lonlat =ds.latu = ds['U'].sel(LEV=850)v = ds['V'].sel(LEV=850)

from meteva.base.tool.plot_tools import add_china_map_2basemap
bounds = [(80.,140., 15., 60.)]
fig = plt.figure(figsize=(17., 12.))ax = plt.subplot(111, projection=ccrs.PlateCarree()) ax.set_extent(*bounds, crs=ccrs.PlateCarree())# 核心代码add_china_map_2basemap(ax, name="river", edgecolor='k', lw=0.5, encoding='gbk', grid0=None)  # "河流"add_china_map_2basemap(ax, name="nation",edgecolor='k', lw=0.5, encoding='gbk', grid0=None)  # "国界"add_china_map_2basemap(ax, name="province", edgecolor='k', lw=0.5, encoding='gbk', grid0=None)  # "省界"
# Plot the surfaceclevmsf = np.arange(0, 4000, 20)cs = ax.contour(lon, lat, msf, clevmsf,                colors='k', linewidths=1.0, linestyles='solid', transform=ccrs.PlateCarree())cs.clabel(fontsize=10, inline=1, inline_spacing=7, fmt='%i', rightside_up=True,          use_clabeltext=True)
# Plot RHcf = ax.contourf(lon, lat, ds['T'].sel(LEV=850),                 range(230, 300, 5), cmap=plt.cm.gist_earth_r, transform=ccrs.PlateCarree())cb = fig.colorbar(cf, orientation='horizontal', aspect=65, shrink=0.5, pad=0.05,                  extendrect='True')cb.set_label('T', size='x-large')
barb_increments = {    'half': 2,   # 短杠 2 m/s    'full': 4,   # 长杠 4 m/s    'flag': 20   # 三角旗 20 m/s}
# Plot wind barbsax.barbs(lon.values, lat.values, u,         v, length=6,         regrid_shape=20, transform=ccrs.PlateCarree(),barb_increments=barb_increments)
fig.tight_layout()plt.show()

"""Compute streamfunction and velocity potential from the long-term-meanflow.
This example uses the standard interface.
Additional requirements for this example:
* netCDF4 (http://unidata.github.io/netcdf4-python/)* matplotlib (http://matplotlib.org/)* cartopy (http://scitools.org.uk/cartopy/)
"""import cartopy.crs as ccrsfrom cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatterfrom cartopy.util import add_cyclic_pointimport matplotlib as mplimport matplotlib.pyplot as pltfrom netCDF4 import Dataset
from windspharm.standard import VectorWindfrom windspharm.tools import prep_data, recover_data, order_latdimfrom windspharm.examples import example_data_path
mpl.rcParams['mathtext.default'] = 'regular'

# Read zonal and meridional wind components from file using the netCDF4# module. The components are defined on pressure levels and are in separate# files.ncu = Dataset(example_data_path('uwnd_mean.nc'), 'r')uwnd = ncu.variables['uwnd'][:]lons = ncu.variables['longitude'][:]lats = ncu.variables['latitude'][:]ncu.close()ncv = Dataset(example_data_path('vwnd_mean.nc'), 'r')vwnd = ncv.variables['vwnd'][:]ncv.close()
# The standard interface requires that latitude and longitude be the leading# dimensions of the input wind components, and that wind components must be# either 2D or 3D arrays. The data read in is 3D and has latitude and# longitude as the last dimensions. The bundled tools can make the process of# re-shaping the data a lot easier to manage.uwnd, uwnd_info = prep_data(uwnd, 'tyx')vwnd, vwnd_info = prep_data(vwnd, 'tyx')
# It is also required that the latitude dimension is north-to-south. Again the# bundled tools make this easy.lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd)
# Create a VectorWind instance to handle the computation of streamfunction and# velocity potential.w = VectorWind(uwnd, vwnd)
# Compute the streamfunction and velocity potential. Also use the bundled# tools to re-shape the outputs to the 4D shape of the wind components as they# were read off files.sf, vp = w.sfvp()sf = recover_data(sf, uwnd_info)vp = recover_data(vp, uwnd_info)
# Pick out the field for December and add a cyclic point (the cyclic point is# for plotting purposes).sf_dec, lons_c = add_cyclic_point(sf[11], lons)vp_dec, lons_c = add_cyclic_point(vp[11], lons)
# Plot streamfunction.ax1 = plt.axes(projection=ccrs.PlateCarree(central_longitude=180))clevs = [-120, -100, -80, -60, -40, -20, 0, 20, 40, 60, 80, 100, 120]sf_fill = ax1.contourf(lons_c, lats, sf_dec * 1e-06, clevs,                       transform=ccrs.PlateCarree(), cmap=plt.cm.RdBu_r,                       extend='both')ax1.coastlines()ax1.gridlines()ax1.set_xticks([0, 60, 120, 180, 240, 300, 359.99], crs=ccrs.PlateCarree())ax1.set_yticks([-90, -60, -30, 0, 30, 60, 90], crs=ccrs.PlateCarree())lon_formatter = LongitudeFormatter(zero_direction_label=True,                                   number_format='.0f')lat_formatter = LatitudeFormatter()ax1.xaxis.set_major_formatter(lon_formatter)ax1.yaxis.set_major_formatter(lat_formatter)plt.colorbar(sf_fill, orientation='horizontal')plt.title('Streamfunction ($10^6$m$^2$s$^{-1}$)', fontsize=16)
# Plot velocity potential.plt.figure()ax2 = plt.axes(projection=ccrs.PlateCarree(central_longitude=180))clevs = [-10, -8, -6, -4, -2, 0, 2, 4, 6, 8, 10]vp_fill = ax2.contourf(lons_c, lats, vp_dec * 1e-06, clevs,                       transform=ccrs.PlateCarree(), cmap=plt.cm.RdBu_r,                       extend='both')ax2.coastlines()ax2.gridlines()ax2.set_xticks([0, 60, 120, 180, 240, 300, 359.99], crs=ccrs.PlateCarree())ax2.set_yticks([-90, -60, -30, 0, 30, 60, 90], crs=ccrs.PlateCarree())lon_formatter = LongitudeFormatter(zero_direction_label=True,                                   number_format='.0f')lat_formatter = LatitudeFormatter()ax2.xaxis.set_major_formatter(lon_formatter)ax2.yaxis.set_major_formatter(lat_formatter)plt.colorbar(vp_fill, orientation='horizontal')plt.title('Velocity Potential ($10^6$m$^2$s$^{-1}$)', fontsize=16)plt.show()

import xarray as xrds  = xr.open_dataset('/home/mw/input/xinvert2128/Helmholtz_atmos.nc')vor = ds.vor
print(ds)

from xinvert import invert_Poisson
iParams = {    'BCs'      : ['extend', 'periodic'],    'mxLoop'   : 1000,    'tolerance': 1e-12,}
sf = invert_Poisson(vor, dims=['lat','lon'], iParams=iParams)

import matplotlib.pyplot as pltimport cartopy.crs as ccrsimport cartopy.feature as cfeaturefrom matplotlib.ticker import FixedLocator, FuncFormatterimport xarray as xrimport numpy as np
# 定义坐标轴格式器def LONGITUDE_FORMATTER(x, pos=None):    return f"{int(x)}°E" if x >= 0 else f"{abs(int(x))}°W"
def LATITUDE_FORMATTER(x, pos=None):    return f"{int(x)}°N" if x >= 0 else f"{abs(int(x))}°S"
# 选择第一个时间步u = ds.u.where(ds.u!=0)[0].load()v = ds.v.where(ds.v!=0)[0].load()m = np.hypot(u, v)
# 创建一个包含两个子图的图形fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(7, 7),                         subplot_kw=dict(projection=ccrs.PlateCarree(central_longitude=180)))
# 设置字体大小fontsize = 13
# 添加海岸线for ax in axes.flat:    ax.add_feature(cfeature.COASTLINE.with_scale('50m'))
# 设置网格线gl = axes[0].gridlines(draw_labels=True, linestyle='--', color='gray', alpha=0.5, linewidth=0.5)gl.xlocator = FixedLocator(np.arange(-180, 181, 60))gl.ylocator = FixedLocator(np.arange(-90, 91, 30))gl.xformatter = FuncFormatter(LONGITUDE_FORMATTER)gl.yformatter = FuncFormatter(LATITUDE_FORMATTER)gl.xlabel_style = {'size': fontsize}gl.ylabel_style = {'size': fontsize}
# 绘制相对涡度p = axes[0].contourf(u.lon, u.lat, vor[0]*1e5, cmap='bwr',                     levels=np.linspace(-10, 10, 22),                     transform=ccrs.PlateCarree())axes[0].set_title('Relative Vorticity', fontsize=fontsize)cb = fig.colorbar(p, ax=axes[0], orientation='vertical', shrink=0.5, pad=0.05)cb.ax.tick_params(labelsize=fontsize)
# 绘制风场和反转流函数p = axes[1].contourf(u.lon, u.lat, sf[0], levels=30, cmap='jet',                     transform=ccrs.PlateCarree())q = axes[1].quiver(u.lon.values, u.lat.values, u.values, v.values,                   transform=ccrs.PlateCarree(),                   width=0.0007, headwidth=12., headlength=15.)axes[1].set_title('Wind Field and Inverted Streamfunction', fontsize=fontsize)cb = fig.colorbar(p, ax=axes[1], orientation='vertical', shrink=0.5, pad=0.05)cb.ax.tick_params(labelsize=fontsize)
plt.tight_layout()plt.show()