使用u,v计算涡度和散度

import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
from cartopy.util import add_cyclic_point
import matplotlib as mpl
import matplotlib.pyplot as plt
from cartopy.io.shapereader import Reader
import os
import numpy as np 
import pdb
import datetime,time
import pandas as pd
import xarray as xr
import math
import cftime
st=datetime.datetime(1900,1,1,0,0)
start=datetime.datetime.now()
from matplotlib.transforms import offset_copy
from matplotlib.font_manager import FontProperties
from math import sin, asin, cos, radians, fabs, sqrt
import warnings
warnings.filterwarnings("ignore")

EARTH_RADIUS=6371           # 地球平均半径，6371km
 
def hav(theta):
    s = sin(theta / 2)
    return s * s
 
def get_distance_hav(lat0, lng0, lat1, lng1):
    "用haversine公式计算球面两点间的距离。"
    # 经纬度转换成弧度
    lat0 = radians(lat0)
    lat1 = radians(lat1)
    lng0 = radians(lng0)
    lng1 = radians(lng1)
 
    dlng = fabs(lng0 - lng1)
    dlat = fabs(lat0 - lat1)
    h = hav(dlat) + cos(lat0) * cos(lat1) * hav(dlng)
    distance = 2 * EARTH_RADIUS * asin(sqrt(h))
 
    return distance
def new_func(data, wind_lons, wind_lats, t,l):
    u=data['u'].sel(time=t,isobaricInhPa=l)
    v=data['v'].sel(time=t,isobaricInhPa=l)
    wind_olon,wind_olat=np.meshgrid(wind_lons,wind_lats)
    nx=len(wind_lons)
    return u,v,wind_olon,wind_olat,nx

if __name__ == '__main__':
    font = FontProperties(fname='/mnt/d/work/数据/SimHei.ttf', size=16)

    SHP = '/mnt/d/work/数据/china_shp'
    filePath='/mnt/d/work_no_backup/ECWMF_data/data/'

    namelist=[]
    file_name_dir=[]
    for i,j,k in os.walk(filePath):
        for o in k:
            if o.endswith('grib'):
                namelist.append(str(i)+str(o))

    for name in namelist:
        print(name)
        data=xr.open_dataset(name,engine='cfgrib')
        time_list=data['time'].values
        lev_list=data['isobaricInhPa'].values
        wind_lons = data['longitude']
        wind_lats = data['latitude']
        for t in time_list:
            t_str=pd.to_datetime(t).strftime('%Y-%m-%d-%H')
            for l in lev_list:
                u, v, wind_olon, wind_olat, nx = new_func(data, wind_lons, wind_lats, t,l)
                ny=len(wind_lats)
                dx=[]
                dy=-111.1949*0.25*1000   #纬距为固定值，一个纬距为111.1949，经距则根据纬距的不同而变化
                for i in range(len(wind_lats)):
                    if i==0:
                        dx.append(get_distance_hav(wind_lats[i],0,wind_lats[i+1],0.25)*1000)
                    elif i==len(wind_lats)-1:
                        dx.append(get_distance_hav(wind_lats[i-1],0,wind_lats[i],0.25)*1000)
                    elif ((i!=0) and (i!=len(wind_lats)-1)):
                        dx.append(get_distance_hav(wind_lats[i-1],0,wind_lats[i+1],0.25)*1000/2)

                ugradient=np.gradient(u)
                vgradient=np.gradient(v)

                dy=np.ones(nx)*dy

                dy,dx=np.meshgrid(dy,dx)

                cau_vort=((vgradient[1]/dx)-(ugradient[0]/dy))*100000     #涡度dv/dx-du/dy
                cau_div=((ugradient[1]/dx)+(vgradient[0]/dy))*100000      #散度du/dy+dv/dx

                #画图

                extent=[117, 123, 27,32]
                proj= ccrs.PlateCarree() 
                font = FontProperties(fname='/mnt/d/work/数据/SimHei.ttf', size=16)


                fig,(ax1,ax2) = plt.subplots(nrows=1,
                                    ncols=2, 
                                    subplot_kw={'projection': proj},
                                    figsize = (12, 8))

                ax1.add_geometries(Reader(os.path.join(SHP, 'cnhimap.shp')).geometries(),ccrs.PlateCarree(),facecolor='none',edgecolor='k', linewidth=0.3)
                ax1.set_xticks(np.arange(100,150,2))#需要显示的经度，一般可用np.arange
                ax1.set_yticks(np.arange(10,50,2))#需要显示的纬度

                ax1.xaxis.set_major_formatter(LongitudeFormatter())#将横坐标转换为经度格式
                ax1.yaxis.set_major_formatter(LatitudeFormatter())#将纵坐标转换为纬度格式
                ax1.set_extent(extent,ccrs.PlateCarree())
                ax1.tick_params(axis='both',labelsize=8,direction='in',length=2.75,width=0.55,right=True,top=True)#修改刻度样式

                vort_levels=np.arange(-26,30,2)
                rh_fill1 = ax1.contourf(wind_olon,wind_olat,cau_vort,levels=vort_levels,cmap='coolwarm',transform=ccrs.PlateCarree())
                cb1=fig.colorbar(rh_fill1,ax=ax1,orientation='horizontal',shrink=0.95,pad=0.05)
                ax1.streamplot(wind_olon,wind_olat,u.values,v.values,
                            transform=ccrs.PlateCarree(),
                            linewidth=0.4, 
                            density=2,
                            arrowsize=0.4,
                            arrowstyle='->',
                            color='blue')

                ax1.set_title('vorticity_'+t_str+'  '+str(l)+'hPa')
                ###############################################################
                ax2.add_geometries(Reader(os.path.join(SHP, 'cnhimap.shp')).geometries(),ccrs.PlateCarree(),facecolor='none',edgecolor='k', linewidth=0.3)
                ax2.set_xticks(np.arange(100,150,2))#需要显示的经度，一般可用np.arange
                ax2.set_yticks(np.arange(10,50,2))#需要显示的纬度
                ax2.xaxis.set_major_formatter(LongitudeFormatter())#将横坐标转换为经度格式
                ax2.yaxis.set_major_formatter(LatitudeFormatter())#将纵坐标转换为纬度格式
                ax2.set_extent(extent,ccrs.PlateCarree())
                ax2.tick_params(axis='both',labelsize=8,direction='in',length=2.75,width=0.55,right=True,top=True)#修改刻度样式
                div_levels=np.arange(-18,20,2)
                rh_fill2 = ax2.contourf(wind_olon,wind_olat,cau_div,levels=div_levels,cmap='coolwarm',transform=ccrs.PlateCarree())
            #    div_position=fig.add_axes([0.92,0.16,0.02,0.68])
            #    cb2=fig.colorbar(rh_fill2,cax=div_position,orientation='horizontal',shrink=0.3,pad=0.01)
                cb2=fig.colorbar(rh_fill2,ax=ax2,orientation='horizontal',shrink=0.95,pad=0.05)

                ax2.streamplot(wind_olon,wind_olat,u.values,v.values,
                            transform=ccrs.PlateCarree(),
                            linewidth=0.4, 
                            density=2,
                            arrowsize=0.4,
                            arrowstyle='->',
                            color='blue')

                ax2.set_title('divergence_'+t_str+'  '+str(l)+'hPa')
                ############################################################### 
                plt.subplots_adjust(wspace=0.1,hspace =0.1)#调整子图间距
                plt.savefig('/mnt/d/work_no_backup/ECWMF_data/out_put/'+str(int(l))+'/'+str(t_str)+'_'+str(int(l))+'_'+'div_vort.png',dpi=500)
                plt.close()
                print('/mnt/d/work_no_backup/ECWMF_data/out_put/'+str(int(l))+'/'+str(t_str)+'_'+str(int(l))+'_'+'div_vort.png','is OK')