def windsfarm(): #Prepare wund for vector w = VectorWind(uwnd, vwnd) #Calculating basic variables eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) planvor = w.planetaryvorticity S = -eta * div - uchi * etax - vchi * etay S = recover_data(S, uwnd_info) return(w, eta, div, uchi, vchi, etax, etay, planvor, S)
def rws(uwnd, vwnd): from windspharm.standard import VectorWind uwnd = np.transpose(uwnd, (1, 2, 0)) vwnd = np.transpose(vwnd, (1, 2, 0)) w = VectorWind(uwnd, vwnd) eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) s1 = -eta * div s2 = - (uchi * etax + vchi * etay) s = s1 + s2 s = np.mean(s, axis=2) s1 = np.mean(s1, axis=2) s2 = np.mean(s2, axis=2) return s, s1, s2
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 computations. w = VectorWind(uwnd, vwnd) # Compute components of rossby wave source: absolute vorticity, divergence, # irrotational (divergent) wind components, gradients of absolute vorticity. eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) # Combine the components to form the Rossby wave source term. Re-shape the # Rossby wave source array to the 4D shape of the wind components as they were # read off files. S = -eta * div - (uchi * etax + vchi * etay) S = recover_data(S, uwnd_info) # Pick out the field for December and add a cyclic point (the cyclic point is # for plotting purposes). S_dec, lons_c = addcyclic(S[11], lons) # Plot Rossby wave source. m = Basemap(projection='cyl', resolution='c', llcrnrlon=0,
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 computations. w = VectorWind(uwnd, vwnd) # Compute components of rossby wave source: absolute vorticity, divergence, # irrotational (divergent) wind components, gradients of absolute vorticity. eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) # Combine the components to form the Rossby wave source term. Re-shape the # Rossby wave source array to the 4D shape of the wind components as they were # read off files. S = -eta * div - (uchi * etax + vchi * etay) S = recover_data(S, uwnd_info) # Pick out the field for December and add a cyclic point (the cyclic point is # for plotting purposes). S_dec, lons_c = add_cyclic_point(S[11], lons) # Plot Rossby wave source. ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=180)) clevs = [-30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30] fill = ax.contourf(lons_c, lats, S_dec * 1e11, clevs, cmap=plt.cm.RdBu_r, transform=ccrs.PlateCarree(), extend="both")
def rws_fprime(ua200, va200): from windspharm.standard import VectorWind from netcdfread import ncread, ncsave from scipy import interpolate lon42 = ncread( '/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc', 'lon') lat42 = ncread( '/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc', 'lat') lon = ncread( '/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/All-Hist/mon/tas/item3236_monthly_mean_a011_2006-01_2016-12.nc', 'longitude0') lat = ncread( '/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/Plus15-Future_LCO2/day/ua/item15201_daily_mean_a00b_2090-01_2100-12.nc', 'latitude1') u = ua200['All-Nat'] #[1] v = va200['All-Nat'] #[1] u_sst = ua200['GHG-Nat'] #[1] v_sst = va200['GHG-Nat'] #[1] u_ghg = ua200['SST-Nat'] #[1] v_ghg = va200['SST-Nat'] #[1] v = w5rem(v) u = w5rem(u) uwnd = np.zeros((64, 128, len(u))) vwnd = np.zeros((64, 128, len(v))) uwnd_sst = np.zeros((64, 128, len(u_sst))) vwnd_sst = np.zeros((64, 128, len(u_sst))) uwnd_ghg = np.zeros((64, 128, len(u_ghg))) vwnd_ghg = np.zeros((64, 128, len(u_ghg))) u = np.transpose(u, (1, 2, 0)) v = np.transpose(v, (1, 2, 0)) u_sst = np.transpose(u_sst, (1, 2, 0)) v_sst = np.transpose(v_sst, (1, 2, 0)) u_ghg = np.transpose(u_ghg, (1, 2, 0)) v_ghg = np.transpose(v_ghg, (1, 2, 0)) for i in range(np.ma.size(uwnd, axis=2)): h = interpolate.interp2d(lon, lat[::-1], u[:, :, i]) uwnd[:, :, i] = h(lon42, lat42[::-1]) g = interpolate.interp2d(lon, lat[::-1], v[:, :, i]) vwnd[:, :, i] = g(lon42, lat42[::-1]) for i in range(np.ma.size(uwnd_sst, axis=2)): h = interpolate.interp2d(lon, lat[::-1], u_sst[:, :, i]) uwnd_sst[:, :, i] = h(lon42, lat42[::-1]) g = interpolate.interp2d(lon, lat[::-1], v_sst[:, :, i]) vwnd_sst[:, :, i] = g(lon42, lat42[::-1]) for i in range(np.ma.size(uwnd_ghg, axis=2)): h = interpolate.interp2d(lon, lat[::-1], u_ghg[:, :, i]) uwnd_ghg[:, :, i] = h(lon42, lat42[::-1]) g = interpolate.interp2d(lon, lat[::-1], v_ghg[:, :, i]) vwnd_ghg[:, :, i] = g(lon42, lat42[::-1]) w = VectorWind(uwnd, vwnd) w_sst = VectorWind(uwnd_sst, vwnd_sst) w_ghg = VectorWind(uwnd_ghg, vwnd_ghg) eta = w.absolutevorticity() eta_sst = w_sst.absolutevorticity() eta_ghg = w_ghg.absolutevorticity() div = w.divergence() div_sst = w_sst.divergence() div_ghg = w_ghg.divergence() uchi, vchi = w.irrotationalcomponent() uchi_sst, vchi_sst = w_sst.irrotationalcomponent() uchi_ghg, vchi_ghg = w_ghg.irrotationalcomponent() etax, etay = w.gradient(eta) # etax_sst, etay_sst = w_sst.gradient(eta_sst) # etax_ghg, etay_ghg = w_ghg.gradient(eta_ghg) eta = np.transpose(eta, (2, 0, 1)) eta_sst = np.transpose(eta_sst, (2, 0, 1)) eta_ghg = np.transpose(eta_ghg, (2, 0, 1)) div = np.transpose(div, (2, 0, 1)) div_sst = np.transpose(div_sst, (2, 0, 1)) div_ghg = np.transpose(div_ghg, (2, 0, 1)) etax = np.transpose(etax, (2, 0, 1)) etay = np.transpose(etay, (2, 0, 1)) uchi = np.transpose(uchi, (2, 0, 1)) uchi_sst = np.transpose(uchi_sst, (2, 0, 1)) uchi_ghg = np.transpose(uchi_ghg, (2, 0, 1)) vchi = np.transpose(vchi, (2, 0, 1)) vchi_sst = np.transpose(vchi_sst, (2, 0, 1)) vchi_ghg = np.transpose(vchi_ghg, (2, 0, 1)) #f_ghg = -eta.mean(axis=0)*(div_ghg-div.mean(axis=0))-(uchi_ghg-uchi.mean(axis=0))*etax.mean(axis=0)-(vchi_ghg-vchi.mean(axis=0))*etay.mean(axis=0) #f_sst = -eta.mean(axis=0)*(div_sst-div.mean(axis=0))-(uchi_sst-uchi.mean(axis=0))*etax.mean(axis=0)-(vchi_sst-vchi.mean(axis=0))*etay.mean(axis=0) f_ghg = -eta * (div_ghg - div) - ((uchi_ghg - uchi) * etax + (vchi_ghg - vchi) * etay) f_sst = -eta * (div_sst - div) - ((uchi_sst - uchi) * etax + (vchi_sst - vchi) * etay) meshlon, meshlat = np.meshgrid(lon42, lat42) ''' ncsave('/home/bakerh/Downloads/vort200_control', lat42, lon42, eta.mean(axis=0)-2*np.sin(meshlat*np.pi/180)*7.2921e-5, 'vorticity') ncsave('/home/bakerh/Downloads/vort200_sst', lat42, lon42, eta_sst.mean(axis=0)-2*np.sin(meshlat*np.pi/180)*7.2921e-5, 'vorticity') ncsave('/home/bakerh/Downloads/vort200_ghg', lat42, lon42, eta_ghg.mean(axis=0)-2*np.sin(meshlat*np.pi/180)*7.2921e-5, 'vorticity') ncsave('/home/bakerh/Downloads/forcing_ghg', lat42, lon42, f_ghg.mean(axis=0), 'forcing') ncsave('/home/bakerh/Downloads/forcing_sst', lat42, lon42, f_sst.mean(axis=0), 'forcing') ''' month = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]) ncsave('/home/bakerh/Downloads/vort200_control', month, lat42, lon42, eta - 2 * np.sin(meshlat * np.pi / 180) * 7.2921e-5, 'vorticity') ncsave('/home/bakerh/Downloads/vort200_sst', month, lat42, lon42, eta_sst - 2 * np.sin(meshlat * np.pi / 180) * 7.2921e-5, 'vorticity') ncsave('/home/bakerh/Downloads/vort200_ghg', month, lat42, lon42, eta_ghg - 2 * np.sin(meshlat * np.pi / 180) * 7.2921e-5, 'vorticity') ncsave('/home/bakerh/Downloads/forcing_ghg', month, lat42, lon42, f_ghg, 'forcing') ncsave('/home/bakerh/Downloads/forcing_sst', month, lat42, lon42, f_sst, 'forcing')
def ming_forcing(ua, va, psl): from windspharm.standard import VectorWind from netcdfread import ncread, ncsave from scipy import interpolate lon42 = ncread( '/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc', 'lon') lat42 = ncread( '/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc', 'lat') lon = ncread( '/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/All-Hist/mon/tas/item3236_monthly_mean_a011_2006-01_2016-12.nc', 'longitude0') lat = ncread( '/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/Plus15-Future_LCO2/day/ua/item15201_daily_mean_a00b_2090-01_2100-12.nc', 'latitude1') lat145 = ncread( '/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/All-Hist/mon/tas/item3236_monthly_mean_a011_2006-01_2016-12.nc', 'latitude0') u = ua['All-Nat'][1] v = va['All-Nat'][1] u_sst = ua['GHG-Nat'][1] v_sst = va['GHG-Nat'][1] u_ghg = ua['SST-Nat'][1] v_ghg = va['SST-Nat'][1] p = psl['All-Nat'][1] p_sst = psl['GHG-Nat'][1] p_ghg = psl['SST-Nat'][1] uwnd = np.zeros((64, 128, len(u))) vwnd = np.zeros((64, 128, len(v))) uwnd_sst = np.zeros((64, 128, len(u_sst))) vwnd_sst = np.zeros((64, 128, len(u_sst))) uwnd_ghg = np.zeros((64, 128, len(u_ghg))) vwnd_ghg = np.zeros((64, 128, len(u_ghg))) ps = np.zeros((64, 128, len(p))) ps_sst = np.zeros((64, 128, len(p_sst))) ps_ghg = np.zeros((64, 128, len(p_ghg))) u = np.transpose(u, (1, 2, 0)) v = np.transpose(v, (1, 2, 0)) u_sst = np.transpose(u_sst, (1, 2, 0)) v_sst = np.transpose(v_sst, (1, 2, 0)) u_ghg = np.transpose(u_ghg, (1, 2, 0)) v_ghg = np.transpose(v_ghg, (1, 2, 0)) p = np.transpose(p, (1, 2, 0)) p_sst = np.transpose(p_sst, (1, 2, 0)) p_ghg = np.transpose(p_ghg, (1, 2, 0)) for i in range(np.ma.size(uwnd, axis=2)): h = interpolate.interp2d(lon, lat[::-1], u[:, :, i]) uwnd[:, :, i] = h(lon42, lat42[::-1]) g = interpolate.interp2d(lon, lat[::-1], v[:, :, i]) vwnd[:, :, i] = g(lon42, lat42[::-1]) j = interpolate.interp2d(lon, lat145[::-1], p[:, :, i]) ps[:, :, i] = j(lon42, lat42[::-1]) for i in range(np.ma.size(uwnd_sst, axis=2)): h = interpolate.interp2d(lon, lat[::-1], u_sst[:, :, i]) uwnd_sst[:, :, i] = h(lon42, lat42[::-1]) g = interpolate.interp2d(lon, lat[::-1], v_sst[:, :, i]) vwnd_sst[:, :, i] = g(lon42, lat42[::-1]) j = interpolate.interp2d(lon, lat145[::-1], p_sst[:, :, i]) ps_sst[:, :, i] = j(lon42, lat42[::-1]) for i in range(np.ma.size(uwnd_ghg, axis=2)): h = interpolate.interp2d(lon, lat[::-1], u_ghg[:, :, i]) uwnd_ghg[:, :, i] = h(lon42, lat42[::-1]) g = interpolate.interp2d(lon, lat[::-1], v_ghg[:, :, i]) vwnd_ghg[:, :, i] = g(lon42, lat42[::-1]) j = interpolate.interp2d(lon, lat145[::-1], p_ghg[:, :, i]) ps_ghg[:, :, i] = j(lon42, lat42[::-1]) w = VectorWind(uwnd, vwnd) w_sst = VectorWind(uwnd_sst, vwnd_sst) w_ghg = VectorWind(uwnd_ghg, vwnd_ghg) psx, psy = w.gradient(ps) psx_sst, psy_sst = w_sst.gradient(ps_sst) psx_ghg, psy_ghg = w_ghg.gradient(ps_ghg) omega = np.mean(uwnd * psx + vwnd * psy, axis=2) omega_sst = uwnd_sst * psx_sst + vwnd_sst * psy_sst omega_ghg = uwnd_ghg * psx_ghg + vwnd_ghg * psy_ghg omega_sst = np.transpose(omega_sst, (2, 0, 1)) omega_ghg = np.transpose(omega_ghg, (2, 0, 1)) oforc_sst = np.mean(omega_sst - omega, axis=0) oforc_ghg = np.mean(omega_ghg - omega, axis=0) return omega, oforc_sst, oforc_ghg
def rws_fprime_co2(ua200, va200): from windspharm.standard import VectorWind from netcdfread import ncread, ncsave from scipy import interpolate lon42 = ncread( '/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc', 'lon') lat42 = ncread( '/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc', 'lat') # HadAM3P lon = ncread( '/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/All-Hist/mon/tas/item3236_monthly_mean_a011_2006-01_2016-12.nc', 'longitude0') lat = ncread( '/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/Plus15-Future_LCO2/day/ua/item15201_daily_mean_a00b_2090-01_2100-12.nc', 'latitude1') # MIROC5 # lat = ncread('/network/aopp/hera/mad/bakerh/HAPPI/MIROC5/All-Hist/day/tas/tas_Aday_MIROC5_All-Hist_est1_v2-0_run001_20060101-20161231.nc','lat') # lon = np.arange(0, 360, 360/256) # CAM4 #lat = np.linspace(-90, 90, 96) #lon = np.linspace(0, 357.5, 144) u = ua200['Plus15-Future_LCO2'] v = va200['Plus15-Future_LCO2'] u_ghg = ua200['Plus15-Future_HCO2'] v_ghg = va200['Plus15-Future_HCO2'] uwnd = np.zeros((64, 128, len(u))) vwnd = np.zeros((64, 128, len(v))) uwnd_ghg = np.zeros((64, 128, len(u_ghg))) vwnd_ghg = np.zeros((64, 128, len(u_ghg))) u = np.transpose(u, (1, 2, 0)) v = np.transpose(v, (1, 2, 0)) u_ghg = np.transpose(u_ghg, (1, 2, 0)) v_ghg = np.transpose(v_ghg, (1, 2, 0)) for i in range(np.ma.size(uwnd, axis=2)): h = interpolate.interp2d(lon, lat[::-1], u[:, :, i]) uwnd[:, :, i] = h(lon42, lat42[::-1]) g = interpolate.interp2d(lon, lat[::-1], v[:, :, i]) vwnd[:, :, i] = g(lon42, lat42[::-1]) for i in range(np.ma.size(uwnd_ghg, axis=2)): h = interpolate.interp2d(lon, lat[::-1], u_ghg[:, :, i]) uwnd_ghg[:, :, i] = h(lon42, lat42[::-1]) g = interpolate.interp2d(lon, lat[::-1], v_ghg[:, :, i]) vwnd_ghg[:, :, i] = g(lon42, lat42[::-1]) w = VectorWind(uwnd, vwnd) w_ghg = VectorWind(uwnd_ghg, vwnd_ghg) eta = w.absolutevorticity() eta_ghg = w_ghg.absolutevorticity() div = w.divergence() div_ghg = w_ghg.divergence() uchi, vchi = w.irrotationalcomponent() uchi_ghg, vchi_ghg = w_ghg.irrotationalcomponent() etax, etay = w.gradient(eta) # etax_ghg, etay_ghg = w_ghg.gradient(eta_ghg) eta = np.transpose(eta, (2, 0, 1)) eta_ghg = np.transpose(eta_ghg, (2, 0, 1)) div = np.transpose(div, (2, 0, 1)) div_ghg = np.transpose(div_ghg, (2, 0, 1)) etax = np.transpose(etax, (2, 0, 1)) etay = np.transpose(etay, (2, 0, 1)) uchi = np.transpose(uchi, (2, 0, 1)) uchi_ghg = np.transpose(uchi_ghg, (2, 0, 1)) vchi = np.transpose(vchi, (2, 0, 1)) vchi_ghg = np.transpose(vchi_ghg, (2, 0, 1)) f_ghg = -eta * (div_ghg - div) - ((uchi_ghg - uchi) * etax + (vchi_ghg - vchi) * etay) meshlon, meshlat = np.meshgrid(lon42, lat42) ''' ncsave('/home/bakerh/Downloads/vort200_control', lat42, lon42, eta.mean(axis=0)-2*np.sin(meshlat*np.pi/180)*7.2921e-5, 'vorticity') ncsave('/home/bakerh/Downloads/forcing_ghg', lat42, lon42, f_ghg.mean(axis=0), 'forcing') ''' month = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]) ncsave('/home/bakerh/Downloads/vort200_control', month, lat42, lon42, eta - 2 * np.sin(meshlat * np.pi / 180) * 7.2921e-5, 'vorticity') ncsave('/home/bakerh/Downloads/forcing_ghg', month, lat42, lon42, f_ghg, 'forcing')
def calc_quantity(uwnd, vwnd, quantity, lat_axis, lon_axis, axis_order): """Calculate a single wind quantity. Args: uwnd (numpy.ndarray): Zonal wind vwnd (numpy.ndarray): Meridional wind quantity (str): Quantity to be calculated lat_axis (list): Latitude axis values lon_axis (list): Longitude axis values axis_order (str): e.g. tyx Design: windsparm requires the input data to be on a global grid (due to the spherical harmonic representation used), latitude and longitude to be the leading axes and the latitude axis must go from 90 to -90. The cdms2 interface is supposed to adjust for these things but I've found that things come back upsidedown if the lat axis isn't right, so I've just used the standard interface here instead. Reference: ajdawson.github.io/windspharm """ check_global(lat_axis, lon_axis) # Make latitude and longitude the leading coordinates uwnd, uwnd_info = prep_data(numpy.array(uwnd), axis_order) vwnd, vwnd_info = prep_data(numpy.array(vwnd), axis_order) # Make sure latitude dimension is north-to-south lats, uwnd, vwnd = order_latdim(lat_axis, uwnd, vwnd) flip_lat = False if lats[0] == lat_axis[0] else True w = VectorWind(uwnd, vwnd) data_out = {} if quantity == 'rossbywavesource': eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) data_out['rws1'] = (-eta * div) / (1.e-11) data_out['rws2'] = (-(uchi * etax + vchi * etay)) / (1.e-11) data_out['rws'] = data_out['rws1'] + data_out['rws2'] elif quantity == 'magnitude': data_out['spd'] = w.magnitude() elif quantity == 'vorticity': data_out['vrt'] = w.vorticity() elif quantity == 'divergence': div = w.divergence() data_out['div'] = div / (1.e-6) elif quantity == 'absolutevorticity': avrt = w.absolutevorticity() data_out['avrt'] = avrt / (1.e-5) elif quantity == 'absolutevorticitygradient': avrt = w.absolutevorticity() ugrad, vgrad = w.gradient(avrt) avrtgrad = numpy.sqrt(numpy.square(ugrad) + numpy.square(vgrad)) data_out['avrtgrad'] = avrtgrad / (1.e-5) elif quantity == 'planetaryvorticity': data_out['pvrt'] = w.planetaryvorticity() elif quantity == 'irrotationalcomponent': data_out['uchi'], data_out['vchi'] = w.irrotationalcomponent() elif quantity == 'nondivergentcomponent': data_out['upsi'], data_out['vpsi'] = w.nondivergentcomponent() elif quantity == 'streamfunction': sf = w.streamfunction() data_out['sf'] = sf / (1.e+6) elif quantity == 'velocitypotential': vp = w.velocitypotential() data_out['vp'] = vp / (1.e+6) else: sys.exit('Wind quantity not recognised') # Return data to its original shape for key in data_out.keys(): data_out[key] = recover_structure(data_out[key], flip_lat, uwnd_info) return data_out
def compute_rws(ds_u, ds_v, lat_coord='lat', lon_coord='lon', time_coord='time'): """ Computation of absolute vorticity, divergence, and Rossby wave source. Outputs xarray datasets of each. Args: ds_u (xarray data array): Zonal (u) wind (m/s). ds_v (xarray data array): Meridional (v) wind (m/s). lat_coord (str): Latitude coordinate. Defaults to ``lat``. lon_coord (str): Longitude coordinate. Defaults to ``lon``. time_coord (str): Time coordinate. Defaults to ``time``. Returns: Xarray datasets for absolute vorticity, divergence, and Rossby wave source. """ from windspharm.standard import VectorWind from windspharm.tools import prep_data, recover_data, order_latdim # grab lat and lon coords lats = ds_u.coords[lat_coord].values lons = ds_u.coords[lon_coord].values time = ds_u.coords[time_coord].values _, wnd_info = prep_data(ds_u.values, 'tyx') # reorder dims into lat, lon, time uwnd = ds_u.transpose(lat_coord, lon_coord, time_coord).values vwnd = ds_v.transpose(lat_coord, lon_coord, time_coord).values # reorder lats to north-south direction lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd) # initialize wind vector instance w = VectorWind(uwnd, vwnd) # Absolute vorticity (sum of relative and planetary vorticity). eta = w.absolutevorticity() # Horizontal divergence. div = w.divergence() # Irrotational (divergent) component of the vector wind. uchi, vchi = w.irrotationalcomponent() # Computes the vector gradient of a scalar field on the sphere. etax, etay = w.gradient(eta) # Compute rossby wave source S = -eta * div - (uchi * etax + vchi * etay) # recover data shape S = recover_data(S, wnd_info) div = recover_data(div, wnd_info) eta = recover_data(eta, wnd_info) # assemble xarray datasets data_rws = xr.Dataset({ 'rws': (['time', 'lat', 'lon'], S), }, coords={ 'time': (['time'], time), 'lat': (['lat'], lats), 'lon': (['lon'], lons) }, attrs={'long_name': 'Rossby wave source'}) data_div = xr.Dataset( { 'div': (['time', 'lat', 'lon'], div), }, coords={ 'time': (['time'], time), 'lat': (['lat'], lats), 'lon': (['lon'], lons) }, attrs={'long_name': 'Horizontal divergence (300-mb)'}) data_eta = xr.Dataset( { 'eta': (['time', 'lat', 'lon'], eta), }, coords={ 'time': (['time'], time), 'lat': (['lat'], lats), 'lon': (['lon'], lons) }, attrs={ 'long_name': 'Absolute vorticity (sum of relative and planetary vorticity)' }) return data_eta, data_div, data_rws